文章总结： 针对大语言模型训练中激活值显存开销成为核心瓶颈的问题，小红书技术团队提出流水线感知的细粒度激活卸载方案，在Megatron-Core框架中实现模块级激活卸载，兼容流水线并行和重计算。实验表明该方案以1%-2%的吞吐损失换取10%-35%的内存节省，或在保证内存收益的同时获取7%-10%的吞吐提升，为千亿级模型长序列训练提供高效内存解决方案。文章提供了诊断内存瓶颈、计算最大可卸载张量大小、配置offload与验证、迭代优化等实施步骤。
综合评分： 92
文章分类： AI安全,安全工具,其他
大模型训练的高效内存解决方案：流水线感知的细粒度激活卸载，实现显存开销与吞吐性能的联合最优
等你加入的
小红书技术REDtech
2025年12月11日 18:04
北京
随着大语言模型（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% 的吞吐收益，为千亿级模型、长序列训练提供高效的内存解决方案。
结合本文的最佳实践，开发者可快速将该方案落地至实际训练任务，进一步突破大模型的内存瓶颈，推动更大规模、更长上下文的模型研发。
本工作由 Rednote HiLab 与英伟达技术专家完成。在此，衷心感谢项目相关同事所提供的支持与帮助，他们包括但不限于 Zhuran（Rednote）、Wenliang（Rednote）、Guangsu（Rednote）、Hongxiao、Zijie、Xin、Yan 及 David。
此外，本工作的开展受到 Kuaishou 与英伟达合作的 GTC Talk 启发，并借鉴了其中相关思路，特此向该分享团队致以谢意。
Yimeng, Rednote HiLab
小红书大模型训练框架开发工程师，目前从事小红书大模型训练框架的开发与优化工作。
Hongbin, NVIDIA Devtech
NVIDIA 加速计算专家，专注于大语言模型训练的功能开发及性能调优，深度参与 Megatron-Core 的相关开发工作，目前主要负责 MoE 模型的通信及显存优化工作。
【hi lab】大模型高性能计算 AI infra 工程师-训练/推理/模型轻量化（社招）
工作职责
【训练推理框架研发】
大模型hi lab AI Infra团队专注于大语言模型领域的前沿技术研究和落地，提供高性能、高可靠、可扩展的机器学习系统、丰富的异构计算资源和极致的端到端的机器学习服务体验，为公司提供核心技术能力和服务：
- 负责机器学习框架的研究与开发，服务于公司各个产品；
- 高效部署，优化NLP/多模态大模型核心业务模型；
【轻量化】
机器学习系统团队需要将传统或者新型的轻量化算法和工程有机结合起来进行加速，提高大语言模型训练或者推理性能的同时，通过算法手段尽可能降低效果损失。候选人将在以下几个方向进行深入探索和落地：
- 量化方向：负责但不限于大语言模型的低精度训练(FP8)、推理（W8A8KV8等）、低精度优化器（量化梯度、优化器状态、参数等）
- 高性能模型结构：大语言模型Finetune或者其他阶段的LoRA系列（熟悉各种变种），训练阶段的MQA/GQA系列等
- 稀疏化方向：大语言模型剪枝、稀疏、蒸馏、Sparse Attention 等
- 新型方向：Medusa、超长文本、Speculative Sampling等
任职资格
【分布式】
- 熟练掌握Linux环境下的C/C++、Python语言
- 具备扎实的计算机科学功底和编程能力，熟悉常见算法和数据结构，具有良好的编程习惯
- 能够熟练使用至少一种主流的机器学习框架(TensorFlow / PyTorch等)，熟悉框架内部实现;
- 熟悉Transformer模型及其应用场景等
【轻量化】
- 拥有上述一个或多个方向的实际落地经验，能独立从头到尾完成上述某个方向
- 熟练使用主流大规模训练或推理引擎开发，例如DeepSpeed/Megatron/Transformer Engine 等
- 两年以上工作经验、工程能力较强者优先，有实际 Cuda 算子开发经验者优先
投递邮箱
往期精彩内容指路
AAAI 2026 | 跨视频推理基准 CrossVid：给多模态大模型出一道“综合题”
EMNLP 2025｜思维链提示虽好，但它可能正在“隐藏”大模型的幻觉
免责声明：
本文所载程序、技术方法仅面向合法合规的安全研究与教学场景，旨在提升网络安全防护能力，具有明确的技术研究属性。
任何单位或个人未经授权，将本文内容用于攻击、破坏等非法用途的，由此引发的全部法律责任、民事赔偿及连带责任，均由行为人独立承担，本站不承担任何连带责任。
本站内容均为技术交流与知识分享目的发布，若存在版权侵权或其他异议，请通过邮件联系处理，具体联系方式可点击页面上方的联系我。
本文转载自：小红书技术REDtech 等你加入的《大模型训练的高效内存解决方案：流水线感知的细粒度激活卸载，实现显存开销与吞吐性能的联合最优》