File size: 10,595 Bytes
0ed6b0e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 | # Offload Training
本文档介绍 DiffSynth-Studio 中的 Offload Training 功能，通过将模型权重逐层在 CPU 和 GPU 间搬运，大幅降低训练时的 GPU 显存占用。
> **注意**：当前 Offload Training 仅支持单卡训练，暂不兼容多卡（DDP）场景。
## 什么是 Offload Training
训练大规模模型（如 Qwen-Image 60 层、Wan2.1-14B 40 层）时，所有层的权重需同时驻留 GPU，仅权重就占用数十 GB 显存。Offload Training 的核心思想是：**任一时刻只将当前正在计算的模块权重加载到 GPU，计算完毕后立即卸载回 CPU**，从而将显存占用从 O(N × 每层参数量) 降低到 O(1 × 每层参数量)。
该功能基于 PyTorch 的 Module Hook 机制实现，不需要修改任何模型代码。
## 方案原理
### 核心机制
`OffloadTrainingManager` 会扫描模型，为每个需要管理的模块注册 4 个 Hook：
```
forward_pre_hook → 将模块权重从 CPU 加载到 GPU (onload)
module.forward() → 正常前向计算
forward_hook → 将模块权重从 GPU 卸载回 CPU (offload)
backward_pre_hook → 将模块权重从 CPU 重新加载到 GPU (onload)
module.backward() → 计算梯度
backward_hook → 将模块权重卸载回 CPU (offload)
```
### 参数与 Buffer 分类管理
根据参数是否可训练以及 buffer 类型，采用不同的 offload 策略：
| 类型 | Offloader 类 | 行为 |
|---------|-------------|------|
| 非可训练参数 (`requires_grad=False`) | `StaticParamOffloader` | 初始化时将权重拷贝到预分配的 pinned memory 中保持永久 CPU 副本，并将 `param.data` 替换为 GPU 端空 placeholder（释放 GPU 显存）；onload 时从 CPU 副本异步拷贝到 GPU，offload 时将 `param.data` 重新指向 placeholder（无需 PCIe 回传） |
| 可训练参数 + `enable_optimizer_cpu_offload=True` | `TrainableParamOffloader` | 权重在训练中会变化，无法保持静态副本；onload/offload 通过 `param.data.to(device)` 实际搬运；backward 后将 `param.grad` 也移到 CPU |
| 可训练参数 + `enable_optimizer_cpu_offload=False` | `AlwaysOnGPUParamOffloader` | 初始化时直接将参数移到 GPU，之后不再搬运；适用于 LoRA 训练（可训练参数量小） |
| 模块 Buffer（如 BatchNorm 的 `running_mean`/`running_var`） | `BufferOffloader` | 与 `StaticParamOffloader` 类似：初始化时将 buffer 拷贝到 pinned memory；onload 时从 CPU 副本异步拷贝到 GPU，offload 时将 `module._buffers[name]` 重新指向 CPU 副本 |
### Pinned Memory Pool
`StaticParamOffloader` 和 `BufferOffloader` 在初始化时需要为每个非可训练参数/buffer 分配一份 CPU 端的 pinned memory 副本（pinned memory 可以实现 CPU→GPU 的异步非阻塞传输，比普通 pageable memory 快得多）。
**问题**：PyTorch 的 `pin_memory()` 底层通过 `CachingHostAllocator` 分配内存，该分配器会将每次申请的大小向上取整到 2 的整数次方。例如一个 17MB 的 tensor 会实际分配 32MB。大模型有数千个参数 tensor，每个都独立 `pin_memory()` 会导致大量内存浪费（实测可能膨胀 50%~100%）。
**解决方案**：`PinnedArenaPool` 预先一次性分配少量大块 pinned memory（即 arena，一块预分配的大内存区域，所有小对象从中切分），然后用 bump-pointer 方式从大块中紧凑地切分出每个 tensor 需要的空间，避免逐 tensor 取整带来的浪费：
- `from_model()` 扫描模型所有非可训练参数和 buffer，计算总大小
- 将总大小分解为若干 power-of-two 大小的 chunk（每个 chunk 是一个 `PinnedBuffer`）
- 分配时顺序查找有剩余空间的 chunk，bump-pointer 前进即完成分配（仅 64 字节对齐，无取整浪费）
- 空间不足时自动 grow 新 chunk
- 异常情况下回退到逐 tensor 的 `pin_memory()`
### Gradient Checkpointing 兼容
Gradient Checkpointing 在 backward 时会重新执行 forward（重算激活），这会再次触发 `forward_hook`。方案通过 `_in_recompute` 集合解决：
- 首次 forward：正常 offload，并将模块加入 `_in_recompute`
- 重算 forward（backward 期间）：检测到模块在 `_in_recompute` 中，跳过 offload，保持权重在 GPU 供 backward 使用
- `after_backward()` 调用时：清空 `_in_recompute`，为下一个 step 做准备
### Hook 注册粒度
`OffloadTrainingManager` 默认以每个叶子模块（`nn.Linear`、`nn.LayerNorm` 等）为单位注册 hook，即每个叶子模块独立进行 onload/offload。此外，未被任何叶子模块管理到的「孤儿参数」和「孤儿 buffer」也会被自动收集并注册 hook。
**实验性功能**：通过 `cpu_offload_split_threshold`（单位 MB）可以调整 hook 的注册粒度。设置后，参数总量超过阈值的模块会被递归拆分为子模块，未超过阈值的模块则作为整体注册 hook。该功能当前版本可能无法兼容所有模型结构，默认不启用。
### 训练流程集成
在 `runner.py` 中的执行流程：
```python
# enable_model_cpu_offload=True 时：
# 1. 模型不调用 model.to(device)，保持在 CPU
# 2. 只 prepare optimizer、dataloader、scheduler（不 prepare model）
# 3. 创建 OffloadTrainingManager，自动为模型注册 hook
# 训练循环：
loss = model(data)
accelerator.backward(loss)
offload_manager.after_backward() # 清空 recompute 标记 + 梯度移到 CPU
optimizer.step()
optimizer.zero_grad()
```
## 如何使用
### 参数说明
| 参数 | 默认值 | 说明 |
|------|--------|------|
| `--enable_model_cpu_offload` | False | 启用逐层 offload 训练 |
| `--enable_optimizer_cpu_offload` | False | 配合 `--enable_model_cpu_offload`，将可训练参数和 optimizer 也放在 CPU |
| `--cpu_offload_split_threshold` | None | 实验性参数（单位 MB），超过此阈值的模块会被递归拆分 |
### 参数组合
| 场景 | `--enable_model_cpu_offload` | `--enable_optimizer_cpu_offload` | 效果 |
|------|:---------------:|:-------------------:|------|
| 默认训练 | ❌ | ❌ | 所有权重和 optimizer 在 GPU |
| 仅 offload 非可训练参数 | ✅ | ❌ | 非可训练参数逐层 offload，可训练参数和 optimizer 留在 GPU |
| offload 所有参数 | ✅ | ✅ | 所有参数逐层 offload，梯度和 optimizer 在 CPU 执行 |
### 使用示例
在现有训练命令中添加 `--enable_model_cpu_offload` 即可启用，以 Qwen-Image LoRA 训练为例：
```bash
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
```
如需完整 offload（optimizer 也在 CPU），添加 `--enable_optimizer_cpu_offload`：
```bash
--enable_model_cpu_offload \
--enable_optimizer_cpu_offload
```
### 兼容性
| 特性 | 兼容 | 说明 |
|------|:----:|------|
| Gradient Checkpointing | ✅ | `_in_recompute` 机制兼容 |
| Accelerate DDP（多卡训练） | ⚠️ | enable_model_cpu_offload 模式下不会对 model 进行 DDP 包装（不调用 `accelerator.prepare(model)`），因此**不会执行梯度 allreduce**。无法保证与多卡训练的兼容性，各卡独立计算梯度而无同步 |
| 拆分训练 | ✅ | `launch_data_process_task` 同样支持 `--enable_model_cpu_offload` |
| DeepSpeed | ❌ | ZeRO 的参数聚集机制与 hook 冲突 |
### 注意事项
- 开启 `--enable_model_cpu_offload` 后，模型不会调用 `model.to(device)`，权重始终由 hook 管理
- 训练速度会因 CPU↔GPU 传输而下降（典型约 2-10 倍），模型越大，速度下降越多，适合显存受限场景
- 建议配合 `--use_gradient_checkpointing` 使用以进一步降低激活值的显存占用
- `--enable_optimizer_cpu_offload` 仅支持梯度累积步数为 1（`--gradient_accumulation_steps 1`）
## 在其他代码库中集成 Offload Training 模块
Offload Training 模块是相对独立的，因此开发者可以将其集成到其他代码库中，以下是一个代码样例，显存占用 4G。
```python
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(), lr=1e-4)
pbar = tqdm(range(100))
for i in pbar:
x = torch.randn((512, 4096), device="cuda")
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(loss=f"{loss.item():.4f}")
```
启用 Offload Training，显存占用降低到 1.4G：
```python
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(), lr=1e-4)
offload_manager = OffloadTrainingManager(model, target_device="cuda", enable_optimizer_cpu_offload=True)
pbar = tqdm(range(100))
for i in pbar:
x = torch.randn((512, 4096), device="cuda")
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(loss=f"{loss.item():.4f}")
```
|