用途:降低 Megatron Bridge 训练时的峰值 GPU 显存,解决 OOM。
主要能力:1) 显存碎片治理:设置 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,多数 Slurm 模板已包含,零吞吐损失。2) 启动前用 estimate_training_memory(cfg, num_microbatches=num_microbatches) 估算最重分片,区分稠密/嵌入、MoE 专家与激活占用,但不含碎片、CUDA/NCCL 工作区、CUDA graph 缓冲等。3) 并行度调整:增加 PP、TP、分布式优化器(use_distributed_optimizer=True)、FSDP、激活重算(recompute_modules=[core_attn])。4) CPU offloading:cfg.model.cpu_offloading = True。
决策顺序:先设 expandable_segments;再加选择性激活重算;避免靠加 TP 省显存;避免牺牲 DP 加 PP;仍 OOM 可考虑 mlp 重算。
限制:expandable_segments:True 与 --use-nccl-ub 不兼容;配合 CUDA graphs 需设 NCCL_GRAPH_REGISTER=0;CPU offloading 要求 pipeline_model_parallel_size = 1,PP>1 报 ValueError;VPP 对峰值显存影响极小,应只用于减少流水线气泡。实测 Llama3 70B 中加 TP 吞吐下降约 28%,加 PP 约 6%。估算器不替代实际显存剖析。