企业官网建设流程全解析

MindSpeed LLM重计算与异步激活卸载：大模型显存优化完全指南

【免费下载链接】MindSpeed-LLM昇腾LLM分布式训练框架项目地址: https://gitcode.com/Ascend/MindSpeed-LLM

训练大模型时，激活值往往是 NPU 显存的头号杀手——序列越长、Batch 越大，前向传播缓存的中间张量就越夸张，OOM（显存溢出）随之而来。本文介绍昇腾分布式训练框架MindSpeed-LLM提供的两大显存优化利器：重计算（Recompute）与异步激活卸载（Async Activation Offload），并给出新手可直接套用的选型顺序与实战案例 🚀。

一、为什么需要这两项技术？

大模型训练的显存消耗主要来自四部分：参数、梯度、优化器状态和激活值。其中激活值占比随序列长度线性甚至平方增长，最先耗尽显存。

MindSpeed-LLM 针对激活值提供了两条互补路线：

官方文档详细阐述了二者各自的背景与瓶颈：重计算 与 异步激活卸载。

二、重计算：三种粒度按需选择

重计算的核心是"用计算换显存"。MindSpeed-LLM 提供三种粒度，按需开启即可 🎛️

2.1 全量重计算（full）

显存非常吃紧时的兜底方案：只保存 Transformer 层（或层组）的输入激活值，其余全部重算。通过参数开启：

--recompute-granularity full --recompute-method uniform --recompute-num-layers 2

• uniform（均匀）：把层均匀分组，每组--recompute-num-layers层按组存储输入
• block（分块）：只对前 N 层重计算，其余层保留

实现逻辑可参考 Transformer 前向调度入口：transformer_block.py。

2.2 选择性重计算（selective）⭐ 推荐

只重算 Attention 中的core attention部分：把占显存少但重算开销大的激活留在显存里，把占显存大但重算便宜的激活丢掉重算。这是显存与性能平衡的最佳选择：

--recompute-granularity selective

2.3 激活函数重计算与细粒度重计算

• 添加--recompute-activation-function可单独开启激活函数重计算，并用--recompute-activation-function-num-layers ${num}指定层数，可与全量重计算叠加（此时仅支持block方法）。
• DeepSeek V4 场景还支持CSA/MHC 细粒度重计算（--recompute-csa-attention、--mhc-recompute），仅重算特定模块的中间结果，详见 DeepSeek V4细粒度重计算。

三、异步激活卸载：把拷贝藏进计算里

重计算的痛点是冗余计算多，尤其在长序列下 Self-Attention 计算量随序列平方增长。异步激活卸载（Async Activation Offload）则走另一条路：把激活值张量从 NPU 卸载到 Host 内存，显著降低峰值显存，同时用异步机制"藏"起拷贝开销。

它由三个机制配合完成：

1. D2H 异步卸载：前向计算 block 时，激活值在独立 stream 上被异步拷到 Host 侧
2. H2D 预取（prefetch）：反向开始之前，提前把后续要用的张量从 Host 拉回 NPU
3. 多流异步执行：D2H/H2D 与计算流并行，拷贝被计算覆盖

3.1 如何接入？

该特性以 Python 上下文管理器的形式提供，核心入口是async_save_on_cpu，按 block 组织张量生命周期：

with async_save_on_cpu( h2d_stream=h2d_stream, d2h_stream=d2h_stream, block_idx=block_idx, depth=depth, custom_check_fn=your_check_fn ): output = layer(input) # 某个 block 的前向计算

关键参数中的custom_check_fn决定哪些张量值得卸载——官方建议筛选"计算量大、激活参数量小"的部分做 offload，"激活大、计算耗时短"的部分交给重计算，否则拷贝开销无法被掩盖。实现代码见 async_offload.py，预训练入口的调用示例可参考 qwen3_moe.py 中对async_save_on_cpu的使用。

3.2 实测收益

• 长序列场景：卸载 Self-Attention 前向激活并在重算时跳过该部分，典型场景端到端性能收益20%+
• FSDP2 场景：短序列下计算掩盖不住通信，卸载激活腾出显存后可增大 micro-batch size 或序列长度，端到端性能收益60%+

四、显存不够了？按这个顺序排查优化

遇到 OOM 时不要盲目开特性，先定位显存花在哪（参数/梯度/优化器状态 vs 激活 vs logits），再按下述优先级评估——优先选对性能影响小的方案：

ChunkLoss / 融合算子 → CP / EP / FSDP 切分 → 重计算 → 异步激活卸载 → Swap Optimizer

配套工具与文档：

• 显存来源分析：Profiling 的 Memory 视图，见 性能数据采集
• logits 峰值优化：chunk_loss.py
• 其他显存优化特性（层输入交换等）位于 features_manager/memory/ 目录
• 完整 FSDP2 瓶颈排查与特性选型表格：fsdp2_backend_performance_optimization.md

五、新手快速上手清单

1. ✅ 显存紧张 → 先开--recompute-granularity selective（成本最低、收益明显）
2. ✅ 仍不够 → 叠加--recompute-granularity full --recompute-method uniform扩大重算范围
3. ✅ 长序列 / 高显存占比激活 + Host 内存充足 → 在 block 前向中接入async_save_on_cpu，用custom_check_fn精挑卸载张量
4. ✅ 调优后用 Profiling 的 Timeline + Memory 视图验证：确认 D2H/H2D 被计算掩盖、峰值显存下降
5. ✅ 每次只动一个开关，用相同口径重新测量 step 耗时与吞吐

重计算解决"存不下"，异步卸载解决"算太慢"——两者组合使用，就是 MindSpeed-LLM 面对超大显存需求时的完整答案。掌握这份指南，你的大模型训练将再无显存之忧 💪

相关资料

• 重计算官方文档
• 异步激活卸载官方文档
• DeepSeek V4 细粒度重计算
• FSDP2 后端性能优化指南
• 异步卸载实现源码：mindspeed_llm/fsdp2/features/async_offload.py

【免费下载链接】MindSpeed-LLM昇腾LLM分布式训练框架项目地址: https://gitcode.com/Ascend/MindSpeed-LLM