SuperAPI AI 工具集

nemo-mbridge-perf-cpu-offloading

通用与其它

查看项目

更新于 10-16 03:49

技能介绍

用途:指导在 Megatron Bridge 中启用两种互相独立的 GPU 到 CPU 内存卸载机制,用于缓解显存压力。

主要能力:激活 offloading 按 transformer 层卸载激活(可含权重),配置在 model.cpu_offloading*;优化器 offloading 通过 HybridDeviceOptimizer 将 Adam 动量与方差状态卸载到 CPU,配置在 optimizer.optimizer_cpu_offload。

选择建议:30B+ 大 MoE 模型且需 PP>1 时用优化器 offloading(激活 offloading 被 PP=1 限制阻断);中小模型 PP=1 且激活内存主导时用激活 offloading;需可调内存速度权衡时用 fractional optimizer_offload_fraction;吞吐优先不要启用;需 CUDA graphs 只用优化器 offloading。

用法:优化器 offloading 可设 cfg.optimizer.optimizer_cpu_offload=True、optimizer_offload_fraction=1.0、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。激活 offloading 可设 cfg.model.cpu_offloading=True、cpu_offloading_num_layers=16、cpu_offloading_activations=True、cpu_offloading_weights=False,并需 pipeline_model_parallel_size=1、recompute_granularity=None、cuda_graph_impl="none"。

限制:激活 offloading 要求 PP=1、禁用重计算、不可与 fine_grained_activation_offloading 及 CUDA graphs 组合,层数须在 [0, num_layers-1)。优化器 offloading 需 use_distributed_optimizer=True,无 PP 等限制,比例须在 [0.0,1.0]。已知吞吐惩罚随比例线性增加,D2H/H2D 重叠仅约 7% 加速。

验证:可运行单元测试与最小训练命令,成功标准为配置校验通过、无 OOM/NCCL 错误、loss 与基线最大偏差小于 0.001、显存按比例下降。