企业官网建设流程全解析

Megatron-LM 细粒度激活卸载（Fine-Grained Activation Offloading）完全指南：原理、配置与实战调优

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM

本文基于 Megatron-LM 官方用户指南 docs/user-guide/features/fine_grained_activation_offloading.md 并结合仓库源码编写。细粒度激活卸载（Fine-Grained Activation Offloading）是 Megatron-LM 与小红书（RedNote）协作贡献的一项显存优化技术：它以单个 Transformer 子模块为粒度，在训练前向过程中把激活异步拷贝到 CPU 内存，从而精确控制"哪些激活被卸载"，在显存节省与 PCIe 带宽开销之间取得可调的折中。读完本文，你将掌握该特性的全部可卸载模块、命令行参数与默认值、CUDA Graph 集成约束、与细粒度重计算的组合方式，以及其三层架构设计与底层 D2H/H2D 异步流程。

背景：为什么需要模块粒度的激活卸载

大模型训练中，前向过程为反向传播保存的中间激活（activations）会占据大量 GPU 显存。传统做法要么全部保留（显存峰值高），要么整层卸载（粒度粗糙、难以精确控制），要么全部重计算（用算力换显存，重计算开销大）。

细粒度激活卸载的思路是：以 Transformer 层内的单个子模块为单位，在子模块前向执行完毕后，把它的输入激活通过专用 CUDA 流异步拷贝（D2H）到 CPU 端固定内存（pinned memory），并在反向传播需要该激活之前再异步拷回（H2D）。相比整层卸载，它可以做到：

• 精确控制：只卸载显存占用大、且重计算成本高的模块（如core_attn、expert_fc1），轻量模块（如 layernorm、激活函数）则留给重计算；
• 可量化折中：通过--activation-offload-fraction、--min-offloaded-tensor-size、--delta-offload-bytes-across-pp-ranks等参数在显存收益与 PCIe 带宽开销之间调节；
• 异步隐藏延迟：D2H/H2D 拷贝在独立 CUDA 流上执行，与计算流重叠，尽可能不阻塞主计算。

从源码看，该特性与 CPU 整层卸载（cpu_offloading）互斥。在 transformer_config.py 的配置校验中明确断言二者不能同时开启：

assert ( not self.cpu_offloading ), "fine_grained_activation_offloading cannot be enabled with cpu_offloading."

快速上手：基本用法

启用细粒度激活卸载只需两个命令行参数。在训练脚本中加入：

# 开启细粒度激活卸载总开关 --fine-grained-activation-offloading # 指定要卸载输入激活的模块（多个模块以空格分隔） # 可选值: "attn_norm", "qkv_linear", "core_attn", "attn_proj", # "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp" --offload-modules core_attn attn_proj expert_fc1

在 transformer_config.py 中，两个开关分别对应TransformerConfig.fine_grained_activation_offloading（布尔开关）与TransformerConfig.offload_modules（模块名列表）。配置校验要求offload_modules非空，且所有模块名必须是上述合法集合之一（见 transformer_config.py），非法模块名会直接报错：

assert self.offload_modules is not None and len(self.offload_modules) > 0 ... invalid_modules = set(self.offload_modules) - allowed_modules assert not invalid_modules, f'Invalid choices for offload_modules: {invalid_modules}. ...'

例如，MoE 训练场景中 megatron/core/transformer/moe/README.md 给出的典型用法是卸载专家网络内部的激活：

--fine-grained-activation-offloading --offload-modules expert_fc1 moe_act

可卸载模块详解

每个模块的输入激活（即该模块前向的输入张量）会被卸载到 CPU，并在该模块反向计算之前重新载回 GPU。8 种可选模块及其语义如下：

模块如何接入卸载框架

从源码看，各模块通过统一的FineGrainedActivationOffloadingInterface接入框架，以"上下文管理器 + 提交点"的范式工作：

• Attention 侧（attention.py）：qkv_linear组的提交点位于get_query_key_value_tensors之后，core_attn组在核心注意力输出处提交，attn_proj组在输出投影处提交，并把core_attn的输出作为强制释放张量（forced_released_tensors），卸载完成后立即释放其 GPU 存储：

qkv_linear_manager = off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") with qkv_linear_manager as hidden_states: qkv_output = self.get_query_key_value_tensors(...) qkv_output = qkv_linear_manager.group_offload(qkv_output, forced_released_tensors=[]) ... attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out])

• MoE 专家侧（experts.py）：fused_group_mlp需要--use-transformer-engine-op-fuser走 Transformer Engine 的融合算子路径，其卸载粒度是整个融合分组 MLP，因此与expert_fc1/moe_act这类"部分卸载"互斥；expert_fc1与moe_act则是串接的两个提交点，先提交 FC1 输入、再提交激活函数输入。配置层面 transformer_config.py 会强制校验该约束：

if "fused_group_mlp" in self.offload_modules: moe_partial_offload = {"expert_fc1", "moe_act"} & set(self.offload_modules) assert not moe_partial_offload, ( "fused_group_mlp offloads the whole fused grouped MLP and cannot be " f"combined with expert_fc1 or moe_act. Remove: {moe_partial_offload}" )

• LayerNorm 侧（transformer_layer.py）：attn_norm/mlp_norm只有在对应 norm 不是IdentityOp时才启用（_set_offload_modules中检查not isinstance(..., IdentityOp)）。

卸载决策与强制释放

group_offload提交时，卸载管理器会针对组内每个被 autograd 保存的张量执行卸载检查（fine_grained_activation_offload.py）：

• 不是Parameter，不是FakeTensor/FunctionalTensor，且位于 CUDA 设备；
• 未被 Transformer Engine 标记为_TE_do_not_offload；
• 元素数不小于min_offloaded_tensor_size。

同时，调用方可以传入forced_released_tensors指定"卸载完成后立即释放 GPU 存储"的张量。释放通过untyped_storage().resize_(0)实现（见 fine_grained_activation_offload.py），适用于那些不会被 PyTorch 垃圾回收自动释放的输出张量。

调优参数

细粒度激活卸载提供 4 个调优旋钮（对应 transformer_config.py 中的TransformerConfig字段）：

# 参与卸载的最小张量元素数。小于该值的张量跳过不卸载。 # 默认: 1048576（约 1M 个元素） --min-offloaded-tensor-size 1048576 # 实际被卸载的"合格卸载组"比例，取值范围 [0, 1]。默认: 1.0 # 当 PCIe 带宽成为瓶颈时，可用于部分卸载 --activation-offload-fraction 0.8 # 高 PP rank 上减少的卸载量（单位: 字节）。默认: 0 # 高 PP rank 在飞行中的 microbatches 更少，少卸载一些既降低开销又不抬高显存峰值 --delta-offload-bytes-across-pp-ranks 1073741824 # 可选: 限制每个卸载组在飞行中的 D2H 卸载数（多数场景可省略或置 None） # 与本地全迭代 CUDA 图（cuda_graph_scope 为 full_iteration）搭配时必须设为非 None 非负整数 --fine-grained-offloading-max-inflight-offloads <N>

参数语义与源码对应

min_offloaded_tensor_size（TransformerConfig.min_offloaded_tensor_size，默认1024 * 1024）过滤掉小张量——对它们执行 D2H/H2D 的开销可能超过收益。它在 fine_grained_activation_offload.py 的张量检查中生效，并在配置校验中要求非负（transformer_config.py）。

activation_offload_fraction（TransformerConfig.activation_offload_fraction，默认1.0）不是"卸载激活字节的比例"，也不决定"启用哪些模块名"，详见下一节专门讲解。

delta_offload_bytes_across_pp_ranks（TransformerConfig.delta_offload_bytes_across_pp_ranks，默认0）用于跨流水线并行（PP）rank 的卸载量均衡。在 1F1B 调度下，高 PP rank 在飞行中的 microbatches 较少、激活峰值更低，因此少卸载一些不会抬高峰值显存，却可以减少 PCIe 流量。源码中该值按keep_on_gpu_bytes = pp_rank * delta_offload_bytes_across_pp_ranks计算每个 rank 需要"留在 GPU 上"的字节数，并从后往前关闭对应的卸载组（fine_grained_activation_offload.py）。配置校验要求其非负（transformer_config.py）。

fine_grained_offloading_max_inflight_offloads（TransformerConfig.fine_grained_offloading_max_inflight_offloads，默认None）为每个卸载组（例如moe_act、qkv_linear）设定"在主线流wait_event之前，最多允许在飞行中的 D2H 拷贝数"：

• 0：每次卸载提交后立即让主计算流等待；
• 更大的值：允许更多卸载与计算重叠，但也意味着主计算流要等更久；
• None：不做这些 join（绝大多数普通训练场景的默认选择）。

该参数仅在与本地全迭代 CUDA 图（cuda_graph_impl='local'且cuda_graph_scope含full_iteration）搭配时才必须设置：那条路径不依赖record_stream，必须显式 join 才能保证内存复用安全。配置校验要求其为非负整数（transformer_config.py）。源码实现上，每个组名维护一个 FIFO 事件队列（_offload_pending_by_name），超过上限时让主计算流等待最旧的事件（fine_grained_activation_offload.py）。

深入理解--activation-offload-fraction的语义

这是一个容易被误解的参数，官方文档专门辟出一节说明。要点如下：

• 它是对"合格卸载组"的比例，既不是字节比例，也不是"启用哪些模块名"的选择器；
• 它与--offload-modules协同工作：--offload-modules中列出的模块名全部照常注册各自的卸载组，然后该比例在所有已配置模块合并后的合格组列表上一次性生效；
• 管理器按前向执行顺序保留前 N% 的组进行卸载，后 (1-N%) 的组留在 GPU 上。

例如：--offload-modules core_attn attn_proj expert_fc1 --activation-offload-fraction 0.5，那么core_attn、attn_proj、expert_fc1三者的组按执行顺序合并成一个列表，前 50% 的组被卸载。这里 0.5 不代表"卸载 50% 的激活字节"，也不代表"只卸载前 50% 的模块名"。

此外，该比例是在其他资格过滤之后应用的。源码中post_warmup_callback的执行顺序是（fine_grained_activation_offload.py）：

1. 先应用 offload margin（为避免反向重载阻塞计算流而保留在 GPU 上的末尾组）；
2. 再应用delta_offload_bytes_across_pp_ranks（PP rank 差异）；
3. 最后应用activation_offload_fraction——此时 N% 是"经过上述过滤后剩余合格组"的比例，且优先保留更早的前向组（因为越早释放激活，显存压力缓解越早）。

eligible_offload_groups = [ group for group in chunk.offload_groups if group.offload and group.total_offload_bytes > 0 ] offloaded_groups_count = len(eligible_offload_groups) disabled_groups_count = int(offloaded_groups_count * (1 - self._activation_offload_fraction)) # 从后往前关闭，优先保留较早的前向组 for group in reversed(eligible_offload_groups): if disabled_groups_count > 0: disabled_groups_count -= 1 group.offload = False else: break

CUDA Graph 集成

细粒度激活卸载与 CUDA Graph 兼容，但有一组严格约束：

• attn_norm和mlp_norm不能被卸载（它们跨越 CUDA Graph 边界）；
• cuda_graph_scope必须包含attn与moe_router；
• cuda_graph_impl必须是transformer_engine；
• 需要torch >= 2.9.0与transformer_engine >= 2.14.0。

约束在 transformer_layer.py 的_set_offload_modules中落地：检测到 CUDA graph 启用时会自动关闭attn_norm/mlp_norm的卸载并打印提示；cuda_graph_impl='local'（本地局部图）下则只支持 MoE 侧模块（expert_fc1、moe_act、fused_group_mlp），见 transformer_config.py。

--delay-offload-until-cuda-graph（TransformerConfig.delay_offload_until_cuda_graph）

这是为"MoE 专家计算位于 CUDA Graph 之外"的典型布局设计的一个优化选项：

图内 vs 图外。位于被捕获的cuda_graph_scope内部的卸载边界（例如attn入图时的qkv_linear、core_attn、attn_proj）属于图捕获与重放的一部分，其卸载相关工作随图一起重放，不会像纯 eager 路径那样每个 step 都产生 CPU launch 开销。而图外的边界（推荐的 MoE 布局下，graphedmoe_router之后的专家计算，如卸载expert_fc1/moe_act）每步仍以普通 eager PyTorch 执行；若不延迟，每个group_offload都会在 forward 到达提交点时立刻从 host 提交 D2H 工作。

该 flag 做了什么。它只影响显式接入"延迟组提交"的卸载提交（当前是 MoE 专家路径：expert_fc1、moe_act）。在每个 layer 的 TransformerEngine CUDA Graph 重放前后，卸载管理器进入replay 模式：延迟提交不再立即 launch D2H，而是把(callback, group name, forced tensors)入队（push_offload_groups），在图重放返回后由flush_delayed_groups按前向顺序发出排队的 D2H 拷贝，且不改变卸载/重载语义。相关代码位于 fine_grained_activation_offload.py 与 transformer_layer.py。

什么时候真正省时间（重放后的 EP A2A）。该收益的前提是图重放后存在一个真实的 CPU/GPU 同步空档——在常见 MoE 训练布局中，专家并行（EP）的 all-to-all 与相关 dispatch 紧跟 graphedmoe_router区域。A2A 路径通常需要 host 协调集合通信并与 GPU 同步（如等待图工作完成或通信 staging），此期间 CPU 无法与有用的 launch 工作完全重叠。在cudaGraphLaunch返回后立即安排flush_delayed_groups，正好利用这个空档从 host 发出 D2H 拷贝：入队成本大部分被 EP A2A 本就要付出的 slack 隐藏。若图重放后没有这样的同步（或专家工作被完整捕获进图内、无 host 可见空档），延迟提交就提供不了这种"免费"的 host 时间。

行为要点

• 不会替换或"延迟"图内 graphedattn区域的 attention 侧卸载——那些不在延迟路径上；
• Warmup 与非重放 forward 仍会立即提交可延迟的组（无 replay 模式的延迟）；
• 必须与细粒度激活卸载 + CUDA Graph 同时使用，并遵守本节的规则（TEcuda_graph_impl、scope 含attn与moe_router等）；
• 图计算路径与d2h_stream之间的流顺序仍沿用既有事件（forward_record/backward_record），该选项只改变"合格 D2H 工作从 host 提交的时机"。

与细粒度重计算组合使用

卸载与重计算是互补的显存优化手段，官方推荐按模块"轻重"分工：

• 重计算用于轻量模块（如 layernorm、激活函数），其重算算力开销可忽略；
• 卸载用于重量模块（如core_attn、expert_fc1），重算它们代价太高。

--recompute-granularity selective --recompute-modules layernorm moe_act --fine-grained-activation-offloading --offload-modules core_attn attn_proj expert_fc1

这样组合后，几乎可以把单个 Transformer 层在设备上的激活全部"赶走"：轻量激活就地重算，重量激活搬到 CPU。

配置校验还包含一条相关约束：当recompute_modules含'mhc'（或对整个 MoE 层做全量重计算）时与细粒度卸载存在互斥/告警逻辑（见 transformer_config.py 与 transformer_config.py）——例如对整个 MoE 层重计算时不能再卸载moe_act/expert_fc1/fused_group_mlp，因为激活根本不会保存，卸载冗余且会引发错误。

兼容性一览

官方文档给出的兼容性矩阵如下：

此外从配置校验可补充：与cpu_offloading互斥；与moe_paged_stash互斥（开启时offload_modules不得包含expert_fc1/moe_act/fused_group_mlp，见 transformer_config.py）；token-drop MoE 暂不支持（transformer_config.py）；cuda_graph_impl='local'下仅支持 MoE 侧模块。

工作原理：三层架构

实现分为三层（核心代码全部位于 megatron/core/pipeline_parallel/fine_grained_activation_offload.py，共约 1600 行）：

1. PipelineOffloadManager（单例）：全局协调者，管理专用 CUDA 流（d2h_stream、h2d_stream、cuda_graph_stream）、共享 CPU 张量池，以及跨流水线 stage 的 chunk 生命周期。通过get_instance()获取单例（fine_grained_activation_offload.py）。
2. ChunkOffloadHandler：每个 microbatch（chunk）一个的处理器，负责张量组的登记、D2H/H2D 传输的执行，以及"哪些组真正卸载"的决策（fine_grained_activation_offload.py 起）。
3. FineGrainedActivationOffloadingInterface：供 Transformer 模块（attention、MoE 等）使用的轻量接口，用于标记卸载边界——group_start进入上下文、group_offload提交卸载组（fine_grained_activation_offload.py 起）。

卸载/重载流程

Forward pass (Layer N): Backward pass (Layer N): ┌─────────────────────┐ ┌───────────────────────┐ │ group_start(input) │─── register ──► │ │ │ │ tensor group │ group_commit_backward │ │ module.forward() │ │ wait H2D complete │ │ │ │ pop tensors from │ │ group_offload(out) │─── D2H async ──► │ CPU → GPU │ │ on d2h_stream │ to pinned CPU │ on h2d_stream │ └─────────────────────┘ └───────────────────────┘

具体步骤：

1. group_start：注册一个新的张量组，并挂接saved_tensors_hooks以拦截 autograd 的save_for_backward。源码中该 hook 由PipelineOffloadManager维护（on_save_for_backward/on_get_saved_tensor，见 fine_grained_activation_offload.py），并通过两个 identity 形式的 autograd Function（FineGrainedOffloadingGroupStartFunction与FineGrainedOffloadingGroupCommitFunction，fine_grained_activation_offload.py）把 forward/backward 事件串进计算图。
2. 前向执行：组内被 autograd 保存的所有张量都被捕获（以(group_index, position)标签登记）。
3. group_offload：在专用 CUDA 流d2h_stream上触发异步 D2H 拷贝，可选地释放输入张量的 GPU 存储（forced_released_tensors的untyped_storage().resize_(0)）。批量卸载前会先让d2h_stream等待计算流（d2h_stream.wait_stream(current_stream)），保证拷贝发生在计算之后（fine_grained_activation_offload.py）。
4. 反向：在组反向开始前，张量从 CPU 经h2d_stream重载回 GPU，计算流等待传输完成事件（wait_reload_event），确认重载就绪后再使用（fine_grained_activation_offload.py）。

一个值得注意的实现细节：对于非连续视图（如transpose/permute产生的张量），若其覆盖底层存储的字节数占比达到BASE_OFFLOAD_MIN_COVERAGE = 0.5，则直接按扁平 storage 做 memcpy（额外携带(size, stride, storage_offset)视图元数据），避免先 gather 成临时连续拷贝带来的额外 GPU 缓冲与 kernel 开销；重载时再通过as_strided还原视图（fine_grained_activation_offload.py）。

Warmup 与自适应卸载决策

第一个训练迭代是 warmup 阶段：管理器记录张量组、大小与执行顺序。warmup 结束后post_warmup_callback（fine_grained_activation_offload.py）按顺序执行：

1. 预留 margin：最后 N 个（按去重组名计数）组保留在 GPU 上，避免反向重载阻塞计算流；
2. 应用 PP rank 差异：高 PP rank 少卸载字节（由delta_offload_bytes_across_pp_ranks控制）；
3. 应用比例：跨所有已配置模块，只卸载剩余合格组的前 N%（由activation_offload_fraction控制）；
4. 打印汇总表：通过print_offload_summary_table输出各 rank 分组的卸载字节数 ASCII 表（单位 MB），并告警同一 storage 在组内被重复拷贝浪费 PCIe 带宽的情况（fine_grained_activation_offload.py）。

CPU 张量池

OffloadTensorPool（fine_grained_activation_offload.py）在 CPU 固定内存上按(shape, dtype)缓存张量，避免反复cudaMallocHost/cudaFreeHost，显著降低首个迭代之后的 D2H 延迟；池内统计（命中/未命中、利用率）可用于诊断。需要留意的是：MoE 卸载组（expert_fc1、moe_act、fused_group_mlp）的张量形状在运行前不可预知，因此不使用 CPU 池，直接分配 pinned 张量（fine_grained_activation_offload.py）。

CUDA Graph 支持

与 CUDA Graph 交互时：

• 专用cuda_graph_stream运行被捕获的计算，d2h_stream与图内区域的 D2H 传输重叠；manager 提供的cuda_graph_stream/cuda_graph_event会被注入 TransformerEngine 的cudagraph_kwargs，让 TE 对图计算与 D2H/H2D 传输做顺序约束（见 module.py 与FineGrainedOffloadingBackwardRecordFunction，后者在图内记录 backward 事件并等待h2d_stream）；
• CUDA Graphwarmup期间卸载被禁用（通过disable_offload/enable_offload与 warmup hooks 配合）；
• delay_offload_until_cuda_graph将 D2H launch 推迟到图重放之后，利用cudaGraphLaunch期间的 CPU 空闲发出卸载命令，几乎零 CPU 开销。

工程落地与验证

在调度层，每个训练迭代开始/结束时通过off_interface.reset(process_group=...)重置管理器状态（schedules.py），校验/推理阶段也会在 forward-only 路径重置卸载状态（schedules.py）。

单元测试 tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py 提供了该特性正确性与显存收益的验证范式：对每个offload_modules组合（attn_norm、qkv_linear、core_attn、core_attn+attn_proj、mlp_norm、expert_fc1、moe_act，覆盖稠密 GPT 与 MoE、MLA 变体），对照"无卸载"基线比较 forward 输出、反向梯度与峰值显存，确认卸载开启后显存下降与记录的卸载字节数大致吻合。测试还演示了一个实操要点：为单元测试的确定性，可将min_offloaded_tensor_size调小（如1024）以强制所有张量参与卸载。

使用建议小结

• 先重计算轻量模块，再卸载重量模块：--recompute-modules layernorm moe_act搭配--offload-modules core_attn attn_proj expert_fc1是文档推荐的组合范式，可在设备上释放几乎全部单层激活；
• PCIe 带宽吃紧时：调低--activation-offload-fraction做部分卸载，优先保住前向靠前、释放收益最大的组；
• PP 规模较大时：用--delta-offload-bytes-across-pp-ranks让高 rank 少卸载，均衡负载；
• 开启 CUDA Graph 时：严格遵守 scope/impl/版本约束，attn_norm、mlp_norm不可卸载；本地全迭代图场景必须为--fine-grained-offloading-max-inflight-offloads设置非 None 非负整数；
• MoE 训练：--delay-offload-until-cuda-graph可利用 EP A2A 的同步空档隐藏 D2H 提交开销，但前提是图重放后确有 host-visible 的同步间隙。

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM