Megatron-LM MIMO 架构深度解析面向多模态理解与生成的可组合 In/Out 模型【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMIMOMultimodal In/Out Model是 Megatron-LM 中一套以「语言模型 模态子模块」为核心的可组合多模态架构支持文本、图像、音频等多种模态的统一理解与生成。本文以 megatron/core/models/mimo/README.md 为主线结合仓库源码与测试系统讲解 MIMO 的架构设计、核心组件、配置方式与使用流程。读完本文你将掌握如何用MimoModelConfigModuleSpec组装出 VLM / 音视频语言模型并理解嵌入对齐、异构并行与多模块优化器的底层实现。MIMO 是什么MIMO 是一种让语言模型能够理解和生成多种模态文本、图像、音频等的模型架构。它由两部分组成一个核心语言模型在统一的嵌入空间中处理所有模态信息若干模态子模块Modality Submodules负责把模态原始输入编码为嵌入如图像 → embeddings、把嵌入解码回模态输出如 embeddings → 图像、并在模态空间与语言模型维度之间做投影Projection。由MimoModel统一负责两类关键工作将各模态嵌入对齐到序列中的特殊 token 位置将合并后的嵌入交由语言模型处理。由于模态子模块是可插拔的MIMO 提供了一条「灵活且规范化」的建模路径可以配置出多种多模态模型例如视觉语言模型VLMVision-Language Models音视频语言模型Audio-Visual Language Models多模态理解与生成模型实验性说明当前MimoModel与MimoModelConfig均标注为 experimental、仍在积极开发中API 可能在未来版本中无通知变更。实际使用前请关注仓库的版本演进见 megatron/core/models/mimo/model/base.py 与 megatron/core/models/mimo/config/base_configs.py 顶部的UserWarning。工作原理与完整数据流MIMO 的模型架构由两大组件构成语言模型Language Model模态子模块Modality Submodules完整的数据流如下Input → Encoder → Projection → Align input embeddings → Language Model → Hidden states for special generation tokens → Output Projection → Decoder → Output编码阶段模态子模块将输入转换为嵌入例如图像 → embeddingsMimoModel按 token 位置将全部模态嵌入与文本嵌入对齐语言模型处理统一的嵌入序列。解码阶段挑选与「模态生成特殊 token」对应的隐藏状态模态子模块将嵌入转换回模态输出例如 embeddings → 图像。注意编码与解码并非必须成对存在。以仓库中的默认实现为例VisionModalitySubmodules 与 AudioModalitySubmodules 的decode目前均抛出NotImplementedErrorNo decoders support yet / Audio decoding not implemented yet即当前主线聚焦于多模态理解encoding 方向生成方向的能力由具体子模块逐步补齐。核心组件详解语言模型统一嵌入空间的中央处理器语言模型是处理所有模态信息的核心组件通过共享词表充当所有模态信息的中央处理器处理同时包含文本 token 与模态 token 的合并序列。在实现上语言模型由MimoModelConfig.language_model_spec指定默认复用 GPT 系列模型。例如 examples/mimo/model_providers/mock.py 中语言模型直接使用GPTModel并传入transformer_layer_spec、vocab_size、max_sequence_length、pre_process/post_process等参数。也就是说任何符合ModuleSpec约定的 Transformer 语言模型都可以作为 MIMO 的中央处理器接入。模态子模块连接原始模态与语言模型的桥ModalitySubmodules抽象基类见 megatron/core/models/mimo/submodules/base.py负责连接原始模态数据与语言模型每个子模块处理编码modality → embeddings与解码embeddings → modality管理模态空间与语言模型维度之间的投影内部由四类部件组成encoders编码器字典、decoders解码器字典、input_projections输入投影编码后 → 语言模型前使用、output_projections输出投影从语言模型接收后 → 解码前使用。基类构造器签名如下源码与 README 一致class ModalitySubmodules(ABC, nn.Module): def __init__( self, encoders: Optional[Dict[str, nn.Module]] None, decoders: Optional[Dict[str, nn.Module]] None, input_projections: Optional[List[nn.Module]] None, output_projections: Optional[List[nn.Module]] None, is_first_stage: bool True, is_last_stage: bool True, pg_collectionNone, **kwargs, ):MIMO 提供了默认实现VisionModalitySubmodules、AudioModalitySubmodules也允许自定义子模块做专用处理# Custom implementation class CustomVisionSubmodules(ModalitySubmodules): def encode(self, inputs): # Specialized encoding logic return projected_embeddings # Use custom submodules when creating the model model MimoModel( mimo_config, modality_submodules{images: ModuleSpec(moduleCustomVisionSubmodules, params{...})} )从源码看子模块的构建逻辑from_specModalitySubmodules.from_specsubmodules/base.py负责从ModuleSpec实例化子模块其关键行为是按流水线阶段条件化地构建部件encoders / decoders所有阶段都会构建供流水线各阶段处理使用input_projections仅在is_last_stageTrue时构建投影发生在编码之后、送往语言模型之前即projection happens after encoding, before sending to language modeloutput_projections仅在is_first_stageTrue时构建投影发生在从语言模型接收之后、解码之前。此外encode方法submodules/base.py对编码器输出做了规范化处理接受(b, s, h)三维或(b*s, h)二维输出若编码器返回(embeddings, aux_state)元组则自动取首个张量combine_embeddings沿第 0 维拼接多个编码器的嵌入project_embeddings则按is_input选择输入或输出投影。forward方法保证encoder_inputs与hidden_states互斥并只在首/末阶段执行相应计算——这套约定正是非共置pipeline布局下跨卡组装嵌入的基础。编码器输入的约束encode要求encoders_data_batch的键与self.encoders的编码器名一一对应缺失会直接抛出ValueError(fNo inputs found for encoder {name})。这意味着modality_inputs的键名必须与初始化时ModuleSpec.submodules[encoders]的键名严格一致这也是 README 强调「modality names and encoder names should match the keys used in mimo config during initialization」的原因。嵌入对齐align_embeddings_by_token_positionsMimoModel通过align_embeddings_by_token_positions方法整合不同模态的嵌入实现见 megatron/core/models/mimo/model/base.py将各模态嵌入放置到输入序列中对应的特殊 token 位置处理维度匹配与位置跟踪确保嵌入放置正确。方法内部的大致行为如下# Inside MimoModels forward method aligned_embeddings self.align_embeddings_by_token_positions( modality_embeddings{text: text_emb, images: image_emb}, input_idstokens, special_token_ids{images: 32000} )源码级实现细节对齐的核心是掩码mask驱动的masked_scatter_具体规则为对text模态构造全 True 掩码再对special_token_ids中每个特殊 token id 做mask input_ids ! token_id即文本位置 非任何特殊 token 的位置对其他模态mask input_ids token_id即该模态的特殊 token 出现的位置校验每个模态的mask.sum()必须等于该模态嵌入的数量否则抛出ValueError防止嵌入与 token 数量错配输出返回(S, B, H)序列优先布局的合并嵌入张量。除掩码路径外该方法还支持预计算 token 索引modality_token_indices的高效路径由调用方提供批内行优先batch-major的扁平索引方法通过index_copy_直接落位。该路径被设计为「trusted producer contract」——索引值不会被重新校验避免在前向热路径引入设备规约与主机同步但_validate_precomputed_token_indices会检查键集合一致、索引为一维、索引总数等于B * S。这一优化对吞吐敏感的大规模训练尤其重要。对应测试见 tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py其中test_basic_alignment用「文本 单个模态」的 batch 验证了嵌入是否正确落到特殊 token 位置。配置与使用MimoModel 参数MimoModel( config: MimoModelConfig, # Required: Configuration for the model )注意MimoModel.__init__的实际签名还接受可选的cp_group、tp_group用于上下文并行 / 张量并行场景下的分区适配器例如 examples/mimo/model_providers/mock.py 中通过pg_collection.cp/pg_collection.tp传入。配置详情MimoModelConfigMIMO 模型使用MimoModelConfig实例化定义见 megatron/core/models/mimo/config/base_configs.py其中包含语言模型的ModuleSpec模态名 → 子模块ModuleSpec的字典映射。MimoModelConfig( language_model: ModuleSpec, # Specification for the language model modality_submodules: Dict[str, ModuleSpec], # Dictionary mapping modality names to their submodule specifications special_token_ids: Dict[str, int] {} # Dictionary mapping modality names to their special token IDs )源码中的完整字段从 base_configs.py 看MimoModelConfig实际包含五个字段README 中的language_model/modality_submodules对应源码的language_model_spec/modality_submodules_spec字段类型默认值说明language_model_specModuleSpecModuleSpec()语言模型规格modality_submodules_specDict[str, ModuleSpec]{}模态名 → 子模块规格special_token_idsDict[str, int]{}模态名 → 特殊 token ID。例如{vision: -200, audio: 32000}这些占位符用于在input_ids中标记插入模态嵌入的位置module_to_grid_mapOptional[Dict[str, HyperCommGrid]]None模块键如 vision、language→HyperCommGrid拓扑。语言模型必须使用固定键MIMO_LANGUAGE_MODULE_KEY即language。网格共享相同 rank → 共置colocatedTP/DP 可不同网格 rank 不相交 → 非共置pipeline 并行为None→ 使用传统全局parallel_state共置kv_formatstrsbhd注意力 KV 格式sbhdseq-batch-head-dim或thdtotal-head-dim其中__post_init__会对module_to_grid_map做键校验网格键集合必须恰好等于modality_submodules_spec的键集合加上语言模型固定键language否则抛出ValueError并提示缺失/多余项见 base_configs.py 与 config/role.py 中MIMO_LANGUAGE_MODULE_KEY language的定义。示例创建视觉语言模型VLM# Language model specification lm_spec ModuleSpec( moduleGPTModel, params{ config: language_config, transformer_layer_spec: get_mock_language_layer_spec(), vocab_size: 50304, } ) # Vision modality specification vision_submodule_spec ModuleSpec( moduleVisionModalitySubmodules, params{ # Any general parameters for the submodule can go here }, submodules{ encoders: { clip_encoder: ModuleSpec( moduleCLIPViTModel, params{ transformer_config: vision_config, transformer_layer_spec: get_mock_vision_layer_spec(), patch_dim: 16, img_h: 224, img_w: 224, } ), }, input_projections: [ ModuleSpec( moduleMultimodalProjector, params{ config: get_mock_projection_config(), submodules: get_mock_projection_layer_spec().submodules, projector_type: mlp, input_size: 128 } ), ], } ) # Instantiate the model vlm MimoModel( MimoModelConfig( language_modellm_spec, modality_submodules{images: vision_submodule_spec}, special_token_ids{images: 32000} ) )这份示例在仓库中有完整可运行的对应用法examples/mimo/model_providers/mock.py 中的model_provider_mock_vlm_single_encoder即为真实生产写法——它额外传入了max_sequence_length、pre_process/post_process并用pg_collection解析cp_group/tp_group。MIMO 前向方法使用# Prepare inputs for multiple modalities and encoders modality_inputs { # modality names and encoder names should match the keys used in mimo config during initialization. images: { clip_encoder: {pixel_values: images}, # Encoder-specific inputs vit_encoder: {images: vit_images} }, audio: { whisper_encoder: {input_features: audio_features} } } # Call forward method outputs, _ mimo_model( input_idsinput_ids, position_idsposition_ids, modality_inputsmodality_inputs, )forward完整参数与返回语义从 model/base.py 的签名看forward还支持attention_mask、loss_mask、labels、packing_kwargs构造 THD 格式PackedSeqParams用于序列打包训练以及modality_token_indices前文所述的高效索引路径。返回值语义取决于当前 rank 的角色仅编码器 rank返回Dict[str, Tensor]各编码器输出供 bridge 送往语言模块 rank语言模块 rank返回语言模型输出logits 或 loss无角色全部共置返回语言模型输出。多模态训练实践真实模型提供者仓库的 examples/mimo/model_providers/ 目录提供了多种可直接参考的组装方式mock.py单视觉编码器的教学用 VLMllava_vlm.pyspecial_token_ids{images: image_special_token_id}的 LLaVA 式 VLMllava_avlm.py同时接入images与audios两种模态、多个编码器的音视频语言模型nemotron_moe_vlm.py通过nemotron_special_token_ids(args)从命令行参数生成特殊 token 映射。对应的训练脚本见 examples/mimo/scripts/run_vlm_train.sh、examples/mimo/scripts/run_avlm_train.sh 等。分布式与异构并行从共置到多模块流水线README 聚焦单机模型组装但仓库源码将 MIMO 扩展到了复杂的分布式场景这也是「In/Out 模型」能支撑大规模训练的关键角色与布局RankRole/ModuleLayoutconfig/role.py 定义了 rank 角色机制ModuleLayout.COLOCATED共置所有模块共享相同 rank。覆盖两种情形——无 grid map沿用全局parallel_state以及异构 TP/DP 但有重叠 rank 的情形此时会构建ColocatedBridgeCommunicator桥接通信器把编码器嵌入从编码器的 TP/DP 布局转换到语言模型的 TP/DP 布局见 model/base.py 与 comm/colocated_communicator.py。该布局走_forward_all_modules路径。ModuleLayout.NON_COLOCATED非共置module_to_grid_map中网格 rank 范围互不重叠每个 rank 要么只跑编码器、要么只跑语言模型通过RankRole分派不同的前向路径_forward_encoders/_forward_language_module本质上是多模块流水线并行。RankRole.build会依据各网格是否共享 rank 自动判定布局非共置模式下还会结合pp维计算每个模块的is_first_stage/is_last_stage如果当前 rank 不属于任何模块网格会直接抛出RuntimeError提示检查module_to_grid_map配置。并行特性与约束上下文并行 / 序列并行当语言模型配置context_parallel_size 1或开启sequence_parallel时MimoModel会创建PartitionAdapterpartition/utils.py对合并后的(S, B, H)嵌入以及labels/loss_mask做 CP/SP 切分。此模式下要求attention_maskNoneCP 切分后的序列无法与稠密 mask 对齐应改用因果attn_mask_type或packed_seq_params掩码model/base.py。MTP多 token 预测支持MimoModel会为语言模型的 MTP 模块准备 token ID、位置 ID 与有效性掩码并保证它们与 CP 本地序列分区一致_prepare_mtp_inputs。嵌入层约束启用 MIMO 序列并行时语言模型嵌入层必须禁用自身的 SP scatter构造时传scatter_embedding_sequence_parallelFalse否则会在对齐前把文本 token 错误切分model/base.py。多模块优化器MimoOptimizer异构并行下不同模块拥有各自的数据并行 / 张量并行组无法共用一个分布式优化器。为此仓库提供了 optimizer.pyget_mimo_optimizer为每个模块语言模型 各模态子模块基于其自身ProcessGroupCollection创建独立的MegatronOptimizerMimoOptimizer聚合所有活跃模块优化器全局梯度范数通过all_reduce(MAX)在模块间取平方和再开方随后用全局范数对各模块梯度统一裁剪step中先用all_reduce(MAX)同步found_inf防止编码器 rank 发现 inf 而语言 rank 未发现时提前 return 导致的死锁再用all_reduce(MIN)同步更新成功标志支持分片检查点param_groups、grad_scaler、param_state_sharding_type以ShardedObject走分布式保存replica_id按(tp_gtp_rank, pp_rank, dp_cp_rank)去重。测试与验证仓库为 MIMO 提供了较完整的单元测试tests/unit_tests/models/mimo/可作为理解行为与自行验证的入口test_mimo_model.py模型构建与前向主流程test_mimo_embedding_alignment.py嵌入对齐的正确性含 mask 路径test_mimo_submodules.py / test_mimo_audio_submodules.py子模块构建、编码、投影行为test_mimo_role.py / test_mimo_hetero_grid_args.pyrank 角色与异构网格配置校验test_mimo_noncolocated_cp_correctness.py / test_mimo_colocated_correctness.py共置 / 非共置布局下的正确性test_mimo_optimizer_consensus.py跨模块优化器的范数 / 成功标志共识test_mimo_checkpoint.py分片 checkpoint 存取test_mimo_forward_step.py / test_mimo_1f1b_schedule.py训练 step 与流水线调度。小结MIMO 把「多模态」问题拆解为「一个共享语言模型 一组可插拔模态子模块」的规范组合编码器负责模态 → 嵌入投影负责维度对齐align_embeddings_by_token_positions负责按特殊 token 位置合并序列语言模型在统一嵌入空间中完成推理解码器与输出投影负责嵌入 → 模态生成。在此基础上仓库通过RankRole/HyperCommGrid/ColocatedBridgeCommunicator/PartitionAdapter/MimoOptimizer等组件把同一套架构平滑扩展到共置与非共置异构并行、上下文并行、序列并行、序列打包THD与多模块分布式优化使其既能作为教学用的单卡 VLM也能支撑工业级的多模态大模型训练。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Megatron-LM MIMO 架构深度解析：面向多模态理解与生成的可组合 In/Out 模型
Megatron-LM MIMO 架构深度解析面向多模态理解与生成的可组合 In/Out 模型【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMIMOMultimodal In/Out Model是 Megatron-LM 中一套以「语言模型 模态子模块」为核心的可组合多模态架构支持文本、图像、音频等多种模态的统一理解与生成。本文以 megatron/core/models/mimo/README.md 为主线结合仓库源码与测试系统讲解 MIMO 的架构设计、核心组件、配置方式与使用流程。读完本文你将掌握如何用MimoModelConfigModuleSpec组装出 VLM / 音视频语言模型并理解嵌入对齐、异构并行与多模块优化器的底层实现。MIMO 是什么MIMO 是一种让语言模型能够理解和生成多种模态文本、图像、音频等的模型架构。它由两部分组成一个核心语言模型在统一的嵌入空间中处理所有模态信息若干模态子模块Modality Submodules负责把模态原始输入编码为嵌入如图像 → embeddings、把嵌入解码回模态输出如 embeddings → 图像、并在模态空间与语言模型维度之间做投影Projection。由MimoModel统一负责两类关键工作将各模态嵌入对齐到序列中的特殊 token 位置将合并后的嵌入交由语言模型处理。由于模态子模块是可插拔的MIMO 提供了一条「灵活且规范化」的建模路径可以配置出多种多模态模型例如视觉语言模型VLMVision-Language Models音视频语言模型Audio-Visual Language Models多模态理解与生成模型实验性说明当前MimoModel与MimoModelConfig均标注为 experimental、仍在积极开发中API 可能在未来版本中无通知变更。实际使用前请关注仓库的版本演进见 megatron/core/models/mimo/model/base.py 与 megatron/core/models/mimo/config/base_configs.py 顶部的UserWarning。工作原理与完整数据流MIMO 的模型架构由两大组件构成语言模型Language Model模态子模块Modality Submodules完整的数据流如下Input → Encoder → Projection → Align input embeddings → Language Model → Hidden states for special generation tokens → Output Projection → Decoder → Output编码阶段模态子模块将输入转换为嵌入例如图像 → embeddingsMimoModel按 token 位置将全部模态嵌入与文本嵌入对齐语言模型处理统一的嵌入序列。解码阶段挑选与「模态生成特殊 token」对应的隐藏状态模态子模块将嵌入转换回模态输出例如 embeddings → 图像。注意编码与解码并非必须成对存在。以仓库中的默认实现为例VisionModalitySubmodules 与 AudioModalitySubmodules 的decode目前均抛出NotImplementedErrorNo decoders support yet / Audio decoding not implemented yet即当前主线聚焦于多模态理解encoding 方向生成方向的能力由具体子模块逐步补齐。核心组件详解语言模型统一嵌入空间的中央处理器语言模型是处理所有模态信息的核心组件通过共享词表充当所有模态信息的中央处理器处理同时包含文本 token 与模态 token 的合并序列。在实现上语言模型由MimoModelConfig.language_model_spec指定默认复用 GPT 系列模型。例如 examples/mimo/model_providers/mock.py 中语言模型直接使用GPTModel并传入transformer_layer_spec、vocab_size、max_sequence_length、pre_process/post_process等参数。也就是说任何符合ModuleSpec约定的 Transformer 语言模型都可以作为 MIMO 的中央处理器接入。模态子模块连接原始模态与语言模型的桥ModalitySubmodules抽象基类见 megatron/core/models/mimo/submodules/base.py负责连接原始模态数据与语言模型每个子模块处理编码modality → embeddings与解码embeddings → modality管理模态空间与语言模型维度之间的投影内部由四类部件组成encoders编码器字典、decoders解码器字典、input_projections输入投影编码后 → 语言模型前使用、output_projections输出投影从语言模型接收后 → 解码前使用。基类构造器签名如下源码与 README 一致class ModalitySubmodules(ABC, nn.Module): def __init__( self, encoders: Optional[Dict[str, nn.Module]] None, decoders: Optional[Dict[str, nn.Module]] None, input_projections: Optional[List[nn.Module]] None, output_projections: Optional[List[nn.Module]] None, is_first_stage: bool True, is_last_stage: bool True, pg_collectionNone, **kwargs, ):MIMO 提供了默认实现VisionModalitySubmodules、AudioModalitySubmodules也允许自定义子模块做专用处理# Custom implementation class CustomVisionSubmodules(ModalitySubmodules): def encode(self, inputs): # Specialized encoding logic return projected_embeddings # Use custom submodules when creating the model model MimoModel( mimo_config, modality_submodules{images: ModuleSpec(moduleCustomVisionSubmodules, params{...})} )从源码看子模块的构建逻辑from_specModalitySubmodules.from_specsubmodules/base.py负责从ModuleSpec实例化子模块其关键行为是按流水线阶段条件化地构建部件encoders / decoders所有阶段都会构建供流水线各阶段处理使用input_projections仅在is_last_stageTrue时构建投影发生在编码之后、送往语言模型之前即projection happens after encoding, before sending to language modeloutput_projections仅在is_first_stageTrue时构建投影发生在从语言模型接收之后、解码之前。此外encode方法submodules/base.py对编码器输出做了规范化处理接受(b, s, h)三维或(b*s, h)二维输出若编码器返回(embeddings, aux_state)元组则自动取首个张量combine_embeddings沿第 0 维拼接多个编码器的嵌入project_embeddings则按is_input选择输入或输出投影。forward方法保证encoder_inputs与hidden_states互斥并只在首/末阶段执行相应计算——这套约定正是非共置pipeline布局下跨卡组装嵌入的基础。编码器输入的约束encode要求encoders_data_batch的键与self.encoders的编码器名一一对应缺失会直接抛出ValueError(fNo inputs found for encoder {name})。这意味着modality_inputs的键名必须与初始化时ModuleSpec.submodules[encoders]的键名严格一致这也是 README 强调「modality names and encoder names should match the keys used in mimo config during initialization」的原因。嵌入对齐align_embeddings_by_token_positionsMimoModel通过align_embeddings_by_token_positions方法整合不同模态的嵌入实现见 megatron/core/models/mimo/model/base.py将各模态嵌入放置到输入序列中对应的特殊 token 位置处理维度匹配与位置跟踪确保嵌入放置正确。方法内部的大致行为如下# Inside MimoModels forward method aligned_embeddings self.align_embeddings_by_token_positions( modality_embeddings{text: text_emb, images: image_emb}, input_idstokens, special_token_ids{images: 32000} )源码级实现细节对齐的核心是掩码mask驱动的masked_scatter_具体规则为对text模态构造全 True 掩码再对special_token_ids中每个特殊 token id 做mask input_ids ! token_id即文本位置 非任何特殊 token 的位置对其他模态mask input_ids token_id即该模态的特殊 token 出现的位置校验每个模态的mask.sum()必须等于该模态嵌入的数量否则抛出ValueError防止嵌入与 token 数量错配输出返回(S, B, H)序列优先布局的合并嵌入张量。除掩码路径外该方法还支持预计算 token 索引modality_token_indices的高效路径由调用方提供批内行优先batch-major的扁平索引方法通过index_copy_直接落位。该路径被设计为「trusted producer contract」——索引值不会被重新校验避免在前向热路径引入设备规约与主机同步但_validate_precomputed_token_indices会检查键集合一致、索引为一维、索引总数等于B * S。这一优化对吞吐敏感的大规模训练尤其重要。对应测试见 tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py其中test_basic_alignment用「文本 单个模态」的 batch 验证了嵌入是否正确落到特殊 token 位置。配置与使用MimoModel 参数MimoModel( config: MimoModelConfig, # Required: Configuration for the model )注意MimoModel.__init__的实际签名还接受可选的cp_group、tp_group用于上下文并行 / 张量并行场景下的分区适配器例如 examples/mimo/model_providers/mock.py 中通过pg_collection.cp/pg_collection.tp传入。配置详情MimoModelConfigMIMO 模型使用MimoModelConfig实例化定义见 megatron/core/models/mimo/config/base_configs.py其中包含语言模型的ModuleSpec模态名 → 子模块ModuleSpec的字典映射。MimoModelConfig( language_model: ModuleSpec, # Specification for the language model modality_submodules: Dict[str, ModuleSpec], # Dictionary mapping modality names to their submodule specifications special_token_ids: Dict[str, int] {} # Dictionary mapping modality names to their special token IDs )源码中的完整字段从 base_configs.py 看MimoModelConfig实际包含五个字段README 中的language_model/modality_submodules对应源码的language_model_spec/modality_submodules_spec字段类型默认值说明language_model_specModuleSpecModuleSpec()语言模型规格modality_submodules_specDict[str, ModuleSpec]{}模态名 → 子模块规格special_token_idsDict[str, int]{}模态名 → 特殊 token ID。例如{vision: -200, audio: 32000}这些占位符用于在input_ids中标记插入模态嵌入的位置module_to_grid_mapOptional[Dict[str, HyperCommGrid]]None模块键如 vision、language→HyperCommGrid拓扑。语言模型必须使用固定键MIMO_LANGUAGE_MODULE_KEY即language。网格共享相同 rank → 共置colocatedTP/DP 可不同网格 rank 不相交 → 非共置pipeline 并行为None→ 使用传统全局parallel_state共置kv_formatstrsbhd注意力 KV 格式sbhdseq-batch-head-dim或thdtotal-head-dim其中__post_init__会对module_to_grid_map做键校验网格键集合必须恰好等于modality_submodules_spec的键集合加上语言模型固定键language否则抛出ValueError并提示缺失/多余项见 base_configs.py 与 config/role.py 中MIMO_LANGUAGE_MODULE_KEY language的定义。示例创建视觉语言模型VLM# Language model specification lm_spec ModuleSpec( moduleGPTModel, params{ config: language_config, transformer_layer_spec: get_mock_language_layer_spec(), vocab_size: 50304, } ) # Vision modality specification vision_submodule_spec ModuleSpec( moduleVisionModalitySubmodules, params{ # Any general parameters for the submodule can go here }, submodules{ encoders: { clip_encoder: ModuleSpec( moduleCLIPViTModel, params{ transformer_config: vision_config, transformer_layer_spec: get_mock_vision_layer_spec(), patch_dim: 16, img_h: 224, img_w: 224, } ), }, input_projections: [ ModuleSpec( moduleMultimodalProjector, params{ config: get_mock_projection_config(), submodules: get_mock_projection_layer_spec().submodules, projector_type: mlp, input_size: 128 } ), ], } ) # Instantiate the model vlm MimoModel( MimoModelConfig( language_modellm_spec, modality_submodules{images: vision_submodule_spec}, special_token_ids{images: 32000} ) )这份示例在仓库中有完整可运行的对应用法examples/mimo/model_providers/mock.py 中的model_provider_mock_vlm_single_encoder即为真实生产写法——它额外传入了max_sequence_length、pre_process/post_process并用pg_collection解析cp_group/tp_group。MIMO 前向方法使用# Prepare inputs for multiple modalities and encoders modality_inputs { # modality names and encoder names should match the keys used in mimo config during initialization. images: { clip_encoder: {pixel_values: images}, # Encoder-specific inputs vit_encoder: {images: vit_images} }, audio: { whisper_encoder: {input_features: audio_features} } } # Call forward method outputs, _ mimo_model( input_idsinput_ids, position_idsposition_ids, modality_inputsmodality_inputs, )forward完整参数与返回语义从 model/base.py 的签名看forward还支持attention_mask、loss_mask、labels、packing_kwargs构造 THD 格式PackedSeqParams用于序列打包训练以及modality_token_indices前文所述的高效索引路径。返回值语义取决于当前 rank 的角色仅编码器 rank返回Dict[str, Tensor]各编码器输出供 bridge 送往语言模块 rank语言模块 rank返回语言模型输出logits 或 loss无角色全部共置返回语言模型输出。多模态训练实践真实模型提供者仓库的 examples/mimo/model_providers/ 目录提供了多种可直接参考的组装方式mock.py单视觉编码器的教学用 VLMllava_vlm.pyspecial_token_ids{images: image_special_token_id}的 LLaVA 式 VLMllava_avlm.py同时接入images与audios两种模态、多个编码器的音视频语言模型nemotron_moe_vlm.py通过nemotron_special_token_ids(args)从命令行参数生成特殊 token 映射。对应的训练脚本见 examples/mimo/scripts/run_vlm_train.sh、examples/mimo/scripts/run_avlm_train.sh 等。分布式与异构并行从共置到多模块流水线README 聚焦单机模型组装但仓库源码将 MIMO 扩展到了复杂的分布式场景这也是「In/Out 模型」能支撑大规模训练的关键角色与布局RankRole/ModuleLayoutconfig/role.py 定义了 rank 角色机制ModuleLayout.COLOCATED共置所有模块共享相同 rank。覆盖两种情形——无 grid map沿用全局parallel_state以及异构 TP/DP 但有重叠 rank 的情形此时会构建ColocatedBridgeCommunicator桥接通信器把编码器嵌入从编码器的 TP/DP 布局转换到语言模型的 TP/DP 布局见 model/base.py 与 comm/colocated_communicator.py。该布局走_forward_all_modules路径。ModuleLayout.NON_COLOCATED非共置module_to_grid_map中网格 rank 范围互不重叠每个 rank 要么只跑编码器、要么只跑语言模型通过RankRole分派不同的前向路径_forward_encoders/_forward_language_module本质上是多模块流水线并行。RankRole.build会依据各网格是否共享 rank 自动判定布局非共置模式下还会结合pp维计算每个模块的is_first_stage/is_last_stage如果当前 rank 不属于任何模块网格会直接抛出RuntimeError提示检查module_to_grid_map配置。并行特性与约束上下文并行 / 序列并行当语言模型配置context_parallel_size 1或开启sequence_parallel时MimoModel会创建PartitionAdapterpartition/utils.py对合并后的(S, B, H)嵌入以及labels/loss_mask做 CP/SP 切分。此模式下要求attention_maskNoneCP 切分后的序列无法与稠密 mask 对齐应改用因果attn_mask_type或packed_seq_params掩码model/base.py。MTP多 token 预测支持MimoModel会为语言模型的 MTP 模块准备 token ID、位置 ID 与有效性掩码并保证它们与 CP 本地序列分区一致_prepare_mtp_inputs。嵌入层约束启用 MIMO 序列并行时语言模型嵌入层必须禁用自身的 SP scatter构造时传scatter_embedding_sequence_parallelFalse否则会在对齐前把文本 token 错误切分model/base.py。多模块优化器MimoOptimizer异构并行下不同模块拥有各自的数据并行 / 张量并行组无法共用一个分布式优化器。为此仓库提供了 optimizer.pyget_mimo_optimizer为每个模块语言模型 各模态子模块基于其自身ProcessGroupCollection创建独立的MegatronOptimizerMimoOptimizer聚合所有活跃模块优化器全局梯度范数通过all_reduce(MAX)在模块间取平方和再开方随后用全局范数对各模块梯度统一裁剪step中先用all_reduce(MAX)同步found_inf防止编码器 rank 发现 inf 而语言 rank 未发现时提前 return 导致的死锁再用all_reduce(MIN)同步更新成功标志支持分片检查点param_groups、grad_scaler、param_state_sharding_type以ShardedObject走分布式保存replica_id按(tp_gtp_rank, pp_rank, dp_cp_rank)去重。测试与验证仓库为 MIMO 提供了较完整的单元测试tests/unit_tests/models/mimo/可作为理解行为与自行验证的入口test_mimo_model.py模型构建与前向主流程test_mimo_embedding_alignment.py嵌入对齐的正确性含 mask 路径test_mimo_submodules.py / test_mimo_audio_submodules.py子模块构建、编码、投影行为test_mimo_role.py / test_mimo_hetero_grid_args.pyrank 角色与异构网格配置校验test_mimo_noncolocated_cp_correctness.py / test_mimo_colocated_correctness.py共置 / 非共置布局下的正确性test_mimo_optimizer_consensus.py跨模块优化器的范数 / 成功标志共识test_mimo_checkpoint.py分片 checkpoint 存取test_mimo_forward_step.py / test_mimo_1f1b_schedule.py训练 step 与流水线调度。小结MIMO 把「多模态」问题拆解为「一个共享语言模型 一组可插拔模态子模块」的规范组合编码器负责模态 → 嵌入投影负责维度对齐align_embeddings_by_token_positions负责按特殊 token 位置合并序列语言模型在统一嵌入空间中完成推理解码器与输出投影负责嵌入 → 模态生成。在此基础上仓库通过RankRole/HyperCommGrid/ColocatedBridgeCommunicator/PartitionAdapter/MimoOptimizer等组件把同一套架构平滑扩展到共置与非共置异构并行、上下文并行、序列并行、序列打包THD与多模块分布式优化使其既能作为教学用的单卡 VLM也能支撑工业级的多模态大模型训练。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考