nemo-mbridge-perf-cpu-offloading

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 图	仅使用优化器卸载——激活值卸载不兼容
内存压力适中	将优化器卸载比例设为 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 必须为 1
• recompute_granularity 必须为 None
• 不能与 fine_grained_activation_offloading 同时使用
• 不能与 CUDA 图同时使用
• cpu_offloading_num_layers 必须在 [0, num_layers-1) 范围内

优化器卸载

• 需要设置 use_distributed_optimizer = True（大多数配方默认如此）
• 对 PP、重计算或 CUDA 图没有限制
• 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 错误
• 损失与未卸载基线一致（最大偏差 < 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 图不兼容

            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:
            # ... setup cpu/gpu optimizer classes ...
            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 图保护逻辑

        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 图 + 激活卸载	检查 cuda_graph_impl	将 cuda_graph_impl="none"
使用激活卸载时发生 OOM	模型对于 PP=1 来说过大	检查已分配内存是否超过 80 GB	在 PP > 1 时使用优化器卸载
性能严重下降（>4x）	100% 优化器卸载，CPU Adam 成为瓶颈	比较不同占比下的迭代时间	降低占比，或启用 overlap_cpu_optimizer_d2h_h2d
部分优化器卸载时发生 OOM	此配置的卸载量不足	检查不同占比下的内存用量	提高占比或增加 PP

已知限制

• 激活卸载要求 PP=1，因此对于需要流水线并行的大型模型（30B+ MoE）并不实用。
• 优化器卸载的吞吐量损失会线性增长（Qwen3-30B-A3B 在 25% 时约为 1.9x，在 100% 时约为 4.2x）。
• D2H/H2D 重叠仅能带来约 7% 的加速，因为 CPU Adam 计算是主要瓶颈。
• fine_grained_activation_offloading 是一种独立的模块级方案，可与 PP > 1 配合使用，但不能与层级的 cpu_offloading 结合使用。