DiffSynth-Studio Offload Training 全解析基于 Hook 机制的逐层 CPU 卸载训练指南【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio导读本文面向需要在有限显存VRAM环境下训练大模型的开发者系统讲解 DiffSynth-Studio 提供的 Offload Training逐层 CPU 卸载训练特性通过 PyTorch Module Hook 机制在任意时刻只把当前计算模块的权重加载到 GPU计算完成立即卸载回 CPU从而把权重显存占用从 O(N × 每层参数量) 降至 O(1 × 每层参数量)全程无需修改模型代码。读完本文你将掌握三个命令行参数--enable_model_cpu_offload、--enable_optimizer_cpu_offload、--cpu_offload_split_threshold的完整用法、底层 Offloader 分类与 Pinned 内存池原理以及如何把 Offload Training 模块集成到自己的训练代码中。注意Offload Training 目前仅支持单卡训练与多卡DDP场景不兼容。相关限制详见 Offload_Training.md 原文。一、为什么需要 Offload Training训练大规模扩散模型时全部层权重必须同时驻留 GPU仅权重一项就会消耗数十 GB 显存。例如 Qwen-Image 有 60 层、Wan2.1-14B 有 40 层权重显存开销巨大使得小显存单卡几乎无法启动训练。Offload Training 的核心思想是任意时刻只把当前正在计算的模块权重放到 GPU 上计算完成后立刻卸载回 CPU将显存占用从 O(N × params_per_layer) 降到 O(1 × params_per_layer)N 为层数。这一特性基于 PyTorch 的 Module Hook 机制实现不需要修改任何模型代码因此可以无缝套用到 DiffSynth-Studio 已支持的各种模型与训练脚本上。二、工作原理2.1 核心机制四个 Hook 的生命周期OffloadTrainingManager会扫描模型并为每个被管理的模块注册 4 个 Hook见 manager.py 中_register_hooks的实现forward_pre_hook → 将模块权重从 CPU 加载到 GPUonload module.forward() → 正常前向计算 forward_hook → 将模块权重从 GPU 卸载回 CPUoffload backward_pre_hook → 重新将模块权重从 CPU 加载到 GPUonload module.backward() → 计算梯度 backward_hook → 将模块权重卸载回 CPUoffload也就是说一个模块在前向 → 卸载 → 反向 → 卸载的循环中反复被搬运GPU 上始终只保留一个模块的权重。需要说明的是非叶子模块的 backward hook 会在子模块 backward 完成之前触发因此 manager.py 会将backward_hook注册到叶子子模块上而不是注册在父模块自身。2.2 参数与 Buffer 的分类卸载策略OffloadTrainingManager根据参数是否可训练以及 Buffer 类型为每个参数选择不同的 Offloader分类逻辑见 manager.py各 Offloader 实现见 offloader.py类型Offloader 类行为不可训练参数requires_gradFalseStaticParamOffloader初始化时把权重拷贝到预分配的 Pinned 内存保留一份永久的 CPU 副本并用空的 GPU 占位符替换param.data释放 GPU 显存onload 时异步从 CPU 拷贝到 GPUoffload 时直接把param.data重新赋值为占位符不产生 PCIe 回传可训练参数 enable_optimizer_cpu_offloadTrueTrainableParamOffloader权重在训练中不断变化因此不保留静态副本onload/offload 通过param.data.to(device)做真实数据搬运backward 之后还会把param.grad移到 CPU可训练参数 enable_optimizer_cpu_offloadFalseAlwaysOnGPUParamOffloader初始化时把参数搬到 GPU 后永不下移适合 LoRA 训练可训练参数数量很少模块 Buffer如 BatchNorm 的running_mean/running_varBufferOffloader与StaticParamOffloader类似初始化时把 buffer 拷贝到 Pinned 内存onload 异步拷贝到 GPUoffload 时把module._buffers[name]重新赋值为 CPU 副本从实现上看StaticParamOffloader在初始化时通过param.data self._placeholder一个空的 GPU 张量替换原始权重从而把权重占据的显存立即归还AlwaysOnGPUParamOffloader则只在__init__里做一次param.data.to(target_device)。这解释了为什么关闭 optimizer offload 时 LoRA 训练依然可行LoRA 注入后的可训练参数极少常驻 GPU 开销可忽略。2.3 Pinned 内存池消灭逐张量分配浪费StaticParamOffloader和BufferOffloader需要为每个不可训练参数/Buffer 在 CPU 侧保留一份 Pinned 内存副本Pinned 内存支持异步非阻塞的 CPU→GPU 传输远快于普通可分页内存。问题PyTorch 的pin_memory()经由CachingHostAllocator分配内存会把每次分配的大小向上取整到下一个 2 的幂。例如一个 17MB 的张量实际会分配 32MB。大模型有数千个参数张量逐个独立pin_memory()会产生巨大浪费实测内存膨胀 50%100%。解决方案PinnedArenaPool见 memory_buffer.py预分配少量大块 Pinned 内存即 arena——预分配的大块内存区域所有小对象从中切分再用 bump-pointer 分配器紧凑地为每个张量切出空间避免逐张量取整浪费from_model()扫描模型中所有不可训练参数和 Buffer计算总大小memory_buffer.py将总大小分解为若干 2 的幂大小的块每个块是一个PinnedBuffermemory_buffer.py分配时顺序探测各块剩余空间bump-pointer 推进完成分配只做 64 字节对齐无取整浪费ALIGNMENT 64空间不足时自动增长新块_grow任何异常时回退到逐张量pin_memory()BaseBufferPool.allocate_like。此外PinnedBuffer采用惰性分配真正的pin_memory在第一次allocate_like时才触发进一步减少初始化开销。2.4 与 Gradient Checkpointing 的兼容Gradient Checkpointing 在反向传播时会重新执行前向重算激活这会再次触发forward_hook。该问题通过_in_recompute集合解决manager.py第一次前向正常 offload模块被加入_in_recompute重算前向反向过程中检测到模块在_in_recompute中跳过 offload让权重留在 GPU 上供反向使用调用after_backward()时清空_in_recompute为下一步做好准备。注意这与 gradient_checkpoint.py 中use_reentrantFalse的 checkpoint 用法相匹配_in_recompute机制正是为这种非重入式 checkpoint 设计的。2.5 Hook 注册粒度与孤儿参数默认情况下OffloadTrainingManager在叶子模块粒度注册 Hooknn.Linear、nn.LayerNorm等每个叶子模块被独立 onload/offloadmanager.py。同时不被任何叶子模块管理的孤儿参数和孤儿 Buffer会被自动收集并单独挂钩manager.py确保模型中的每个参数都不会漏管。实验性cpu_offload_split_threshold参数单位 MB可调整 Hook 注册粒度。设置后总参数量超过阈值的模块会被递归拆分到子模块低于阈值的模块整体挂钩。从 manager.py 的_should_force_recurse看除参数量超过阈值外未定义forward的容器模块、以及同时具备encode/decode的模块也会被强制递归。该特性在当前版本可能与部分模型架构不兼容默认关闭。2.6 训练循环中的集成执行流程见 runner.py# 当 enable_model_cpu_offloadTrue 时 # 1. 模型不调用 model.to(device)保持在 CPU 上 # 2. 只 prepare optimizer、dataloader、scheduler模型不参与 prepare # 3. 创建 OffloadTrainingManager自动在模型上注册 Hook # 训练循环 loss model(data) accelerator.backward(loss) offload_manager.after_backward() # 清空重算标记 将梯度移到 CPU optimizer.step() optimizer.zero_grad()OffloadTrainingManager(model, accelerator.device, enable_optimizer_cpu_offload, cpu_offload_split_threshold)的三个参数与命令行一一对应其中cpu_offload_split_threshold会在构造函数中乘以1024*1024转换为字节数manager.py。此外launch_data_process_taskSplit Training 的数据处理任务也支持同样的三个参数runner.py。三、使用方法3.1 参数说明命令行参数在 parsers.py 的add_offload_training_config中定义参数默认值说明--enable_model_cpu_offloadFalse开启逐层卸载训练。权重保存在 CPU逐层加载到 GPU--enable_optimizer_cpu_offloadFalse需与--enable_model_cpu_offload配合使用把可训练参数与优化器放到 CPU--cpu_offload_split_thresholdNone实验性单位 MB总参数量超过该阈值的模块被递归拆分launch_training_task会从args中读取这三个参数并传给OffloadTrainingManagerrunner.py。3.2 参数组合效果场景--enable_model_cpu_offload--enable_optimizer_cpu_offload效果默认训练❌❌全部权重与优化器都在 GPU 上仅卸载不可训练参数✅❌不可训练参数逐层卸载可训练参数与优化器留在 GPU卸载全部参数✅✅所有参数逐层卸载梯度与优化器运行在 CPU3.3 实战示例Qwen-Image LoRA 训练只需在原有训练命令中追加--enable_model_cpu_offload。以 Qwen-Image LoRA 训练为例基于 examples/qwen_image/model_training/train.py 的命令行接口accelerate launch examples/qwen_image/model_training/train.py \ --dataset_base_path data/example_dataset \ --dataset_metadata_path data/example_dataset/metadata.json \ --max_pixels 1048576 \ --dataset_repeat 50 \ --model_id_with_origin_paths Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors,Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors \ --learning_rate 1e-4 \ --num_epochs 5 \ --remove_prefix_in_ckpt pipe.dit. \ --output_path ./models/train/Qwen-Image_lora \ --lora_base_model dit \ --lora_target_modules to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_out.0,to_add_out,img_mlp.net.2,img_mod.1,txt_mlp.net.2,txt_mod.1 \ --lora_rank 32 \ --use_gradient_checkpointing \ --dataset_num_workers 8 \ --find_unused_parameters \ --enable_model_cpu_offload若需要完整卸载优化器也放到 CPU追加--enable_model_cpu_offload \ --enable_optimizer_cpu_offload值得说明的是开启--enable_model_cpu_offload后模型初始化设备会被切换为 CPU——在 train.py 中可见devicecpu if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device这与权重常驻 CPU、由 Hook 管理的设计一致。3.4 兼容性特性兼容性说明Gradient Checkpointing✅_in_recompute机制处理重算Accelerate DDP多卡⚠️enable_model_cpu_offload模式下模型不被 DDP 包装无accelerator.prepare(model)不执行梯度 allreduce。多卡训练兼容性无法保证每张卡独立训练、不做梯度同步Split Training✅launch_data_process_task同样支持--enable_model_cpu_offloadDeepSpeed❌ZeRO 的参数收集与 Hook 冲突3.5 使用注意事项开启--enable_model_cpu_offload后模型永远不会调用model.to(device)权重完全由 Hook 管理由于存在 CPU↔GPU 搬运训练速度会下降通常慢 210 倍模型越大减速越明显适合显存受限的场景建议与--use_gradient_checkpointing配合使用进一步降低激活显存--enable_optimizer_cpu_offload只支持梯度累积步数为 1--gradient_accumulation_steps 1--enable_optimizer_cpu_offload生效时TrainableParamOffloader会在 backward 后把param.grad移到 CPUoffload_grad所以after_backward()必须在optimizer.step()之前调用runner.py。四、在其他代码库中集成 Offload Training 模块Offload Training 模块相对独立diffsynth/core/offload_training 下仅有manager.py、offloader.py、memory_buffer.py三个核心文件对外只暴露OffloadTrainingManager可以很方便地集成到其他训练代码中。下面是一个仅约 4GB 显存的普通训练示例import torch from tqdm import tqdm class ToyModel(torch.nn.Module): def __init__(self): super().__init__() self.layers torch.nn.ModuleList(torch.nn.Linear(4096, 4096) for _ in range(10)) def forward(self, x): for layer in self.layers: x x layer(torch.nn.functional.layer_norm(x, (4096,))) return x model ToyModel().to(cuda) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) pbar tqdm(range(100)) for i in pbar: x torch.randn((512, 4096), devicecuda) y x 1 y_pred model(x) loss torch.nn.functional.mse_loss(y_pred, y) loss.backward() optimizer.step() optimizer.zero_grad() pbar.set_postfix(lossf{loss.item():.4f})启用 Offload Training 后显存占用降到约 1.4GB改动极小import torch from tqdm import tqdm from diffsynth.core import OffloadTrainingManager class ToyModel(torch.nn.Module): def __init__(self): super().__init__() self.layers torch.nn.ModuleList(torch.nn.Linear(4096, 4096) for _ in range(10)) def forward(self, x): for layer in self.layers: x x layer(torch.nn.functional.layer_norm(x, (4096,))) return x model ToyModel().to(cpu) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) offload_manager OffloadTrainingManager(model, target_devicecuda, enable_optimizer_cpu_offloadTrue) pbar tqdm(range(100)) for i in pbar: x torch.randn((512, 4096), devicecuda) y x 1 y_pred model(x) loss torch.nn.functional.mse_loss(y_pred, y) loss.backward() offload_manager.after_backward() optimizer.step() optimizer.zero_grad() pbar.set_postfix(lossf{loss.item():.4f})对比两段代码可以看出集成要点模型保持在 CPUmodel.to(cpu)用OffloadTrainingManager替代model.to(cuda)并在loss.backward()之后、optimizer.step()之前调用offload_manager.after_backward()。其余训练逻辑数据、优化器、进度条完全不变。五、相关文档与源码导航官方文档Offload_Training.md中文版见 docs/zh/Training/Offload_Training.md核心实现manager.pyOffloadTrainingManager、UnitWiseHookManager、UnitWiseParamManager、offloader.py四类 Offloader、memory_buffer.pyPinnedArenaPool、PinnedBuffer、BaseBufferPool训练入口runner.pylaunch_training_task、launch_data_process_task参数定义parsers.pyadd_offload_training_config训练脚本示例examples/qwen_image/model_training/train.py更多模型Wan、FLUX、LTX-2、MiniMax 等的训练脚本位于 examples 对应目录下【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 DiffSynth-Studio Offload Training 全解析：基于 Hook 机制的逐层 CPU 卸载训练指南
DiffSynth-Studio Offload Training 全解析基于 Hook 机制的逐层 CPU 卸载训练指南【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio导读本文面向需要在有限显存VRAM环境下训练大模型的开发者系统讲解 DiffSynth-Studio 提供的 Offload Training逐层 CPU 卸载训练特性通过 PyTorch Module Hook 机制在任意时刻只把当前计算模块的权重加载到 GPU计算完成立即卸载回 CPU从而把权重显存占用从 O(N × 每层参数量) 降至 O(1 × 每层参数量)全程无需修改模型代码。读完本文你将掌握三个命令行参数--enable_model_cpu_offload、--enable_optimizer_cpu_offload、--cpu_offload_split_threshold的完整用法、底层 Offloader 分类与 Pinned 内存池原理以及如何把 Offload Training 模块集成到自己的训练代码中。注意Offload Training 目前仅支持单卡训练与多卡DDP场景不兼容。相关限制详见 Offload_Training.md 原文。一、为什么需要 Offload Training训练大规模扩散模型时全部层权重必须同时驻留 GPU仅权重一项就会消耗数十 GB 显存。例如 Qwen-Image 有 60 层、Wan2.1-14B 有 40 层权重显存开销巨大使得小显存单卡几乎无法启动训练。Offload Training 的核心思想是任意时刻只把当前正在计算的模块权重放到 GPU 上计算完成后立刻卸载回 CPU将显存占用从 O(N × params_per_layer) 降到 O(1 × params_per_layer)N 为层数。这一特性基于 PyTorch 的 Module Hook 机制实现不需要修改任何模型代码因此可以无缝套用到 DiffSynth-Studio 已支持的各种模型与训练脚本上。二、工作原理2.1 核心机制四个 Hook 的生命周期OffloadTrainingManager会扫描模型并为每个被管理的模块注册 4 个 Hook见 manager.py 中_register_hooks的实现forward_pre_hook → 将模块权重从 CPU 加载到 GPUonload module.forward() → 正常前向计算 forward_hook → 将模块权重从 GPU 卸载回 CPUoffload backward_pre_hook → 重新将模块权重从 CPU 加载到 GPUonload module.backward() → 计算梯度 backward_hook → 将模块权重卸载回 CPUoffload也就是说一个模块在前向 → 卸载 → 反向 → 卸载的循环中反复被搬运GPU 上始终只保留一个模块的权重。需要说明的是非叶子模块的 backward hook 会在子模块 backward 完成之前触发因此 manager.py 会将backward_hook注册到叶子子模块上而不是注册在父模块自身。2.2 参数与 Buffer 的分类卸载策略OffloadTrainingManager根据参数是否可训练以及 Buffer 类型为每个参数选择不同的 Offloader分类逻辑见 manager.py各 Offloader 实现见 offloader.py类型Offloader 类行为不可训练参数requires_gradFalseStaticParamOffloader初始化时把权重拷贝到预分配的 Pinned 内存保留一份永久的 CPU 副本并用空的 GPU 占位符替换param.data释放 GPU 显存onload 时异步从 CPU 拷贝到 GPUoffload 时直接把param.data重新赋值为占位符不产生 PCIe 回传可训练参数 enable_optimizer_cpu_offloadTrueTrainableParamOffloader权重在训练中不断变化因此不保留静态副本onload/offload 通过param.data.to(device)做真实数据搬运backward 之后还会把param.grad移到 CPU可训练参数 enable_optimizer_cpu_offloadFalseAlwaysOnGPUParamOffloader初始化时把参数搬到 GPU 后永不下移适合 LoRA 训练可训练参数数量很少模块 Buffer如 BatchNorm 的running_mean/running_varBufferOffloader与StaticParamOffloader类似初始化时把 buffer 拷贝到 Pinned 内存onload 异步拷贝到 GPUoffload 时把module._buffers[name]重新赋值为 CPU 副本从实现上看StaticParamOffloader在初始化时通过param.data self._placeholder一个空的 GPU 张量替换原始权重从而把权重占据的显存立即归还AlwaysOnGPUParamOffloader则只在__init__里做一次param.data.to(target_device)。这解释了为什么关闭 optimizer offload 时 LoRA 训练依然可行LoRA 注入后的可训练参数极少常驻 GPU 开销可忽略。2.3 Pinned 内存池消灭逐张量分配浪费StaticParamOffloader和BufferOffloader需要为每个不可训练参数/Buffer 在 CPU 侧保留一份 Pinned 内存副本Pinned 内存支持异步非阻塞的 CPU→GPU 传输远快于普通可分页内存。问题PyTorch 的pin_memory()经由CachingHostAllocator分配内存会把每次分配的大小向上取整到下一个 2 的幂。例如一个 17MB 的张量实际会分配 32MB。大模型有数千个参数张量逐个独立pin_memory()会产生巨大浪费实测内存膨胀 50%100%。解决方案PinnedArenaPool见 memory_buffer.py预分配少量大块 Pinned 内存即 arena——预分配的大块内存区域所有小对象从中切分再用 bump-pointer 分配器紧凑地为每个张量切出空间避免逐张量取整浪费from_model()扫描模型中所有不可训练参数和 Buffer计算总大小memory_buffer.py将总大小分解为若干 2 的幂大小的块每个块是一个PinnedBuffermemory_buffer.py分配时顺序探测各块剩余空间bump-pointer 推进完成分配只做 64 字节对齐无取整浪费ALIGNMENT 64空间不足时自动增长新块_grow任何异常时回退到逐张量pin_memory()BaseBufferPool.allocate_like。此外PinnedBuffer采用惰性分配真正的pin_memory在第一次allocate_like时才触发进一步减少初始化开销。2.4 与 Gradient Checkpointing 的兼容Gradient Checkpointing 在反向传播时会重新执行前向重算激活这会再次触发forward_hook。该问题通过_in_recompute集合解决manager.py第一次前向正常 offload模块被加入_in_recompute重算前向反向过程中检测到模块在_in_recompute中跳过 offload让权重留在 GPU 上供反向使用调用after_backward()时清空_in_recompute为下一步做好准备。注意这与 gradient_checkpoint.py 中use_reentrantFalse的 checkpoint 用法相匹配_in_recompute机制正是为这种非重入式 checkpoint 设计的。2.5 Hook 注册粒度与孤儿参数默认情况下OffloadTrainingManager在叶子模块粒度注册 Hooknn.Linear、nn.LayerNorm等每个叶子模块被独立 onload/offloadmanager.py。同时不被任何叶子模块管理的孤儿参数和孤儿 Buffer会被自动收集并单独挂钩manager.py确保模型中的每个参数都不会漏管。实验性cpu_offload_split_threshold参数单位 MB可调整 Hook 注册粒度。设置后总参数量超过阈值的模块会被递归拆分到子模块低于阈值的模块整体挂钩。从 manager.py 的_should_force_recurse看除参数量超过阈值外未定义forward的容器模块、以及同时具备encode/decode的模块也会被强制递归。该特性在当前版本可能与部分模型架构不兼容默认关闭。2.6 训练循环中的集成执行流程见 runner.py# 当 enable_model_cpu_offloadTrue 时 # 1. 模型不调用 model.to(device)保持在 CPU 上 # 2. 只 prepare optimizer、dataloader、scheduler模型不参与 prepare # 3. 创建 OffloadTrainingManager自动在模型上注册 Hook # 训练循环 loss model(data) accelerator.backward(loss) offload_manager.after_backward() # 清空重算标记 将梯度移到 CPU optimizer.step() optimizer.zero_grad()OffloadTrainingManager(model, accelerator.device, enable_optimizer_cpu_offload, cpu_offload_split_threshold)的三个参数与命令行一一对应其中cpu_offload_split_threshold会在构造函数中乘以1024*1024转换为字节数manager.py。此外launch_data_process_taskSplit Training 的数据处理任务也支持同样的三个参数runner.py。三、使用方法3.1 参数说明命令行参数在 parsers.py 的add_offload_training_config中定义参数默认值说明--enable_model_cpu_offloadFalse开启逐层卸载训练。权重保存在 CPU逐层加载到 GPU--enable_optimizer_cpu_offloadFalse需与--enable_model_cpu_offload配合使用把可训练参数与优化器放到 CPU--cpu_offload_split_thresholdNone实验性单位 MB总参数量超过该阈值的模块被递归拆分launch_training_task会从args中读取这三个参数并传给OffloadTrainingManagerrunner.py。3.2 参数组合效果场景--enable_model_cpu_offload--enable_optimizer_cpu_offload效果默认训练❌❌全部权重与优化器都在 GPU 上仅卸载不可训练参数✅❌不可训练参数逐层卸载可训练参数与优化器留在 GPU卸载全部参数✅✅所有参数逐层卸载梯度与优化器运行在 CPU3.3 实战示例Qwen-Image LoRA 训练只需在原有训练命令中追加--enable_model_cpu_offload。以 Qwen-Image LoRA 训练为例基于 examples/qwen_image/model_training/train.py 的命令行接口accelerate launch examples/qwen_image/model_training/train.py \ --dataset_base_path data/example_dataset \ --dataset_metadata_path data/example_dataset/metadata.json \ --max_pixels 1048576 \ --dataset_repeat 50 \ --model_id_with_origin_paths Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors,Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors \ --learning_rate 1e-4 \ --num_epochs 5 \ --remove_prefix_in_ckpt pipe.dit. \ --output_path ./models/train/Qwen-Image_lora \ --lora_base_model dit \ --lora_target_modules to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_out.0,to_add_out,img_mlp.net.2,img_mod.1,txt_mlp.net.2,txt_mod.1 \ --lora_rank 32 \ --use_gradient_checkpointing \ --dataset_num_workers 8 \ --find_unused_parameters \ --enable_model_cpu_offload若需要完整卸载优化器也放到 CPU追加--enable_model_cpu_offload \ --enable_optimizer_cpu_offload值得说明的是开启--enable_model_cpu_offload后模型初始化设备会被切换为 CPU——在 train.py 中可见devicecpu if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device这与权重常驻 CPU、由 Hook 管理的设计一致。3.4 兼容性特性兼容性说明Gradient Checkpointing✅_in_recompute机制处理重算Accelerate DDP多卡⚠️enable_model_cpu_offload模式下模型不被 DDP 包装无accelerator.prepare(model)不执行梯度 allreduce。多卡训练兼容性无法保证每张卡独立训练、不做梯度同步Split Training✅launch_data_process_task同样支持--enable_model_cpu_offloadDeepSpeed❌ZeRO 的参数收集与 Hook 冲突3.5 使用注意事项开启--enable_model_cpu_offload后模型永远不会调用model.to(device)权重完全由 Hook 管理由于存在 CPU↔GPU 搬运训练速度会下降通常慢 210 倍模型越大减速越明显适合显存受限的场景建议与--use_gradient_checkpointing配合使用进一步降低激活显存--enable_optimizer_cpu_offload只支持梯度累积步数为 1--gradient_accumulation_steps 1--enable_optimizer_cpu_offload生效时TrainableParamOffloader会在 backward 后把param.grad移到 CPUoffload_grad所以after_backward()必须在optimizer.step()之前调用runner.py。四、在其他代码库中集成 Offload Training 模块Offload Training 模块相对独立diffsynth/core/offload_training 下仅有manager.py、offloader.py、memory_buffer.py三个核心文件对外只暴露OffloadTrainingManager可以很方便地集成到其他训练代码中。下面是一个仅约 4GB 显存的普通训练示例import torch from tqdm import tqdm class ToyModel(torch.nn.Module): def __init__(self): super().__init__() self.layers torch.nn.ModuleList(torch.nn.Linear(4096, 4096) for _ in range(10)) def forward(self, x): for layer in self.layers: x x layer(torch.nn.functional.layer_norm(x, (4096,))) return x model ToyModel().to(cuda) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) pbar tqdm(range(100)) for i in pbar: x torch.randn((512, 4096), devicecuda) y x 1 y_pred model(x) loss torch.nn.functional.mse_loss(y_pred, y) loss.backward() optimizer.step() optimizer.zero_grad() pbar.set_postfix(lossf{loss.item():.4f})启用 Offload Training 后显存占用降到约 1.4GB改动极小import torch from tqdm import tqdm from diffsynth.core import OffloadTrainingManager class ToyModel(torch.nn.Module): def __init__(self): super().__init__() self.layers torch.nn.ModuleList(torch.nn.Linear(4096, 4096) for _ in range(10)) def forward(self, x): for layer in self.layers: x x layer(torch.nn.functional.layer_norm(x, (4096,))) return x model ToyModel().to(cpu) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) offload_manager OffloadTrainingManager(model, target_devicecuda, enable_optimizer_cpu_offloadTrue) pbar tqdm(range(100)) for i in pbar: x torch.randn((512, 4096), devicecuda) y x 1 y_pred model(x) loss torch.nn.functional.mse_loss(y_pred, y) loss.backward() offload_manager.after_backward() optimizer.step() optimizer.zero_grad() pbar.set_postfix(lossf{loss.item():.4f})对比两段代码可以看出集成要点模型保持在 CPUmodel.to(cpu)用OffloadTrainingManager替代model.to(cuda)并在loss.backward()之后、optimizer.step()之前调用offload_manager.after_backward()。其余训练逻辑数据、优化器、进度条完全不变。五、相关文档与源码导航官方文档Offload_Training.md中文版见 docs/zh/Training/Offload_Training.md核心实现manager.pyOffloadTrainingManager、UnitWiseHookManager、UnitWiseParamManager、offloader.py四类 Offloader、memory_buffer.pyPinnedArenaPool、PinnedBuffer、BaseBufferPool训练入口runner.pylaunch_training_task、launch_data_process_task参数定义parsers.pyadd_offload_training_config训练脚本示例examples/qwen_image/model_training/train.py更多模型Wan、FLUX、LTX-2、MiniMax 等的训练脚本位于 examples 对应目录下【免费下载链接】DiffSynth-StudioEnjoy the magic of Diffusion models!项目地址: https://gitcode.com/GitHub_Trending/dif/DiffSynth-Studio创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考