我要投稿

大模型训练的高效内存解决方案：流水线感知的细粒度激活卸载，实现显存开销与吞吐性能的联合最优

发布日期：2025-12-11 18:10:33 浏览次数： 1543

作者：小红书技术REDtech

微信搜一搜，关注“小红书技术REDtech”

随着大语言模型（LLM）参数规模突破千亿级（如 Kimi-K2 达 1000B 参数）、上下文长度扩展至 32K 甚至更长，激活值巨大的显存开销已经成为大语言模型训练过程中的核心瓶颈 —— 激活值数据量随序列长度增长呈二次方增长，极易超出 GPU 内存容量，该问题在多模态、强化学习（RL）等场景中尤为尖锐。传统解决方案（如完全重计算、增大并行度）或带来显著计算开销，或导致通信成本激增，最终导致模型训练效率的次优。

为兼顾内存占用与训练效率，技术团队提出细粒度激活卸载方案（Fine-grained Activation Offloading），在 Megatron-Core 框架中实现模块 / 算子级别的激活卸载，并兼容流水线并行（PP）、虚拟流水线并行（VPP）及细粒度重计算，最终在内存节省与性能损失间取得最优平衡。本文将从背景、实现原理、兼容性设计、性能验证到最佳实践，全面解析该方案。

本文介绍的工作已 release 到 Megatron-LM 仓库，欢迎试用：https://github.com/NVIDIA/Megatron-LM/tree/dev。

在常规大语言模型训练中，激活数据（模型前向传播过程中产生的、且将用于反向传播阶段梯度计算的中间结果）的内存开销通常与序列长度、特征维度正相关。当序列长度或特征维度持续增加时，模型训练的整体动态内存开销显著增加，最终可能超出 GPU 内存上限。

这类问题常见于以下两类场景：

在长序列训练场景中，例如多模态模型训练及强化学习训练，序列长度相比较预训练来说有显著的增长，通常可达几十 k 甚至上百 k，造成激活动态显存开销急剧增加，而模型本身的显存占用保持不变，那么此时训练过程中激活动态显存的开销占比就会非常大。该场景下，相比于增加并行度这一同时针对模型参数和激活值的传统方案，我们更需要一种类似重计算的、专门针对激活值的显存优化技术，并期望获得性能上的优势。
第二类场景是 fine-grained MoE 模型的训练，目前这类模型牢牢占据了大语言模型排行榜的前列。它们的特点是总参数量和激活参数量的比值非常悬殊，典型的模型如下表所示，DeepSeek-V3 的总参数量和激活量的比值是 18，而 Kimi-K2 的这一比值达到了 31。这意味着和同等参数量的 Dense 模型相比，fine-grained MoE 模型的计算量是前者的 1/18 或 1/31。若不对显存做额外的优化，那每张卡上分到的计算量将远远少于 Dense 模型，这种计算密度的降低不仅会造成 kernel efficiency 降低，还会放大通信开销，同时使 Host 端的开销暴露出来，对性能优化提出了较大挑战。

针对以上问题，本文提出了 Fine-grained Activation Offloading，支持 module-level/operator-level 的 offloading，类似于 Fine-grained Recomputing，用户可指定 offload 某个特定的 module 的 activation，并与 Fine-grained Recomputing 功能组合使用，实现精准、高效、全面的内存管理。

本文介绍的工作已 release 到 Megatron-LM 仓库，欢迎试用：https://github.com/NVIDIA/Megatron-LM/tree/dev。

2.1 方案核心思想

Fine-grained Activation Offloading 的核心是“按需卸载、重叠计算、兼容并行”，具体包含三大设计原则：

粒度精细化：以模型模块（如 qkv_linear、core_attn、moe_act）为单位卸载激活，而非层或块。
计算 - 卸载重叠：前向卸载、反向重载与后续计算并行执行，隐藏数据传输开销。
全场景兼容：支持 PP=1/PP>1/VPP>1，兼容 BF16/Blockwise-FP8/MXFP8/NVFP4 精度、1F1B A2A 重叠、CUDA Graph 等 Megatron-Core 最新优化，同时适配 MoE、MLA 等复杂模型结构。

2.2 前向传播（Forward）的 offload 逻辑

前向传播过程中，激活卸载的关键是”模块计算完成后立即卸载激活，且不影响后续计算“，具体流程如下：

Offload 触发时机：当一个模块（如 core_attn）完成计算后，立即将其输入激活和中间激活 offload 至 CPU 内存，张量是否被 offload 取决于其是否通过 save_for_backward 存储至反向阶段使用。
计算 - offload 重叠：offload 操作与下一个模块（如 attn_proj）的计算并行执行，通过独立 CUDA 流（D2H 流）实现异步传输。

特殊规则：对于在 timeline 上 forward pass 与 backward pass 相邻的 chunk，其内部的最后一层的激活将不进行 activation offloading（如下图中标红的 chunk 所示），因为该层激活将在反向传播阶段优先使用，offload 后需立刻进行 reload，否则存在增加延迟的风险。

2.3 反向传播（Backward）的 reload 逻辑

反向传播过程中，reload 激活需与前向 offload 操作保持对称且错峰，避免出现局部时刻的内存开销翻倍，具体流程如下：

隔层 reload：当前模块（如 expert_fc2）完成反向计算后，再 reload 下一层对应模块（如 expert_fc2）的激活。
避免内存翻倍：通过“计算完成后再 reload ”的时序设计，确保同一时刻仅保留一层激活，无需维护 2 倍的激活数据量。
计算 - reload 重叠：reload 操作与下一个模块的反向计算并行执行，通过独立 CUDA 流（H2D 流）隐藏传输开销。

2.4 同步机制

激活的 offload/reload 操作和计算 kernel 存在数据依赖，因此必须添加同步机制以确保数据的正确性。本工作设计了四个同步规则：

启动 offload 时需确保当前 module 完成了前向计算。
启动 reload 时需确保当前 module 完成了反向计算。
启动 reload 时需确保 reload 对应的 offload 已经完成，即 cpu tensor 已经准备好被 reload 回 GPU。
启动反向计算前需确保其使用的所有激活值已经完成 reload 操作，即 gpu tensor 已经准备好被用于反向计算。

需要说明的是，上图中相邻的 forward chunk 与 backward chunk 是不同的两个 model chunk，因此 layer2 的 offload 可以与反向时 layer2 的 reload 操作进行 overlap，否则依据第 3 条同步规则，这两个操作应当是串行执行的。

3.1 兼容流水线并行（PP/VPP）

Fine-grained Activation Offloading 在 PP/VPP 场景下的核心挑战是“块间数据依赖与顺序管理”，解决方案如下：

Model chunk 内原生兼容：对于单个 PP/VPP chunk 内的模块，activation offloading 逻辑无需特别修改，与 PP=1 场景一致。
ChunkOffloadHandler 队列管理：通过单例管理器PipelineOffloadManager 维护 ChunkOffloadHandler 双端队列，每个 PP/VPP chunk 的每个 micro batch 分别对应一个独立的 ChunkOffloadHandler，其负责管理本 chunk 的所有 offload/reload 逻辑。
双端队列维护逻辑：ChunkOffloadHandler 以“ VPP stage 逆序、micro batch 正序”顺序入队，此时出队顺序即为反向传播时 VPP chunk 的执行顺序。
重叠优化：VPP chunk 间的 activation offload 与 reload 可并行进行，以减少传输延迟。
特殊规则：对于在 timeline 上 forward pass 与 backward pass 相邻的 chunk，其内部的最后一层的激活将不进行 activation offloading，因为该层激活将在反向传播阶段优先使用，offload 后需立刻进行 reload，否则存在增加延迟的风险。

3.2 兼容 Fine-grained Recomputing

Fine-grained Activation Offloading 与 Fine-grained Recomputing 可以互补增效。重计算减少轻量算子的内存占用，offloading 降低重计算开销较大的算子的内存压力，具体协同逻辑如下：

算子/模块差异化策略

轻量算子/模块：轻量算子/模块的重计算开销相对较低，而激活显存开销与高计算量算子/模块的显存开销在同一数量级，使用 offloading 会带来较前者更多的 overhead，因此进行重计算性价比更高。轻量算子/模块通常包括 layernorm、moe_act 等。
高计算量算子/模块：高计算量算子/模块进行重计算会引入很重的计算开销，使用 offloading 并通过与计算进行 overlap，可以在节省内存的同时达到吞吐性能上的无损或微损。高计算量算子/模块通常包括 core_attn、expert_fc1 等。

内存完全释放：对同一模块，结合“重计算释放中间变量+输出”和“ offloading 释放输入和中间变量”，可以实现对单模块动态显存占用的全覆盖，最终达到模块级别的 memory-free。

在现有的代码实现里，重计算和 offloading 可以覆盖一个 transformer layer 的所有 activation，用户可按需选择，极限情况下可以释放一个 transformer layer 的所有 activation。

3.3 与 Full Recomputation 的对比

Full Recomputation 通常会存储每个 transformer layer 的 input，释放掉 transformer layer 前向传播时的中间变量，因此最终的显存峰值会包括一层 layer 的临时中间变量以及前面所有 layer 的 input。

而 Fine-grained Activation Offloading + Recomputing 可以释放掉一个 transformer layer 的所有显存，最终的显存峰值为一层 layer 的临时中间变量，所以极限情况下，这种方案有着比 Full Recomputation 更高的显存优化潜力。

4.1 实验配置与核心指标

实验基于 Megatron-LM 框架，选取 DeepSeek-V3、Qwen3-235B 和 Dots2.llm 等典型大模型，围绕吞吐量和峰值内存开销两大指标，考察：

在使能了部分重计算的 baseline 基础上，对比开关 offload 功能对两大指标产生的影响，预期以极少的吞吐量损失换取较大内存收益。
在全层重计算的基础上，将部分子模块替换为 offload，其余模块采用部分重计算，对比 Fine-grained Activation Offloading + Fine-grained Recomputing 的组合策略与全层重计算在两大指标上的差异，预期在保证极致内存收益基本不变的同时获取吞吐量收益。

4.2 实验 1: DeepSeek-V3-proxy （64卡）

基本配置：14 层 moe layer、num-experts=64、采用 TP1PP4EP6VPP1 策略、MBS1GBS512、BF16 精度
offload 策略：baseline 不使用 offload，对比实验对 expert_fc1、moe_act 和 layernorm 使能 offload
以极少的吞吐量损失换取较大内存收益：

4.3 实验 2: DeepSeek-V3 （256卡）

基本配置：不包含 MTP 的完整 DeepSeek-V3 模型、采用 TP1PP8EP32VPP4 策略、MBS1GBS2048、MXFP8 精度
offload 策略：baseline 不使用 offload，对比实验对 moe_act 使能 offload
以极少的吞吐量损失换取较大内存收益：

4.4 实验 3: Qwen3-235B 长序列训练（256卡）

基本配置：Qwen3-235B 模型、64k 序列长度、BF16 精度

offload 策略：

bassline：TP4PP8EP8VPP4CP4，recompute layernorm、moe
offload: 在 baseline 基础上 CP4->CP1，offload qkv、layernorm、core_attn、attn_proj

在保证极致内存收益基本不变的同时获取吞吐量收益：

4.5 实验 4: Dots2.llm （512卡）

基本配置：完整 Dots2.llm 模型（Dots2 是小红书的最新一代自研模型，采用 MoE 架构，本实验采用该系列的文本模型 Dots2.llm，更多模型具体信息会在未来工作中逐步披露）、MBS1GBS4096

offload + recompute 组合策略：

offload：moe_act
recompute：除 moe_act 以外的所有 submodule，包括 attention module、router、moe_fc2 等

在保证极致内存收益基本不变的同时获取吞吐量收益：

4.6 开销分析与优化

实验中观察到两类潜在开销，具体实践中通过针对性优化将影响降至最低：

PCIe 带宽竞争：H2D/D2H 传输与跨节点 A2A 通信竞争 PCIe 带宽，导致 A2A 通信变慢。这个问题首先可以通过 Megatron-Core 提供的A2A Overlapping 解决，A2A 通信被计算 kernel 所隐藏，它的变慢对整体的性能影响不大；其次在具有比较大的 NVLink Domain的集群，A2A 通信可以控制在 NVLink domain 内，不会受 H2D/D2H 的影响
CPU 开销：张量打包 / 解包、钩子回调带来额外 CPU 负担，可以通过开启 CUDA Graph 来降低 CPU 开销对整体性能的影响；另外 Offloading 节省的显存可以用来增加 GPU 的计算密度，进一步降低训练中 CPU 开销带来的影响。

5.1 步骤 1：诊断内存瓶颈

通过手动计算或者 torch 的 memory snapshot 找出一个 transformer layer 占用显存量比较大的 topK 个 activation
Dump memory snapshot 的方法：

https://pytorch.org/blog/understanding-gpu-memory-1/

5.2 步骤 2：计算最大可卸载张量大小

需确保 offload/reload 的传输时间可被计算完全掩盖，关键公式：最大可卸载字节 = （单模块计算时间） × （H2D/D2H带宽）

带宽测试：使用 https://github.com/NVIDIA/nvbandwidth工具测试硬件带宽
计算时间：通过 NSys report 计算一个 transformer layer 的耗时，例如单层计算耗时 2ms，则可卸载最大字节数为 2ms×40GBps=80MB。
前反向计算不均衡：通常情况下我们只需考虑前向即可，因为反向计算时间是前向的 2 倍，而且 nvbandwidth 给出的数据显示 H2D 要略快于 D2H，故反向传播时的 reload 相较于前向的 offload 更容易被隐藏，offloading 的 overlap 瓶颈集中在前向阶段。

5.3 步骤 3：配置 offload 与后验证

参数配置：通过 --offload-modules 指定卸载模块，例如 --offload-modules=core_attn,expert_fc1
验证指标：

NSys report 时序图：检查 D2H/H2D 是否与计算完全重叠（如下图）
Memory snapshot：确认目标模块的 activation 已按预期释放。具体地，开发者可以仔细查看 memory snapshot 里的每次显存分配操作，来确认相应的 activation 是否被及时释放掉，查看显存节省量是否已达到预期。

5.4 步骤 4：迭代优化

调整 offload 模块：若内存仍不足，增加 offload 模块（如 attn_proj）；若性能损失过大，将部分 offload 模块换成重计算（如保留 layernorm 重计算）
并行度适配：结合上一步的显存节省量，通过显存估计器（https://huggingface.co/spaces/ISEEKYAN/megatron_memory_estimator）调整 TP/CP/PP 并行度（如将 TP 从 4 降低至 1，提升计算强度）
重复上述步骤

当前方案仍有进一步优化空间，后续将重点推进五大方向：

支持 CUDA Graph 捕获卸载模块：解决当前 CUDA Graph 无法捕获卸载操作的限制，进一步降低 CPU 开销；
优化 FP8 张量卸载：减少 FP8 张量在 CPU-GPU 传输中的类型转换与打包开销；
兼容 Megatron-FSDP：实现细粒度卸载与 FSDP（完全共享数据并行）的协同，适配更大规模模型训练。
支持不同层的卸载策略的差异化配置：当前卸载策略为全局配置，所有层采用统一的卸载策略，后续将允许不同层采取不同卸载策略，以适应流水线并行在 warm-up 阶段与 steady 阶段可掩盖的数据搬运量不一样的特点。
支持跨 pp rank 的显存均衡：我们可以让 Offload 策略感知 PP 里不同 rank 的显存消耗不一致的问题，从而控制较大的 pp rank 减少非必要的 offload 操作，使它们的显存占用尽可能均衡，最大程度利用显存和降低对性能的影响。

Megatron Core 中的细粒度激活卸载方案，通过 “模块级粒度、计算 - 传输重叠、全并行兼容” 的设计，在长上下文训练中实现了 “峰值内存收益 + 吞吐性能收益” 的联合最优。实验表明，该方案可使 DeepSeek-V3、Dots2.llm 等模型以吞吐损失 1%-2% 的代价换取峰值内存开销 10%-35% 的收益，或在保证极致内存收益的同时获取 7%-10% 的吞吐收益，为千亿级模型、长序列训练提供高效的内存解决方案。

结合本文的最佳实践，开发者可快速将该方案落地至实际训练任务，进一步突破大模型的内存瓶颈，推动更大规模、更长上下文的模型研发。