当前位置:首页 > 技术类 > 树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

微信用户2年前 (2023-07-24)技术类1252

编辑:梓文、张倩

Stable Diffusion 能在树莓派上运行了!

11 个月前 Stable Diffusion 诞生,它能够在消费级 GPU 上运行的消息让不少研究者备受鼓舞。不仅如此,苹果官方很快下场,将 Stable Diffusion「塞进」iPhone、iPad 和 Mac 中运行。这大大降低了 Stable Diffusion 对硬件设备的要求,让其逐渐成为人人都能使用的「黑科技」。现在,它甚至已经可以在 Raspberry Pi Zero 2 上运行了。

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技 提供


Raspberry Pi Zero 2 「Just as small. Five times as fast.」这是怎样一个概念?运行 Stable Diffusion 并不是一件容易的事,它包含一个 10 亿参数的大型 Transformer 模型,建议使用的最低 RAM/VRAM 通常为 8GB。而 RPI Zero 2 只是内存为 512MB 的微型计算机。这意味着在 RPI Zero 2 上运行 Stable Diffusion 是一个巨大的挑战。而且,在运行过程中,作者没有增加存储空间,也没有将中间结果卸载到磁盘上。一般而言,主要的机器学习框架和库都专注于最小化推理延迟和 / 或最大化吞吐量,但以上这些都以内存使用为代价。因此,作者决定写一个超小的、可破解的推理库,致力于将内存消耗最小化。OnnxStream 做到了。

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技 提供


OnnxStream 基于将推理引擎与负责提供模型权重的组件解耦的思路,后者是派生自 WeightsProvider 的一个类。一个 WeightsProvider 的专门化可以实现任何类型的模型参数加载、缓存和预取。例如,一个自定义的 WeightsProvider 可以决定直接从 HTTP 服务器下载数据,而不加载或写入任何内容到磁盘(这也是 OnnxStream 命名中有 Stream 的原因)。有两个默认的 WeightsProviders 可用:DiskNoCache 和 DiskPrefetch与微软的推理框架 OnnxStream 相比,OnnxStream 只需要消耗 1/55 的内存就可以达到同样的效果,但(在 CPU 上的)速度只比前者慢 0.5-2 倍。接下来你将看到 Stable Diffusion 在 RPI Zero 2 上运行的效果,以及背后的方法。需要注意的是,虽然运行速度较慢,但是它是大模型在更小、更有限的设备上运行的崭新尝试。

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技 提供


 网友们认为这个项目很酷将 Stable Diffusion 在 Raspberry Pi Zero 2 上运行VAE 解码器是 Stable Diffusion 中唯一无法以单精度或半精度放入 RPI Zero 2 RAM 的模型。这是因为模型中存在残差连接、非常大的张量和卷积。唯一的解决办法就是静态量化(8 bit)。以下这些图像是由作者 repo 中包含的 Stable Diffusion 示例实现在不同精度的 VAE 解码器下使用 OnnxStream 生成的。第一张图像是在作者的 PC 上生成的,使用了由 RPI Zero 2 生成的相同的 latent。

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技 提供


精度为 W16A16 的 VAE 解码器的生成效果

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技 提供


精度为 W8A32 的 VAE 解码器的生成效果第三张图由 RPI Zero 2 在大约 3 小时内生成。

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技 提供


图注:精度为 W8A8 的 VAE 解码器的生成效果OnnxStream 的特点

  • 推理引擎与 WeightsProvider 解耦

  • WeightsProvider 可以是 DiskNoCache、DiskPrefetch 或自定义

  • 注意力切片

  • 动态量化(8 bit 无符号、非对称、百分位数)

  • 静态量化(W8A8 无符号、非对称、百分位数)

  • 轻松校准量化模型

  • 支持 FP16(使用或不使用 FP16 运算)

  • 实现了 24 个 ONNX 算子(最常用的算子)

  • 运算按顺序执行,但所有算子都是多线程的

  • 单一实现文件 + header 文件

  • XNNPACK 调用被封装在 XnnPack 类中 (用于将来的替换)

并且需要注意的是,OnnxStream 依赖 XNNPACK 来加速某些原语:MatMul、Convolution、element-wise Add/Sub/Mul/Div、Sigmoid 和 Softmax。性能对比Stable Diffusion 由三个模型组成:文本编码器(672 次运算和 1.23 亿个参数)、UNET 模型(2050 次运算和 8.54 亿个参数)和 VAE 解码器(276 次运算和 4900 万个参数。假设批大小等于 1,生成完整图像则需要 10 步,这需要运行 2 次文本编码器、运行 20 次(即 2*10)UNET 模型和运行 1 次 VAE 解码器,才能获得良好效果(使用 Euler Ancestral 调度器)。该表显示了 Stable Diffusion 的三个模型不同的推理时间,以及内存消耗(即 Windows 中的 Peak Working Set Size 或 Linux 中的 Maximum Resident Set Size)。可以发现,在 UNET 模型中(以 FP16 精度运行时,OnnxStream 中启用了 FP16 算术),OnnxStream 的内存消耗量仅为 OnnxRuntime 的 1/55,但速度只慢 0.5-2 倍。这次测试需要注明的几点是:

  • OnnxRuntime 的第一次运行是预热推理,因为它的 InferenceSession 是在第一次运行前创建的,并在随后的所有运行中重复使用。而 OnnxStream 没有预热推理,因为它的设计是纯粹「eager」的(不过,后续运行可以受益于操作系统对权重文件的缓存)。

  • 目前 OnnxStream 不支持 batch size ! = 1 的输入,这与 OnnxRuntime 不同,后者在运行 UNET 模型时使用 batch size = 2 可以大大加快整个扩散过程。

  • 在测试中,改变 OnnxRuntime 的 SessionOptions(如 EnableCpuMemArena 和 ExecutionMode)对结果没有产生明显影响。

  • 在内存消耗和推理时间方面,OnnxRuntime 的性能与 NCNN(另一个框架)非常相似。

  • 测试的运行条件:Windows Server 2019、16GB 内存、8750H CPU (AVX2)、970 EVO Plus SSD, VMWare 上的 8 个虚拟内核。

注意力切片与量化在运行 UNET 模型时,采用「注意力切片」技术,并对 VAE 解码器使用 W8A8 量化,这对于将模型内存消耗降低到适合在 RPI Zero 2 上运行的水平至关重要。虽然互联网上有很多关于量化神经网络的信息,但关于「注意力切片」的却很少。这里的想法很简单:目标是在计算 UNET 模型中各种多头注意力的缩放点积注意力时,避免生成完整的 Q @ K^T 矩阵。在 UNET 模型中,注意力头数为 8 时,Q 的形状为 (8,4096,40),同时 K^T 为 (8,40,4096)。因此,第一个 MatMul 的最终形状为 (8,4096,4096),这是一个 512MB 的张量(FP32 精度)。

树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型

树莓派上运行 Stable Diffusion,260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技 提供


解决方案是垂直分割 Q,然后在每个 Q 块上正常进行注意力操作。Q_sliced 形状为 (1,x,40),其中 x 为 4096(在本例中),除以 onnxstream::Model::m_attention_fused_ops_parts(默认值为 2,但可以自定义。这个简单的技巧可以将 UNET 模型以 FP32 精度运行时的整体内存消耗从 1.1GB 降低到 300MB。一个更高效的替代方案是使用 FlashAttention,但是 FlashAttention 需要为每个支持的架构(AVX, NEON)等编写自定义内核,在作者给出的例子中绕过 XnnPack。


免责声明
    本网站在国家相关法律法规规定的范围内,只按现有状况提供文章发布第三方网络平台服务,本网站及其所有者非交易一方,也非交易任何一方之代理人或代表;同时,本网站及其所有者也未授权任何人代表或代理本网站及其所有者从事任何网络交易行为或做出任何承诺、保证或其他类似行为,除非有明确的书面授权。
    鉴于互联网的特殊性,本网站无法鉴别和判断相关交易各主体之民事权利和行为能力、资质、信用等状况,也无法鉴别和判断虚拟交易或正在交易或已交易之虚拟物品来源、权属、真伪、性能、规格、质量、数量等权利属性、自然属性及其他各种状况。因此,交易各方在交易前应加以仔细辨明,并慎重考虑和评估交易可能产生的各项风险。

扫描二维码推送至手机访问。

版权声明:本文由点度点度金讯时代-BLOG发布,如需转载请注明出处。

本文链接:https://lmwmm.com/post/1373.html

标签: 树莓派
分享给朋友:

“树莓派上运行 Stable Diffusion,260MB 的 RAM“hold”住 10 亿参数大模型” 的相关文章

批处理文件扩展名(扩展名为.bat)

批处理文件扩展名(扩展名为.bat)

批处理的介绍 扩展名是bat(在Windows NT/2000/xp/2003/win 7下也可,cmd)的文件就是批处理文件。首先批处理文件是一个文本文件,这个文件的每一行都是一条DOS命令(大部分时候就好象我们在DOS提示符下...

为什么AutoGPT这么火爆?手把手带你安装和体验AI自主机器人

为什么AutoGPT这么火爆?手把手带你安装和体验AI自主机器人

最近几天,一直在忙自己的AIHub.cn网站,目前已经上线了几百款优质AI产品和教程近日,AI界貌似出现了一种新的趋势:自主人工智能。继ChatGPT后,一个叫AutoGPT的开源项目火了,短短几天,在GitHub star达到近4万了,可...

用ChatGPT写假新闻!深圳一自媒体公司法人代表被抓

用ChatGPT写假新闻!深圳一自媒体公司法人代表被抓

ChatGPT问世以来,引发了不小的争议。最让人担心的事发生了——有人用它制作假新闻!此前,一篇题为《今晨甘肃一火车撞上修路工人致9人死亡》的“新闻”吸引了甘肃省平凉市警方的注意,并引发网友关注。经判定,这是一篇假新闻,警方进一步侦查发现,...

chatGPT解析人生八字 2

chatGPT解析人生八字 2

留言区留下你的出生、性别、地址, 隔天,公式分析你的八字人生全系,每日解析按顺序,编排。 由于都是AI机器人智能回答,大家可以作为参考,务必全信,古话说的好,信而不迷,则人生。 命主姓名...

AI孙燕姿翻唱爆火,多亏这个开源项目!广西老表带头打造,上手指南已出

AI孙燕姿翻唱爆火,多亏这个开源项目!广西老表带头打造,上手指南已出

AI孙燕姿这么快翻唱了这么多首歌,到底是咋实现的?关键在于一个开源项目。最近,这波AI翻唱趋势大火,不仅是AI孙燕姿唱的歌越来越多,AI歌手的范围也在扩大,就连制作教程都层出不穷了。而如果在各大教程中溜达一圈后就会发现,其中的关键秘诀,还是...

Windows 12 系统概念版 iSO 镜像下载

Windows 12 系统概念版 iSO 镜像下载

[Download]资源名称:(不包含PE)Windows 12 by sun12文件大小:3.96G下载地址1:https://pan.baidu.com/s/1sWristlACMEpMcQz67zrww?pwd=szpv 下载地址2:...