Megatron Core 多模态模型实战指南：从 LLaVA、NVLM 到 MIMO 框架-平芜编程栈

Megatron Core 多模态模型实战指南：从 LLaVA、NVLM 到 MIMO 框架

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM

Megatron Core（本仓库megatron/）原生支持语言与视觉、音频等多模态的组合理解，本文基于仓库 docs/models/multimodal.md 系统梳理其多模态技术栈：从可任意组合视觉/音频/文本的 MIMO（Multimodal In/Out）实验框架，到 LLaVA、NVLM、LLaMA 3.1 Nemotron Nano VL 等视觉语言模型的完整训练链路（权重转换、数据预处理、预训练、SFT、评测）。读完本文，你将掌握如何在本仓库中搭建一套可运行的多模态模型训练与评测流程，并理解底层并行与架构设计。

MIMO：Multimodal In/Out 实验框架

MIMO（Multimodal 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 prefetch）：EncoderPrefetchLoader在数据加载的同时执行冻结编码器前向（prefetch_frozen_features），隐藏编码器计算延迟，详见 examples/mimo/training/encoder_prefetch.py。
训练流程复用：MIMO 直接复用标准megatron.training.training.pretrain循环与mimo_forward_step前向步骤，仅要求使用分布式优化器（--use-distributed-optimizer），并约定dataloader_type=external。

MIMO 的配套训练脚本与数据示例位于 examples/mimo/scripts/（如run_mock_train.sh、run_avlm_train.sh）。

视觉语言模型（VLM）矩阵

本仓库支持以下视觉语言模型的训练、微调与评测：

模型	描述	视觉编码器	语言模型
LLaVA	视觉指令微调	CLIP ViT-L/14	Mistral-7B / LLaMA
NVLM	NVIDIA 视觉语言模型	CLIP / 自定义 ViT（InternViT）	LLaMA 系（Yi-34B / Qwen2-72B）
LLaMA 3.1 Nemotron Nano VL	高效多模态模型	Vision Transformer（RADIO）	LLaMA 3.1 8B

各模型的完整训练示例分别位于：

examples/multimodal/ — LLaVA 风格训练（Mistral + CLIP）
examples/multimodal/nvlm/ — NVLM 训练脚本
examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/ — Nemotron VL 训练脚本
examples/multimodal/radio/ — RADIO 视觉编码器集成

视觉编码器

编码器	描述	关键特性
CLIP ViT	OpenAI 的 CLIP 视觉 Transformer	图像-文本对齐，支持多尺度（L/14@336px）
RADIO	Resolution-Agnostic Dynamic Image Optimization	灵活分辨率处理，高效视觉编码

从 examples/multimodal/config.py 的get_vision_model_config可以看到各视觉编码器在 mcore 中的默认超参：

clip：24 层、16 头、hidden_size 1024、ffn_hidden_size 4096、quick_gelu激活、LayerNorm、class_token_len 1；
siglip：27 层、16 头、hidden_size 1152、fast_gelu、class_token_len 0；
internvit：45 层、hidden_size 3200、ffn_hidden_size 12800、RMSNorm（NVLM 使用）；
radio：32 层、16 头、hidden_size 1280、class_token_len 8（Nemotron VL 使用）。

扩散模型

对于多模态扩散模型（图像生成、文生图等），可参考 NVIDIA NeMo 生态中的 Diffusion Models 仓库（DFM），其中提供生产级实现，包括：Stable Diffusion 系列变体、文生图、图生图转换、ControlNet 及其他条件控制机制。本仓库主体聚焦自回归多模态理解模型，扩散类生成模型的示例不在本仓库范围内。

多模态核心特性

图像-文本对齐：在图像-描述（image-caption）对上进行预训练；
视觉指令微调：在指令跟随数据集上进行 SFT；
灵活视觉编码器：支持不同 ViT 架构与分辨率（如 336px / 448px 输入）；
组合式检查点（Combined Checkpointing）：将视觉与语言模型统一合并为一个多模态检查点；
高效训练：视觉与语言组件均支持完整并行（TP、PP、DP）。

端到端实战：以 LLaVA（Mistral + 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/14@336px，运行转换脚本：

python examples/multimodal/model_converter/clip_converter.py --download-root /some/download/folder --output /some/output/folder --tensor-parallel-size 4 --use-te

4. 合并多模态检查点

将 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_LOAD=1。仅可对可信检查点使用此变量，因为它允许加载时执行任意代码。

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_7b，config 中自动填充--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 256（DEBUG 模式为 32）；
tokenizer：--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.sh

SFT 脚本同样支持重复执行以恢复训练。

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.py：

python 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 8

34B 语言模型：从 Nous-Hermes-2-Yi-34B 出发，用 tools/checkpoint/convert.py 转换（--model-size yi-34B）；
72B 语言模型：从 Qwen2-72B-Instruct 出发（--model-size qwen2.5-72Bf）；
合并检查点：examples/multimodal/combine_lm_vision_checkpoints.sh <lm dir> <vision dir> <out dir> nvlm（TP=8 模式）；
训练脚本：pretrain_yi_34b_internvit_6b.sh/sft_34b_internvit.sh（34B），pretrain_qwen20_72b_internvit_6b.sh/sft_qwen20_72b_internvit_6b.sh（72B）；
PP 重分片：72B 在 SFT 前需用 examples/multimodal/nvlm/pp_checkpoint_converter.py 将预训练检查点从 PP=1 转为 PP=4，评测时再转回 PP=1；
评测任务：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 llama3，TP=4）；
视觉编码器：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：语言部分支持 TE（get_layer_spec_te）、混合架构（get_hybrid_layer_spec_te）、MoE（get_gpt_decoder_block_spec）甚至hf://前缀直接加载 HuggingFace 模型（此时要求 TP=1、PP=1、无 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_model=args.freeze_LM, freeze_vision_model=args.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 等模型的预训练与 SFT；MIMO 框架进一步把组合范围扩展到音频等更多模态与异构并行。任何新模型接入的关键都落在 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），仅供参考