尧图网站建设 
Megatron-LM MoE Paged Stash 完整指南HybridEP 同步免等待调度与分页激活暂存实战【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM导读本文围绕 Megatron-LM 中用于 MoEMixture of Experts训练的Paged Stash 实验特性展开讲解如何通过moe_expert_rank_capacity_factor驱动的sync-free免同步专家执行与paged stashing分页激活暂存在规避逐 step 设备查询/重分配循环的同时显著降低反向传播所需的激活显存占用。读完本文你将掌握 Paged Stash 的完整配置参数、容量因子调优方法论、溢出回退rerun机制与源码级工作原理能够直接在你的 MoE 训练脚本中启用并调优该特性。请注意该特性当前标注为experimental实验性接口与行为可能随版本演进而变化。什么是 Paged Stash两个正交机制的组合Paged stash 本质上由两个可以独立开关、又协同工作的机制组成关联文档给出的核心公式是Paged stash sync-free 专家执行 paged stashing将反向传播所需的 routed-expert 激活分页打包进分页缓冲区Sync-free免同步的静态预分配调度传统 HybridEP 调度下每个 step 都需要通过设备查询device query/ 重分配循环来确定 dispatch 与 grouped expert 缓冲区的大小这一循环引入了同步开销。Sync-free 通过以下三个开关预先根据用户控制的容量来固定缓冲区大小配置项作用--moe-flex-dispatcher-backend hybridep使用 HybridEP 灵活调度后端--use-transformer-engine-op-fuser启用 Transformer Engine 算子融合支撑 fused grouped experts--moe-expert-rank-capacity-factor float预设每个 expert rank 的 token 预算据此静态确定 dispatch / 融合分组专家缓冲区大小从 transformer_config.py 的源码注释可见moe_expert_rank_capacity_factor即每个 expert rank 的 token 预算None表示不丢弃任何 token默认值在hybridep后端下超出预算的 token 会被丢弃在ncclep后端下设置该值会启用静态形状固定接收缓冲区可捕获 CUDA Graph置None则回到 eager 模式每个 step 按实际收到的 token 数动态确定接收缓冲区大小。Paged stashing分页激活暂存--moe-paged-stash将反向传播所需的 routed-expert 激活存储到分页的 CUDA 缓冲区可选配 pinned host 溢出区。它帮助节省激活显存即使没有它sync-free 依然可以工作代价是更高的激活显存占用——因为那时激活必须完整保留在显存中以备反向传播使用。启用前提Prerequisites只要设置了moe_expert_rank_capacity_factor就必须同时满足 HybridEP TE fused grouped experts 这两个前提。具体约束从 transformer_config.py 的__post_init__校验中可以得到印证moe_expert_rank_capacity_factor非空时moe_flex_dispatcher_backend必须是hybridep或ncclep使用hybridep后端时必须同时开启use_transformer_engine_op_fuser或moe_use_grouped_tensor。当启用moe_paged_stash时还有额外限制见 transformer_config.pycapacity factor 必须设置moe_paged_stash requires moe_expert_rank_capacity_factor to be set否则直接抛ValueError不能启用cpu_offloadingmoe_paged_stash cannot be enabled with cpu_offloadingoffload_modules不能包含expert_fc1、moe_act、fused_group_mlp这三类激活本身正是 paged stash 要接管的对象不能同时被细粒度 offloading 接管。关于 Runner 的激活条件文档明确指出只要设置了 capacity factorRunner 就会激活即便不开--moe-paged-stash用于处理 over-budgettoken 超预算重跑而 stash overflow暂存溢出检查仅在开启 paged stashing 时才进行。在 training.py 中可以看到训练循环正是以args.moe_expert_rank_capacity_factor is not None为条件将forward_backward_func包装为PagedStashRunner与文档描述完全一致。完整配置示例将以下参数添加到训练命令行即可启用完整能力# Sync-free免同步调度 --moe-flex-dispatcher-backend hybridep --use-transformer-engine-op-fuser --moe-expert-rank-capacity-factor float # Paged stashing避免因碎片化导致的显存浪费 --moe-paged-stashPaged stashing 专属调优参数# 暂存页大小每个 page 容纳的 token 数 --moe-paged-stash-page-size 64 # CUDA 暂存缓冲区的缩放因子默认 1.10 --moe-paged-stash-buffer-size-factor-cuda 1.10 # Host 溢出0 关闭符号规则与 CUDA 相同 --moe-paged-stash-buffer-size-factor-cpu 0.0这些参数在 transformer_config.py 中的默认值与语义为moe_paged_stash默认False为所有 routed-expert 反向所需激活启用分页暂存moe_paged_stash_page_size默认64分页存储管理的每页 token 数moe_paged_stash_buffer_size_factor_cuda默认1.10CUDA 暂存缓冲区分配比例。符号选择尺寸基准正数 基于 avg tokens 推导的最大值负数 基于实际 num tokens 推导的最大值旧行为数值大小取绝对值表示余量headroom如1.10即 10% 余量moe_paged_stash_buffer_size_factor_cpu默认0.0Host 缓冲区比例0表示不分配 host 缓冲区符号约定与 CUDA 相同。moe_expert_rank_capacity_factor与moe_paged_stash_buffer_size_factor_cuda的含义两者本质上是相对于完美均衡情形perfectly balanced case的缓冲区大小乘数——所谓完美均衡即路由 token 恰好均匀分布在各个 expert rank 上时所需的容量。真实世界的路由存在偏斜skew因此更大的因子为偏斜预留了额外余量。从 paged_stash.py 的allocate_stash_buffers实现可以看到该语义的落地点cuda_factor 0时以max_avg_tokens_across_vp_stages基于平均 token 数推导的跨 VP 阶段最大值为基准并乘以cuda_factor为负时改用max_tokens_across_vp_stages实际 token 数最大值缩放取-cuda_factor若基于 avg 的字典尚未填充会自动回退到基于实际最大值的字典保证安全每个(dtype, hidden_size)组合独立分配一个PagedStashBufferFP8 dtypetorch.float8_e4m3fn等的缓冲区以torch.uint8存储1 字节元素进一步降低显存占用。PagedStashBufferpaged_stash.py的内部组织为[num_pages, page_size, hidden_size]形态使用**基于页的 free list循环缓冲区**管理free list 状态以二元素张量表示[0]CUDA、[1]host。CUDA 页与可选 pinned host 页各自独立成环reset()通过预分配的 reset 张量完成 free list 重置不在 reset 路径上做任何新的内存分配——这是为了 CUDA Graph 安全graph 回放时不允许分配。如何选择moe_expert_rank_capacity_factor与暂存缓冲区缩放调优的核心方法论是先 profile 真实路由与均衡基准的偏离程度再选择因子使偏斜尖峰skew spikes很少越过余量边界从而避免频繁触发 rerun。具体到三个参数moe_expert_rank_capacity_factor根据 profile 结果选择使over-budget token 丢弃不常见建议略高于profile 得到的峰值让 rerun 保持罕见。一旦发生 over-budgetRunner 会记录当 step 实际需要的峰值接收容量并给出建议因子见下节rerun 机制。moe_paged_stash_buffer_size_factor_cuda基于**同一份统计峰值 vs 均值**来定尺寸使stash overflow不常见尺寸偏小会像 over-budget 一样触发 rerun。moe_paged_stash_buffer_size_factor_cpu设置 0允许在 CUDA 页用尽时溢出到 pinned host——这通常能以 host 内存和 paged stashing 额外开销为代价避免 overflow / rerun。开启后若发生 host spillRunner 会输出提示建议考虑增大 CUDA 因子以获得更好性能。Runner 与一次失败自动重跑rerun机制触发条件与执行流程只要设置了moe_expert_rank_capacity_factor训练循环中的forward_backward_func就会被包装为PagedStashRunnertraining.py。Runner 包装前向-反向计算每次 pass 结束后检查两类失败信号——stash overflow仅当--moe-paged-stash开启时检查token over-budgetHybridEP dispatch 丢弃了 token。若任一 rank触发任一信号该 step 将以不填充容量、不启用 paged stashing 的方式重跑一次。从 paged_stash.py 的__call__实现可见其完整逻辑每次尝试前先通过data_read预取全部 microbatch 输入保证重跑时数据一致首次尝试以静态 HybridEP 接收预算capacity factor paged stashing 运行check_moe_overflow通过一次 all_reduce汇总三种标志stash_overflow_ranks、overbudget_ranks、host_spill_ranks仅在失败路径上再做第二次 all_reduce 查询所需容量Happy path 保持单一集合通信全部为 0 则成功恢复各 MoE 层的 capacity factor、恢复moe_paged_stash配置若发生过 host spill 则打印提示否则调用prepare_for_rerun清空 capacity factorHybridEP 转为 dropless 模式、关闭 paged stashing、清空梯度、重置 CUDA Graph然后重试一次。为什么一次重试足够代码中的_MAX_RERUN_ATTEMPTS 2paged_stash.py注释给出了严谨的理由prepare_for_rerun会清空 capacity factor——dropless 模式下没有接收预算可供超出同时禁用 paged stashing——因此重试不可能以同样的方式再次失败。若重试仍失败理论上不可能会触发断言保护。重跑日志与后续动作over-budget 触发重跑时Runner 会打印峰值接收容量需求并给出建议Set moe_expert_rank_capacity_factor X.XX to avoid the rerun该建议不会跨重启保留需要手动更新配置对 NCCL EP 后端还会自动增长接收容量grow_ep_recv_capacity作为运行时自适应。训练场景的 fallback 会释放大页缓冲区release_stash_buffers在下一个paged_stash_reset时按相同 capture 基准与当前因子重新分配验证场景forward_only则不释放避免训练 CUDA Graph 记录期指针失效。源码级工作原理解析三阶段状态机begin → capture → capturedpaged_stash.py 的paged_stash_reset驱动一个三阶段状态机begin首个迭代将状态置为capturecapture第二迭代以实际运行统计 token 峰值与均值构建跨 VP 阶段的max_tokens_across_vp_stages与max_avg_tokens_across_vp_stages字典随后转入captured并调用allocate_stash_buffers按因子分配所有缓冲captured稳态每个迭代重置 free list、清零overflow/host_spill标志并断言待暂存/暂存中的张量队列为空。capture 阶段同时记录真实的 PP 调度序列_pp_schedule用vp_stage * 1000000 layer_no * 1000 microbatch_no编码正号表示前向、负号表示反向供反向时精确定位需要 reload 的层。与 autograd 的集成saved_tensors_hooksPaged stash 通过torch.autograd.graph.saved_tensors_hookspaged_stash.py透明接管反向传播所需的保存张量on_save_for_backward前向中 autograd 保存激活时将其扁平化并封装为PagedTensor携带实际 token 数张量、VP stage、调度层号、原始形状、页记录等元数据加入待暂存列表并把原始张量内存置空on_get_saved_tensor反向需要激活时从暂存缓冲区 reload必要时从 host 溢出区取回按实际 token 数截断/填充回原始形状后返回。PagedTensor.offload_to_stash/reload_from_stashpaged_stash.py分别调用paged_stash_copy_kernel与paged_stash_pop_kernel自定义 CUDA 核函数完成页级拷贝核函数通过 free list head/tail 以 O(1) 分配/回收页SCALE_INV_BLOCK_SIZE 32用于处理 columnwise scale-inv 分组张量FP8 场景的 token 数换算。流式流水线pack/unpack 流与 PP 调度钩子PagedStashManager单例维护专门的pack 流当前实现中 unpack 复用同一流因为暂存尚未做到流安全并利用PipelinePreScheduleFunction/PipelinePostScheduleFunction两个 autograd.Function 在 PP 调度的层边界插入存/取钩子前向经过层组时发起 stash在 pack 流上执行拷贝wait_for_stash_to_complete确保下一层开始前暂存完成并释放原始张量反向时在 unpack 流上按调度顺序预取reload下一层所需激活并与主流同步实现暂存/重载与前向/反向计算的流水线重叠。MoE 层中的接入点在 experts.py 的 fused grouped MLP 前向中config.moe_paged_stash开启时permute 后的本地 hidden states 先经paged_stash_group_start标记层组起点随后以get_paged_stash_context包裹专家计算最后paged_stash_group_commit标记层组终点并触发暂存调度。也就是说被暂存的是 routed-expertpermute 之后的激活而非 attention 等其他模块的激活。与 CUDA Graph 的兼容整个设计对 CUDA Graph 友好PagedStashBuffer.reset()不做任何分配所有 free list、overflow、host_spill 状态均为预分配的设备内存张量通过copy_复位reload 张量在主流上分配、由 unpack 流填充先wait_stream再写入。这使 Paged Stash 可以在FullCudaGraphWrapper捕获的训练图上运行而不会因图回放中的分配操作报错。测试验证行为如何被守护仓库中的单元测试覆盖了 Paged Stash 的核心行为可作为理解与调试的参考test_paged_stashing.py覆盖 4 层前向-反向全流程test_forward_backward_4_layers、容量因子过载与 over-budgettest_overload_factor_and_over_budget、over-budget 恢复test_over_budget_recovery验证prepare_for_rerun清空 capacity factor 转为 eager 后不再有预算可超以及 FP8/MXFP8 场景下的溢出核算test_over_budgetwire_dtypemxfp8。test_paged_stash_runner.py用 fake 模型/优化器验证 Runner 的重试逻辑包括共享 root config 的行为保持test_retry_preserves_shared_root_config_behavior与逐模块 config 的禁用/恢复test_retry_disables_and_restores_per_module_configs以及 MXFP8 复用梯度缓冲区场景下的 staged 参数拷贝test_retry_stages_only_overlapped_reused_mxfp8_buffers。小结与最佳实践清单先 profile 后设参统计真实路由中各 expert rank 的 token 峰值/均值分布以均衡基准为 1.0将moe_expert_rank_capacity_factor与moe_paged_stash_buffer_size_factor_cuda设在略高于 profile 峰值的位置善用 host spill 兜底将moe_paged_stash_buffer_size_factor_cpu设为 0例如0.51.0可在不牺牲 CUDA 余量的情况下大幅降低 rerun 概率留意 rerun 日志日志会明确告知 over-budget 的峰值容量与建议因子、stash overflow 建议增大 CUDA/CPU 因子、host spill 建议增大 CUDA 因子牢记约束moe_paged_stash必须与moe_expert_rank_capacity_factor、HybridEP、TE fuser 共存且不能与cpu_offloading或offload_modules中包含expert_fc1/moe_act/fused_group_mlp的配置混用知晓实验性边界该特性处于实验阶段行为可能随版本变化NCCL EP 后端下容量自适应不跨重启保留需要手动将建议因子写入训练脚本。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Megatron-LM MoE Paged Stash 完整指南：HybridEP 同步免等待调度与分页激活暂存实战
Megatron-LM MoE Paged Stash 完整指南HybridEP 同步免等待调度与分页激活暂存实战【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM导读本文围绕 Megatron-LM 中用于 MoEMixture of Experts训练的Paged Stash 实验特性展开讲解如何通过moe_expert_rank_capacity_factor驱动的sync-free免同步专家执行与paged stashing分页激活暂存在规避逐 step 设备查询/重分配循环的同时显著降低反向传播所需的激活显存占用。读完本文你将掌握 Paged Stash 的完整配置参数、容量因子调优方法论、溢出回退rerun机制与源码级工作原理能够直接在你的 MoE 训练脚本中启用并调优该特性。请注意该特性当前标注为experimental实验性接口与行为可能随版本演进而变化。什么是 Paged Stash两个正交机制的组合Paged stash 本质上由两个可以独立开关、又协同工作的机制组成关联文档给出的核心公式是Paged stash sync-free 专家执行 paged stashing将反向传播所需的 routed-expert 激活分页打包进分页缓冲区Sync-free免同步的静态预分配调度传统 HybridEP 调度下每个 step 都需要通过设备查询device query/ 重分配循环来确定 dispatch 与 grouped expert 缓冲区的大小这一循环引入了同步开销。Sync-free 通过以下三个开关预先根据用户控制的容量来固定缓冲区大小配置项作用--moe-flex-dispatcher-backend hybridep使用 HybridEP 灵活调度后端--use-transformer-engine-op-fuser启用 Transformer Engine 算子融合支撑 fused grouped experts--moe-expert-rank-capacity-factor float预设每个 expert rank 的 token 预算据此静态确定 dispatch / 融合分组专家缓冲区大小从 transformer_config.py 的源码注释可见moe_expert_rank_capacity_factor即每个 expert rank 的 token 预算None表示不丢弃任何 token默认值在hybridep后端下超出预算的 token 会被丢弃在ncclep后端下设置该值会启用静态形状固定接收缓冲区可捕获 CUDA Graph置None则回到 eager 模式每个 step 按实际收到的 token 数动态确定接收缓冲区大小。Paged stashing分页激活暂存--moe-paged-stash将反向传播所需的 routed-expert 激活存储到分页的 CUDA 缓冲区可选配 pinned host 溢出区。它帮助节省激活显存即使没有它sync-free 依然可以工作代价是更高的激活显存占用——因为那时激活必须完整保留在显存中以备反向传播使用。启用前提Prerequisites只要设置了moe_expert_rank_capacity_factor就必须同时满足 HybridEP TE fused grouped experts 这两个前提。具体约束从 transformer_config.py 的__post_init__校验中可以得到印证moe_expert_rank_capacity_factor非空时moe_flex_dispatcher_backend必须是hybridep或ncclep使用hybridep后端时必须同时开启use_transformer_engine_op_fuser或moe_use_grouped_tensor。当启用moe_paged_stash时还有额外限制见 transformer_config.pycapacity factor 必须设置moe_paged_stash requires moe_expert_rank_capacity_factor to be set否则直接抛ValueError不能启用cpu_offloadingmoe_paged_stash cannot be enabled with cpu_offloadingoffload_modules不能包含expert_fc1、moe_act、fused_group_mlp这三类激活本身正是 paged stash 要接管的对象不能同时被细粒度 offloading 接管。关于 Runner 的激活条件文档明确指出只要设置了 capacity factorRunner 就会激活即便不开--moe-paged-stash用于处理 over-budgettoken 超预算重跑而 stash overflow暂存溢出检查仅在开启 paged stashing 时才进行。在 training.py 中可以看到训练循环正是以args.moe_expert_rank_capacity_factor is not None为条件将forward_backward_func包装为PagedStashRunner与文档描述完全一致。完整配置示例将以下参数添加到训练命令行即可启用完整能力# Sync-free免同步调度 --moe-flex-dispatcher-backend hybridep --use-transformer-engine-op-fuser --moe-expert-rank-capacity-factor float # Paged stashing避免因碎片化导致的显存浪费 --moe-paged-stashPaged stashing 专属调优参数# 暂存页大小每个 page 容纳的 token 数 --moe-paged-stash-page-size 64 # CUDA 暂存缓冲区的缩放因子默认 1.10 --moe-paged-stash-buffer-size-factor-cuda 1.10 # Host 溢出0 关闭符号规则与 CUDA 相同 --moe-paged-stash-buffer-size-factor-cpu 0.0这些参数在 transformer_config.py 中的默认值与语义为moe_paged_stash默认False为所有 routed-expert 反向所需激活启用分页暂存moe_paged_stash_page_size默认64分页存储管理的每页 token 数moe_paged_stash_buffer_size_factor_cuda默认1.10CUDA 暂存缓冲区分配比例。符号选择尺寸基准正数 基于 avg tokens 推导的最大值负数 基于实际 num tokens 推导的最大值旧行为数值大小取绝对值表示余量headroom如1.10即 10% 余量moe_paged_stash_buffer_size_factor_cpu默认0.0Host 缓冲区比例0表示不分配 host 缓冲区符号约定与 CUDA 相同。moe_expert_rank_capacity_factor与moe_paged_stash_buffer_size_factor_cuda的含义两者本质上是相对于完美均衡情形perfectly balanced case的缓冲区大小乘数——所谓完美均衡即路由 token 恰好均匀分布在各个 expert rank 上时所需的容量。真实世界的路由存在偏斜skew因此更大的因子为偏斜预留了额外余量。从 paged_stash.py 的allocate_stash_buffers实现可以看到该语义的落地点cuda_factor 0时以max_avg_tokens_across_vp_stages基于平均 token 数推导的跨 VP 阶段最大值为基准并乘以cuda_factor为负时改用max_tokens_across_vp_stages实际 token 数最大值缩放取-cuda_factor若基于 avg 的字典尚未填充会自动回退到基于实际最大值的字典保证安全每个(dtype, hidden_size)组合独立分配一个PagedStashBufferFP8 dtypetorch.float8_e4m3fn等的缓冲区以torch.uint8存储1 字节元素进一步降低显存占用。PagedStashBufferpaged_stash.py的内部组织为[num_pages, page_size, hidden_size]形态使用**基于页的 free list循环缓冲区**管理free list 状态以二元素张量表示[0]CUDA、[1]host。CUDA 页与可选 pinned host 页各自独立成环reset()通过预分配的 reset 张量完成 free list 重置不在 reset 路径上做任何新的内存分配——这是为了 CUDA Graph 安全graph 回放时不允许分配。如何选择moe_expert_rank_capacity_factor与暂存缓冲区缩放调优的核心方法论是先 profile 真实路由与均衡基准的偏离程度再选择因子使偏斜尖峰skew spikes很少越过余量边界从而避免频繁触发 rerun。具体到三个参数moe_expert_rank_capacity_factor根据 profile 结果选择使over-budget token 丢弃不常见建议略高于profile 得到的峰值让 rerun 保持罕见。一旦发生 over-budgetRunner 会记录当 step 实际需要的峰值接收容量并给出建议因子见下节rerun 机制。moe_paged_stash_buffer_size_factor_cuda基于**同一份统计峰值 vs 均值**来定尺寸使stash overflow不常见尺寸偏小会像 over-budget 一样触发 rerun。moe_paged_stash_buffer_size_factor_cpu设置 0允许在 CUDA 页用尽时溢出到 pinned host——这通常能以 host 内存和 paged stashing 额外开销为代价避免 overflow / rerun。开启后若发生 host spillRunner 会输出提示建议考虑增大 CUDA 因子以获得更好性能。Runner 与一次失败自动重跑rerun机制触发条件与执行流程只要设置了moe_expert_rank_capacity_factor训练循环中的forward_backward_func就会被包装为PagedStashRunnertraining.py。Runner 包装前向-反向计算每次 pass 结束后检查两类失败信号——stash overflow仅当--moe-paged-stash开启时检查token over-budgetHybridEP dispatch 丢弃了 token。若任一 rank触发任一信号该 step 将以不填充容量、不启用 paged stashing 的方式重跑一次。从 paged_stash.py 的__call__实现可见其完整逻辑每次尝试前先通过data_read预取全部 microbatch 输入保证重跑时数据一致首次尝试以静态 HybridEP 接收预算capacity factor paged stashing 运行check_moe_overflow通过一次 all_reduce汇总三种标志stash_overflow_ranks、overbudget_ranks、host_spill_ranks仅在失败路径上再做第二次 all_reduce 查询所需容量Happy path 保持单一集合通信全部为 0 则成功恢复各 MoE 层的 capacity factor、恢复moe_paged_stash配置若发生过 host spill 则打印提示否则调用prepare_for_rerun清空 capacity factorHybridEP 转为 dropless 模式、关闭 paged stashing、清空梯度、重置 CUDA Graph然后重试一次。为什么一次重试足够代码中的_MAX_RERUN_ATTEMPTS 2paged_stash.py注释给出了严谨的理由prepare_for_rerun会清空 capacity factor——dropless 模式下没有接收预算可供超出同时禁用 paged stashing——因此重试不可能以同样的方式再次失败。若重试仍失败理论上不可能会触发断言保护。重跑日志与后续动作over-budget 触发重跑时Runner 会打印峰值接收容量需求并给出建议Set moe_expert_rank_capacity_factor X.XX to avoid the rerun该建议不会跨重启保留需要手动更新配置对 NCCL EP 后端还会自动增长接收容量grow_ep_recv_capacity作为运行时自适应。训练场景的 fallback 会释放大页缓冲区release_stash_buffers在下一个paged_stash_reset时按相同 capture 基准与当前因子重新分配验证场景forward_only则不释放避免训练 CUDA Graph 记录期指针失效。源码级工作原理解析三阶段状态机begin → capture → capturedpaged_stash.py 的paged_stash_reset驱动一个三阶段状态机begin首个迭代将状态置为capturecapture第二迭代以实际运行统计 token 峰值与均值构建跨 VP 阶段的max_tokens_across_vp_stages与max_avg_tokens_across_vp_stages字典随后转入captured并调用allocate_stash_buffers按因子分配所有缓冲captured稳态每个迭代重置 free list、清零overflow/host_spill标志并断言待暂存/暂存中的张量队列为空。capture 阶段同时记录真实的 PP 调度序列_pp_schedule用vp_stage * 1000000 layer_no * 1000 microbatch_no编码正号表示前向、负号表示反向供反向时精确定位需要 reload 的层。与 autograd 的集成saved_tensors_hooksPaged stash 通过torch.autograd.graph.saved_tensors_hookspaged_stash.py透明接管反向传播所需的保存张量on_save_for_backward前向中 autograd 保存激活时将其扁平化并封装为PagedTensor携带实际 token 数张量、VP stage、调度层号、原始形状、页记录等元数据加入待暂存列表并把原始张量内存置空on_get_saved_tensor反向需要激活时从暂存缓冲区 reload必要时从 host 溢出区取回按实际 token 数截断/填充回原始形状后返回。PagedTensor.offload_to_stash/reload_from_stashpaged_stash.py分别调用paged_stash_copy_kernel与paged_stash_pop_kernel自定义 CUDA 核函数完成页级拷贝核函数通过 free list head/tail 以 O(1) 分配/回收页SCALE_INV_BLOCK_SIZE 32用于处理 columnwise scale-inv 分组张量FP8 场景的 token 数换算。流式流水线pack/unpack 流与 PP 调度钩子PagedStashManager单例维护专门的pack 流当前实现中 unpack 复用同一流因为暂存尚未做到流安全并利用PipelinePreScheduleFunction/PipelinePostScheduleFunction两个 autograd.Function 在 PP 调度的层边界插入存/取钩子前向经过层组时发起 stash在 pack 流上执行拷贝wait_for_stash_to_complete确保下一层开始前暂存完成并释放原始张量反向时在 unpack 流上按调度顺序预取reload下一层所需激活并与主流同步实现暂存/重载与前向/反向计算的流水线重叠。MoE 层中的接入点在 experts.py 的 fused grouped MLP 前向中config.moe_paged_stash开启时permute 后的本地 hidden states 先经paged_stash_group_start标记层组起点随后以get_paged_stash_context包裹专家计算最后paged_stash_group_commit标记层组终点并触发暂存调度。也就是说被暂存的是 routed-expertpermute 之后的激活而非 attention 等其他模块的激活。与 CUDA Graph 的兼容整个设计对 CUDA Graph 友好PagedStashBuffer.reset()不做任何分配所有 free list、overflow、host_spill 状态均为预分配的设备内存张量通过copy_复位reload 张量在主流上分配、由 unpack 流填充先wait_stream再写入。这使 Paged Stash 可以在FullCudaGraphWrapper捕获的训练图上运行而不会因图回放中的分配操作报错。测试验证行为如何被守护仓库中的单元测试覆盖了 Paged Stash 的核心行为可作为理解与调试的参考test_paged_stashing.py覆盖 4 层前向-反向全流程test_forward_backward_4_layers、容量因子过载与 over-budgettest_overload_factor_and_over_budget、over-budget 恢复test_over_budget_recovery验证prepare_for_rerun清空 capacity factor 转为 eager 后不再有预算可超以及 FP8/MXFP8 场景下的溢出核算test_over_budgetwire_dtypemxfp8。test_paged_stash_runner.py用 fake 模型/优化器验证 Runner 的重试逻辑包括共享 root config 的行为保持test_retry_preserves_shared_root_config_behavior与逐模块 config 的禁用/恢复test_retry_disables_and_restores_per_module_configs以及 MXFP8 复用梯度缓冲区场景下的 staged 参数拷贝test_retry_stages_only_overlapped_reused_mxfp8_buffers。小结与最佳实践清单先 profile 后设参统计真实路由中各 expert rank 的 token 峰值/均值分布以均衡基准为 1.0将moe_expert_rank_capacity_factor与moe_paged_stash_buffer_size_factor_cuda设在略高于 profile 峰值的位置善用 host spill 兜底将moe_paged_stash_buffer_size_factor_cpu设为 0例如0.51.0可在不牺牲 CUDA 余量的情况下大幅降低 rerun 概率留意 rerun 日志日志会明确告知 over-budget 的峰值容量与建议因子、stash overflow 建议增大 CUDA/CPU 因子、host spill 建议增大 CUDA 因子牢记约束moe_paged_stash必须与moe_expert_rank_capacity_factor、HybridEP、TE fuser 共存且不能与cpu_offloading或offload_modules中包含expert_fc1/moe_act/fused_group_mlp的配置混用知晓实验性边界该特性处于实验阶段行为可能随版本变化NCCL EP 后端下容量自适应不跨重启保留需要手动将建议因子写入训练脚本。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考