尧图网站建设 
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, fInvalid 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 种可选模块及其语义如下模块说明备注attn_normAttention 输入 LayerNorm 的输入使用IdentityOp如 QK layernorm 被融合/替换时自动跳过qkv_linearQKV 线性投影的输入core_attn核心注意力softmax matmul的输入attn_projAttention 输出投影的输入必须与core_attn搭配使用mlp_normMLP 前 LayerNorm 的输入使用IdentityOp时自动跳过expert_fc1MoE 专家第一个 FC 层的输入仅 MoE 模型moe_actMoE 专家激活函数的输入仅 MoE 模型fused_group_mlp整个融合分组 MLP 的输入需要--use-transformer-engine-op-fuser不能与expert_fc1、moe_act同时使用模块如何接入卸载框架从源码看各模块通过统一的FineGrainedActivationOffloadingInterface接入框架以上下文管理器 提交点的范式工作Attention 侧attention.pyqkv_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.pyfused_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 fcombined with expert_fc1 or moe_act. Remove: {moe_partial_offload} )LayerNorm 侧transformer_layer.pyattn_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_sizeTransformerConfig.min_offloaded_tensor_size默认1024 * 1024过滤掉小张量——对它们执行 D2H/H2D 的开销可能超过收益。它在 fine_grained_activation_offload.py 的张量检查中生效并在配置校验中要求非负transformer_config.py。activation_offload_fractionTransformerConfig.activation_offload_fraction默认1.0不是卸载激活字节的比例也不决定启用哪些模块名详见下一节专门讲解。delta_offload_bytes_across_pp_ranksTransformerConfig.delta_offload_bytes_across_pp_ranks默认0用于跨流水线并行PPrank 的卸载量均衡。在 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_offloadsTransformerConfig.fine_grained_offloading_max_inflight_offloads默认None为每个卸载组例如moe_act、qkv_linear设定在主线流wait_event之前最多允许在飞行中的 D2H 拷贝数0每次卸载提交后立即让主计算流等待更大的值允许更多卸载与计算重叠但也意味着主计算流要等更久None不做这些 join绝大多数普通训练场景的默认选择。该参数仅在与本地全迭代 CUDA 图cuda_graph_impllocal且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先应用 offload margin为避免反向重载阻塞计算流而保留在 GPU 上的末尾组再应用delta_offload_bytes_across_pp_ranksPP rank 差异最后应用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: breakCUDA Graph 集成细粒度激活卸载与 CUDA Graph 兼容但有一组严格约束attn_norm和mlp_norm不能被卸载它们跨越 CUDA Graph 边界cuda_graph_scope必须包含attn与moe_routercuda_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_impllocal本地局部图下则只支持 MoE 侧模块expert_fc1、moe_act、fused_group_mlp见 transformer_config.py。--delay-offload-until-cuda-graphTransformerConfig.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因为激活根本不会保存卸载冗余且会引发错误。兼容性一览官方文档给出的兼容性矩阵如下特性支持情况PP / Interleaved PP / PP1支持细粒度重计算Fine-grained recomputation支持FP8 训练支持MTP多 Token 预测支持混合稠密与 MoE 层支持A2A overlapEP支持CUDA GraphTE 实现支持此外从配置校验可补充与cpu_offloading互斥与moe_paged_stash互斥开启时offload_modules不得包含expert_fc1/moe_act/fused_group_mlp见 transformer_config.pytoken-drop MoE 暂不支持transformer_config.pycuda_graph_impllocal下仅支持 MoE 侧模块。工作原理三层架构实现分为三层核心代码全部位于 megatron/core/pipeline_parallel/fine_grained_activation_offload.py共约 1600 行PipelineOffloadManager单例全局协调者管理专用 CUDA 流d2h_stream、h2d_stream、cuda_graph_stream、共享 CPU 张量池以及跨流水线 stage 的 chunk 生命周期。通过get_instance()获取单例fine_grained_activation_offload.py。ChunkOffloadHandler每个 microbatchchunk一个的处理器负责张量组的登记、D2H/H2D 传输的执行以及哪些组真正卸载的决策fine_grained_activation_offload.py 起。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 │ └─────────────────────┘ └───────────────────────┘具体步骤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 FunctionFineGrainedOffloadingGroupStartFunction与FineGrainedOffloadingGroupCommitFunctionfine_grained_activation_offload.py把 forward/backward 事件串进计算图。前向执行组内被 autograd 保存的所有张量都被捕获以(group_index, position)标签登记。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。反向在组反向开始前张量从 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_callbackfine_grained_activation_offload.py按顺序执行预留 margin最后 N 个按去重组名计数组保留在 GPU 上避免反向重载阻塞计算流应用 PP rank 差异高 PP rank 少卸载字节由delta_offload_bytes_across_pp_ranks控制应用比例跨所有已配置模块只卸载剩余合格组的前 N%由activation_offload_fraction控制打印汇总表通过print_offload_summary_table输出各 rank 分组的卸载字节数 ASCII 表单位 MB并告警同一 storage 在组内被重复拷贝浪费 PCIe 带宽的情况fine_grained_activation_offload.py。CPU 张量池OffloadTensorPoolfine_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_streamCUDA 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_attnattn_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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Megatron-LM 细粒度激活卸载（Fine-Grained Activation Offloading）完全指南：原理、配置与实战调优
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, fInvalid 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 种可选模块及其语义如下模块说明备注attn_normAttention 输入 LayerNorm 的输入使用IdentityOp如 QK layernorm 被融合/替换时自动跳过qkv_linearQKV 线性投影的输入core_attn核心注意力softmax matmul的输入attn_projAttention 输出投影的输入必须与core_attn搭配使用mlp_normMLP 前 LayerNorm 的输入使用IdentityOp时自动跳过expert_fc1MoE 专家第一个 FC 层的输入仅 MoE 模型moe_actMoE 专家激活函数的输入仅 MoE 模型fused_group_mlp整个融合分组 MLP 的输入需要--use-transformer-engine-op-fuser不能与expert_fc1、moe_act同时使用模块如何接入卸载框架从源码看各模块通过统一的FineGrainedActivationOffloadingInterface接入框架以上下文管理器 提交点的范式工作Attention 侧attention.pyqkv_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.pyfused_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 fcombined with expert_fc1 or moe_act. Remove: {moe_partial_offload} )LayerNorm 侧transformer_layer.pyattn_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_sizeTransformerConfig.min_offloaded_tensor_size默认1024 * 1024过滤掉小张量——对它们执行 D2H/H2D 的开销可能超过收益。它在 fine_grained_activation_offload.py 的张量检查中生效并在配置校验中要求非负transformer_config.py。activation_offload_fractionTransformerConfig.activation_offload_fraction默认1.0不是卸载激活字节的比例也不决定启用哪些模块名详见下一节专门讲解。delta_offload_bytes_across_pp_ranksTransformerConfig.delta_offload_bytes_across_pp_ranks默认0用于跨流水线并行PPrank 的卸载量均衡。在 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_offloadsTransformerConfig.fine_grained_offloading_max_inflight_offloads默认None为每个卸载组例如moe_act、qkv_linear设定在主线流wait_event之前最多允许在飞行中的 D2H 拷贝数0每次卸载提交后立即让主计算流等待更大的值允许更多卸载与计算重叠但也意味着主计算流要等更久None不做这些 join绝大多数普通训练场景的默认选择。该参数仅在与本地全迭代 CUDA 图cuda_graph_impllocal且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先应用 offload margin为避免反向重载阻塞计算流而保留在 GPU 上的末尾组再应用delta_offload_bytes_across_pp_ranksPP rank 差异最后应用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: breakCUDA Graph 集成细粒度激活卸载与 CUDA Graph 兼容但有一组严格约束attn_norm和mlp_norm不能被卸载它们跨越 CUDA Graph 边界cuda_graph_scope必须包含attn与moe_routercuda_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_impllocal本地局部图下则只支持 MoE 侧模块expert_fc1、moe_act、fused_group_mlp见 transformer_config.py。--delay-offload-until-cuda-graphTransformerConfig.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因为激活根本不会保存卸载冗余且会引发错误。兼容性一览官方文档给出的兼容性矩阵如下特性支持情况PP / Interleaved PP / PP1支持细粒度重计算Fine-grained recomputation支持FP8 训练支持MTP多 Token 预测支持混合稠密与 MoE 层支持A2A overlapEP支持CUDA GraphTE 实现支持此外从配置校验可补充与cpu_offloading互斥与moe_paged_stash互斥开启时offload_modules不得包含expert_fc1/moe_act/fused_group_mlp见 transformer_config.pytoken-drop MoE 暂不支持transformer_config.pycuda_graph_impllocal下仅支持 MoE 侧模块。工作原理三层架构实现分为三层核心代码全部位于 megatron/core/pipeline_parallel/fine_grained_activation_offload.py共约 1600 行PipelineOffloadManager单例全局协调者管理专用 CUDA 流d2h_stream、h2d_stream、cuda_graph_stream、共享 CPU 张量池以及跨流水线 stage 的 chunk 生命周期。通过get_instance()获取单例fine_grained_activation_offload.py。ChunkOffloadHandler每个 microbatchchunk一个的处理器负责张量组的登记、D2H/H2D 传输的执行以及哪些组真正卸载的决策fine_grained_activation_offload.py 起。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 │ └─────────────────────┘ └───────────────────────┘具体步骤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 FunctionFineGrainedOffloadingGroupStartFunction与FineGrainedOffloadingGroupCommitFunctionfine_grained_activation_offload.py把 forward/backward 事件串进计算图。前向执行组内被 autograd 保存的所有张量都被捕获以(group_index, position)标签登记。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。反向在组反向开始前张量从 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_callbackfine_grained_activation_offload.py按顺序执行预留 margin最后 N 个按去重组名计数组保留在 GPU 上避免反向重载阻塞计算流应用 PP rank 差异高 PP rank 少卸载字节由delta_offload_bytes_across_pp_ranks控制应用比例跨所有已配置模块只卸载剩余合格组的前 N%由activation_offload_fraction控制打印汇总表通过print_offload_summary_table输出各 rank 分组的卸载字节数 ASCII 表单位 MB并告警同一 storage 在组内被重复拷贝浪费 PCIe 带宽的情况fine_grained_activation_offload.py。CPU 张量池OffloadTensorPoolfine_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_streamCUDA 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_attnattn_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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考