Megatron Core 多模态模型实战指南从 LLaVA、NVLM 到 MIMO 框架【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMegatron Core本仓库megatron/原生支持语言与视觉、音频等多模态的组合理解本文基于仓库 docs/models/multimodal.md 系统梳理其多模态技术栈从可任意组合视觉/音频/文本的 MIMOMultimodal In/Out实验框架到 LLaVA、NVLM、LLaMA 3.1 Nemotron Nano VL 等视觉语言模型的完整训练链路权重转换、数据预处理、预训练、SFT、评测。读完本文你将掌握如何在本仓库中搭建一套可运行的多模态模型训练与评测流程并理解底层并行与架构设计。MIMOMultimodal In/Out 实验框架MIMOMultimodal In/Out Model是 Megatron Core 中的实验性框架支持视觉、音频、文本等模态的任意组合为构建自定义多模态模型提供灵活架构。注意MIMO 目前处于实验阶段正在积极开发中API 可能在后续版本发生变化。核心特性支持任意模态组合视觉、音频、文本针对不同输入模态的灵活编码器架构跨模态统一嵌入空间同时支持视觉-语言模型与音频-视觉-语言模型MIMO 的入口脚本位于 examples/mimo/pretrain_mimo.py从源码结构看其设计要点包括异构并行网格通过--mimo-llm-tp/pp/cp/ep等参数见 examples/mimo/training/args.py 的add_hetero_grid_args为语言模型单独配置张量/流水线/上下文/专家并行度再由build_module_grid_specs与create_topology构建多模块异构拓扑实现编码器与 LLM 各自独立的并行布局。模型提供器model provider通过resolve_provider按参数解析具体的多模态模型实现视觉、音频编码器等模块以modality_submodules形式挂载模型提供器示例见 examples/mimo/model_providers/如 Nemotron MoE VLM。编码器预取encoder prefetchEncoderPrefetchLoader在数据加载的同时执行冻结编码器前向prefetch_frozen_features隐藏编码器计算延迟详见 examples/mimo/training/encoder_prefetch.py。训练流程复用MIMO 直接复用标准megatron.training.training.pretrain循环与mimo_forward_step前向步骤仅要求使用分布式优化器--use-distributed-optimizer并约定dataloader_typeexternal。MIMO 的配套训练脚本与数据示例位于 examples/mimo/scripts/如run_mock_train.sh、run_avlm_train.sh。视觉语言模型VLM矩阵本仓库支持以下视觉语言模型的训练、微调与评测模型描述视觉编码器语言模型LLaVA视觉指令微调CLIP ViT-L/14Mistral-7B / LLaMANVLMNVIDIA 视觉语言模型CLIP / 自定义 ViTInternViTLLaMA 系Yi-34B / Qwen2-72BLLaMA 3.1 Nemotron Nano VL高效多模态模型Vision TransformerRADIOLLaMA 3.1 8B各模型的完整训练示例分别位于examples/multimodal/ — LLaVA 风格训练Mistral CLIPexamples/multimodal/nvlm/ — NVLM 训练脚本examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/ — Nemotron VL 训练脚本examples/multimodal/radio/ — RADIO 视觉编码器集成视觉编码器编码器描述关键特性CLIP ViTOpenAI 的 CLIP 视觉 Transformer图像-文本对齐支持多尺度L/14336pxRADIOResolution-Agnostic Dynamic Image Optimization灵活分辨率处理高效视觉编码从 examples/multimodal/config.py 的get_vision_model_config可以看到各视觉编码器在 mcore 中的默认超参clip24 层、16 头、hidden_size 1024、ffn_hidden_size 4096、quick_gelu激活、LayerNorm、class_token_len 1siglip27 层、16 头、hidden_size 1152、fast_gelu、class_token_len 0internvit45 层、hidden_size 3200、ffn_hidden_size 12800、RMSNormNVLM 使用radio32 层、16 头、hidden_size 1280、class_token_len 8Nemotron VL 使用。扩散模型对于多模态扩散模型图像生成、文生图等可参考 NVIDIA NeMo 生态中的 Diffusion Models 仓库DFM其中提供生产级实现包括Stable Diffusion 系列变体、文生图、图生图转换、ControlNet 及其他条件控制机制。本仓库主体聚焦自回归多模态理解模型扩散类生成模型的示例不在本仓库范围内。多模态核心特性图像-文本对齐在图像-描述image-caption对上进行预训练视觉指令微调在指令跟随数据集上进行 SFT灵活视觉编码器支持不同 ViT 架构与分辨率如 336px / 448px 输入组合式检查点Combined Checkpointing将视觉与语言模型统一合并为一个多模态检查点高效训练视觉与语言组件均支持完整并行TP、PP、DP。端到端实战以 LLaVAMistral CLIP为例以下流程基于 examples/multimodal/README.md完整复现预训练 指令微调 评测一条链路。该示例已在 A100 DGX 集群验证64 张 GPU、tensor parallel size 4 时预训练约 1 天、SFT 约 11 小时训练速度随 GPU 数量近似线性扩展。1. 环境搭建在 Megatron-LM 根目录构建多模态 Docker 镜像docker build -t megatron-multimodal -f examples/multimodal/Dockerfile .2. 准备语言模型权重参考 docs/llama_mistral.md 中 Mistral-7B 章节从 HuggingFace 下载Mistral-7B-Instruct-v0.3权重并转换为 mcore 格式tensor parallel size 4。请使用 HuggingFace 自带的 tokenizer。3. 转换视觉模型权重本示例使用 OpenAI CLIPViT-L/14336px运行转换脚本python examples/multimodal/model_converter/clip_converter.py --download-root /some/download/folder --output /some/output/folder --tensor-parallel-size 4 --use-te4. 合并多模态检查点将 mcore 格式的 CLIP 与 Mistral 权重合并为单一多模态检查点目录examples/multimodal/combine_lm_vision_checkpoints.sh /path/to/mistral/model /path/to/clip/model /output/dir该脚本内部调用 examples/multimodal/combine_state_dicts.py按language_model/vision_model前缀交叉排列各 TP 分片的model_optim_rng.pt并在输出目录写入latest_checkpointed_iteration.txt内容为 1其细节见 examples/multimodal/combine_lm_vision_checkpoints.sh。注意若加载时报错可尝试设置TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD1。仅可对可信检查点使用此变量因为它允许加载时执行任意代码。5. 预训练数据准备从 HuggingFace 下载 LLaVA-Pretrain 数据集并解压 images 文件夹需约 79GB 磁盘空间运行 examples/multimodal/convert_llava_pretrain_to_wds.py 转换为 webdataset 格式在 wds 目录执行energon prepare ./转换为 Megatron-Energon 格式交互选项如下 Please enter a desired train/val/test split like 0.5, 0.2, 0.3 or 8,1,1: 9,1,0 Do you want to create a dataset.yaml interactively? [Y/n]: Y Please enter a number to choose a class: 9 (VQASample) Do you want to set a simple field_map[Y] (or write your own sample_loader [n])? [Y/n]: Y Please enter a webdataset field name for image (class torch.Tensor): jpg Please enter a webdataset field name for context (class str): json[0][value] Please enter a webdataset field name for answers (typing.Optional[typing.List[str]], default: None): json[1][value] Please enter a webdataset field name for answer_weights (typing.Optional[torch.Tensor], default: None):修改 examples/multimodal/pretrain_dataset.yaml将两个path变量指向LLaVA-Pretrain/wds。6. 预训练运行cd megatron-lm dir examples/multimodal/pretrain_mistral_clip.sh从 examples/multimodal/pretrain_mistral_clip.sh 可见关键训练配置并行--tensor-model-parallel-size 4、--pipeline-model-parallel-size 1模型尺寸32 层、hidden_size 4096、32 头、ffn_hidden_size 14336对应--language-model-type mistral_7bconfig 中自动填充--ffn-hidden-size 14336、SwiGLU、RMSNorm 等序列长度--seq-length 576视觉编码器序列长度即图像 token 数、--decoder-seq-length 1024语言模型序列长度视觉输入--patch-dim 14 --img-h 336 --img-w 336 --disable-vision-class-token冻结策略--freeze-LM --freeze-ViT预训练阶段冻结语言与视觉骨干仅训练投影层优化器--use-distributed-optimizer、bf16、cosine 学习率 1.5e-4、global batch size 256DEBUG 模式为 32tokenizer--tokenizer-type MultimodalTokenizer--tokenizer-prompt-format mistral并配套 examples/multimodal/manual_prompts.json 提示模板。在 examples/multimodal/model.py 的model_provider中模型整体由megatron.core.models.multimodal.llava_model.LLaVAModel构建视觉 Transformer、视觉投影层与 GPT 解码器使用独立 config视觉部分强制pipeline_model_parallel_size 1、关闭 SP/CP--seq-length会被自动改写为视觉嵌入数语言模型序列长度由--decoder-seq-length单独控制须大于最大图像嵌入数。训练完成后可观察到类似如下的训练/验证损失曲线以上曲线在 global batch size 256 下获得改动该值会改变曲线形状。注意对 LLaVA 而言损失曲线并非下游任务表现的有效预测指标仍需通过文本生成与多指标评测判断模型质量。预训练脚本可重复执行以断点续训恢复时会加载最新的模型、优化器与 dataloader 状态。7. 指令微调SFT准备指令微调数据集Megatron-Energon 格式修改 examples/multimodal/sft_dataset.yaml 中path指向训练与验证 split运行examples/multimodal/sft_mistral_clip.shSFT 脚本同样支持重复执行以恢复训练。8. 评测文本生成examples/multimodal/text_generation_mistral_clip.sh --input-image-path /path/to/input/images --output-path /some/output/directory \ --model-path /path/to/model.pt --gt-path /path/to/groundtruth/file --task generation-task-name其中--task generation-task-name为评测基准名如captioning或MMMU。预训练后评测 — COCO captioning下载 COCO 2014 测试集test2014.zip下载 COCO 测试标注coco_karpathy_test.json以--task captioning运行文本生成运行 examples/multimodal/evaluate_coco.pypython examples/multimodal/evaluate_coco.py --input-path /output/directory/from/generation --groundtruth-path /path/to/groundtruth/file对 mistral-7b-instruct clip 的 LLaVA 模型COCO CIDEr 分数约为 94。SFT 后评测 — MMMU将官方 MMMU 代码 clone 到examples/multimodal目录下以--task MMMU运行文本生成数据集由代码自动从 HuggingFace 加载随后运行python examples/multimodal/evaluate_mmmu.py --input-path /output/directory/from/generation对指令微调后的 mistral-7b-instruct clip LLaVA 模型MMMU 分数约为 38。NVLM 训练要点NVLM 1.0 的完整流程权重转换、预训练、SFT、PP 重分片、评测记录在 examples/multimodal/nvlm/README.md要点如下视觉编码器使用 InternViT-6B-448px-V1-5转换命令python examples/multimodal/model_converter/internvit_converter.py --output-dir some output dir --use-te --tensor-parallel-size 834B 语言模型从 Nous-Hermes-2-Yi-34B 出发用 tools/checkpoint/convert.py 转换--model-size yi-34B72B 语言模型从 Qwen2-72B-Instruct 出发--model-size qwen2.5-72Bf合并检查点examples/multimodal/combine_lm_vision_checkpoints.sh lm dir vision dir out dir nvlmTP8 模式训练脚本pretrain_yi_34b_internvit_6b.sh/sft_34b_internvit.sh34Bpretrain_qwen20_72b_internvit_6b.sh/sft_qwen20_72b_internvit_6b.sh72BPP 重分片72B 在 SFT 前需用 examples/multimodal/nvlm/pp_checkpoint_converter.py 将预训练检查点从 PP1 转为 PP4评测时再转回 PP1评测任务captioning、MMMU、TextVQA运行文本生成脚本时需加--use-tiling。NVLM 训练脚本中的关键多模态参数见 examples/multimodal/nvlm/pretrain_yi_34b_internvit_6b.sh--seq-length 256图像嵌入序列、--decoder-seq-length 512、--img-h 448 --img-w 448 --patch-dim 14、--pixel-shuffle、--image-tag-type nvlm、--tokenizer-prompt-format nvlm-yi-34b。其中--pixel-shuffle会将图像嵌入数压缩 4 倍在 examples/multimodal/model.py 中体现为num_image_embeddings // 4--image-tag-type nvlm则会按_get_tile_tags生成tile_i、tile_global_thumbnail等 tile 标签包围图像分块。LLaMA 3.1 Nemotron Nano VL 8B 要点该模型的完整流程记录在 examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/README.md语言模型从 meta-llama/Llama-3.1-8B-Instruct 转换--saver-transformer-impl transformer_engine --model-size llama3TP4视觉编码器RADIO 转换命令python examples/multimodal/model_converter/radio_converter.py --output radio_tp_4 --tensor-parallel-size 4 --use-te \ --version c-radio_v2-vlm-h --model-type radio_v2.5-h合并examples/multimodal/combine_lm_vision_checkpoints.sh lm dir vision dir out dir训练预训练脚本pretraining_llama_3p1_nemotron_nano_vl_8b_v1.sh、SFT 脚本sft_llama_3p1_nemotron_nano_vl_8b_v1.sh推理评测text_generation.sh --task inference或改为MMMU/TextVQA需--tensor-model-parallel-size 4。从源码看多模态模型构建原理examples/multimodal/model.py 的model_provider完整展示了 mcore 多模态模型的组装逻辑独立子配置基于core_transformer_config_from_args深拷贝出language_config、vision_config、vision_projection_config三套独立配置分别由get_language_model_config、get_vision_model_config、get_vision_projection_config均在 examples/multimodal/config.py按--language-model-type/--vision-model-type注入默认超参可插拔 layer spec语言部分支持 TEget_layer_spec_te、混合架构get_hybrid_layer_spec_te、MoEget_gpt_decoder_block_spec甚至hf://前缀直接加载 HuggingFace 模型此时要求 TP1、PP1、无 SP/CP视觉部分按vision_model_type选择 clip/siglip/radio/internvit/pixtral 等 spec图像 token 数量计算静态分辨率下由get_num_image_embeddings(img_h, img_w, patch_dim, ...)计算并自动改写seq_length动态分辨率--dynamic-resolution下直接以--seq-length为上限冻结控制构建完成后调用model.freeze(freeze_language_modelargs.freeze_LM, freeze_vision_modelargs.freeze_ViT, ...)对应训练脚本中的--freeze-LM/--freeze-ViT重计算隔离--recompute*参数对视觉与语言部分分开生效--recompute-vision控制视觉骨干的重计算。小结本仓库的多模态能力以megatron.core.models.multimodal.llava_model.LLaVAModel为统一载体向下兼容多种视觉编码器CLIP、SigLIP、InternViT、RADIO、Pixtral 等与语言骨干Mistral、Yi、Qwen、Llama、Nemotron 系列向上支撑 LLaVA、NVLM、Nemotron VL 等模型的预训练与 SFTMIMO 框架进一步把组合范围扩展到音频等更多模态与异构并行。任何新模型接入的关键都落在 examples/multimodal/config.py 的配置函数与 examples/multimodal/model.py 的 spec 选择逻辑上配合 tools/checkpoint/convert.py 完成权重转换后即可复用同一套训练、评测与断点续训管线。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考