Megatron SFT 长上下文训练显存优化指南

针对 Qwen3-Omni-30B-A3B-Instruct 模型，60K tokens 长上下文 Megatron SFT 训练场景。基于 ms-swift 源码分析，所有参数引用均标注源码位置。

一、激活值重计算（Activation Recomputation）

激活值重计算是长上下文训练中最关键的显存优化手段，用计算时间换取显存空间。

1.1 核心参数

参数	类型	默认值	说明	源码位置
`recompute_granularity`	`'selective'` / `'full'`	`'selective'`	重计算粒度。`full` = 整层重计算（最省显存），`selective` = 只重计算选定模块	`megatron_args.py:402`
`recompute_method`	`'uniform'` / `'block'`	`None`	重计算方法。仅在 `recompute_granularity='full'` 时生效。`uniform` = 均匀分布重计算层，`block` = 块状分布	`megatron_args.py:403`
`recompute_num_layers`	`int`	`None`	每次重计算的层数。仅在 `recompute_granularity='full'` 时生效。The larger the recompute_num_layers, the smaller the memory usage but higher computation cost. Default is None.	`megatron_args.py:404`
`recompute_modules`	`List[str]`	`['core_attn']`	仅在 `recompute_granularity='selective'` 时生效。可选模块：`core_attn`, `moe_act`, `layernorm`, `mla_up_proj`, `mlp`, `moe`	`megatron_args.py:405`
`moe_layer_recompute`	`bool`	`False`	MoE 层额外重计算，对 MoE 模型（如 30B-A3B）有显著效果	`megatron_args.py:551`
`vit_gradient_checkpointing`	`bool`	`True`	对视觉编码器（ViT）启用梯度检查点（HuggingFace 风格）	`megatron_args.py:362`
`gradient_checkpointing_kwargs`	`dict`	`None`	传递给 `torch.utils.checkpoint` 的额外参数，如 `{"use_reentrant": false}`	`megatron_args.py:365`

1.2 重计算的两种模式详解

Full 模式（推荐用于长上下文）：

前向传播时不保存中间激活值，反向传播时从 checkpoint 点重新计算。设 recompute_num_layers=1 意味着每一层都重新计算，最大化节省显存。

yaml
展开代码
recompute_granularity: full
recompute_method: uniform
recompute_num_layers: 1   # 最省显存，每层都重计算

实现位置：swift/megatron/model/mm_gpt/qwen3_vl.py:251-280（_checkpointed_forward 方法），使用 tensor_parallel.checkpoint() 对 transformer 层做分段 checkpoint。

Selective 模式（默认）：

只对 recompute_modules 中指定的模块做 checkpoint，其余模块正常保存激活值。core_attn、mlp、moe 使用标准 checkpoint；moe_act、layernorm、mla_up_proj 使用 output-discarding checkpoint（CheckpointWithoutOutput）。

1.3 当前配置问题

你当前 recompute_num_layers: 2，表示每 2 层为一组做 checkpoint。

二、Context Parallel（长序列最重要的参数）

Context Parallel (CP) 将序列按长度维度切分到多个 GPU，是处理超长序列的核心手段。

2.1 参数

参数	类型	默认值	说明	源码位置
`context_parallel_size`	`int`	`1`	CP 大小。将序列切分到 N 个 GPU，每 GPU 处理 `seq_length / N` 个 tokens	`megatron_args.py:481`
`sequence_parallel`	`bool`	`False`	序列并行，在 TP 组内切分 LayerNorm/Dropout 的激活值。需要 `tensor_model_parallel_size > 1`	`megatron_args.py:480`

2.2 关键澄清：SP 和 CP 可以同时使用

源码中唯一的互斥约束是 mlp_padding_free 与 SP/CP（megatron_args.py:608-609）：

python
展开代码
if self.mlp_padding_free and (self.sequence_parallel or self.context_parallel_size > 1):
    raise ValueError('mlp_padding_free is not compatible with sequence parallel or context parallel.')

这说明：

不兼容：mlp_padding_free + SP 或 CP
兼容：SP + CP（可以同时开启）
兼容：padding_free（默认 True）+ SP + CP

SP 和 CP 是互补关系：SP 切分 TP 组内的激活值序列维度，CP 在更大范围切分整个序列。

2.3 当前配置问题

你当前 context_parallel_size: 1，意味着 60K tokens 全部放在单 GPU 上处理。这是 OOM 的最主要原因之一。

建议：

CP 值	每 GPU tokens	DP 数（64卡, TP=4, PP=2）	说明
`1`（当前）	60,000	8	极易 OOM
`2`	30,000	4	推荐首选
`4`	15,000	2	如果 CP=2 仍 OOM

2.4 Padding 对齐规则

开启 SP/CP 后，序列长度会被自动 padding 到特定倍数（源码 swift/megatron/trainers/utils.py:318-332）：

条件	Padding 对齐到
SP 开启	`tensor_model_parallel_size` 的倍数
CP 开启	`TP × CP` 的倍数
FP8 blockwise	`TP × CP × 128` 的倍数
FP8 其他	`max(TP × CP × 8, 16)` 的倍数

三、优化器显存优化

全参数训练时，优化器状态（Adam 的一阶矩、二阶矩、fp32 主参数副本）通常占总显存的 50% 以上。

3.1 核心参数

参数	类型	默认值	说明	源码位置
`use_distributed_optimizer`	`bool`	`True`	ZeRO-1：将优化器状态分片到各 DP rank，显存减少为 `1/DP_size`	`megatron_args.py:472`
`optimizer_cpu_offload`	`bool`	`False`	将优化器状态卸载到 CPU，大幅释放 GPU 显存	`megatron_args.py:422`
`use_precision_aware_optimizer`	`bool`	`False`	使用 TransformerEngine 的精度感知优化器，支持低精度优化器状态	`megatron_args.py:424`
`optimizer_offload_fraction`	`float`	`1.0`	卸载到 CPU 的优化器状态比例（0.0 ~ 1.0），配合 `optimizer_cpu_offload` 使用	`megatron_args.py:423`

3.2 低精度优化器状态参数

参数	类型	默认值	可选值	说明	源码位置
`main_grads_dtype`	`str`	`'fp32'`	`'fp32'`, `'bf16'`	主梯度精度，`bf16` 节省一半梯度显存	`megatron_args.py:425`
`main_params_dtype`	`str`	`'fp32'`	`'fp32'`, `'fp16'`	主参数副本精度	`megatron_args.py:426`
`exp_avg_dtype`	`str`	`'fp32'`	`'fp32'`, `'fp16'`, `'bf16'`, `'fp8'`	Adam 一阶矩精度，`fp8` 最省	`megatron_args.py:427`
`exp_avg_sq_dtype`	`str`	`'fp32'`	`'fp32'`, `'fp16'`, `'bf16'`, `'fp8'`	Adam 二阶矩精度，`fp8` 最省	`megatron_args.py:428`

3.3 官方 Offload 示例

examples/megatron/moe/qwen3_moe_offload.sh 针对同款 Qwen3-30B-A3B 模型的配置：

bash
展开代码
--optimizer_cpu_offload true \
--use_precision_aware_optimizer true \
--optimizer_offload_fraction 1

该示例在 4 × A100 上，每卡 75GiB 显存。

四、并行策略

4.1 并行参数总览

参数	类型	默认值	说明	源码位置
`tensor_model_parallel_size`	`int`	`1`	TP：将权重矩阵按列切分到多个 GPU	`megatron_args.py:473`
`pipeline_model_parallel_size`	`int`	`1`	PP：将模型按层切分到多个 GPU 组	`megatron_args.py:474`
`context_parallel_size`	`int`	`1`	CP：将序列按长度切分到多个 GPU	`megatron_args.py:481`
`sequence_parallel`	`bool`	`False`	SP：TP 组内切分序列维度的激活值	`megatron_args.py:480`
`expert_model_parallel_size`	`int`	`1`	EP：将不同 expert 分配到不同 GPU	`megatron_args.py:542`
`expert_tensor_parallel_size`	`int`	`1`	Expert TP：单个 expert 内部的张量并行	`megatron_args.py:543`

4.2 DP 计算公式


展开代码
DP = total_GPUs / (TP × PP × CP)

源码位置：megatron_args.py:227-228

注意：EP 不参与 DP 公式计算（EP 在 TP 组内部再分），但会影响每个 GPU 上的 expert 数量。

4.3 约束条件

约束	说明	源码位置
`num_query_groups` 必须是 TP 的倍数	Qwen3-Omni-30B-A3B 的 `num_query_groups=4`，所以 TP 只能是 1, 2, 4	Megatron 内部校验
`mlp_padding_free` 与 SP/CP 互斥	不能同时开启 `mlp_padding_free` 和 SP 或 CP	`megatron_args.py:608-609`
`freeze_parameters_ratio`(0~1) 与 PP>1 互斥	部分冻结不能与流水线并行同时使用	`megatron_args.py:308-309`
`decoder_first/last_pipeline_num_layers` 需要 PP>1	PP 首尾层数自定义仅在 PP>1 时生效	`megatron_args.py:713-716`
`global_batch_size` 必须被 `micro_batch_size × DP` 整除	框架自动计算梯度累积步数	Megatron 内部校验

4.4 Virtual Pipeline Parallelism（VPP）

VPP 通过增加 virtual stages 减少 PP bubble，适合 PP 较大的场景。

参数	类型	默认值	说明	源码位置
`num_layers_per_virtual_pipeline_stage`	`int`	`None`	每个虚拟 PP stage 的层数	`megatron_args.py:486`
`num_virtual_stages_per_pipeline_rank`	`int`	`None`	每个 PP rank 的虚拟 stage 数	`megatron_args.py:487`
`decoder_first_pipeline_num_layers`	`int`	`None`	第一个 PP stage 的层数（平衡 embedding 层的显存）	`megatron_args.py:475`
`decoder_last_pipeline_num_layers`	`int`	`None`	最后一个 PP stage 的层数（平衡 LM head 的显存）	`megatron_args.py:476`

4.5 64 卡并行方案对比

方案	TP	EP	PP	SP	CP	DP	每 GPU tokens	特点
当前配置	4	4	4	✓	1	1	60,000	DP=1 吞吐极低，单卡 60K tokens 易 OOM
推荐方案 A	4	4	2	✓	2	4	30,000	平衡显存和吞吐
推荐方案 B	4	4	2	✓	4	2	15,000	最省激活显存
极限方案	4	4	4	✓	2	1	30,000	PP=4 减少模型权重显存，但 bubble 大

五、其他显存优化参数

5.1 注意力后端

参数	类型	默认值	说明	源码位置
`attention_backend`	`str`	`'flash'`	`'flash'` = Flash Attention（O(n) 显存），`'fused'`, `'unfused'`, `'local'`, `'auto'`	`megatron_args.py:420`

Flash Attention 对长序列至关重要，将注意力显存从 O(n²) 降为 O(n)。

5.2 融合优化

参数	类型	默认值	说明	源码位置
`cross_entropy_loss_fusion`	`bool`	`False`	融合交叉熵损失计算，减少 logits 张量的显存占用	`megatron_args.py:416`
`no_gradient_accumulation_fusion`	`bool`	`False`	禁用梯度累积融合（需要 apex）	`megatron_args.py:415`
`moe_permute_fusion`	`bool`	`False`	MoE token permute 操作融合	`megatron_args.py:547`
`moe_grouped_gemm`	`bool`	`True`	MoE grouped GEMM 优化	`megatron_args.py:546`
`moe_shared_expert_overlap`	`bool`	`False`	共享专家计算与通信重叠	`megatron_args.py:550`

5.3 数据效率

参数	类型	默认值	说明	源码位置
`padding_free`	`bool`	`True`	去除 batch 内 padding，减少无效计算和显存浪费	`megatron_args.py:321`
`packing`	`bool`	`False`	将多个短序列打包到一个 batch，减少 padding	继承自 `BaseArguments`

5.4 FP8 训练（如果硬件支持）

参数	类型	默认值	说明	源码位置
`fp8_format`	`str`	`None`	FP8 格式：`'e4m3'` 或 `'hybrid'`，大幅减少权重和激活显存	`megatron_args.py:568`
`fp8_recipe`	`str`	`'delayed'`	FP8 算法：`'tensorwise'`, `'delayed'`, `'mxfp8'`, `'blockwise'`	`megatron_args.py:569`
`fp8_param_gather`	`bool`	`False`	all-gather 时保持 FP8 参数，节省通信和显存	`megatron_args.py:572`

5.5 MoE 专用

参数	类型	默认值	说明	源码位置
`moe_expert_capacity_factor`	`float`	`None`	专家容量因子，限制每个 expert 处理的 token 数，防止显存峰值	`megatron_args.py:552`
`moe_token_dispatcher_type`	`str`	`'alltoall'`	Token 分发方式：`'allgather'`, `'alltoall'`, `'flex'`, `'alltoall_seq'`，不同方式显存特性不同	`megatron_args.py:544`

5.6 其他

参数	类型	默认值	说明	源码位置
`micro_batch_size`	`int`	`1`	每 GPU 微批次大小，长上下文必须设为 1	`megatron_args.py:400`
`use_cpu_initialization`	`bool`	`False`	CPU 初始化模型权重，节省初始化阶段 GPU 显存	`megatron_args.py:406`
`manual_gc`	`bool`	`False`	手动控制垃圾回收，减少显存碎片	`megatron_args.py:430`
`manual_gc_interval`	`int`	`0`	手动 GC 间隔	`megatron_args.py:431`
`no_save_optim`	`bool`	`False`	不保存优化器状态到 checkpoint（节省存储，不影响训练显存）	`megatron_args.py:454`

5.7 环境变量

bash
展开代码
# 必须设置，启用 PyTorch 可扩展内存段，减少显存碎片
PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True'

所有官方 Megatron 示例脚本均使用此环境变量。

六、针对当前配置的具体修改建议

6.1 当前值 vs 推荐值对比

参数	当前值	推荐值	改动原因
`recompute_num_layers`	`2`	`1`	每层都重计算，最大化节省激活显存
`moe_layer_recompute`	未设	`true`	MoE 层额外重计算，对 30B-A3B 有效
`context_parallel_size`	`1`	`2`（或 `4`）	核心改动：将 60K 序列切分，每 GPU 处理 30K（或 15K）
`pipeline_model_parallel_size`	`4`	`2`	PP=4 时 DP=1 吞吐太低，降为 2 配合 CP=2 更平衡
`optimizer_cpu_offload`	未设	`true`	优化器状态卸载到 CPU
`use_precision_aware_optimizer`	未设	`true`	配合 CPU offload
`optimizer_offload_fraction`	未设	`1.0`	全部卸载
`moe_shared_expert_overlap`	未设	`true`	共享专家计算重叠，提升性能

6.2 推荐的 YAML 配置片段

yaml
展开代码
# ============ Memory Optimization（关键改动）============
recompute_granularity: full
recompute_method: uniform
recompute_num_layers: 1            # 从 2 改为 1，最省显存
moe_layer_recompute: true          # 新增：MoE 层重计算
vit_gradient_checkpointing: true

# ============ 优化器 CPU Offload（新增）============
optimizer_cpu_offload: true
use_precision_aware_optimizer: true
optimizer_offload_fraction: 1.0

# ============ 并行策略（调整）============
tensor_model_parallel_size: 4      # TP=4（受 num_query_groups=4 约束）
expert_model_parallel_size: 4      # EP=4
pipeline_model_parallel_size: 2    # PP: 从 4 降为 2，减少 pipeline bubble
sequence_parallel: true            # SP: TP 组内序列并行
context_parallel_size: 2           # CP=2: 60K ÷ 2 = 每 GPU 30K tokens

# DP = 64 / (TP=4 × PP=2 × CP=2) = 4

# ============ 注意力 & 融合优化 ============
attention_backend: flash
cross_entropy_loss_fusion: true
moe_permute_fusion: true
moe_grouped_gemm: true
moe_shared_expert_overlap: true

# ============ Batch ============
micro_batch_size: 1
global_batch_size: 64              # 需确保能被 micro_batch_size × DP=4 整除

6.3 如果 CP=2 仍然 OOM

逐步升级策略：


展开代码
第 1 步：CP=2, PP=2              → DP=4, 每 GPU 30K tokens
第 2 步：CP=4, PP=2              → DP=2, 每 GPU 15K tokens
第 3 步：CP=4, PP=2 + optimizer_cpu_offload → 进一步释放优化器显存
第 4 步：CP=4, PP=4              → DP=1, 吞吐最低但最省显存
第 5 步：在以上基础上加 FP8      → fp8_format: e4m3（需要硬件支持）

七、显存优化优先级排序

从最有效到最不有效：

优先级	手段	预期效果	代价
⭐⭐⭐⭐⭐	`context_parallel_size` 增大	激活显存按比例线性切分	增加通信开销
⭐⭐⭐⭐⭐	`optimizer_cpu_offload: true`	优化器状态全部卸载到 CPU，释放大量 GPU 显存	训练速度略降（CPU-GPU 数据传输）
⭐⭐⭐⭐	`recompute_num_layers: 1`	最大化激活重计算节省	约增加 30% 计算时间
⭐⭐⭐⭐	`sequence_parallel: true`	TP 组内激活值序列维度切分	几乎无额外代价（需 TP>1）
⭐⭐⭐	`moe_layer_recompute: true`	MoE 层额外重计算	略增加计算时间
⭐⭐⭐	`attention_backend: flash`	注意力显存 O(n²) → O(n)	无（纯优化）
⭐⭐⭐	`cross_entropy_loss_fusion: true`	减少 logits 显存	无（纯优化）
⭐⭐	调整 PP/TP 比例	平衡权重和激活显存分布	PP 增大 → pipeline bubble 增大
⭐⭐	`fp8_format: e4m3`	权重和激活显存减半	需要 H100/H800 等硬件支持
⭐	`use_cpu_initialization: true`	节省初始化阶段显存	初始化慢一些
⭐	`manual_gc: true`	减少显存碎片	可忽略

附录：参考示例脚本

脚本路径	场景	关键配置
`examples/megatron/long_text.sh`	32K 长文本，4×A100	TP=4, SP, recompute full/uniform/1
`examples/megatron/moe/qwen3_moe_offload.sh`	30B-A3B + CPU offload，4×A100	EP=4, optimizer_cpu_offload, recompute full/uniform/1
`examples/megatron/moe/qwen3_moe.sh`	30B-A3B 多节点，16卡	PP=2, EP=8, recompute full/uniform/1
`examples/megatron/multimodal/omni/moe.sh`	Qwen3-Omni-30B LoRA，2卡	EP=2, LoRA, vit_gradient_checkpointing

目录