人工智能大模型强化学习AI Agent微调【免费下载链接】OpenClaw-RLOpenClaw-RL: Train any agent simply by talking项目地址https://gitcode.com/gh_mirrors/op/OpenClaw-RL点击查看免费下载导读megatron.core.distributed是 Megatron-LM 中负责优化器更新前跨所有 rank 完成模型梯度同步的分布式基础组件它既提供了将模型包装起来、以连续内存 buffer 存储梯度并支持通信与反向传播重叠的DistributedDataParallel也提供了跨数据并行、张量并行序列并行、流水线并行与专家并行等全部并行模式统一收口梯度的finalize_model_grads。读完本文你将理解 Megatron-LM 梯度同步的分桶机制、all-reduce/reduce-scatter 两种路径的选择逻辑、各类配置参数的作用以及训练循环中梯度从 backward 到 optimizer step 之间的完整调用链。一、distributed包的整体定位根据官方 API 文档 Megatron-LM/docs/source/api-guide/distributed.rst该包的目标非常聚焦在 optimizer step 之前在每个 rank 上完成模型权重梯度的收口finalize。它由两部分核心能力构成分布式数据并行包装器DistributedDataParallel在数据并行副本之间对梯度执行 all-reduce 或 reduce-scatterfinalize_model_grads方法跨不同并行模式同步梯度例如流水线不同 stage 上共享tied层的梯度、因专家并行expert parallelism而分散在不同 rank 上的 MoE 专家梯度等。包的对外导出定义在 megatron/core/distributed/init.pyDistributedDataParallel、DistributedDataParallelConfig、finalize_model_grads以及FullyShardedDataParallel、TorchFullyShardedDataParallel等 FSDP 相关实现。可见该包不仅覆盖传统 DDP也延伸到了 FSDP 路径。二、DistributedDataParallel梯度连续缓冲与通信重叠2.1 设计目标DistributedDataParallel源码见 megatron/core/distributed/distributed_data_parallel.py是一个模型包装器其类注释明确了两个设计要点梯度存储在连续的 buffer 中contiguous buffers避免大量小张量的分散通信支持通信与反向传播重叠把完整模型的梯度拆分为更小的 bucket对每个 bucket 异步发起 all-reduce / reduce-scatter从而让通信与后续 backward 计算并行进行。它继承自 megatron/core/distributed/data_parallel_base.py 中的_BaseDataParallel模板类。模板类定义了forward、no_sync、start_grad_sync、finish_grad_sync、zero_grad_buffer、scale_gradients、broadcast_params等协议方法DistributedDataParallel在这些接口上给出真正的分布式实现而 Megatron FSDP 路径同样遵循这套接口约定。2.2 构造流程与核心机制构造函数签名distributed_data_parallel.pyDistributedDataParallel( config: TransformerConfig, ddp_config: DistributedDataParallelConfig, module: torch.nn.Module, disable_bucketing: bool False, pg_collection: Optional[ProcessGroupCollection] None, )初始化时按如下顺序组织内部结构1进程组分配。通过ProcessGroupCollection.setup_process_groups_for_ddp解析出dp_group、dp_cp_group含上下文并行的数据并行组、intra_dp_cp_group、expt_dp_group、intra_expt_dp_group专家数据并行相关、tp_group、pp_group、ep_group等一系列进程组供后续不同类参数的梯度使用不同组做通信。2bucket 尺寸决策。若用户未显式指定bucket_size默认取max(40000000, 1000000 * dp_group.size())——DP 规模越大bucket 越大以保证 NCCL ring-reduce 每次通信的 chunk 足够大处于带宽受限bandwidth-bound而非延迟受限latency-bound状态若关闭overlap_grad_reduce则bucket_size置为None即不切分整段梯度一次通信。此外流水线 stage 非首段pp_rank 0或传入disable_bucketingTrue时也会关闭分桶因为此时数据并行通信不在关键路径上distributed_data_parallel.py。3参数分类与缓冲分配。按参数的allreduce属性把参数分为两类需要跨 DP 组 all-reduce 的稠密参数dense params与不参与常规 all-reduce 的专家并行参数expert parallel params分别为它们分配独立的_ParamAndGradBuffer见 megatron/core/distributed/param_and_grad_buffer.py。缓冲分配还会按(param_dtype, grad_dtype)分组若开启grad_reduce_in_fp32梯度以 fp32 累加对 bf16 模型尤其重要对 fp8 参数则使用torch.uint8作为实际存储 dtype。4梯度缩放因子的计算。为了最终得到除以 DP 规模的效果代码提供两种等价路径distributed_data_parallel.pyaverage_in_collectiveTrue非专家参数预缩放因子为1.0在 DP 组内做平均归约专家参数预缩放edp_size/dp_size再在专家 DP 组内做平均归约最终缩放仍为1/dp_sizeaverage_in_collectiveFalse所有参数先预缩放1/dp_size再在集体通信中做求和归约。5注册 backward hook。为每个requires_grad的参数注册反向 post-hook_make_backward_post_hook在反向传播中把param.grad累加到param.main_grad并置空param.grad若开启overlap_grad_reduce则通知所属 bucket group 登记该参数梯度已就绪distributed_data_parallel.py。当 bucket group 内所有参数的梯度都就绪时异步通信all-reduce 或 reduce-scatter会自动发起实现通信与反向传播的重叠。2.3 对外核心方法DistributedDataParallel暴露的方法与_BaseDataParallel模板一一对应是训练循环直接调用的入口方法作用zero_grad_buffer()每个训练迭代开始时清零所有梯度缓冲并重置grad_added_to_main_grad标记与 bucket group 元数据distributed_data_parallel.pystart_grad_sync()对全部 bucket group 发起梯度同步all-reduce 或 reduce-scatteroverlap_grad_reduceTrue时为异步派发否则为同步调用L510-L520finish_grad_sync()等待/执行梯度同步完成L522-L532scale_gradients(factor)按因子缩放缓冲内全部梯度供 per-token loss 归一化使用start_param_sync()发起模型参数的 all-gather 同步分布式优化器路径下overlap_param_gatherTrue时异步派发L460-L472broadcast_params()从 DP 源 rank 向其他 DP rank 广播参数用于随机初始化的数据并行对齐no_sync()上下文管理器临时关闭梯度同步梯度累积场景见 L447-L458值得注意的是start_param_sync与overlap_param_gather的配合当启用分布式优化器且overlap_param_gatherTrue时参数 all-gather 通过 forward pre-hook 与 forward 计算重叠enable_forward_pre_hookL354-L364并在finish_param_sync中等待对应 bucket 的 all-gather 句柄align_param_gatherTrue时所有流水线 stage 同时发起 all-gather避免各自按需发起造成的错峰。2.4 分桶与通信聚合的底层实现bucket 与 bucket group 的具体实现位于 megatron/core/distributed/param_and_grad_buffer.py_ParamAndGradBucket负责一个参数子集的 param/grad 视图param_data、grad_data记录param_to_index映射与未填充元素数numel_unpadded_ParamAndGradBucketGroup将多个 bucket 编组使它们的通信可以聚合coalesce为少数几个内核调用。当组内全部参数梯度就绪时自动发起异步通信在use_distributed_optimizer下通信组为分布式优化器实例内部组否则为数据并行组。通信内核的合并利用了 PyTorch 的 coalescing manager文件开头会探测 NCCL 后端是否支持reduce_scatter_tensor_coalescedPyTorch 2.2 才具备支持时用_coalescing_manager把同一 bucket group 内多个 bucket 的 all-gather/reduce-scatter 合并派发否则回退到顺序执行param_and_grad_buffer.py。此外_ParamAndGradBucketGroup还提供check_grads方法在通信前检查各 bucket 梯度的 NaN、Inf 与异常大梯度对应check_for_nan_in_grad/check_for_large_grads配置配合重跑状态机rerun_state_machine实现故障定位。三、DistributedDataParallelConfig全部配置项详解配置类定义在 megatron/core/distributed/distributed_data_parallel_config.py以下按功能分组梳理默认值均以当前仓库为准3.1 通信重叠与数据并行策略参数默认值说明overlap_grad_reduceFalse是否让梯度 all-reduce / reduce-scatter 与反向计算重叠按 bucket 异步派发overlap_param_gatherFalse是否让参数 all-gather 与 forward 计算重叠分布式优化器路径align_param_gatherFalse为 True 时所有 PP stage 同时发起参数 all-gather否则各 stage 独立按需发起use_distributed_optimizerFalse为 True 时用 reduce-scatter 聚合梯度并清理原始模型参数DistOpt 路径为 False 时用 all-reducenum_distributed_optimizer_instances1将 DP 域分片为多个分布式优化器实例partial DistOpt默认 1 即 DistOpt 覆盖整个 DP 域average_in_collectiveFalse在集体通信内直接做平均而非先除以 dp_size 再做求和3.2 数值精度与安全校验参数默认值说明grad_reduce_in_fp32False梯度通信/累加以 fp32 进行bf16 模型常用reduce_scatter_with_fp32_accumulationFalse用线上低精度、本地 fp32 累加的 reduce-scatter 实现通过 all-to-all 保持与 ring 实现相当的通信量且不支持num_distributed_optimizer_instances 1check_for_nan_in_gradFalse通信前检查梯度中的 NaN 与 Infcheck_for_large_gradsFalse通信前检查异常大的梯度gradient_reduce_div_fusionTrue梯度归约与除法融合执行3.3 分桶调优参数默认值说明bucket_sizeNone每个 bucket 最大参数数。未指定时用默认max(40000000, 1000000 * dp_size)overlap_grad_reduceFalse时强制为None不分桶pad_buckets_for_high_nccl_busbwFalse将 bucket 大小对齐到 2^16 的大次幂保证大 DP 规模下 NCCL 消息长度ring 算法下为bucket_size / dp_size可被 2 的幂整除以获得高总线带宽3.4 FP8 / MXFP8 支持参数默认值说明fp8_param_gatherFalse计算参数保持 fp8参数 all-gather 也以 fp8 进行reuse_grad_buf_for_mxfp8_param_agFalse复用梯度缓冲做 MXFP8 参数 all-gather仅在fp8_recipemxfp8且fp8_param_gatherTrue时可用校验逻辑要求同时开启overlap_param_gather与overlap_grad_reducekeep_fp8_transpose_cacheFalseMegatron FSDP 路径下是否保留 fp8 transpose 缓存3.5 FSDP 相关参数默认值说明use_megatron_fsdpFalse使用 Megatron FSDP 代码路径做 DDP取代已废弃的use_custom_fsdpdata_parallel_sharding_strategyno_shardFSDP 分片策略合法值为no_shard、optim、optim_grads、optim_grads_paramsouter_dp_sharding_strategyno_shard混合分片数据并行HSDP下外层 DP 组的分片策略suggested_communication_unit_sizeNoneFSDP 单次通信的元素数影响 all-gather 预取行为preserve_fp32_weightsTrueMegatron FSDP 的 ParamAndGradBuffer 中保留 fp32 权重nccl_ubFalse为 param/grad buffer 分配并注册 NCCL userbuffer启用 SM 高效的 NCCL 算法配合 SHARP 效果更佳__post_init__会拒绝与PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True同时使用fsdp_double_bufferFalse为 Megatron FSDP 通信分配持久双缓冲nccl_ubTrue时会自动开启disable_symmetric_registrationFalse禁用对称窗口userbuffer 注册回退到常规本地注册3.6 其他参数默认值说明delay_wgrad_computeFalse延迟权重梯度计算以提升批级通信重叠3.7 从命令行参数到配置对象的组装在训练入口 megatron/training/training.py 中DistributedDataParallelConfig由命令行参数构造而来bucket_size可来自--ddp-bucket-size或由--ddp-num-buckets换算num_parameters // args.ddp_num_buckets另有--ddp-pad-buckets-for-high-nccl-busbw、--ddp-reduce-scatter-with-fp32-accumulation、--ddp-average-in-collective等开关。随后用DistributedDataParallel(config..., ddp_config..., modulemodel_chunk, disable_bucketing...)包装每个模型 chunk其中disable_bucketing对第 2 个及以后的 model chunkinterleaved pipeline 的多 chunk 场景传True因为它们的通信已与计算重叠无需再分桶。四、finalize_model_grads跨并行模式统一收口梯度finalize_model_grads定义于 megatron/core/distributed/finalize_model_grads.py签名如下finalize_model_grads( model: List[torch.nn.Module], num_tokens: Optional[torch.Tensor] None, pg_collection: Optional[ProcessGroupCollection] None, )它应在每个 optimizer step 之前、backward 之后被调用依次完成以下同步与文档描述一一对应4.1 第一步DP 副本间的 all-reduce / reduce-scatter对每个模型 chunk 调用model_chunk.finish_grad_sync()收口DistributedDataParallel中按 bucket 异步发起的梯度通信确保数据并行副本间梯度一致finalize_model_grads.py。该步骤用config.timers中的all-grads-sync计时。4.2 第二步条件嵌入器梯度Diffusion 模型_allreduce_conditional_embedding_grads针对在各 PP/VPP rank 上复制、带pipeline_parallel标记的条件嵌入器如 DiT 的时间步嵌入器、FPS 嵌入器、标签嵌入器梯度跨 PP 组做 all-reduce并先累加虚拟 PP rank 的梯度再广播回各副本finalize_model_grads.py。4.3 第三步非张量并行梯度序列并行等_allreduce_non_tensor_model_parallel_grads在 TP 组内按两种归约语义处理finalize_model_grads.py求和SUM序列并行config.sequence_parallel且参数带sequence_parallel标记下的 LayerNorm 梯度以及qk_layernorm下q_layernorm/k_layernorm参数的梯度平均AVG带average_gradients_across_tp_domainTrue标记的模块梯度。_allreduce_layernorm_grads是该函数的向后兼容别名计划在 mcore 0.14 移除。4.4 第四步Embedding 梯度跨流水线 stage_allreduce_word_embedding_grads当词嵌入权重在首尾 PP stage 间共享share_embeddings_and_output_weights或 MTP 场景下通过 embedding 进程组 all-reduce 词嵌入梯度finalize_model_grads.py_allreduce_position_embedding_grads跨编码器/解码器 stage 同步位置嵌入梯度L262-L275。两者共用_allreduce_embedding_grad辅助函数并处理 Megatron FSDP 下_local_tensor与DTensor的 unshard/reshard 转换。4.5 MoE 专家偏置更新与临时张量重置当moe_router_enable_expert_bias开启时_update_router_expert_bias会跨 TP×CP×DP rank all-reducelocal_tokens_per_expert并按moe_router_bias_update_rate更新 router 的 expert biasfinalize_model_grads.py。随后reset_model_temporary_tensors清零临时张量并重置全局 aux loss 追踪器L278-L290。这部分逻辑同样被 megatron/core/transformer/cuda_graphs.py 的 CUDA Graph 路径复用。4.6 可选per-token loss 归一化若传入num_tokensper-token loss 模式该值只在流水线末段存在因此先从末段 broadcast 到整个 PP 组再跨dp_cp_groupall-reduce 得到全局非 padding token 总数最后调用model_chunk.scale_gradients(1.0 / num_tokens)对全部梯度归一化finalize_model_grads.py。注意此模式下average_in_collective必须关闭构造 DDP 时已有断言。4.7 进程组来源finalize_model_grads支持两种进程组获取方式默认从parallel_state取tp、pp、embd、pos_embd、dp_cp组或通过pg_collectionmegatron/core/process_groups_config.py 中的ProcessGroupCollection显式传入统一组织好的进程组集合。文档中的示例用法为pgs ProcessGroupCollection() pgs.tp tp_group pgs.pp pp_group pgs.dp dp_group model TransformerModel(..., pg_collectionpgs) ddp_model DistributedDataParallel(..., pg_collectionpgs) finalize_model_grads(..., pg_collectionpgs)使用pg_collection时finalize_model_grads会强制校验其包含tp、pp、embd、pos_embd、dp_cp属性避免静默回退到全局默认进程组。五、从 backward 到 optimizer step 的完整调用链综合源码可梳理出典型 Megatron-LM 训练迭代中的梯度流迭代开始zero_grad_buffer()清零梯度缓冲与元数据forward/backward各 bucket 的梯度随反向传播逐步就绪overlap_grad_reduceTrue时异步发起 all-reduceDDP 路径或 reduce-scatterDistOpt 路径backward 结束后调用finalize_model_grads(model, num_tokens...)依次收口 DP 梯度同步、同步条件嵌入器/序列并行 LayerNorm/embedding 梯度、更新 MoE expert bias、重置临时张量并按需用num_tokens归一化梯度详见 finalize_model_grads.py优化器 stepDistOpt 下优化器消费 reduce-scatter 产生的本地梯度分片下一个迭代前再次zero_grad_buffer()。六、FSDP 扩展Megatron FSDP 与 Torch FSDP2distributed包还包含两条 FSDP 路径Megatron FSDPmegatron/core/distributed/fsdp/通过use_megatron_fsdpTrue启用复用_ParamAndGradBuffer、支持混合分片HSDP、NCCL userbuffernccl_ub与 fp8 等能力Torch FSDP2megatron/core/distributed/torch_fully_sharded_data_parallel.py根据 megatron/core/distributed/README.md启用方式为--use-torch-fsdp2 --no-gradient-accumulation-fusion --ckpt-format torch_dist并特别注明此时不应设置CUDA_MAX_CONNECTIONS1以保证 FSDP 通信与主流的计算能充分并行。对应配置类为TorchFullyShardedDataParallelConfig。七、源码导航与进一步阅读核心 DDP 实现megatron/core/distributed/distributed_data_parallel.pyDistributedDataParallel类在 L22 起配置类megatron/core/distributed/distributed_data_parallel_config.pyDistributedDataParallelConfig在 L8 起梯度收口megatron/core/distributed/finalize_model_grads.pyfinalize_model_grads在 L396 起分桶与通信底层megatron/core/distributed/param_and_grad_buffer.py模板基类megatron/core/distributed/data_parallel_base.py统一进程组megatron/core/process_groups_config.py训练入口组装megatron/training/training.pyDDP 包装见 L1088-L1115FSDP 说明megatron/core/distributed/README.md 与 megatron/core/distributed/fsdp/src/README.md单元测试tests/unit_tests/distributed/ 目录集中覆盖了本包的进程组、分布式状态与梯度同步逻辑结语megatron.core.distributed是 Megatron-LM 多维度并行训练中承上启下的关键组件DistributedDataParallel以连续梯度缓冲、动态分桶与通信/计算重叠为手段把数据并行的梯度同步做得高效且可配置finalize_model_grads则在 optimizer step 前把数据并行、序列并行、流水线并行与专家并行产生的各类梯度统一收口并兼顾 FP8、DistOpt、FSDP 与 per-token loss 等现代训练需求。理解这两者的配置语义与调用顺序是排查梯度同步问题、调优大集群训练通信性能的起点。赞分享人工智能大模型强化学习AI Agent微调【免费下载链接】OpenClaw-RLOpenClaw-RL: Train any agent simply by talking项目地址https://gitcode.com/gh_mirrors/op/OpenClaw-RL点击查看免费下载相关推荐Megatron-LM 分布式优化器Distributed Optimizer深度解析ZeRO 式优化器状态分片与数据流原理Megatron LM 分布式优化器Distributed Optimizer深度解析ZeRO 式优化器状态分片与数据流原理 导读 本篇技术指南聚焦 Me人工智能大模型强化学习AI Agent微调PyTorch DistributedDataParallelDDP原理详解从 DistributedSampler 到 Ring All-Reduce 梯度同步PyTorch DistributedDataParallelDDP原理详解从 DistributedSampler 到 Ring All Reduce示例工程Megatron-LM内存优化激活重计算与梯度检查点Megatron LM内存优化激活重计算与梯度检查点 引言大模型训练的内存挑战 在训练数十亿甚至数千亿参数的大型语言模型时GPU内存成为最关键的瓶颈资源。人工智能大模型预训练分布式训练深度学习强化学习上一篇终极免费方案3步解决Mac NTFS读写难题告别Windows文件交换烦恼下一篇终极免费在线3D模型浏览器从零开始构建你的专业可视化平台创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
