
diffusers 如何用 group offloading 在生成视频时减少显存占用【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers用 CogVideoX、Wan 这类视频扩散模型做推理时模型参数动辄数十亿到数十十亿级别单块 GPU 的显存常常放不下。diffusers 的 memory 优化指南 提供了多种减显存手段其中group offloading分组卸载的思路是把模型内部的层组torch.nn.ModuleList或torch.nn.Sequential整体移到 CPU 存放需要计算时才把对应层组加载回 GPU。相比按整块模型搬运的 model offloading它占用的显存更少相比逐叶子层搬运的 CPU offloading它因为减少了设备间通信次数而更快。官方文档同时指出视频生成模型更偏计算密集compute-bound量化叠加 group offloading 的组合在视频生成上通常比图片生成效能更好详见 量化与 offloading 组合指南。先弄清 offload_typeblock_level 还是 leaf_level启用 group offloading 时通过offload_type参数选择粒度block_level按num_blocks_per_group参数把层分组卸载。例如一个 40 层的模型设num_blocks_per_group2就是每次 onload/offload 2 层全程共 20 次往返能大幅降低显存需求。使用block_level时num_blocks_per_group是必填参数见 apply_group_offloading 的参数说明。leaf_level按最细粒度逐层卸载效果等价于 CPU offloading但配合 CUDA stream 后可以在不牺牲推理速度的前提下变快。group offloading 可以施加在整个 pipeline 上也可以只施加到单个模型组件上。整个 pipeline 上启用是最省事的方案按组件启用则允许给不同组件transformer、VAE、text encoder分别配置不同的 offloading 策略。主路径在 pipeline 上启用 group offloading在DiffusionPipeline上调用enable_group_offload即可。下面的例子来自官方文档使用CogVideoXPipeline加载 CogVideoX-5b 视频模型并生成一段视频import torch from diffusers import CogVideoXPipeline from diffusers.utils import export_to_video onload_device torch.device(cuda) offload_device torch.device(cpu) pipeline CogVideoXPipeline.from_pretrained(THUDM/CogVideoX-5b, dtypetorch.bfloat16) pipeline.enable_group_offload( onload_deviceonload_device, offload_deviceoffload_device, offload_typeleaf_level, use_streamTrue ) prompt ( A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The pandas fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The pandas face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance. ) video pipeline(promptprompt, guidance_scale6, num_inference_steps50).frames[0] print(fMax memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB) export_to_video(video, output.mp4, fps8)其中onload_device是层组被加载到的计算设备这里是 CUDAoffload_device是空闲时停放层组的目标设备默认就是 CPU。生成完成后用export_to_video把帧序列导出为output.mp4脚本结尾打印的Max memory reserved一行用于查看本次推理的显存峰值。可选只给个别组件启用 group offloading如果只想让 transformer 走 group offloading、text encoder 用别的策略可以分别调用。继承自ModelMixin的标准 diffusers 组件用enable_group_offload方法不是ModelMixin的普通torch.nn.Module比如很多 text encoder则用apply_group_offloading函数# Use the enable_group_offload method for Diffusers model implementations pipeline.transformer.enable_group_offload(onload_deviceonload_device, offload_deviceoffload_device, offload_typeleaf_level) pipeline.vae.enable_group_offload(onload_deviceonload_device, offload_typeleaf_level) # Use the apply_group_offloading method for other model components apply_group_offloading(pipeline.text_encoder, onload_deviceonload_device, offload_typeblock_level, num_blocks_per_group2)注意block_level分支必须给出num_blocks_per_group上面的示例值为文档示例可按自己的模型层数调整分组大小。用 CUDA stream 重叠传输与计算use_stream参数面向支持异步数据流传输的 CUDA 设备。它通过层预取让传输和计算重叠当前层在 GPU 上执行的同时下一个要执行的层已经开始加载整体耗时低于纯 CPU offloading。代价是 CPU 内存占用会明显上升文档要求你预留约为模型大小 2 倍的 CPU 内存。与 stream 相关的其他要点record_streamTrue可以进一步提速代价是显存占用略增需与use_streamTrue同时设置。使用block_level且开启use_stream时num_blocks_per_group应设为1否则实现会打印警告并自动改为1见 group_offloading.py 中的处理。low_cpu_mem_usageTrue可在 stream 场景下减少 CPU 内存它改为按需创建 pinned tensor 而不预先 pin最适合leaf_level且 CPU 内存吃紧的环境但可能增加整体执行时间。对启用 tiling 的 VAE 使用use_streamTrue时文档建议在正式推理前先用 dummy 输入做一次 forward避免 device mismatch 错误并非所有实现都支持遇到问题可向项目提交 issue。系统内存不够时offload 到磁盘group offloading 会把不活跃的层组存放在系统内存里模型越大占用的 CPU 内存越多。如果机器内存有限可以在enable_group_offload或apply_group_offloading中设置offload_to_disk_path把权重卸载到磁盘作为二级存储文档给出的示例path/to/disk需要替换为你自己的磁盘目录pipeline.transformer.enable_group_offload(onload_deviceonload_device, offload_deviceoffload_device, offload_typeleaf_level, offload_to_disk_pathpath/to/disk) apply_group_offloading(pipeline.text_encoder, onload_deviceonload_device, offload_typeblock_level, num_blocks_per_group2, offload_to_disk_pathpath/to/disk)官方文档还链接了两张速度/内存权衡对比表见 memory 指南中 Offloading to disk 一节可用于评估该方案在你机器上的取舍。验证效果与已知限制验证方式沿用上面的脚本推理结束后打印torch.cuda.max_memory_allocated() / 1024**3得到的 GB 数值即为本次运行的显存峰值导出后的output.mp4用于确认生成质量未受影响。该数值随模型和硬件变化文档未给出固定预期值只作为示例流程展示。使用前的两个必要检查兼容性警告如果某个模型的 forward 实现中包含依赖权重的设备/输入转换可能与 group offloading 的设备转换机制冲突导致不工作。遇到异常时先检查模型 forward 里是否有这类写法。场景匹配官方文档明确 group offloading 的优势在视频生成对图片生成由于计算核往往结束得更快、难以完全重叠数据传输它会带来一些 CPU-GPU 通信开销不如量化加 model offloading 划算。如果显存还需要进一步压缩memory 指南 建议把 group offloading 与 layerwise casting 组合权重以小精度格式如torch.float8_e4m3fn存储、计算时上转到torch.float16或torch.bfloat16同样以CogVideoXPipeline为例官方给出了transformer.enable_layerwise_casting(storage_dtypetorch.float8_e4m3fn, compute_dtypetorch.bfloat16)的完整用法layerwise casting 一节也提示归一化和 modulation 相关层会被跳过存储压缩。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考