nemo-mbridge-perf-memory-tuning

内存调优

稳定版文档：@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 或接近内存上限时：

1. 首先设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True。 这可以修复由碎片化引起的 OOM，且不会产生性能开销。大多数 Slurm 启动模板已包含此设置。
2. 对于使用序列并行的 LoRA，启用输入重新收集 （LoRA(sequence_parallel_input_regather=True)）。这样可以避免在每个符合条件的层中保留完整的已收集 LoRA-A 输入；禁用 SP 时此设置不会产生影响。
3. 如果尚未启用，请添加选择性激活重计算（recompute_modules=[core_attn]）。参见 @skills/nemo-mbridge-perf-activation-recompute/SKILL.md。
4. 不要通过提高 TP 来解决内存问题——TP 加倍会大幅增加 NVLink 全归约流量，且往往会严重损害吞吐量（Llama3 70B 上降低 28%）。
5. 不要以牺牲 DP 为代价提高 PP——DP 减半会让梯度累积步数加倍，损害吞吐量（约 6%）。
6. 如果仍然 OOM，可考虑使用 mlp 重计算。它可节省约 3 GB 内存，但在大型稠密模型（Llama3 70B）上会使 GPU 利用率降低约 16%。
7. 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 + sequence-parallel 输入重新收集

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

构建 PEFT 配置时启用输入重新收集：

from megatron.bridge.peft.lora import LoRA

cfg.peft = LoRA(
    # Keep the recipe's existing LoRA settings here.
    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，而非 58.8 GB），但由于 expandable segments 防止了碎片，因此没有发生 OOM。

调整 VPP 应以减少流水线气泡为目标（参见 @docs/parallelisms.md）， 而不应将其作为内存修复方法。

兼容性和限制

• expandable_segments:True 与 --use-nccl-ub（NCCL 用户缓冲区注册）不兼容。 参见 Megatron-FSDP 文档。
• 使用 CUDA 图且启用了 expandable_segments:True 时，设置 NCCL_GRAPH_REGISTER=0（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（当前缩放）：

• 基线：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	相对黄金值	峰值内存（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	性能严重回退
基线 + expandable_segments	4	4	5	2	~704	-0.8%	~59	通过

关键要点：

• expandable_segments:True 是最佳方案。 基线 OOM 是由内存碎片导致， 而非容量不足。设置此环境变量即可消除 OOM，不影响吞吐量，也无需更改并行度。
• PP=8 能解决内存问题，但会失去 DP（2→1），这意味着每个批次需要进行 32 步梯度累积，吞吐量会降低约 6%。
• TP=8 会造成灾难性影响（-28%），因为 TP 翻倍会使 NVLink 上的 all-reduce 通信量按比例增加，而 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 峰值内存，但在此工作负载中会使 GPU 利用率降低约 16%。完整结果请参阅 @skills/nemo-mbridge-perf-activation-recompute/SKILL.md。

LoRA + SP 输入重新收集

使用真实检查点在 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 梯度，以及双 microbatch 融合的 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}"
                    )

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

To minimize the pipeline bubble, the computation on each GPU can be divided into multiple subsets of layers (referred to as model chunks), rather than a single contiguous block. Enable this by setting `virtual_pipeline_model_parallel_size`:

model_config = GPTModelProvider(
    pipeline_model_parallel_size=4,
    virtual_pipeline_model_parallel_size=2,  # 2 model chunks per pipeline stage
    # ... other model parameters
)

故障诊断

症状	原因	确认方法	修复方法
尽管其他 rank 仍有余量，单个 rank 却发生 OOM	内存碎片化	检查是否设置了 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	使用 cpu_offloading 且 PP > 1	检查 PP 配置	禁用 CPU 卸载或设置 PP=1
使用 --use-nccl-ub + 可扩展内存段时出现 RuntimeError	NCCL UB 与可扩展分配器不兼容	检查环境变量	移除 expandable_segments:True 或禁用 --use-nccl-ub

已知限制

• PP > 1 时无法使用 CPU 卸载
• 调整并行度（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