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）
• 激活内存 — 由激活重计算、序列长度、 微批次大小控制
• 临时／工作区内存 — 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. 添加选择性激活重计算（recompute_modules=[core_attn]），前提是尚未启用。请参阅 @skills/nemo-mbridge-perf-activation-recompute/SKILL.md。
3. 避免通过增加 TP 来解决内存问题 — 将 TP 翻倍会大幅增加 NVLink 全归约流量，并且往往会严重影响吞吐量（Llama3 70B 上下降 28%）。
4. 避免以牺牲 DP 为代价增加 PP — DP 减半会使梯度累积步数翻倍，并影响吞吐量（约 6%）。
5. 如果仍然 OOM，可考虑 mlp 重计算。它可节省约 3 GB，但会使大型稠密模型（Llama3 70B）的 GPU 利用率下降约 16%。
6. 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。

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。

测量结果

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。

代码锚点

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 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、使用分布式优化器或重计算
ValueError: PP + CPU offloading	PP > 1 时使用了 cpu_offloading	检查 PP 配置	禁用 CPU offloading 或设置 PP=1
--use-nccl-ub 与可扩展分段同时使用时出现 RuntimeError	NCCL UB 与可扩展分配器不兼容	检查环境变量	移除 expandable_segments:True 或禁用 --use-nccl-ub

已知限制

• PP > 1 时会禁用 CPU offloading
• 调整并行度（TP/PP）通常会显著降低吞吐量
• 理论估算器基于公式，不能替代运行时性能分析或 CUDA 内存报告

验证

快速检查 expandable_segments:True 是否已启用：

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

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