拓冰网站定制 
Megatron-LM MIMO 训练示例全解析从单元测试命令到多模态异构并行实现【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMIMOMultimodal In/Out Model是 Megatron-LM 中用于让语言模型同时理解与生成文本、图像、音频等多种模态的统一模型架构。本文以 tests/unit_tests/models/mimo/README.md 中给出的两条训练验证命令为主线深入讲解同址Colocated与非同址Non-colocated两种编码器 LLM 部署形态下的测试命令、验证目标与底层并行机制并结合 megatron/core/models/mimo/README.md、test_mimo_colocated_correctness.py、test_mimo_1f1b_schedule.py 及examples/mimo下的训练脚本给出可复制、可运行、可深入研究的完整指南。读完本文你将掌握如何运行 MIMO 的多模态训练验证、如何理解异构 TP/DP/PP 拓扑下的梯度缩放正确性以及 MIMO 的桥接通信bridge communication与梯度同步grad sync在生产训练中的真实调用方式。一、MIMO 训练示例的定位从 README 出发tests/unit_tests/models/mimo/README.md全文极其精炼但信息密度很高它直接指向仓库根目录下的两类可运行验证同址Colocated编码器encoder与 LLM 部署在同一组 rank 上运行test_mimo_colocated_correctness.py验证异构 DP 下的梯度缩放正确性非同址Non-colocated编码器与 LLM 部署在互斥的 rank 集合上运行test_mimo_1f1b_schedule.py验证 1F1B 流水线调度在多模块异构拓扑上的端到端训练。两条命令都要求从仓库根目录执行并使用uv作为 Python 包管理器配合torch.distributed.run启动多进程分布式测试# 同址8 卡编码器与 LLM 共享同一批 GPU uv run python -m torch.distributed.run --nproc_per_node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py -v -s # 非同址8 卡运行 1F1B 调度下的一个代表性用例编码器 TP2LLM TP2 PP3 uv run python -m torch.distributed.run --nproc-per-node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py::TestMimo1F1BSchedule::test_encoder_tp2_llm_tp2_pp3_8gpu -v -s注意两条命令在参数上的细微差异同址示例使用--nproc_per_node下划线非同址示例使用--nproc-per-node连字符二者在torch.distributed.run中是等价的写法。README 还提示test_mimo_1f1b_schedule.py中还包含更多 8 卡非同址示例稍后我们会逐一展开。二、MIMO 模型架构速览这些测试在验证什么在深入测试之前先理解 MIMO 本身。根据 megatron/core/models/mimo/README.mdMIMO 由两大组件构成语言模型Language Model处理统一嵌入空间的“中央处理器”通过共享词表处理同时包含文本与模态 token 的合并序列模态子模块Modality Submodules负责模态数据与语言模型之间的双向转换——编码modality → embeddings、解码embeddings → modality以及模态空间与语言模型维度的投影projection。完整数据流为Input → Encoder → Projection → Align input embeddings → Language Model → Hidden states for special generation tokens → Output Projection → Decoder → Output其核心类是MimoModelmegatron/core/models/mimo/model/base.py它通过align_embeddings_by_token_positions把各模态嵌入按特殊 token 位置对齐进输入序列解码时则选取特殊模态生成 token 对应的隐藏状态交由子模块转换回模态输出。默认提供VisionModalitySubmodules与AudioModalitySubmodules见 megatron/core/models/mimo/submodules/vision.py 与 megatron/core/models/mimo/submodules/audio.py并支持自定义子模块。因此tests/unit_tests/models/mimo/下的这些测试并非普通的“冒烟测试”而是对 MIMO 在多模态异构并行heterogeneous parallel场景下正确性的系统性验证。目录下共 23 个测试文件覆盖了模型构建、前向/反向、梯度同步、检查点、量化、角色划分、嵌入对齐、编码器预取、优化器共识、重叠生命周期、非同址 CP 正确性等方方面面。三、同址训练示例异构 DP 梯度缩放正确性验证3.1 运行方式与前置条件uv run python -m torch.distributed.run --nproc_per_node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py -v -s该文件test_mimo_colocated_correctness.py中的TestColocatedGradientScalingCorrectness要求8 张 GPU不满足时测试会以pytest.skip跳过并要求 PyTorch 2.3。它通过环境变量固定确定性算法NVTE_ALLOW_NONDETERMINISTIC_ALGO0 CUDA_DEVICE_MAX_CONNECTIONS1 CUBLAS_WORKSPACE_CONFIG:4096:83.2 测试目标异构 DP 下的梯度必须等于 DP1 的每 token 均值同址意味着编码器与 LLM 在同一组 rank 上但它们可以拥有不同的 TP/DP 拓扑。测试参数化的核心组合为参数组合含义拓扑形态enc_tp2, enc_dp4, llm_tp4, llm_dp2编码器 DP4 大于 LLM DP2Fan-in收敛enc_tp4, enc_dp2, llm_tp2, llm_dp4编码器 DP2 小于 LLM DP4Fan-out发散此外还在mask_patternuniform/asymmetric与num_microbatches1 / 4两个维度上做交叉验证。其核心断言是异构 DP 的 MIMO 模型在一步 Adam 更新后编码器的分片后权重必须与一个基于相同编码器 TP/DP 布局的等 DPequal-DP参考模型完全一致fp32 精度内。为什么这是正确的验证方式关键机制如下等 DP 参考模型的编码器分片与待测模型完全对齐enc_tpdist_enc_tp, enc_dpdist_enc_dp分片一一对应无需 gather-and-slice参考侧enc_dp llm_dp因此桥接器退化为恒等直通BridgeDirection.EQUAL编码器每个 rank 直接把输出喂给同址的 LLM rank不产生任何重分布集合通信两侧都设置calculate_per_token_lossTrue将 DDP 的gradient_scaling_factor钉在1.0纯 SUM 归约随后通过_wire_training_hooks中挂接的自定义finalize_grads_func在 LLM DP 组上 all-reducetotal_num_tokens得到全局有效 token 数N_global再对编码器与 LLM 统一执行scale_gradients(1/N_global)。也就是说每个编码器分片上落地的梯度就是 DP1 的每 token 均值梯度。如果缩放错误——例如某个 DDP 在 per-token 均值之上又应用了默认的1/dp_size或 finalize hook 用编码器 DP 组内的局部 token 计数替代了全局N_global——那么仅仅一步 Adam 更新就足以让待测模型与参考模型的编码器权重产生可检测的分歧。测试同时在前向信号进入 LLM decoder 的桥接后隐藏状态、logits与编码器首层梯度三个 oracle 层面进行对比从而把“桥接是否数学等价”与“LLM TP 前向数值漂移”两个问题隔离诊断。测试还通过_generate_and_broadcast_global_batches在 rank 0 生成全局 batch 并广播给所有 rank支持asymmetric掩码模式——每个样本的尾部文本 token 被按样本编号确定性清零使得不同 DP 分片携带不同的有效 token 数专门检验“num den 全局均值 CE”路径旧的局部均值配方在这种场景下只是近似正确。3.3 生产路径的验证configure_grad_sync值得强调的是测试并没有“自造”一套缩放逻辑而是直接委托 examples/mimo/training/grad_sync.py 中的configure_grad_sync来挂接 no-sync 与 grad-sync hook——这正是生产 MIMO 训练使用的路径。该模块实现的双梯度终结dual gradient finalization逻辑包括通过mark_modality_participation/reset_modality_participation按 step 标记每个模态子模块是否在本步接收到了对应模态输入batch 中缺失模态 key 即视为未参与在 LLM 最后一个 PP 阶段、TP rank 0 这一单一坐标上汇总所有数据通路的全局 token 数避免跨 TP/PP 副本重复计数依据参与模态的 DP rank 数做归一化缩放_vision_participation_count在 vision DP 组上 all-reduce 参与指示量。因此通过test_mimo_colocated_correctness.py也就直接验证了生产configure_grad_sync在异构 DP 场景下的正确性。四、非同址训练示例1F1B 流水线调度下的多模块异构拓扑4.1 代表性用例uv run python -m torch.distributed.run --nproc-per-node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py::TestMimo1F1BSchedule::test_encoder_tp2_llm_tp2_pp3_8gpu -v -s该用例test_encoder_tp2_llm_tp2_pp3_8gpu的拓扑为编码器 TP2 PP1LLM TP2 PP38 卡。非同址的含义是编码器与 LLM 位于互斥的 rank 区间上——本用例中编码器占用 rank 0-1encoder_offset0LLM 占用 rank 2-7llm_offset2。测试通过run_mimo_1f1b_test驱动schedule.forward_backward_pipelining_without_interleaving完成一次含 4 个 micro-batch 的前向/反向随后执行真实的优化器 step并断言 grad_norm 为正且有限、loss 有限。4.2 8 卡与非 8 卡用例全景TestMimo1F1BScheduletest_mimo_1f1b_schedule.py中的每个测试方法都会按当前 world size 做pytest.skip守卫因此既可以在 2 卡、4 卡环境跑小规模用例也可以在 8 卡环境跑完整矩阵测试方法GPU 数拓扑验证要点test_baseline_2gpu2编码器 PP1LLM PP1最小非同址基线test_lm_pp3_4gpu4编码器 PP1LLM PP3深 LLM 流水线test_encoder_tp2_llm_tp2_pp3_8gpu8编码器 TP2 PP1LLM TP2 PP3TP 深 PP 组合README 指向的代表用例test_full_pp_8gpu8编码器 TP2 PP2LLM TP2 PP2双模块完整流水线test_fan_in_dp4_to_dp1_llm_pp2_8gpu8编码器 DP4 → LLM TP2 或 CP2 PP2 DP1高 fan-in 比 4:1CP2 场景额外验证目标端 CP 梯度重建test_fan_out_dp1_to_dp4_enc_tp2_pp2_8gpu8编码器 TP2 PP2 DP1 → LLM DP4fan-out 1:4编码器输出被拆成 4 份test_fan_in_dp2_to_dp1_llm_pp3_8gpu8编码器 DP2 → LLM TP2 PP3 DP1深 LLM 流水线下的 fan-in参数化use_layer_wise_distributed_optimizer为 False/True 两种路径其中 fan-in / fan-out 的命名与桥接方向直接相关Fan-in收敛源端编码器DP 副本数大于目标端LLM。前向时沿 batch 维 all-gather 源端输出反向时把目标端梯度收窄narrow回源 rank 的槽位Fan-out发散目标端 DP 副本更多。前向时收窄反向时在目标端兄弟 DP rank 之间 all-gathernarrow 的伴随算子而非零填充 scatter因为每个目标 rank 消费的是同一源激活的不同切片。4.3 非对称 DP 下的 micro-batch 守恒非同址 异构 DP 带来的一个关键工程细节是每个 rank 的 micro-batch size 不同。在run_mimo_1f1b_test中llm_mbs micro_batch_size encoder_mbs micro_batch_size * llm_dp // encoder_dpLLM 的 MBS 即调度级 MBS而编码器 MBS 按 DP 比例调整保证穿过桥接器的总工作量守恒。数据迭代器只在需要的 rank 上创建编码器需要数据的条件是“位于编码器网格内且处于其 PP 首级”LLM 则需要首级或末级末级用于 loss。同址colocated情况下二者都持有数据此时以驱动调度的 LLM MBS 为准。4.4 LayerWise 优化器路径test_fan_in_dp2_to_dp1_llm_pp3_8gpu通过参数化同时覆盖了两条优化路径use_layer_wise_distributed_optimizerFalse标准DistributedDataParallel 分布式 AdamOptimizerConfig中optimizeradamuse_layer_wise_distributed_optimizerTrue走wrap_active_modules_with_ddpexamples/mimo/training/runtime.py的 MIMO LayerWise 参数布局路径配optimizermuon通过get_mimo_optimizermegatron/core/models/mimo/optimizer.py返回LayerWiseDistributedOptimizermegatron/core/optimizer/layer_wise_optimizer.py。测试对优化器类型做了显式断言当启用 LayerWise 时optimizer.chained_optimizers中必须存在LayerWiseDistributedOptimizer实例。五、底层机制桥接通信与超通信网格5.1 BridgeCommunicator跨网格的流水线通信非同址拓扑下编码器输出需要穿过 TP/DP/PP/CP 均可能不同的两个网格。这正是 megatron/core/pipeline_parallel/bridge_communicator.py 中BridgeCommunicator的职责初始化一对源/目标网格之间的通信器为每个 rank 构建显式通信计划并提供send_forward、recv_forward、send_forward_recv_backward、send_backward_recv_forward等被流水线调度调用的公开方法。每个 rank 通过CommRole枚举扮演三种角色之一SENDER源网格每个 DP 副本内的 tp-cp 领头 rank向目标网格接收者发送数据RECEIVER目标网格每个 DP 副本内的领头 rank从源网格发送者接收数据MEMBERDP 副本内的非领头 rank参与本地领头 rank 的广播。BridgeCommunicator内部维护两类进程组缓存_broadcast_pg_cache避免为相同 rank 集合重复创建 NCCL 通信器与_bridge_pg_cache跨网格领头 rank 间的 P2P按排序后的领头 rank 集合为 key并提供destroy_broadcast_pgs/destroy_bridge_pgs供测试清理。其限制与约束也在 docstring 中明确源网格 CP 大于 1 不支持目标端 CP 梯度在反向发送前于目标领头 rank 的 TP lane 上归约。dim_mapping默认{s: 1, b: 0, h: 2}tensor_ndim支持 2 或 3——对 MIMO 预展平的(s*b, h)编码器输出应传dim_mapping{b: 0, h: 1}。5.2 ColocatedBridgeCommunicator同址异构布局桥接同址场景则使用 megatron/core/models/mimo/comm/colocated_communicator.py 中的ColocatedBridgeCommunicator。它以BridgeDirection枚举区分三种桥接方向FAN_IN源端 DP 副本更多前向沿 batch 维 all-gather反向收窄FAN_OUT目标端 DP 副本更多前向收窄反向 all-gatherEQUALDP 匹配桥接为纯直通identity passthrough。其前置条件是输入必须在源 TP 组内沿 batch 维复制TP-replicated——桥接器从不沿 TP 维做 gather违反该约定会静默产生错误结果。同址正确性测试正是利用EQUAL直通特性构造“零重分布”参考模型作为 oracle。5.3 HyperCommGrid统一网格抽象两类测试都通过 megatron/core/hyper_comm_grid.py 的HyperCommGrid构造拓扑例如create_hypercomm_grid(offset0, tpenc_tp, cp1, ppenc_pp, dpenc_dp)。HyperCommGrid以shape[tp, cp, pp, dp]dim_names描述一个逻辑网格支持register_view注册专家视图dense 场景下ep1专家视图将expt_tptp, epcp1, expt_dpdp重标定到相同 rank 集合并通过create_pg预创建 TP/CP/PP/DP 及组合进程组。测试工具函数get_pg_collection会把网格解析为ProcessGroupCollection含dp_cp、expt_dp、tp_ep_pp、intra_dist_opt等add_embedding_groups则进一步为 PP 各段的首/末 rank 补建 embedding 同步进程组。需要特别注意的是dist.new_group是集合操作所有 rank 都必须调用因此测试用create_all_embedding_groups提前为每个网格枚举全部 PP 子组、去重后统一创建。六、从单元测试到真实训练examples/mimo 的完整入口tests/unit_tests/models/mimo/验证的是 MIMO 的“骨架正确性”而 examples/mimo 提供了真实训练入口与生产级组件训练入口examples/mimo/train.py入口脚本与 examples/mimo/pretrain_mimo.py启动脚本examples/mimo/scripts/run_mock_train.sh2 卡 mock 数据快速验证、examples/mimo/scripts/run_vlm_train.sh、examples/mimo/scripts/run_avlm_train.sh、examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh异构 nemotron MoE 训练与 examples/mimo/scripts/run_video_vlm_train.sh训练运行时examples/mimo/training 下的builder.py模型构建、step.py训练 step、runtime.pyDDP 包装与激活模块管理、grad_sync.py双梯度终结即测试验证的生产 hook、topology.py异构拓扑、encoder_prefetch.py编码器预取、distributed.py与args.py模型提供器examples/mimo/model_providers 提供llava_vlm、llava_avlm、mock、nemotron_moe_vlm、radio_encoder、hf_clip_encoder、hf_whisper_encoder等数据提供器examples/mimo/data 提供 VLM/AVLM 的 energon 任务编码器与样本加载器。以run_mock_train.sh为例其核心训练参数包括TP2、micro-batch-size2、global-batch-size4、train-iters100、cosine LR初始 6.0e-5min 6.0e-6、--dataset-provider mock、--model-provider mock并使用llava-hf/llava-1.5-7b-hf的 HuggingFace tokenizer。该脚本支持-d参数启用 debugpy 远程调试默认端口 5678并预留DRY_RUN开关查看实际命令。模型结构参数--num-layers 1、--hidden-size 128、--num-attention-heads 4、--max-position-embeddings 512、--encoder-seq-length 512在脚本注释中说明最终由 model provider 设置 LLM 的 config 与 specCLI 覆盖能力仍在演进中。七、测试中的工程实践要点可直接复用的模式最后梳理这些测试代码中值得复用到你自己的 MIMO 训练或验证脚本中的工程模式确定性环境固定_set_deterministic_env设置NVTE_ALLOW_NONDETERMINISTIC_ALGO0、CUDA_DEVICE_MAX_CONNECTIONS1、CUBLAS_WORKSPACE_CONFIG:4096:8并显式弹出NVTE_FLASH_ATTN/NVTE_FUSED_ATTN/NVTE_UNFUSED_ATTNGPTModel 要求这些变量未设置或与 attention 后端一致fp32 add_bias_linearFalse 关闭 dropout 可把数值噪声压到最低per-token loss 三元组约定calculate_per_token_lossTrue时 loss 函数必须返回(local_sum, local_num_tokens, log_dict)三元组其中num_tokens必须是int 张量调度器会跨 micro-batch 累加得到total_num_tokens外部 finalize hook 再统一除以全局N_global跨 TP 重分片参数拷贝参考模型与待测模型 TP 不同时先沿ref_tp_groupall-gather 还原完整权重再按partition_dim与目标 TP size 用torch.tensor_split切回分片——且必须在构造分布式优化器其会在__init__时克隆 fp32 master 权重之前完成拷贝全局 diff 统计跨进程组对 min/max/mean/分位数做归约同时报告参考张量量级ref_max、ref_mean与相对 diff避免“绝对 diff 很小但相对灾难”的误判资源清理顺序测试在 teardown 中先销毁 MimoModel释放ColocatedBridgeCommunicator子组再调用destroy_all_grids依次销毁网格、embedding PG 缓存、bridge 广播/P2P PG 缓存避免 NCCL 通信器泄漏。八、小结tests/unit_tests/models/mimo/README.md虽然只有两条命令却精确指向了 Megatron-LM MIMO 训练验证的两个最关键入口同址异构 DP 梯度缩放正确性test_mimo_colocated_correctness.py与非同址 1F1B 流水线多拓扑端到端训练test_mimo_1f1b_schedule.py。前者验证了生产configure_grad_sync在 fan-in/fan-out 下落地 DP1 每 token 均值梯度的正确性后者覆盖了从 2 卡基线到 8 卡 TP/PP/DP/CP 全组合、从标准分布式优化器到 LayerWise 优化器的完整矩阵。理解这些测试就等于拿到了在真实多模态大模型训练中正确配置 MIMO 异构并行、调试梯度缩放与桥接通信的完整地图。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Megatron-LM MIMO 训练示例全解析：从单元测试命令到多模态异构并行实现
Megatron-LM MIMO 训练示例全解析从单元测试命令到多模态异构并行实现【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMIMOMultimodal In/Out Model是 Megatron-LM 中用于让语言模型同时理解与生成文本、图像、音频等多种模态的统一模型架构。本文以 tests/unit_tests/models/mimo/README.md 中给出的两条训练验证命令为主线深入讲解同址Colocated与非同址Non-colocated两种编码器 LLM 部署形态下的测试命令、验证目标与底层并行机制并结合 megatron/core/models/mimo/README.md、test_mimo_colocated_correctness.py、test_mimo_1f1b_schedule.py 及examples/mimo下的训练脚本给出可复制、可运行、可深入研究的完整指南。读完本文你将掌握如何运行 MIMO 的多模态训练验证、如何理解异构 TP/DP/PP 拓扑下的梯度缩放正确性以及 MIMO 的桥接通信bridge communication与梯度同步grad sync在生产训练中的真实调用方式。一、MIMO 训练示例的定位从 README 出发tests/unit_tests/models/mimo/README.md全文极其精炼但信息密度很高它直接指向仓库根目录下的两类可运行验证同址Colocated编码器encoder与 LLM 部署在同一组 rank 上运行test_mimo_colocated_correctness.py验证异构 DP 下的梯度缩放正确性非同址Non-colocated编码器与 LLM 部署在互斥的 rank 集合上运行test_mimo_1f1b_schedule.py验证 1F1B 流水线调度在多模块异构拓扑上的端到端训练。两条命令都要求从仓库根目录执行并使用uv作为 Python 包管理器配合torch.distributed.run启动多进程分布式测试# 同址8 卡编码器与 LLM 共享同一批 GPU uv run python -m torch.distributed.run --nproc_per_node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py -v -s # 非同址8 卡运行 1F1B 调度下的一个代表性用例编码器 TP2LLM TP2 PP3 uv run python -m torch.distributed.run --nproc-per-node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py::TestMimo1F1BSchedule::test_encoder_tp2_llm_tp2_pp3_8gpu -v -s注意两条命令在参数上的细微差异同址示例使用--nproc_per_node下划线非同址示例使用--nproc-per-node连字符二者在torch.distributed.run中是等价的写法。README 还提示test_mimo_1f1b_schedule.py中还包含更多 8 卡非同址示例稍后我们会逐一展开。二、MIMO 模型架构速览这些测试在验证什么在深入测试之前先理解 MIMO 本身。根据 megatron/core/models/mimo/README.mdMIMO 由两大组件构成语言模型Language Model处理统一嵌入空间的“中央处理器”通过共享词表处理同时包含文本与模态 token 的合并序列模态子模块Modality Submodules负责模态数据与语言模型之间的双向转换——编码modality → embeddings、解码embeddings → modality以及模态空间与语言模型维度的投影projection。完整数据流为Input → Encoder → Projection → Align input embeddings → Language Model → Hidden states for special generation tokens → Output Projection → Decoder → Output其核心类是MimoModelmegatron/core/models/mimo/model/base.py它通过align_embeddings_by_token_positions把各模态嵌入按特殊 token 位置对齐进输入序列解码时则选取特殊模态生成 token 对应的隐藏状态交由子模块转换回模态输出。默认提供VisionModalitySubmodules与AudioModalitySubmodules见 megatron/core/models/mimo/submodules/vision.py 与 megatron/core/models/mimo/submodules/audio.py并支持自定义子模块。因此tests/unit_tests/models/mimo/下的这些测试并非普通的“冒烟测试”而是对 MIMO 在多模态异构并行heterogeneous parallel场景下正确性的系统性验证。目录下共 23 个测试文件覆盖了模型构建、前向/反向、梯度同步、检查点、量化、角色划分、嵌入对齐、编码器预取、优化器共识、重叠生命周期、非同址 CP 正确性等方方面面。三、同址训练示例异构 DP 梯度缩放正确性验证3.1 运行方式与前置条件uv run python -m torch.distributed.run --nproc_per_node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py -v -s该文件test_mimo_colocated_correctness.py中的TestColocatedGradientScalingCorrectness要求8 张 GPU不满足时测试会以pytest.skip跳过并要求 PyTorch 2.3。它通过环境变量固定确定性算法NVTE_ALLOW_NONDETERMINISTIC_ALGO0 CUDA_DEVICE_MAX_CONNECTIONS1 CUBLAS_WORKSPACE_CONFIG:4096:83.2 测试目标异构 DP 下的梯度必须等于 DP1 的每 token 均值同址意味着编码器与 LLM 在同一组 rank 上但它们可以拥有不同的 TP/DP 拓扑。测试参数化的核心组合为参数组合含义拓扑形态enc_tp2, enc_dp4, llm_tp4, llm_dp2编码器 DP4 大于 LLM DP2Fan-in收敛enc_tp4, enc_dp2, llm_tp2, llm_dp4编码器 DP2 小于 LLM DP4Fan-out发散此外还在mask_patternuniform/asymmetric与num_microbatches1 / 4两个维度上做交叉验证。其核心断言是异构 DP 的 MIMO 模型在一步 Adam 更新后编码器的分片后权重必须与一个基于相同编码器 TP/DP 布局的等 DPequal-DP参考模型完全一致fp32 精度内。为什么这是正确的验证方式关键机制如下等 DP 参考模型的编码器分片与待测模型完全对齐enc_tpdist_enc_tp, enc_dpdist_enc_dp分片一一对应无需 gather-and-slice参考侧enc_dp llm_dp因此桥接器退化为恒等直通BridgeDirection.EQUAL编码器每个 rank 直接把输出喂给同址的 LLM rank不产生任何重分布集合通信两侧都设置calculate_per_token_lossTrue将 DDP 的gradient_scaling_factor钉在1.0纯 SUM 归约随后通过_wire_training_hooks中挂接的自定义finalize_grads_func在 LLM DP 组上 all-reducetotal_num_tokens得到全局有效 token 数N_global再对编码器与 LLM 统一执行scale_gradients(1/N_global)。也就是说每个编码器分片上落地的梯度就是 DP1 的每 token 均值梯度。如果缩放错误——例如某个 DDP 在 per-token 均值之上又应用了默认的1/dp_size或 finalize hook 用编码器 DP 组内的局部 token 计数替代了全局N_global——那么仅仅一步 Adam 更新就足以让待测模型与参考模型的编码器权重产生可检测的分歧。测试同时在前向信号进入 LLM decoder 的桥接后隐藏状态、logits与编码器首层梯度三个 oracle 层面进行对比从而把“桥接是否数学等价”与“LLM TP 前向数值漂移”两个问题隔离诊断。测试还通过_generate_and_broadcast_global_batches在 rank 0 生成全局 batch 并广播给所有 rank支持asymmetric掩码模式——每个样本的尾部文本 token 被按样本编号确定性清零使得不同 DP 分片携带不同的有效 token 数专门检验“num den 全局均值 CE”路径旧的局部均值配方在这种场景下只是近似正确。3.3 生产路径的验证configure_grad_sync值得强调的是测试并没有“自造”一套缩放逻辑而是直接委托 examples/mimo/training/grad_sync.py 中的configure_grad_sync来挂接 no-sync 与 grad-sync hook——这正是生产 MIMO 训练使用的路径。该模块实现的双梯度终结dual gradient finalization逻辑包括通过mark_modality_participation/reset_modality_participation按 step 标记每个模态子模块是否在本步接收到了对应模态输入batch 中缺失模态 key 即视为未参与在 LLM 最后一个 PP 阶段、TP rank 0 这一单一坐标上汇总所有数据通路的全局 token 数避免跨 TP/PP 副本重复计数依据参与模态的 DP rank 数做归一化缩放_vision_participation_count在 vision DP 组上 all-reduce 参与指示量。因此通过test_mimo_colocated_correctness.py也就直接验证了生产configure_grad_sync在异构 DP 场景下的正确性。四、非同址训练示例1F1B 流水线调度下的多模块异构拓扑4.1 代表性用例uv run python -m torch.distributed.run --nproc-per-node8 \ -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py::TestMimo1F1BSchedule::test_encoder_tp2_llm_tp2_pp3_8gpu -v -s该用例test_encoder_tp2_llm_tp2_pp3_8gpu的拓扑为编码器 TP2 PP1LLM TP2 PP38 卡。非同址的含义是编码器与 LLM 位于互斥的 rank 区间上——本用例中编码器占用 rank 0-1encoder_offset0LLM 占用 rank 2-7llm_offset2。测试通过run_mimo_1f1b_test驱动schedule.forward_backward_pipelining_without_interleaving完成一次含 4 个 micro-batch 的前向/反向随后执行真实的优化器 step并断言 grad_norm 为正且有限、loss 有限。4.2 8 卡与非 8 卡用例全景TestMimo1F1BScheduletest_mimo_1f1b_schedule.py中的每个测试方法都会按当前 world size 做pytest.skip守卫因此既可以在 2 卡、4 卡环境跑小规模用例也可以在 8 卡环境跑完整矩阵测试方法GPU 数拓扑验证要点test_baseline_2gpu2编码器 PP1LLM PP1最小非同址基线test_lm_pp3_4gpu4编码器 PP1LLM PP3深 LLM 流水线test_encoder_tp2_llm_tp2_pp3_8gpu8编码器 TP2 PP1LLM TP2 PP3TP 深 PP 组合README 指向的代表用例test_full_pp_8gpu8编码器 TP2 PP2LLM TP2 PP2双模块完整流水线test_fan_in_dp4_to_dp1_llm_pp2_8gpu8编码器 DP4 → LLM TP2 或 CP2 PP2 DP1高 fan-in 比 4:1CP2 场景额外验证目标端 CP 梯度重建test_fan_out_dp1_to_dp4_enc_tp2_pp2_8gpu8编码器 TP2 PP2 DP1 → LLM DP4fan-out 1:4编码器输出被拆成 4 份test_fan_in_dp2_to_dp1_llm_pp3_8gpu8编码器 DP2 → LLM TP2 PP3 DP1深 LLM 流水线下的 fan-in参数化use_layer_wise_distributed_optimizer为 False/True 两种路径其中 fan-in / fan-out 的命名与桥接方向直接相关Fan-in收敛源端编码器DP 副本数大于目标端LLM。前向时沿 batch 维 all-gather 源端输出反向时把目标端梯度收窄narrow回源 rank 的槽位Fan-out发散目标端 DP 副本更多。前向时收窄反向时在目标端兄弟 DP rank 之间 all-gathernarrow 的伴随算子而非零填充 scatter因为每个目标 rank 消费的是同一源激活的不同切片。4.3 非对称 DP 下的 micro-batch 守恒非同址 异构 DP 带来的一个关键工程细节是每个 rank 的 micro-batch size 不同。在run_mimo_1f1b_test中llm_mbs micro_batch_size encoder_mbs micro_batch_size * llm_dp // encoder_dpLLM 的 MBS 即调度级 MBS而编码器 MBS 按 DP 比例调整保证穿过桥接器的总工作量守恒。数据迭代器只在需要的 rank 上创建编码器需要数据的条件是“位于编码器网格内且处于其 PP 首级”LLM 则需要首级或末级末级用于 loss。同址colocated情况下二者都持有数据此时以驱动调度的 LLM MBS 为准。4.4 LayerWise 优化器路径test_fan_in_dp2_to_dp1_llm_pp3_8gpu通过参数化同时覆盖了两条优化路径use_layer_wise_distributed_optimizerFalse标准DistributedDataParallel 分布式 AdamOptimizerConfig中optimizeradamuse_layer_wise_distributed_optimizerTrue走wrap_active_modules_with_ddpexamples/mimo/training/runtime.py的 MIMO LayerWise 参数布局路径配optimizermuon通过get_mimo_optimizermegatron/core/models/mimo/optimizer.py返回LayerWiseDistributedOptimizermegatron/core/optimizer/layer_wise_optimizer.py。测试对优化器类型做了显式断言当启用 LayerWise 时optimizer.chained_optimizers中必须存在LayerWiseDistributedOptimizer实例。五、底层机制桥接通信与超通信网格5.1 BridgeCommunicator跨网格的流水线通信非同址拓扑下编码器输出需要穿过 TP/DP/PP/CP 均可能不同的两个网格。这正是 megatron/core/pipeline_parallel/bridge_communicator.py 中BridgeCommunicator的职责初始化一对源/目标网格之间的通信器为每个 rank 构建显式通信计划并提供send_forward、recv_forward、send_forward_recv_backward、send_backward_recv_forward等被流水线调度调用的公开方法。每个 rank 通过CommRole枚举扮演三种角色之一SENDER源网格每个 DP 副本内的 tp-cp 领头 rank向目标网格接收者发送数据RECEIVER目标网格每个 DP 副本内的领头 rank从源网格发送者接收数据MEMBERDP 副本内的非领头 rank参与本地领头 rank 的广播。BridgeCommunicator内部维护两类进程组缓存_broadcast_pg_cache避免为相同 rank 集合重复创建 NCCL 通信器与_bridge_pg_cache跨网格领头 rank 间的 P2P按排序后的领头 rank 集合为 key并提供destroy_broadcast_pgs/destroy_bridge_pgs供测试清理。其限制与约束也在 docstring 中明确源网格 CP 大于 1 不支持目标端 CP 梯度在反向发送前于目标领头 rank 的 TP lane 上归约。dim_mapping默认{s: 1, b: 0, h: 2}tensor_ndim支持 2 或 3——对 MIMO 预展平的(s*b, h)编码器输出应传dim_mapping{b: 0, h: 1}。5.2 ColocatedBridgeCommunicator同址异构布局桥接同址场景则使用 megatron/core/models/mimo/comm/colocated_communicator.py 中的ColocatedBridgeCommunicator。它以BridgeDirection枚举区分三种桥接方向FAN_IN源端 DP 副本更多前向沿 batch 维 all-gather反向收窄FAN_OUT目标端 DP 副本更多前向收窄反向 all-gatherEQUALDP 匹配桥接为纯直通identity passthrough。其前置条件是输入必须在源 TP 组内沿 batch 维复制TP-replicated——桥接器从不沿 TP 维做 gather违反该约定会静默产生错误结果。同址正确性测试正是利用EQUAL直通特性构造“零重分布”参考模型作为 oracle。5.3 HyperCommGrid统一网格抽象两类测试都通过 megatron/core/hyper_comm_grid.py 的HyperCommGrid构造拓扑例如create_hypercomm_grid(offset0, tpenc_tp, cp1, ppenc_pp, dpenc_dp)。HyperCommGrid以shape[tp, cp, pp, dp]dim_names描述一个逻辑网格支持register_view注册专家视图dense 场景下ep1专家视图将expt_tptp, epcp1, expt_dpdp重标定到相同 rank 集合并通过create_pg预创建 TP/CP/PP/DP 及组合进程组。测试工具函数get_pg_collection会把网格解析为ProcessGroupCollection含dp_cp、expt_dp、tp_ep_pp、intra_dist_opt等add_embedding_groups则进一步为 PP 各段的首/末 rank 补建 embedding 同步进程组。需要特别注意的是dist.new_group是集合操作所有 rank 都必须调用因此测试用create_all_embedding_groups提前为每个网格枚举全部 PP 子组、去重后统一创建。六、从单元测试到真实训练examples/mimo 的完整入口tests/unit_tests/models/mimo/验证的是 MIMO 的“骨架正确性”而 examples/mimo 提供了真实训练入口与生产级组件训练入口examples/mimo/train.py入口脚本与 examples/mimo/pretrain_mimo.py启动脚本examples/mimo/scripts/run_mock_train.sh2 卡 mock 数据快速验证、examples/mimo/scripts/run_vlm_train.sh、examples/mimo/scripts/run_avlm_train.sh、examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh异构 nemotron MoE 训练与 examples/mimo/scripts/run_video_vlm_train.sh训练运行时examples/mimo/training 下的builder.py模型构建、step.py训练 step、runtime.pyDDP 包装与激活模块管理、grad_sync.py双梯度终结即测试验证的生产 hook、topology.py异构拓扑、encoder_prefetch.py编码器预取、distributed.py与args.py模型提供器examples/mimo/model_providers 提供llava_vlm、llava_avlm、mock、nemotron_moe_vlm、radio_encoder、hf_clip_encoder、hf_whisper_encoder等数据提供器examples/mimo/data 提供 VLM/AVLM 的 energon 任务编码器与样本加载器。以run_mock_train.sh为例其核心训练参数包括TP2、micro-batch-size2、global-batch-size4、train-iters100、cosine LR初始 6.0e-5min 6.0e-6、--dataset-provider mock、--model-provider mock并使用llava-hf/llava-1.5-7b-hf的 HuggingFace tokenizer。该脚本支持-d参数启用 debugpy 远程调试默认端口 5678并预留DRY_RUN开关查看实际命令。模型结构参数--num-layers 1、--hidden-size 128、--num-attention-heads 4、--max-position-embeddings 512、--encoder-seq-length 512在脚本注释中说明最终由 model provider 设置 LLM 的 config 与 specCLI 覆盖能力仍在演进中。七、测试中的工程实践要点可直接复用的模式最后梳理这些测试代码中值得复用到你自己的 MIMO 训练或验证脚本中的工程模式确定性环境固定_set_deterministic_env设置NVTE_ALLOW_NONDETERMINISTIC_ALGO0、CUDA_DEVICE_MAX_CONNECTIONS1、CUBLAS_WORKSPACE_CONFIG:4096:8并显式弹出NVTE_FLASH_ATTN/NVTE_FUSED_ATTN/NVTE_UNFUSED_ATTNGPTModel 要求这些变量未设置或与 attention 后端一致fp32 add_bias_linearFalse 关闭 dropout 可把数值噪声压到最低per-token loss 三元组约定calculate_per_token_lossTrue时 loss 函数必须返回(local_sum, local_num_tokens, log_dict)三元组其中num_tokens必须是int 张量调度器会跨 micro-batch 累加得到total_num_tokens外部 finalize hook 再统一除以全局N_global跨 TP 重分片参数拷贝参考模型与待测模型 TP 不同时先沿ref_tp_groupall-gather 还原完整权重再按partition_dim与目标 TP size 用torch.tensor_split切回分片——且必须在构造分布式优化器其会在__init__时克隆 fp32 master 权重之前完成拷贝全局 diff 统计跨进程组对 min/max/mean/分位数做归约同时报告参考张量量级ref_max、ref_mean与相对 diff避免“绝对 diff 很小但相对灾难”的误判资源清理顺序测试在 teardown 中先销毁 MimoModel释放ColocatedBridgeCommunicator子组再调用destroy_all_grids依次销毁网格、embedding PG 缓存、bridge 广播/P2P PG 缓存避免 NCCL 通信器泄漏。八、小结tests/unit_tests/models/mimo/README.md虽然只有两条命令却精确指向了 Megatron-LM MIMO 训练验证的两个最关键入口同址异构 DP 梯度缩放正确性test_mimo_colocated_correctness.py与非同址 1F1B 流水线多拓扑端到端训练test_mimo_1f1b_schedule.py。前者验证了生产configure_grad_sync在 fan-in/fan-out 下落地 DP1 每 token 均值梯度的正确性后者覆盖了从 2 卡基线到 8 卡 TP/PP/DP/CP 全组合、从标准分布式优化器到 LayerWise 优化器的完整矩阵。理解这些测试就等于拿到了在真实多模态大模型训练中正确配置 MIMO 异构并行、调试梯度缩放与桥接通信的完整地图。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考