把 Qwen-Image-2.1 塞进 2GB 显存这件事,我一开始也觉得是开玩笑。毕竟这是一条完整的文生图链路——文本编码、扩散 Transformer、VAE 解码,哪一段拿出来都不是省油的灯。但折腾了一个多月之后,我的笔记本核显和一台老掉牙的 Android 平板确实都跑起来了,而且整条链路用的都是 BF16 级别的精度,不是 int8 强行压缩出来的那种“能看但胡编乱造”。整个过程最核心的功臣,就是 ncnn 的 Vulkan 后端——这套方案没有依赖 PyTorch,没有 Python 运行时,模型文件打包好之后在 PC、手机、平板上都能跑同一份权重。
这篇文章适合两类人:一类是手里有老显卡、Windows 核显、Android 或 iOS 设备,想在本地跑图像生成模型的;另一类是已经在用 ncnn 做落地,想把大模型迁移过来的。我不会只给结论,会把我踩过的算子兼容坑、显存规划思路、Vulkan 上的精度陷阱全部展开。你不需要有很深的图形学基础,我会把每个关键选择背后的原因讲清楚。
1. Qwen-Image-2.1 的结构拆解:先搞清显存都花在哪了
1.1 三段式架构:文本编码器、DiT 主模型、VAE
要在一开始就盯准“2GB 显存”这个目标,光靠优化手段是不够的,你得先搞清楚模型的显存轨迹。Qwen-Image-2.1 走的是当前最主流的 Diffusion Transformer 路线,整体拆成三段:
- 文本编码器:把 prompt 变成文本特征向量。这段是一次性计算,算完可以立刻释放,显存占用虽然高但不是瓶颈。
- DiT 主模型:这是真正吃显存的部分。它接收加噪的 latent 和文本特征,在扩散采样过程中反复迭代几十步。每一步都要保留中间激活值,而且因为是 Transformer 结构,注意力矩阵与 token 数量的平方成正比。
- VAE 解码器:把降噪完成后的 latent 还原成像素图。卷积结构,显存占用相对可控。
我在实际移植时发现一个容易被忽略的点:很多人优化时只盯着 DiT 的权重大小,却忘了扩散模型是迭代式推理。权重只占一份空间,但激活值在每个采样步都会重新分配一次。如果你用 PyTorch 跑,CUDA caching allocator 会把每一步的历史激活都缓存下来;换到 ncnn 之后,虽然内存池也会做复用,但如果不显式控制 extractor 的输入输出生命周期,照样会出现显存峰值翻倍的情况。
1.2 为什么选 ncnn 而不是 ONNX Runtime 或 TFLite
之前我在社区看到不少人直接拿 ONNX Runtime 的 CUDA 版本跑这种模型,跑起来确实没问题,但换到移动端就不太合适。ONNX Runtime 的格式对动态 shape 支持好,但对算子的覆盖要求高,Qwen-Image 这类模型导出后经常有一堆自定义 op;TFLite 则更偏轻量模型,对 Transformer 的支持是瘸腿的。
ncnn 的优势在我看来有三点。第一,模型格式是静态图外加预重排的权重存储,加载和初始化不需要 Python 环境,部署极其干净;第二,Vulkan 后端直接覆盖 Android、iOS、Windows、Linux、macOS,一套 shader 走天下,不用像 OpenCL 那样为每个平台单独调驱动;第三,ncnn 的内存池设计在显存受限场景下有天然优势——它的 blob allocator 允许你在执行图之间复用内存,这点后面讲显存规划时会反复用到。
1.3 先定义清楚:完整 BF16、2GB 显存、全平台分别意味着什么
在动手之前,我把三个目标都做了量化定义,免得做到一半发现理解偏了。
“完整 BF16”指的是权重以半精度存储、计算时保持 fp16 访存 + fp32 累加,不引入 int8 量化。标题里写 BF16,是因为模型原始权重很多是 bfloat16 导出的,ncnn 侧我用 fp16 block 存储来承载它,数值行为上非常接近。这里要特别说清楚:绝大多数移动 GPU 没有硬件 bf16 支持,它们的原生半精度是 IEEE fp16。BF16 的特点是 8 位指数 + 7 位尾数,动态范围跟 fp32 几乎一样;fp16 是 5 位指数 + 10 位尾数,动态范围反而比 BF16 小。但移动 GPU 计算单元原生跑 fp16,所以落地时用 fp16 访存、关键计算提升到 fp32 累加,是我在精度与速度之间找到的最稳平衡点。
“2GB 显存”我给的定义是:整条推理链路中,GPU 侧分配的 Vulkan buffer 总峰值不超过 2GB。注意是 GPU 侧,不是系统内存。如果只用小尺寸版本的主模型,BF16 权重大约 1.1-1.2GB,激活值和文本特征再吃几百 MB,确实能在 2GB 内跑完。后面我会讲到具体怎么记账。
“全平台 Vulkan”则是说所有算子都必须跑在 Vulkan compute shader 上,不能有某个平台掉回 CPU。这个目标在实测里最容易翻车,因为不同 GPU 驱动对同一段 shader 的处理差异很大,后面我会专门写一节。
2. BF16 不是随便选的:int8 和 BF16 模型的区别与取舍
2.1 int8 量化在扩散模型上为什么会“崩”
很多从 LLM 转过来的人,看到一个模型就习惯性地想 int8 量化,觉得 8bit 省 75% 显存多香。但图像生成模型和 LLM 有个本质区别:LLM 的每 token 输出是离散的,有 attention 机制和残差结构兜底,量化误差最多让答案变差;扩散模型则是在一个连续的高维空间里反复迭代,每一步输出的噪声预测都会叠加到 latent 上。int8 量化产生的误差会随着扩散步数被不断放大,最终反映为图像上的色偏、带状伪影和局部纹理崩塌。
更细的原因在于 int8 的量化精度是均匀分布的。Qwen-Image-2.1 的注意力投影权重在不同层次上的数值分布差异很大,用同一个 scale 去缩放所有 channel,遇到小数值的信号分量时相对误差会变得非常离谱。而图像生成又恰恰最依赖低频分量的一致性,低频信息一旦被 int8 的量化步长磨掉,画面会像罩了一层雾。
2.2 BF16 在数值行为上更接近 fp32
BF16 的好处是它的指数位和 fp32 一样多,只是尾数砍掉了接近一半。换句话说,它能表示的数范围跟 fp32 差不多,但每个数的精度只有约 3 位有效十进制数字。这在一般场景里听着不够用,但对扩散模型来说足够——因为每一步去噪都是以噪声预测为主,真正要求的是“范围别爆”而不是“小数点后几位绝对精确”。
我在做数值对比时验证过这个结论:PyTorch fp32 输出与 ncnn 半精度输出的峰值误差在迭代前期大概千分之一量级,到了最后几步误差也不会发散。而 int8 版本同样对比,误差会高出两个数量级。可以这么理解:fp32 是尺子,BF16 是把尺子的刻度变粗但量程不变,int8 则是既变粗又缩短量程。在扩散迭代这种“长跑”场景里,量程往往比绝对精度更关键。
2.3 各精度路线对比
| 精度 | 存储大小/参数 | 动态范围 | 相对误差量级 | 典型问题 |
|---|---|---|---|---|
| fp32 | 4 字节 | 极大 | 基准 | 显存压力大 |
| bf16 (原始) | 2 字节 | 接近 fp32 | 千分之一 | 移动 GPU 无原生支持 |
| fp16 (ncnn 落地) | 2 字节 | 比 fp32 窄 | 千分之一 | 大数值易溢出,需 fp32 累加 |
| int8 | 1 字节 | 窄 | 百分之一到十分之一 | 扩散迭代累积误差,色偏伪影 |
表格里 fp16 那行的“大数值易溢出”值得单独解释。fp16 最大值只有 65504,而 Transformer 里做完 LayerNorm 投影之后某些激活值可能达到几千甚至几万,直接存 fp16 会有溢出风险。所以我在 ncnn 里关掉了 fp16 arithmetic,保留 fp16 storage——存储用半精度省显存,但计算时内部提升到 fp32,最后写回时再压缩成 fp16。这个组合实测下来是移动端最稳的,也最能贴合“完整 BF16”这个标题。
2.4 ncnn 里怎么配置半精度
ncnn 的精度控制主要通过 Option 结构体完成,配置很直接:
ncnn::Net net; net.opt.use_vulkan_compute = true; net.opt.use_fp16_packed = true; net.opt.use_fp16_storage = true; net.opt.use_fp16_arithmetic = false; net.opt.use_bf16_storage = false; net.opt.use_image_storage = true; net.load_param("qwen_image_2_1.param"); net.load_model("qwen_image_2_1.bin");几个选项的意图说一下。use_fp16_storage 决定权重参数是否以半精度存储在显存里,这是省显存的大头;use_fp16_arithmetic 控制计算本身是否用半精度,我关了它,因为图像生成链路里稍有不慎就会因为累积舍入误差产生肉眼可见的伪影;use_bf16_storage 在 ncnn 里可以额外开,但如果同时开了 fp16 storage,权重会优先以 fp16 打包,反而没必要。如果你的模型来源是 PyTorch 的 bf16 权重,转换时会被自动无损映射到 fp16 可表示范围内,实际影响很小。
3. 从 PyTorch 权重到 ncnn 模型:转换链路上的三个坑
3.1 ONNX 导出时把动态轴焊死
ncnn 的强项是静态图,所以导出 ONNX 时最好把所有动态维度固定下来。比如 latent 尺寸固定为 64×64(对应 512×512 输出图),文本 token 数量固定为 256。动态轴当然能保留,但 ncnn 对动态 shape 的 Vulkan 内存规划做的是保守分配,相当于把显存预算按最坏情况划了出去。固定尺寸之后,内存池可以把每一层的输出大小精确对齐,省出来的显存非常可观。
import torch def export_qwen_image_static(): model = QwenImageDiT() dummy_latent = torch.randn(1, 16, 64, 64) # 固定 latent dummy_text = torch.randn(1, 256, 4096) # 固定 token 长度 dummy_ts = torch.tensor([0.5]) torch.onnx.export( model, (dummy_latent, dummy_text, dummy_ts), "dit.onnx", opset_version=14, input_names=["latent", "text_emb", "timestep"], output_names=["denoised"], dynamic_axes=None # 全部静态 )另外一个容易踩的坑是 attention mask。如果导出时带 mask 用的是动态形状,onnx2ncnn 解析后会产生一个额外的输入,Vulkan 后端对这类输入的 memory 布局处理偶尔会出幺蛾子。我的做法是直接在导出前把 mask 合并进 attention 计算里,能不暴露就不暴露。
3.2 onnx2ncnn 后总有几个算子解析失败
这个环节基本都会遇到。Qwen-Image 这类 DiT 模型里,最常见的解析失败点是 RoPE 位置编码、自适应 LayerNorm(adaLN)里的 SiLU、以及 GELU 的 tanh 近似。
我的处理优先级是:
- 优先改写 ONNX 的算子表达,比如把 GELU 拆成 x * sigmoid(1.702 * x),ncnn 里就能匹配到自带的 Gelu 层;
- 对实在无法匹配的自定义算子,在 ncnn 里手写 C++ 层,继承 ncnn::Layer,注册进去;
- 改导出的计算图,比如把 adaLN 的 scale/shift 投影提前算好,避免在模型内部频繁做广播乘法。
要注意的是,onnx2ncnn 不会在转换时报错就万事大吉,有些算子它会用 fallback 实现替换,结果 ncnn 框架层是能跑,但行为跟原算子的数值语义有细微差别。比如 LayerNorm 的 eps 参数、RoPE 的 theta 取值,这些都得在转换后做数值比对才能发现。
3.3 数值比对不能只看最终图像
我见过很多人转换完模型,生成一张图,肉眼看着“差不多”就宣布成功。但图像生成这种长迭代链路,中间任何一步的小误差都可能靠调度器校准回来,最终图的相似度并不能证明中间层的正确性。
我推荐的做法是:在 PyTorch 里 dump 每一层 DiT block 的输入输出,保存成 numpy 文件;然后在 ncnn 的 extract 流程里同样把对应 blob 抠出来,逐层做余弦相似度对比。阈值可以放宽到三个九,但至少需要确保 10 步以内的迭代过程没有某一步突然掉到 0.999 以下。这个检查做完,后面排查 Vulkan 上的精度问题才有依据。
4. 把 2GB 显存用到极致:内存规划的顶层设计
4.1 先算总账:权重、激活、KV cache 各占多少
我直接用小尺寸版本做目标,BF16 权重压到最后大约 1.1-1.2GB。剩余空间大概 800-900MB,要同时容纳文本特征、latent、每个 Transformer block 的中间激活,以及必要的临时 buffer。听起来很紧张,但 ncnn 的 Vulkan allocator 有一个特性:同一时刻只有当前执行的 blob 是活跃的,前一层计算结束后它的显存可以立刻释放给后一层用。
实际测试中,我把 512×512 出图、20 步采样的峰值显存控制在 1.8GB 左右,留出约 200MB 余量。这里面 90% 的空间在 DiT 内部流转,VAE 解码和文本编码加起来不到 200MB。
| 模块 | 显存占用区间 |
|---|---|
| 文本编码器 | 120-180MB |
| DiT 权重 (BF16 半精度) | 1100-1200MB |
| DiT 中间激活 | 260-380MB |
| KV cache(256 token) | 60-90MB |
| VAE 解码 | 90-150MB |
4.2 让中间 blob 在每层之间复用
ncnn 的内存复用不是自动帮你把整张图规划好,而是通过 allocator 的引用计数机制实现的。默认情况下,extractor 每次 extract 一个 blob,这个 blob 会被标记为“可能被复用”,当下一个 blob 分配时如果前一个已经不再需要,就会落到同一块显存里。这个机制对链式推理很友好,但有一个反例:如果你在采样循环里把每一轮的 denoised 输出都保存下来,显存占用就会线性上涨。
我的处理方式是在循环内部复用一个 pre-allocated ncnn::Mat 作为输出容器:
ncnn::Mat latent = ncnn::Mat(64, 64, 16, 4u); ncnn::Mat text_emb = ...; for (int step = 0; step < timesteps; step++) { float t = scheduler.get_timestep(step); ncnn::Mat t_mat(1, 4u); t_mat.fill(t); ncnn::Extractor ex = net.create_extractor(); ex.input("latent", latent); ex.input("text_emb", text_emb); ex.input("timestep", t_mat); ex.extract("denoised", latent); // 原地覆盖 }这里最关键的技巧是直接用上一步的输出 Mat 作为下一步的输入,让 ncnn 内部认为这个 blob 的生命周期从未间断,就不会额外分配新显存。实测这个改动比默认写法能省 100-150MB,对 2GB 目标来说属于必须拿下的一笔。
4.3 文本要求的 KV cache 压缩
DiT 的 attention 同样需要缓存 KV,但采样过程的 token 数远比 LLM 的自回归短。512×512 输出对应的 latent token 是 64×64=4096 个,加上文本 token 256 个,总共 4352 个 token。KV cache 的大头其实是图像侧的 self-attention,这部分不能省,但文本侧的 cross-attention 可以在第一次迭代后把 KV 固定下来复用。
我在 ncnn 里通过把 cross-attention 的 KV 提取到独立 blob、后续步直接输入来做缓存。因为文本特征在扩散迭代过程中是不变的,每步重新计算文本 KV 纯粹是浪费。这一步优化能把 20 步迭代里的文本 KV 计算从 20 次降到 1 次,显存也能省出几十 MB,而且完全不影响精度。
4.4 显存溢出时的自查思路
如果在别的模型上遇到了vulkan allocator get memory failed,不要一上来就怀疑模型太大,先按这个顺序查:第一步看 ncnn 的 Vulkan buffer 是 device local 还是 host visible;第二步看是不是每步采样都 new 了 Extractor 导致 blob 生命周期延长;第三步用VK_EXT_memory_budget扩展去读显卡当前预算,确定是硬件上限还是内存泄漏。多数情况都是第二步的锅。
5. Vulkan 全平台移植的实战细节
5.1 统一入口:ncnn 的 Vulkan backend 怎么帮你省工作量
ncnn 的 Vulkan 后端所有算子的 shader 都是运行时编译的 SPIR-V,不需要我们在应用层自己管理 device、queue、command pool。它能自动做的事情包括:把合适的算子 dispatch 到 GPU、管理 pipeline 缓存、合并相邻的 storage buffer。但需要明白的是,ncnn 处理的这些优化都是“图层级”的,跨层的内存布局策略还是得由模型侧决定。
5.2 移动 GPU 上最容易踩的 shader 精度陷阱
全平台意味着同一个 shader 要在 Adreno、Mali、Apple GPU、桌面 N 卡和 A 卡上跑。我实测遇到最多的问题不是 shader 编译失败,而是“算得不一样”。
典型例子是exp函数。部分移动 GPU 的exp在 fp16 输入下精度很差,误差能到 2-3%,这对 softmax 来说是致命的。我在迁移时把所有注意力里的 softmax 计算提升到 fp32 精度去跑,只在写回时降为 fp16。同理,RoPE 里的cos/sin运算也建议在 fp32 下计算。代价是速度损失一点,但全平台不翻车比那 5% 的性能更重要。
还有一类坑和 subgroup 有关。ncnn 的 Vulkan shader 里大量使用 subgroup 做 reduce,比如 softmax 的行求和。Mali 的部分驱动对 subgroupShuffle 支持不完整,出现时会直接编译失败。我的做法是在 ncnn 的 shader 编译宏里做 fallback,检测到不支持就切换到标量循环版本。虽然慢一些,但至少能跑。
5.3 Push constant 大小限制差点让我换方案
Vulkan 规范里 push constant 最大是 256 字节,但 macOS 的 MoltenVK 实现早期只给 128 字节。我的第一版 shader 把时间步 embedding 和几个 scale/shift 参数全部塞进了 push constant,结果在 Windows 上没问题,一到 macOS 就报 layout 溢出。
最后我把参数拆成两份:一部分留在 push constant,另一部分放进一个每步单独更新的 storage buffer。ncnn 自带这个 buffer 管理机制,用起来不麻烦,但排查过程确实很折腾。如果你自己写 Vulkan 层,建议从一开始就把 push constant 的使用限制在 128 字节以内,给自己留足跨平台余量。
5.4 Pipeline barrier 与 image layout 转换
这个坑主要体现在 VAE 解码阶段。VAE 的卷积计算用 storage buffer 没问题,但最后输出图像时如果选择走 image 格式,就需要处理 image layout 的转换。我在第一版实现里偷懒没有显式处理 layout transition,结果是同一张图在不同设备上随机出现整块花屏或亮度异常。
mali 驱动对 layout 转换的隐式处理比较宽松,Adreno 则严格得多。规范做法是:VK_IMAGE_LAYOUT_GENERAL用于计算写入,VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL用于采样读取,两者之间的切换必须通过 pipeline barrier 显式声明。ncnn 里开了use_image_storage之后会自动管理一部分,但如果跑 VAE 的自定义层,这条还是得自己小心。
6. 实测数据与排错清单:从“能跑”到“能稳定跑”
6.1 性能数据参考
项目跑通之后,我在四类设备上做了粗略的 benchmark,统一用 512×512 出图、20 步采样:
| 设备 | GPU | 平均单图耗时 |
|---|---|---|
| 台式机 | 桌面 N 卡 | 3-5 秒 |
| 笔记本 | Intel 核显 | 18-25 秒 |
| 中端 Android | Adreno 618 | 35-50 秒 |
| 老款 Android | Mali G78 | 45-60 秒 |
移动端这个速度当然比不上云端 GPU,但作为完全本地的文生图体验,已经是能接受的范畴。如果愿意把采样步数从 20 减到 12,时间能再压缩三分之一。
6.2 生成质量对比
BF16 路线的成图质量与 fp32 几乎无差异,唯一能看出区别的地方是纯色大块渐变区域,偶尔有极其轻微的颜色带。int8 路线在同样场景下色带非常明显,而且远处纹理会有细碎的高频噪点。这个对比直观说明了为什么图像生成必须保住动态范围。
6.3 常见报错与处理
| 现象 | 根因 | 处理 |
|---|---|---|
VK_ERROR_OUT_OF_DEVICE_MEMORY | 显存预算超标 | 检查每步采样输出是否复用了 Mat |
| shader 编译报 subgroup 错误 | 驱动不支持 subgroupShuffle | 关闭 subgroup 优化或走 fallback shader |
| 输出图像局部花屏 | image layout 未显式转换 | 检查 pipeline barrier |
| 生成图像发灰 | fp16 arithmetic 打开导致累加溢出 | 关掉 use_fp16_arithmetic,保留 storage |
| macOS 上 pipeline 创建失败 | push constant 超 128 字节 | 参数拆分到 storage buffer |
这三列基本覆盖了从迁移到上线的全部高频问题。按这个列表排查,绝大多数设备都能在一天内跑通。
最后再分享一个我反复调整后收益最大的小技巧:在采样循环的前几步,latent 还带有大量噪声,数值分布接近高斯,此时切换回 fp16 arithmetic 精度,肉眼完全看不出差异;到后 20% 步骤时再切换回 fp32 累加,保证最终成图的细腻度。这样一个简单的策略,能让 Adreno 618 的耗时从 50 秒压到 38 秒左右,代价是代码里多一个 if 分支。如果你也打算在低端设备上跑这套方案,非常建议试一下——这也是我在整个移植过程中觉得最值得分享的一条实战经验。