用途:在 Megatron Bridge 中验证并使用 CPU Offloading,把数据从 GPU 显存移到 CPU 内存,以缓解显存压力。包含两种相互独立的机制。
主要能力:一是激活卸载,配置命名空间为 model.cpu_offloading*,按 transformer 层卸载激活值(可选权重),要求流水线并行 PP 必须为 1。二是优化器卸载,配置项为 optimizer.optimizer_cpu_offload,通过 HybridDeviceOptimizer 把 Adam 优化器状态(动量与方差)放到 CPU,无 PP 限制。
选择建议:30B+ 大型 MoE 模型需要 PP>1 时用优化器卸载;中小模型 PP=1 且激活显存占主导时用激活卸载;需要可调的内存与速度权衡时用分数式 optimizer_offload_fraction;吞吐优先则不要启用;需要 CUDA graphs 时只能用优化器卸载;显存压力中等时用 25%–50% 的优化器卸载比例效率最佳。
启用方式:优化器卸载可设 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。激活卸载可设 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,取值 0.0–1.0)、overlap_cpu_optimizer_d2h_h2d(默认 False)、use_torch_optimizer_for_cpu_offload(默认 False)。激活侧有 cpu_offloading(默认 False)、cpu_offloading_num_layers(默认 0)、cpu_offloading_activations(默认 True)、cpu_offloading_weights(默认 False)、cpu_offloading_double_buffering(默认 False)。
限制:激活卸载要求 PP=1、recompute_granularity 为 None,不能与 fine_grained_activation_offloading 或 CUDA graphs 组合,cpu_offloading_num_layers 需在 [0, num_layers-1) 内。优化器卸载要求 use_distributed_optimizer = True,无 PP、重计算或 CUDA graph 限制。激活卸载因 PP=1 约束对 30B+ MoE 不实用;优化器卸载吞吐损失随比例线性增长(Qwen3-30B-A3B 约 25% 时 1.9 倍、100% 时 4.2 倍),D2H/H2D 重叠仅约 7% 加速,因为 CPU Adam 计算是主要瓶颈。
验证:可运行单元测试 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)、显存按卸载比例下降。