Name: nemo-mbridge-perf-memory-tuningSkill
Rating: 5 (1 reviews)
Author: nemo-mbridge-perf-memory-tuning


名称	nemo-mbridge-perf-memory-tuning
描述	降低 Megatron Bridge 中 GPU 峰值内存的技术 — 可扩展段、PEFT + 序列并行输入重新收集、并行度调整、激活重算、CPU 卸载约束及常见 OOM 修复。
开源协议	Apache-2.0 when_to_use: GPU OOM 错误、降低峰值内存、使用序列并行降低 LoRA 或 PEFT 激活内存、或将 OOM 回归追踪到特定提交或配置更改的场景；关键词包括 ‘out of memory’、‘OOM’、‘memory fragmentation’、‘expandable_segments’、‘reduce GPU memory’、‘LoRA memory’、‘PEFT memory’、‘sequence_parallel_input_regather’、‘PYTORCH_CUDA_ALLOC_CONF’。

内存调优

稳定文档: @docs/parallelisms.md 技能卡: @skills/nemo-mbridge-perf-memory-tuning/card.yaml

它是什么

GPU 训练期间的 OOM（显存不足）失败往往源于内存碎片化，而不是原始容量不足。PyTorch 的默认 CUDA 分配器可能会在分配之间留下不可用的间隙。最有效的单一修复是：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

这告诉 PyTorch 使用可扩展（非固定大小）的内存段，从而大幅减少碎片化，并在无需任何模型或并行度改动的情况下消除临界 OOM。

除了碎片化，实际峰值内存由以下因素决定：

参数 + 优化器状态内存——由 TP、PP、DP 分片控制（分布式优化器、FSDP）
激活内存——由激活重算、序列长度、微批大小以及 PEFT 特有的保留的已聚合输入控制
临时/工作区内存——CUDA 内核、NCCL 缓冲区、CUDA 图

对于配置规划，请在进行大规模任务启动前使用 Bridge 理论估算器：

from megatron.bridge.training.utils.theoretical_memory_utils import estimate_training_memory

estimate = estimate_training_memory(cfg, num_microbatches=num_microbatches)

估算器会报告负载最高的 GPU 分片，并区分稠密/嵌入、路由 MoE 专家和激活组件。它不包含分配器碎片化、CUDA/NCCL 工作区、CUDA 图缓冲区、token 不平衡或调度器工作区，因此需要结合运行时内存指标验证最终配置。

快速决策

当训练运行出现 OOM 或接近内存限制时：

首先设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True。 这会以零性能开销修复碎片化导致的 OOM。大多数 Slurm 启动模板已包含该设置。
对于启用序列并行的 LoRA，启用输入重新聚合 (LoRA(sequence_parallel_input_regather=True))。这可避免在每个符合条件的层中保留完整的聚合 LoRA-A 输入；当 SP 关闭时无效果。
如果尚未启用，添加选择性激活重算 (recompute_modules=[core_attn])。参见 @skills/nemo-mbridge-perf-activation-recompute/SKILL.md。
避免为了减少内存而增加 TP —— 加倍 TP 会大幅增加 NVLink 全归约通信量，并常常导致吞吐量下降（Llama3 70B 上 -28%）。
避免以牺牲 DP 为代价增加 PP —— DP 减半会使梯度累积步数加倍并损害吞吐量（约 6%）。
如果仍然 OOM，可考虑重算 mlp。可在大型稠密模型（Llama3 70B）上节省约 3 GB，但会损失约 16% 的 GPU 利用率。
当 PP > 1 时，CPU 卸载被阻止。

启用方法

可扩展段（推荐首选）

在启动前设置作业环境变量：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

在 Slurm 脚本中通常与其他环境变量一起放置：

export CUDA_DEVICE_MAX_CONNECTIONS=1
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

无需修改模型配置。零吞吐开销。

并行度调整

如果模型确实放不下（不是碎片化问题），调整并行度：

策略	内存影响	吞吐成本	备注
增加 PP（保持 DP）	每个阶段层数更少	中等（如果 DP 减半约 6%）	仅当 GPU 数量允许时
增加 TP	每个 GPU 参数量更少	严重（70B 上 -28%）	最后手段
分布式优化器	在 DP rank 间分片优化器状态	约 1-2%	推荐用于大模型
FSDP	分片参数+梯度+优化器	各异	参见 @skills/nemo-mbridge-perf-megatron-fsdp/SKILL.md

激活重算

完整细节参见 @skills/nemo-mbridge-perf-activation-recompute/SKILL.md。

PEFT + 序列并行输入重新聚合

对于启用序列并行的 LoRA 训练，符合条件的列并行 linear_qkv 和 linear_fc1 适配器会消费聚合的 LayerNorm 输出。由于 LoRA-A 是可训练的，默认路径会保留该完整的聚合输入直到反向计算 LoRA-A 权重梯度。

在构造 PEFT 配置时启用输入重新聚合：

from megatron.bridge.peft.lora import LoRA

cfg.peft = LoRA(
    # 保持此处已有的 LoRA 设置。
    sequence_parallel_input_regather=True,
)

开启此选项后，前向仍会临时为 LoRA-A GEMM 物化完整输入，但 MCore 自动求导只保留其序列本地分片。反向会异步地再次聚合完整输入，在可能的情况下与 dgrad 重叠该集合通信，计算 LoRA-A 权重梯度，然后重用临时通信缓冲区。

这是一种内存换通信的权衡，而不是传统的激活检查点：不会重跑 LayerNorm、注意力、MLP 或 LoRA GEMM。预计会有一定吞吐下降，收益随所保留的合格 LoRA-A 激活量增长。当序列并行关闭时，此选项无效。

CPU 卸载

cfg.model.cpu_offloading = True

与 PP > 1 不兼容。 只能在 pipeline_model_parallel_size = 1 时使用。

关于 VPP 的说明

虚拟流水线并行 (VPP) 主要是一种吞吐量优化，通过交错更小的模型块来减少流水线气泡开销。其对峰值内存的影响很小 —— 更改 VPP 并不会显著改变 GPU 上的总激活、参数或优化器内存。

在早期实验中，我们错误地把一次 OOM 修复归因于 VPP 调整（VPP 5→10）。真正的修复是 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，它消除了内存碎片。VPP=10 的运行实际上使用的峰值内存略高（60.2 GB vs 58.8 GB），但由于可扩展段防止了碎片化而没有 OOM。

VPP 应面向流水线气泡减少来调整（参见 @docs/parallelisms.md），而不是作为内存修复手段。

兼容性和限制

expandable_segments:True 与 --use-nccl-ub（NCCL 用户缓冲区注册）不兼容。参见 Megatron-FSDP 文档。
在使用 CUDA 图配合 expandable_segments:True 时，设置 NCCL_GRAPH_REGISTER=0（在 pre-Blackwell GPU 上必需，由 MCore CudaGraphManager 强制执行）。
CPU 卸载要求 pipeline_model_parallel_size = 1。
分布式优化器要求在优化器配置中 use_distributed_optimizer = True。
sequence_parallel_input_regather 仅适用于符合条件的非专家列并行 LoRA-A 投影。行并行适配器、专家适配器、TP=1、CUDA 图、CPU 激活卸载以及重叠的全层或选择性 MLP 激活重算都会回退到既有路径。

实测结果

Llama3 70B SFT，32x H100 80GB，FP8（当前 Scaling）：

基线：TP=4, PP=4, VPP=5, DP=2, MBS=1, GBS=32, seq_len=4096
黄金 GPU 利用率：709.93 TFLOP/s/GPU
回归阈值：5%

用于降低内存的并行度变更策略对比

实验	TP	PP	VPP	DP	TFLOP/s/GPU	vs 黄金	峰值内存 (GB)	结果
基线	4	4	5	2	~704	-0.8%	58.8	OOM（碎片化）
更多 PP	4	8	5	1	668.0	-5.9%	53.2	性能临界
更多 TP	8	4	5	1	508.7	-28.4%	50.2	严重回归
基线 + 可扩展段	4	4	5	2	~704	-0.8%	~59	通过

关键结论：

expandable_segments:True 是赢家。 基线的 OOM 由内存碎片导致，而非容量不足。设置该环境变量消除了 OOM，实现零吞吐成本且无需更改并行度。
PP=8 虽能降低内存但导致 DP 丢失（2→1），意味着每个 batch 需要 32 个梯度累积步，这使吞吐量下降约 6%。
TP=8 是灾难性的 (-28%)，因为加倍 TP 使 NVLink 上的全归约通信量成比例增加，且 DP=1 意味着没有微批重叠。

CPU 卸载：被阻止

实验	offload_layers	结果
Exp 4	2	不兼容（PP > 1）
Exp 5	4	不兼容（PP > 1）
Exp 6	6	不兼容（PP > 1）

ValueError: Currently there is no support for Pipeline parallelism with CPU offloading. 该方式对任何使用 PP > 1 的模型都不可用。

激活重算：昂贵的备选方案

在此工作负载上，使用 mlp 的选择性激活重算节省了约 3 GB 峰值内存，但代价是约 16% 的 GPU 利用率。完整结果参见 @skills/nemo-mbridge-perf-activation-recompute/SKILL.md。

LoRA + SP 输入重新聚合

基于真实 checkpoint 的 H100 训练配合 SQuAD 显示，在所有测试配置中峰值内存均较低，吞吐成本因负载而异：

模型/配置	基线峰值	输入重新聚合峰值	节省内存	吞吐变化
Qwen3-8B, TP2, seq 8192	47.545 GB	42.814 GB	4.731 GB (10.0%)	-6.74%
Qwen3-30B-A3B, TP4/EP4	29.890 GB	28.321 GB	1.569 GB (5.2%)	-2.89%
GPT-OSS-120B, TP2/EP8	52.185 GB	51.371 GB	0.814 GB (1.6%)	-0.34%

所有运行的损失均有限，零跳过或 NaN 迭代。两 rank 的 BF16 和 FP32 检查在输出、输入梯度、LoRA-A 和 LoRA-B 梯度以及两微批融合的 main_grad 累积方面与基线一致。

代码锚点

LoRA 序列并行输入重新聚合

src/megatron/bridge/peft/lora.py
    LoRA.sequence_parallel_input_regather

src/megatron/bridge/peft/utils.py
    ParallelLinearAdapter._sequence_parallel_input_regather_eligibility()
    ParallelLinearAdapter.forward()

CPU 卸载 PP 不兼容（MCore）

        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
            raise ValueError('Currently there is no support for Pipeline parallelism with CPU offloading')

VPP 配置和层数整除性校验（MCore）

            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:
                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size
                if (
                    not num_layers_per_middle_pipeline_rank
                    % self.virtual_pipeline_model_parallel_size
                    == 0
                ):
                    raise ValueError(
                        f'number of layers on each middle pipeline rank:'
                        f'{num_layers_per_middle_pipeline_rank} must be divisible by virtual'
                        f'pipeline parallel degree {self.virtual_pipeline_model_parallel_size}'
                    )

关于交错流水线调度的并行度文档

为了最小化流水线气泡，每块 GPU 上的计算可被分成多个层子集（称为模型块），而不是单个连续块。通过设置 `virtual_pipeline_model_parallel_size` 启用：

model_config = GPTModelProvider(
    pipeline_model_parallel_size=4,
    virtual_pipeline_model_parallel_size=2,  # 每个流水线阶段 2 个模型块
    # ... 其他模型参数
)

故障诊断

症状	原因	确认	修复
单个 rank OOM 而其他 rank 有余量	内存碎片	检查是否设置了 `expandable_segments:True`	设置 `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`
已启用 `expandable_segments` 仍 OOM	真正的容量限制	检查 `nvidia-smi` 中的参数/优化器内存	增加 PP、使用分布式优化器或添加重算
启动前估算内存超过 GPU 容量	模型状态或激活确实过大	运行 `estimate_training_memory` 并检查最大组件	启动前调整 PP/TP/CP/EP、分布式优化器或重算
LoRA + SP 保留异常高的激活内存	完整的聚合 LoRA-A 输入被保留到反向	检查 `cfg.peft.sequence_parallel_input_regather` 是否启用，目标是否合格	设置 `LoRA(sequence_parallel_input_regather=True)`；验证回退约束
`ValueError: PP + CPU offloading`	在 PP > 1 时使用 cpu_offloading	检查 PP 配置	禁用 CPU 卸载或将 PP 设为 1
`--use-nccl-ub` + 可扩展段出现 `RuntimeError`	NCCL UB 与可扩展分配器不兼容	检查环境变量	移除 `expandable_segments:True` 或禁用 `--use-nccl-ub`

已知限制

CPU 卸载在 PP > 1 时被阻止
并行度调整（TP/PP）往往有显著的吞吐成本
理论估算器基于公式，不能替代运行时性能分析或 CUDA 内存报告
LoRA 输入重新聚合不覆盖行并行或专家适配器，并且当少量符合条件的 LoRA-A 激活占主导时可能收益可忽略

验证

快速检查 expandable_segments:True 是否生效：

import os
assert 'expandable_segments:True' in os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '')

对于 Slurm 作业，确认启动脚本中的训练命令之前已导出该环境变量。

对于 LoRA + SP 输入重新聚合，运行聚焦的配置测试和真实两 rank MCore 反向等价性测试：

uv run python -m pytest tests/unit_tests/peft/test_utils.py -k 'sequence_parallel_input_regather' tests/unit_tests/peft/test_lora.py -k 'sequence_parallel_input_regather'
uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest tests/unit_tests/peft/test_lora_sp_input_regather_distributed.py