← 返回实验室

按阶段运行 Qwen-Image-2.1:stageload 的第二条流水线

pytorchapple-siliconmemorybenchmark

stageload 的第一篇笔记只测量了一条流水线,即 Pixal3D。文生图流水线的结构和它一样:文本编码器、去噪器和图像解码器依次运行。所以10月6日,我测量了 diffusers 里的 Qwen-Image-2.1,用的是同一台 48 GB 的 Mac。它的文本编码器占 16.3 GB,Transformer 占 13.3 GB,VAE 占 1.3 GB。每次运行生成一张 1024 × 1024 的图片,共 20 步,每次都在内存防护程序之下的独立进程里运行。

eager 是常见的 from_pretrained(...).to("mps")。在 staged 模式下,流水线的两个大模型换成了替身。替身从检查点回答 .config,这样流水线在去噪开始之前就能读取配置;其他任何用法都会让替身加载真正的模型。三个钩子在 encode_prompt、prepare_latents 和 _unpack_latents 处标出各阶段的开始。流水线的代码没有改动。

Qwen-Image-2.1 各阶段的峰值占用(48 GB 的 Mac,eager 与 staged)

在编码和去噪阶段,eager 一直在 31 到 33 GB 之间。VAE 解码在权重之外还需要大约 11 GB,这让它达到了 43.6 GB;swap 增加了 8 GB,防护程序在它写出图片之前就把它停了。staged 在解码之前释放了 Transformer。它在编码时峰值为 19.0 GB,在解码时为 14.4 GB,两次运行都没有增加 swap。

staged 加载两个模型用了 13 秒。去噪在它的第一次运行中用了 109 秒,第二次用了 79 秒,eager 用了 77.5 秒;从第一次运行的跟踪记录里看不出它为什么更慢。

我没能证明的

staged 的两次运行生成的图片逐位相同。eager 在这台 Mac 上没有生成图片,所以这些运行无法说明 staged 加载和 eager 生成的图片相同。更早的一次 eager 运行连加载都没有完成:它在一个视频任务结束一分钟后启动,swap 增加了 10 GB,只用了 18 秒。

代码在 bench/qwen_image.py,跟踪记录和汇总在仓库里。