Megatron-Core 中 NVLM 1.0 多模态模型训练实战：从检查点转换、预训练/SFT 到评估的完整指南

Megatron-Core 中 NVLM 1.0 多模态模型训练实战：从检查点转换、预训练/SFT 到评估的完整指南

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

导读

本文基于 Megatron-LM 仓库中 examples/multimodal/nvlm/README.md 及配套脚本，系统讲解如何在 Megatron-Core 框架下复现并训练 NVLM 1.0（34B 与 72B 两个尺寸）多模态大模型。文章覆盖从 HuggingFace 权重到 Megatron 格式的视觉/语言模型转换、视觉与语言检查点合并、基于 Megatron Energon 的数据准备、预训练与 SFT 全流程，以及文本生成与评测脚本的使用方法。读完本文，你将掌握一套完整可复现的 NVLM 训练流水线，并理解其中涉及的张量并行（TP=8）、流水线并行（PP 重切分）、InternViT 视觉编码器接入等关键实现细节。

一、NVLM 1.0 与 Megatron 中的 VLM 支持概览

NVLM 1.0 是由 NVIDIA 发布的多模态大语言模型（VLM），Megatron-LM 仓库为 34B 与 72B 两个规格提供了完整的训练与推理脚本。仓库明确指出：Megatron 中的 VLM（视觉语言模型）仍处于积极开发阶段，相关接口与行为预期会持续变化，因此使用时应以当前仓库版本为准。

NVLM 1.0 模型权重公开提供两种格式：

• HuggingFace 格式：NVLM-1.0-D 72B（HuggingFace 上对应nvidia/NVLM-D-72B）；
• Megatron-Core 格式：NVLM-1.0-D 72B（HuggingFace 上对应nvidia/NVLM-D-72B-mcore）。

从源码结构看，NVLM 在 Megatron 中由三个核心部分组成：InternViT 视觉编码器（vision_model）、语言模型主干（language_model，Yi-34B 或 Qwen2-72B-Instruct）以及vision projection 投影层。相关的模型组装与训练入口位于 examples/multimodal/train.py，构建逻辑可参考 examples/multimodal/model.py，所有 NVLM 专属脚本集中在 examples/multimodal/nvlm/ 目录下。

二、环境搭建与数据集准备

2.1 Docker 镜像

NVLM 训练脚本依赖 Transformer Engine（--use-te）等环境，官方要求使用仓库提供的多模态 Dockerfile 构建镜像：

examples/multimodal/Dockerfile

该镜像包含了运行 examples/multimodal/train.py 所需的依赖。训练脚本同时支持两种启动方式（见各.sh脚本末尾）：

• 交互式：torchrun --nproc_per_node 8 examples/multimodal/train.py ${OPTIONS}；
• 批处理/SLURM：通过srun在容器内执行同样的 Python 入口，并输出日志到${LOGS_DIR}。

2.2 数据集准备（Megatron Energon 格式）

NVLM 的预训练与 SFT 均使用 Megatron Energon 数据加载器，对应的数据描述文件为：

• 预训练混合数据集：examples/multimodal/nvlm/pretrain_blend.yaml；
• SFT 混合数据集：examples/multimodal/nvlm/sft_blend.yaml。

两个文件均为 Energon 的Metadataset配置，结构如下：

__module__: megatron.energon __class__: Metadataset splits: train: datasets: - weight: 0.579 # 数据集按其规模加权，所有权重之和为 1 path: <path to laion dataset> subflavors: augmentation: False - weight: 0.02 path: <path to coco> subflavors: augmentation: False # 完整预训练数据集清单请参考 NVLM 论文 Table 4 val: datasets: - weight: 1. path: <path to validation dataset> subflavors: augmentation: False

其中weight决定各数据集在混合采样中的占比；subflavors.augmentation用于控制是否做数据增强。SFT 混合文件结构相同，训练数据覆盖 COCO、CLEVR-Math 等（论文 Table 6）。读者需按 Megatron Energon 的数据准备规范将原始数据（如 LAION、COCO、VQAv2）预处理为 Energon 格式后，再把对应路径填入上述 yaml。此外，训练时通过--dataloader-type external启用外部数据加载器，并将数据文件通过--data-path传入。

三、模型权重转换：从 HuggingFace 到 Megatron-Core

训练前需要把 NVLM 的两个组成部分——视觉模型与语言模型——分别从 HuggingFace 格式转换为 Megatron 格式，再合并成多模态检查点。

3.1 视觉模型：InternViT 转换

NVLM 1.0 使用的视觉编码器是 HuggingFace 上的OpenGVLab/InternViT-6B-448px-V1-5（448px 分辨率、6B 参数）。下载后运行 examples/multimodal/model_converter/internvit_converter.py 转换：

python examples/multimodal/model_converter/internvit_converter.py --output-dir <some output dir> --use-te --tensor-parallel-size 8

从转换器源码可以看到几个关键技术点：

• 参数名映射：脚本把 HF 权重逐项改名为 Megatron 的命名，例如embeddings.patch_embedding.weight→conv1.weight、encoder.layers.*.attn.qkv.weight→decoder.layers.*.self_attention.linear_qkv.weight、mlp.fc1.weight→mlp.linear_fc1.weight等；
• 注意力头填充：InternViT 有 25 个注意力头，当--tensor-parallel-size 8时，由于 25 不能被 8 整除，脚本会用全零的虚拟头把 25 个头补齐到 32 个头（num_padded_heads = 32），以支持张量并行的均匀切分——这也是为什么仓库建议自行设计模型时使用偶数注意力头；
• QKV 权重重排：脚本通过order张量将 HF 的 Q/K/V 拼接顺序重排为 Megatron 期望的布局，并对linear_qkv等张量沿 dim 0 做torch.chunk切分到各 TP rank；
• TE 兼容：当--use-te开启时，会为linear_qkv、linear_proj、linear_fc1、linear_fc2写入空的_extra_state，与 Transformer Engine 的 FP8 状态机制兼容；
• 输出布局：每个 TP rank 生成iter_0000001/mp_rank_0{i}/model_optim_rng.pt。

注意：该脚本只支持--tensor-parallel-size 1或8两种取值，其它取值会抛出NotImplementedError。

3.2 34B 语言模型（Yi-34B）转换

NVLM 1.0 34B 以 HuggingFace 上的NousResearch/Nous-Hermes-2-Yi-34B为起点，使用通用转换工具 tools/checkpoint/convert.py 转换：

python tools/checkpoint/convert.py --bf16 --model-type GPT --loader llama_mistral --saver mcore --target-tensor-parallel-size 8 --checkpoint-type hf \ --load-dir <hf model directory> --save-dir <output dir> --tokenizer-model <hf model name/directory> \ --saver-transformer-impl transformer_engine --model-size yi-34B --make-vocab-size-divisible-by 1

关键参数说明：

3.3 72B 语言模型（Qwen2-72B-Instruct）转换

NVLM 1.0 72B 以Qwen/Qwen2-72B-Instruct为起点，转换命令几乎相同：

python tools/checkpoint/convert.py --bf16 --model-type GPT --loader llama_mistral --saver mcore --target-tensor-parallel-size 8 --checkpoint-type hf \ --load-dir <hf model directory> --save-dir <output directory> --tokenizer-model <hf model name/directory> \ --saver-transformer-impl transformer_engine --model-size qwen2.5-72Bf

区别仅在于--model-size qwen2.5-72Bf。转换时同样输出 TP=8 的 Megatron 检查点布局，供后续合并与训练使用。

3.4 合并视觉与语言检查点

转换完成后，用 examples/multimodal/combine_lm_vision_checkpoints.sh 把 InternViT 视觉检查点与 34B/72B 语言检查点合并为多模态检查点：

examples/multimodal/combine_lm_vision_checkpoints.sh <language model directory> <vision model directory> <output directory> nvlm

该脚本对 NVLM（TP=8）会调用 examples/multimodal/combine_state_dicts.py，将 8 个 TP rank 的语言模型与视觉模型状态字典按language_model/vision_model前缀交替配对合并：

python examples/multimodal/combine_state_dicts.py \ --input ${MCORE_LM}/iter_0000001/mp_rank_00/model_optim_rng.pt \ ${MCORE_VISION}/iter_0000001/mp_rank_00/model_optim_rng.pt \ ...（共 16 个文件，LM 与 ViT 交替） \ --prefixes language_model vision_model language_model vision_model ... \ --output ${OUTPUT_DIR}/iter_0000001/mp_rank_00/model_optim_rng.pt ...（共 8 个文件）

其原理（见 combine_state_dicts.py）：对每个输出文件，取同一 TP rank 的语言与视觉状态字典，把视觉模型的所有参数键加上vision_model.前缀、语言模型参数键加上language_model.前缀后合并到一个字典中；最后脚本会写入latest_checkpointed_iteration.txt。合并结果即为训练脚本中--pretrained-checkpoint指向的初始权重。

四、预训练（Pretraining）

4.1 34B 预训练

运行 examples/multimodal/nvlm/pretrain_yi_34b_internvit_6b.sh，需使用前面生成的 InternViT + 34B 合并检查点，tokenizer 直接使用 HuggingFace 的NousResearch/Nous-Hermes-2-Yi-34B。该脚本的关键配置如下：

• 模型规模：60 层、hidden 7168、FFN 20480、56 注意力头、GQA（8 个 query group）、SwiGLU、RMSNorm、RoPE（--rotary-base 5000000），--language-model-type yi-34b；
• 图像编码：--vision-model-type internvit，--img-h/--img-w 448、--patch-dim 14，开启--pixel-shuffle（像素重排下采样）、--image-tag-type nvlm、--disable-vision-class-token；
• 序列长度：--seq-length 256（图像 embedding 序列长度）、--decoder-seq-length 512（语言模型序列长度）、--max-position-embeddings 512；
• 训练规模：--global-batch-size 2048、--micro-batch-size 1、--train-samples 122880000、--lr 1e-4、cosine 衰减、--weight-decay 0.1、--clip-grad 10.0；
• 冻结策略：同时--freeze-LM --freeze-ViT（预训练阶段冻结视觉与语言主干，只训练 vision projection 与多模态对齐部分）；
• 检查点：--pretrained-checkpoint指向合并检查点，并配合--allow-missing-vision-projection-checkpoint（允许缺失视觉投影权重）与--use-checkpoint-args、--ckpt-format torch；
• 环境变量：export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1、export NVTE_APPLY_QK_LAYER_SCALING=0；
• 并行度：--tensor-model-parallel-size 8，交互模式下通过torchrun --nproc_per_node 8启动。

脚本还内置了DEBUG=1的快速调试分支（micro-batch 1、dropout 0、log-interval 1），以及BATCH变量切换 torchrun / srun 两种运行模式。

4.2 72B 预训练

运行 examples/multimodal/nvlm/pretrain_qwen20_72b_internvit_6b.sh，使用 InternViT + 72B 合并检查点与 HuggingFace 的Qwen/Qwen2-72B-Instructtokenizer。与 34B 的主要差异：

• 模型规模：80 层、hidden 8192、FFN 29568、64 注意力头、GQA（8 query groups）、--add-qkv-bias、--norm-epsilon 1e-06、--rotary-base 1000000，--language-model-type qwen2.0_72B；
• 并行配置：显式给出--tensor-model-parallel-size 8 --pipeline-model-parallel-size 1；
• Attention 数值稳定性：增加--no-masked-softmax-fusion与--attention-softmax-in-fp32；
• tokenizer 提示词格式：--tokenizer-prompt-format qwen2p0（34B 脚本则为nvlm-yi-34b）；
• 其余（冻结策略、序列长度、--pixel-shuffle、--allow-missing-vision-projection-checkpoint等）与 34B 一致。

五、流水线并行检查点转换（72B SFT 前置步骤）

72B 的预训练在 PP=1 下进行，而 SFT 需要 PP=4。为此仓库提供了专用工具 examples/multimodal/nvlm/pp_checkpoint_converter.py，用于在 PP=1 与 PP=N 之间互相转换。

5.1 从 PP=1 切分为 PP=4（预训练 → SFT）

examples/multimodal/nvlm/pp_checkpoint_converter.py --input <pretrained checkpoint directory> \ --input-pipeline-parallel 1 --output <some output dir> --output-pipeline-parallel 4 \ --tensor-parallel 8

5.2 从 PP=4 合并回 PP=1（SFT → 评测）

examples/multimodal/nvlm/pp_checkpoint_converter.py --input <sft checkpoint directory> \ --input-pipeline-parallel 4 --output <some output dir> --output-pipeline-parallel 1 \ --tensor-parallel 8

从 pp_checkpoint_converter.py 源码可以看出其处理规则：

• split 方向（PP 1→N）：逐 TP rank 读取mp_rank_0{tp}/model_optim_rng.pt，按层号区间把language_model.decoder.layers分发到各 PP rank；视觉模型（vision_model）与vision_projection只保留在第一个 PP rank；word_embeddings仅归第一个 PP rank，output_layer与final_layernorm仅归最后一个 PP rank；每层在所有 PP rank 上从 0 开始重新编号；
• combine 方向（PP N→1）：反向合并各 PP rank 的层，层号按 offset 递增回填，视觉/投影权重取自第一个 PP rank，并把args.pipeline_model_parallel_size改回 1；
• 不均等切分：当num_layers % output_pp != 0时，可用--num-layers-per-pp-rank指定每个 PP rank 的层数（每个 rank 至少 1 层）；
• 迭代号：可用--iteration指定要处理的检查点迭代，输出目录结构为iter_{iter:0>7}/mp_rank_0{tp}_00{pp}/model_optim_rng.pt，并自动写出latest_checkpointed_iteration.txt。

该脚本仅支持 PP 1↔N 两种方向，其他组合会抛出NotImplementedError。

六、监督微调（SFT）

6.1 34B SFT

运行 examples/multimodal/nvlm/sft_34b_internvit.sh，--pretrained-checkpoint指向 34B 预训练产出的检查点。关键差异：

• --global-batch-size 128、--lr 2e-6（远低于预训练）、--min-lr 2.5e-7；
• 序列长度放大：--seq-length 261（256 个图像 embedding + 5 个 tile tag embedding）、--decoder-seq-length 3200、--max-position-embeddings 3200；
• 开启动态 tile 分块：--use-tiling --max-num-tiles 6 --use-thumbnail --use-tile-tags --pixel-shuffle，支持把大图切分为最多 6 个 tile 并附带缩略图与 tile 标签；
• 仅--freeze-ViT（语言模型参与微调），数据切换到 sft_blend.yaml；
• 环境变量额外设置export NCCL_ALGO=^NVLS；DEBUG 分支为节省显存额外追加--freeze-LM（仅用于交互调试，正式训练不要使用）。

6.2 72B SFT

运行 examples/multimodal/nvlm/sft_qwen20_72b_internvit_6b.sh，使用第 5.1 节切分好的 PP=4 检查点（LOAD_NAME="mcore-qwen20-72b-internvit-pp4"）。关键差异：

• --pipeline-model-parallel-size 4（配合--tensor-model-parallel-size 8）；
• --global-batch-size 256、--lr 2e-6、--decoder-seq-length 3200、--max-position-embeddings 8192；
• 同样开启--use-tiling --max-num-tiles 6 --use-thumbnail --use-tile-tags，仅--freeze-ViT。

脚本注释明确提示：预训练检查点必须先手动切分到 4 个 PP stage（见 README 与pp_checkpoint_converter.py），SFT 完成后如需评测，再按第 5.2 节合并回 PP=1。

七、文本生成与评测

7.1 文本生成

评测前先用对应尺寸的推理脚本生成模型输出：

• 34B：examples/multimodal/nvlm/run_text_generation_yi_34b_internvit_6b.sh
• 72B：examples/multimodal/nvlm/run_text_generation_qwen20_72b_internvit_6b.sh

调用方式（两者一致）：

examples/multimodal/nvlm/run_text_generation_yi_34b_internvit_6b.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 --use-tiling

参数说明：--task为评测基准名（如captioning、MMMU、TextVQA）；--use-tiling开启动态 tile 分块；此外脚本还支持--use-pixel-shuffle-only（仅像素重排不分块）。内部调用 examples/multimodal/run_text_generation.py，默认推理配置为--micro-batch-size 1、--out-seq-length 16、--temperature 1.0、--top_k 1（贪心解码）、--no-load-rng --no-load-optim，并可通过NUM_PARTITIONS/START/END将评测集分片并行执行。

注意：脚本按生成模式动态调整序列长度——开启 tiling 时--seq-length 261，仅 pixel shuffle 时--seq-length 256，否则为 1024；语言模型解码长度固定为--decoder-seq-length 8192、--max-position-embeddings 8192。

7.2 运行评测脚本

生成完成后，使用 examples/multimodal 下的评测脚本对输出打分，例如 MMMU：

python examples/multimodal/evaluate_mmmu.py --input-path /output/directory/from/generation

仓库在 examples/multimodal/evaluation/ 目录提供了覆盖 AI2D、ChartQA、COCO、InfoVQA、MathVista、MMMU、OCRBench、TextVQA、VQAv2、视频类（MotionBench、MVBench）等十余个基准的评测实现（如 evaluate_ai2d.py、evaluate_textvqa.py 等），可配合--task选择对应任务。

八、全流程速查与注意事项

实战注意事项总结：

1. 并行度贯穿一致：从转换（--target-tensor-parallel-size 8）到训练（--tensor-model-parallel-size 8）再到推理，TP=8 必须保持一致；72B SFT 引入 PP=4 时需先用pp_checkpoint_converter.py切分，评测前再合并回 PP=1；
2. 冻结策略有讲究：预训练阶段--freeze-LM --freeze-ViT（只训练投影对齐），SFT 阶段仅--freeze-ViT，让语言模型参与指令微调；DEBUG 模式为省显存临时加--freeze-LM，仅限交互测试；
3. 图像预处理标志：--pixel-shuffle与--use-tiling/--use-thumbnail/--use-tile-tags必须在训练/推理间保持一致，且 tiling 会改变--seq-length（261 vs 256 vs 1024）；
4. tokenizer 配套：--tokenizer-type MultimodalTokenizer配合--tokenizer-prompt-format（34B 用nvlm-yi-34b，72B 用qwen2p0），提示词模板来自 examples/multimodal/nvlm/nvlm_prompts.json；
5. 视觉头数限制：InternViT 的 25 个注意力头在 TP=8 下会填充为 32 个虚拟头，这是当前张量并行实现不支持不均匀切分所致；
6. VLM 处于活跃开发中：Megatron 的 VLM 相关接口后续可能变化，请以当前仓库 examples/multimodal/nvlm/README.md 与脚本为准。

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