|
|
| 名称 |
nemo-mbridge-perf-cpu-offloading |
| 描述 |
验证并使用 Megatron Bridge 中的 CPU 卸载功能，包括层级别激活卸载，以及使用 HybridDeviceOptimizer 的部分优化器状态卸载。 |
| 开源协议 |
Apache-2.0 when_to_use: 启用 CPU 卸载以降低 GPU 内存占用，或调查因 CPU 卸载配置变更导致 OOM 或崩溃的提交；‘cpu_offloading’、‘optimizer_cpu_offload’、‘optimizer_offload_fraction’、‘HybridDeviceOptimizer’、‘将优化器移至 CPU’。 |
CPU 卸载
参考文献
- 稳定版文档：@docs/training/cpu-offloading.md
- 结构化元数据：@skills/nemo-mbridge-perf-cpu-offloading/card.yaml
简介
有两种独立机制可将数据从 GPU 内存移动到 CPU 内存：
| 机制 |
配置命名空间 |
卸载内容 |
PP 限制 |
| 激活卸载 |
model.cpu_offloading* |
每个 transformer 层的激活（可选权重） |
PP 必须为 1 |
| 优化器卸载 |
optimizer.optimizer_cpu_offload |
通过 HybridDeviceOptimizer 卸载 Adam 优化器状态（动量 + 方差） |
无限制 |
快速决策
| 情况 |
建议 |
| 大型 MoE 模型（30B+），需要 PP > 1 |
优化器卸载 —— 激活卸载受 PP=1 限制 |
| 中小型模型，PP=1 可容纳，激活内存占主导 |
激活卸载 |
| 需要可调的内存/速度权衡 |
使用带分数优化器卸载（optimizer_offload_fraction） |
| 吞吐量优先 |
不启用 —— 卸载总会增加开销 |
| 需要使用 CUDA graphs |
只能优化器卸载 —— 激活卸载不兼容 |
| 内存压力适中 |
优化器卸载分数设为 25–50% 效率最佳 |
启用方法
优化器 CPU 卸载（大型模型推荐）
cfg.optimizer.optimizer_cpu_offload = True
cfg.optimizer.optimizer_offload_fraction = 1.0
cfg.optimizer.overlap_cpu_optimizer_d2h_h2d = True
CLI 覆盖：
optimizer.optimizer_cpu_offload=True \
optimizer.optimizer_offload_fraction=0.5 \
optimizer.overlap_cpu_optimizer_d2h_h2d=True
激活 CPU 卸载（仅限中小型模型）
cfg.model.cpu_offloading = True
cfg.model.cpu_offloading_num_layers = 16
cfg.model.cpu_offloading_activations = True
cfg.model.cpu_offloading_weights = False
cfg.model.pipeline_model_parallel_size = 1
cfg.model.recompute_granularity = None
cfg.model.cuda_graph_impl = "none"
配置参数参考
优化器卸载
| 参数 |
默认值 |
说明 |
optimizer_cpu_offload |
False |
总开关 |
optimizer_offload_fraction |
0.0 |
优化器状态卸载到 CPU 的比例（0.0–1.0） |
overlap_cpu_optimizer_d2h_h2d |
False |
GPU↔CPU 传输与计算重叠 |
use_torch_optimizer_for_cpu_offload |
False |
CPU 部分使用 torch.optim 而非融合优化器 |
激活卸载
| 参数 |
默认值 |
说明 |
cpu_offloading |
False |
总开关 |
cpu_offloading_num_layers |
0 |
要卸载的 transformer 层数（0 至 num_layers-1） |
cpu_offloading_activations |
True |
卸载激活 |
cpu_offloading_weights |
False |
卸载权重 |
cpu_offloading_double_buffering |
False |
重新加载时跨层双缓冲 |
兼容性与约束
激活卸载
pipeline_model_parallel_size 必须为 1recompute_granularity 必须为 None- 不能与
fine_grained_activation_offloading 同时使用
- 不能与 CUDA graphs 同时使用
cpu_offloading_num_layers 必须在 [0, num_layers-1) 内
优化器卸载
- 需要
use_distributed_optimizer = True（大多数 recipe 中为默认）
- 无 PP、recompute 或 CUDA graph 限制
optimizer_offload_fraction 必须在 [0.0, 1.0] 内
实际：大型 MoE 模型
Qwen3-30B-A3B 等大型 MoE 模型无法使用激活卸载。PP=1 约束意味着每个 GPU 需要容纳全部 48 层；仅模型权重 + 优化器状态（约 70 GB）就已超过 H100 80 GB 容量。
最小可运行命令
uv run python scripts/training/run_recipe.py \
--recipe qwen3_30b_a3b_pretrain_config \
optimizer.optimizer_cpu_offload=True \
optimizer.optimizer_offload_fraction=0.5 \
train.train_iters=20 \
train.global_batch_size=8 \
train.micro_batch_size=1
验证
单元测试
uv run python -m pytest \
tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py -k "cpu_offload" \
tests/unit_tests/peft/test_utils.py -k "cpu_offload" -q
成功标准
- 所选卸载模式的配置验证通过
- 训练在无 OOM 或 NCCL 错误的情况下完成
- 损失与非卸载 baseline 匹配（最大差值 < 0.001）
- 内存使用量随卸载比例成比例下降
代码锚点
MCore 激活卸载约束
if self.cpu_offloading and (
self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers
):
raise ValueError(...)
if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
raise ValueError("Currently there is no support for Pipeline parallelism with CPU offloading")
if self.cpu_offloading and self.recompute_granularity is not None:
raise ValueError("CPU offloading does not work when activation recomputation is enabled")
MCore CUDA graph 不兼容
if self.cpu_offloading:
raise ValueError("CUDA graphs not supported with CPU offloading.")
MCore 细粒度卸载互斥
if self.fine_grained_activation_offloading:
assert (
not self.cpu_offloading
), "fine_grained_activation_offloading cannot be enabled with cpu_offloading."
MCore HybridDeviceOptimizer 实例化
if config.optimizer_cpu_offload:
# ... 设置 cpu/gpu 优化器类 ...
optimizer = HybridDeviceOptimizer(
param_groups,
offload_fraction=config.optimizer_offload_fraction,
cpu_optimizer_cls=cpu_optimizer_cls,
gpu_optimizer_cls=gpu_optimizer_cls,
overlap_cpu_optimizer_d2h_h2d=config.overlap_cpu_optimizer_d2h_h2d,
pin_cpu_grads=config.pin_cpu_grads,
pin_cpu_params=config.pin_cpu_params,
)
Bridge CUDA graph 保护
assert not config.cpu_offloading and config.recompute_granularity is None, "Cudagraphs not supported"
Bridge PEFT 中的激活卸载
if self.config.cpu_offloading and self.config.cpu_offloading_activations:
x.activation_offloading = True
x, _ = self.linear_in(x)
x = self.activation(x)
if self.config.cpu_offloading and self.config.cpu_offloading_activations:
x.activation_offloading = True
x, _ = self.linear_out(x)
故障诊断
| 症状 |
可能原因 |
如何确认 |
修复 |
Currently there is no support for Pipeline parallelism with CPU offloading |
激活卸载 + PP > 1 |
检查 pipeline_model_parallel_size |
设置 PP=1 或使用优化器卸载 |
CPU offloading does not work when activation recomputation is enabled |
激活卸载 + 重计算 |
检查 recompute_granularity |
设置 recompute_granularity=null |
fine_grained_activation_offloading cannot be enabled with cpu_offloading |
两种卸载模式同时启用 |
检查两个标志 |
只使用其中一种 |
CUDA graphs not supported with CPU offloading |
CUDA graphs + 激活卸载 |
检查 cuda_graph_impl |
设置 cuda_graph_impl="none" |
| 激活卸载时 OOM |
模型过大，PP=1 装不下 |
检查已分配内存与 80 GB |
使用优化器卸载并配合 PP > 1 |
| 严重变慢（>4 倍） |
100% 优化器卸载，CPU Adam 成为瓶颈 |
比较不同分数下的单次迭代时间 |
降低分数或启用 overlap_cpu_optimizer_d2h_h2d |
| 部分优化器卸载时 OOM |
该配置卸载不足 |
检查不同分数下的内存 |
增加分数或增加 PP |
已知限制
- 激活卸载要求 PP=1，导致对需要 PP 的大型模型（30B+ MoE）不实用。
- 优化器卸载吞吐量损失近似线性（Qwen3-30B-A3B：25% 约 1.9 倍，100% 约 4.2 倍）。
- D2H/H2D 重叠仅带来约 7% 加速，因为 CPU Adam 计算是主要瓶颈。
fine_grained_activation_offloading 是独立的模块级方法，可与 PP > 1 配合，但不能与层级的 cpu_offloading 同时使用。