Megatron-LM 多模态示例完全指南：从零预训练与指令微调 LLaVA 架构视觉语言模型（Mistral-7B + CLIP）

Megatron-LM 多模态示例完全指南：从零预训练与指令微调 LLaVA 架构视觉语言模型（Mistral-7B + CLIP）

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

本文基于 examples/multimodal/README.md 编写，结合仓库中的训练脚本、模型转换器与评估代码进行源码级扩充。文中所有命令、参数与文件路径均以当前 Megatron-LM 仓库实际内容为准。

导读

本指南完整演示了在 Megatron-LM 中预训练（Pretrain）与指令微调（Instruction Tune）一个LLaVA 架构的视觉语言模型（VLM）：以 Mistral-7B-Instruct-v0.3 作为语言主干、OpenAI CLIPViT-L/14@336px作为视觉编码器，覆盖从环境搭建、模型权重转换、checkpoint 合并、数据格式转换，到预训练、SFT 及 COCO/MMMU 下游评测的全流程。读完本文，你将能够独立复现仓库中的多模态训练管线，并理解每一步背后的源码实现。

需要提前说明的是：多模态支持在 Megatron-LM 中仍处于积极开发阶段（README 明确标注under active development and is expected to change），本示例的目的并非产出 SOTA 模型质量，而是演示 megatron 的多模态功能。以下流程已在A100 基座的 DGX 集群上验证：在 64 块 GPU（tp=4，四路张量并行）上，预训练约耗时 1 天，指令微调约 11 小时，训练速度随 GPU 数量近似线性扩展。

一、环境搭建：Docker 容器

仓库提供了开箱即用的镜像定义 examples/multimodal/Dockerfile。在 Megatron-LM 仓库根目录执行：

docker build -t megatron-multimodal -f examples/multimodal/Dockerfile .

构建完成后即可在该容器内完成本文后续全部步骤。注意，容器本身只负责提供运行时环境（依赖、TE 等），模型权重与数据集仍需按下列小节自行下载与转换。

二、模型权重准备

多模态模型由"语言模型 + 视觉模型"两部分拼装而成，二者需要分别下载并转换为 Megatron-Core（mcore）格式。

2.1 语言模型：Mistral-7B-Instruct-v0.3

按文档 docs/llama_mistral.md 中 Mistral-7B 一节的说明，从 HuggingFace 下载 Mistral-7B-Instruct-v0.3 权重，并转换为mcore 格式、张量并行度为 4（tensor parallel size 4）的 checkpoint。转换后请继续使用 HuggingFace 的 tokenizer（而非自行训练的词表），因为后续训练脚本通过--tokenizer-model mistralai/Mistral-7B-Instruct-v0.3直接引用该 tokenizer 名称。

2.2 视觉模型：OpenAI CLIP ViT-L/14@336px

本示例使用 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

从源码 examples/multimodal/model_converter/clip_converter.py 可以清晰看到转换器做了什么：

• 加载 OpenAI 原始 CLIP 权重（clip.load("ViT-L/14@336px", ...)），仅保留visual前缀（即文本塔被丢弃），跳过visual.proj与ln_post等本模型不使用的位置；
• 将 PyTorch 多头注意力的参数布局重排为 Megatron 的linear_qkv布局：通过kv_channels=64、num_heads=16、hidden_dim=1024计算 head 索引，把 Q、K、V 拼接权重重排成 mcore 期望的[q0,k0,v0,q1,k1,v1,...]顺序；
• 将参数名映射到 mcore 命名空间，例如transformer.resblocks.N.attn.in_proj_weight→decoder.layers.N.self_attention.linear_qkv.weight，mlp.c_fc.weight→mlp.linear_fc1.weight；
• 在指定维度（chunk_dim）上对 QKV/FFN 等张量按--tensor-parallel-size切分，同时把class_embedding展开为class_token（扩展至输入维度），并将 fp16 张量转回 fp32。

2.3 组合多模态 checkpoint

将上两步得到的 mcore CLIP 与 Mistral checkpoint 合并为单一多模态 checkpoint 文件夹：

examples/multimodal/combine_lm_vision_checkpoints.sh \ /path/to/mistral/model \ /path/to/clip/model \ /output/dir

脚本内部调用 examples/multimodal/combine_state_dicts.py：以语言模型和视觉模型的iter_0000001/mp_rank_XX/model_optim_rng.pt为输入，按 rank 一一配对，分别冠以language_model与vision_model前缀后合并写入输出目录，最后写入latest_checkpointed_iteration.txt。该脚本同时支持 NVLM（TP=8）与 Mistral-CLIP（TP=4）两种模型类型，默认走 TP=4 分支。

注意：若加载 checkpoint 时遇到错误，可尝试设置环境变量TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1。该变量仅在处理可信 checkpoint 时使用，因为它允许加载过程中的任意代码执行（潜在安全风险）。

三、预训练（Pretraining）

预训练的目标是让模型学会图片描述（image captioning），数据来自 LLaVA-Pretrain 数据集。

3.1 下载并解压 LLaVA-Pretrain

从 HuggingFace 克隆 LLaVA-Pretrain 数据集并解压图片（注意：需要约 79GB 磁盘空间）：

git clone https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain cd LLaVA-Pretrain unzip images.zip

3.2 转换为 webdataset 格式

回到 Megatron-LM 根目录，运行数据转换脚本：

cd <megatron-lm dir> python examples/multimodal/convert_llava_pretrain_to_wds.py

该脚本（examples/multimodal/convert_llava_pretrain_to_wds.py）读取blip_laion_cc_sbu_558k.json，逐条将图片二进制与对话 JSON 打包为 webdataset tar 分片：每条样本的__key__为图片 id，jpg字段存图片字节，json字段存entry['conversations']的 JSON 编码，每个分片最多 10000 条，输出到LLaVA-Pretrain/wds目录。使用前请将脚本顶部的llava_pretrain_dir = '<path_to_LLaVA-Pretrain>'替换为实际路径。

3.3 转换为 megatron-energon 格式

进入wds目录后运行 energon 的交互式准备命令：

cd <LLaVA-Pretrain dir>/wds energon prepare ./

按如下选项交互回答（对应VQASample样本类型）：

> 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):

要点：context取对话 JSON 的第一个元素（即用户提问），answers取第二个元素（即助手回答），answer_weights留空（可选字段，默认 None）。切分比例选择9,1,0，即 90% 训练、10% 验证、0 测试。

3.4 配置数据路径

编辑 examples/multimodal/pretrain_dataset.yaml，将其中train与val两个 split 下的path变量都指向LLaVA-Pretrain/wds。该文件是一个megatron.energon.Metadataset定义，subflavors.augmentation: false表示不做数据增强。

3.5 启动预训练

cd <megatron-lm dir> examples/multimodal/pretrain_mistral_clip.sh

脚本 examples/multimodal/pretrain_mistral_clip.sh 在启动前要求两个环境变量：

• WORKSPACE：模型 checkpoint 输出根目录；
• LOAD_NAME：上一节组合好的多模态 checkpoint 所在目录名（实际从${WORKSPACE}/${LOAD_NAME}/checkpoints加载）。

该脚本的核心训练配置如下（可在脚本中调整）：

其他值得注意的点：

• 脚本使用--use-te+--transformer-impl transformer_engine、--use-flash-attn、--use-distributed-optimizer、--apply-layernorm-1p、--attention-softmax-in-fp32、--untie-embeddings-and-output-weights、--disable-bias-linear、--position-embedding-type rope（rotary-base 1000000）等 Megatron-Core 标准配置；
• --pretrained-checkpoint ${CHECKPOINT_DIR}指向组合好的多模态 checkpoint，配合--load/--save ${FINETUNE_DIR}完成训练过程断点保存；
• --allow-missing-vision-projection-checkpoint允许视觉投影层权重缺失时继续加载；
• 环境变量NVTE_ALLOW_NONDETERMINISTIC_ALGO=1允许非确定性注意力算法，NVTE_APPLY_QK_LAYER_SCALING=0关闭 QK 层缩放；
• 训练入口为 examples/multimodal/train.py，它通过megatron.core.models.multimodal.llava_model.LLaVAModel构建模型，并仅在流水线首末级 stage 上运行 dataloader（train_valid_test_dataloaders_provider）。

一切正常时，你会观察到与下图类似的训练/验证 loss 曲线：

该曲线在global batch size 256下获得；更改该值曲线会相应变化。对 LLaVA 类模型而言，loss 曲线并不能可靠预测下游任务表现，因此必须通过后续的测试生成与多指标评测来判断模型质量（文档表示未来版本会加入训练期零样本评测）。

断点续训：可多次执行预训练脚本；恢复时脚本会自动加载最新的 model、optimizer 与 dataloader 状态。

四、指令微调（SFT）

SFT 阶段让预训练模型学会遵循指令完成任务。

1. 准备一份指令微调数据集，同样需要 megatron-energon 格式（README 明确说明仓库不提供该数据集的制作指引，需自行准备）；
2. 编辑 examples/multimodal/sft_dataset.yaml，将train与val两个 split 的path分别指向微调数据集的训练与验证部分；
3. 运行：

examples/multimodal/sft_mistral_clip.sh

SFT 脚本 examples/multimodal/sft_mistral_clip.sh 需要三个环境变量：WORKSPACE、LOAD_NAME（预训练输出目录名）、LOAD_ITER（要加载的预训练迭代号，用于定位${WORKSPACE}/${LOAD_NAME}/checkpoints下的具体 iter 目录）。

与预训练相比，SFT 配置的关键差异（其余架构参数保持一致）：

SFT 同样支持多次执行以断点续训，恢复时加载最新的 model、optimizer 与 dataloader 状态。

五、评测（Evaluation）

评测分为"生成（Generation）"与"下游指标评估"两个环节，前者由统一的文本生成脚本完成。

5.1 文本生成

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取值是评测基准名称，例如captioning（COCO 标题生成）或MMMU。脚本 examples/multimodal/text_generation_mistral_clip.sh 还支持--num-frames（视频帧数，默认 1）、--gt-path等参数，并内置了NUM_PARTITIONS/START/END分区变量用于大评测集分片并行生成；生成时通过torchrun --nproc_per_node 8拉起 examples/multimodal/run_text_generation.py，推理侧采用--top_k 1、--temperature 1.0、--seq-length 2048、--out-seq-length 12等贪心解码配置，并加载与训练一致的多模态 checkpoint（--ckpt-format torch）。

5.2 预训练模型评估：COCO Captioning

1. 下载 COCO 2014 测试图片集：

   wget http://images.cocodataset.org/zips/test2014.zip

2. 下载 COCO 测试图片标注（Karpathy 测试标注）：

   https://storage.googleapis.com/sfr-vision-language-research/datasets/coco_karpathy_test.json

3. 先用--task captioning运行文本生成；
4. 计算 CIDEr 分数：

   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。

5.3 SFT 模型评估：MMMU

MMMU 官方仓库目前无法通过 pip 安装，需在examples/multimodal目录下手动克隆：

cd examples/multimodal git clone https://github.com/MMMU-Benchmark/MMMU.git

MMMU 数据集由代码自动从 HuggingFace 加载。先用--task MMMU运行文本生成，然后：

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

对mistral-7b-instruct + clip的指令微调 LLaVA 模型，MMMU 分数应约为38。

六、评估脚本一览

评测环节用到的脚本均位于 examples/multimodal 目录，除本文使用的 evaluate_coco.py 与 evaluate_mmmu.py 外，仓库还提供了覆盖多基准的评估器：evaluate_ai2d.py、evaluate_chartqa.py、evaluate_mathvista.py、evaluate_ocrbench.py、evaluate_textvqa.py、evaluate_vqav2.py 等，数据集定义统一收敛在 evaluation_datasets.py 中，可参照 COCO/MMMU 的用法扩展更多评测基准。

七、常见问题与注意事项

• checkpoint 加载报错：设置TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1重试（仅限可信文件，存在任意代码执行风险）；
• TP 一致性：CLIP 转换（--tensor-parallel-size 4）、Mistral 转换、checkpoint 合并脚本（TP=4 分支）与训练脚本（--tensor-model-parallel-size 4）必须保持张量并行度一致；
• 磁盘空间：LLaVA-Pretrain 解压后约 79GB，请提前规划存储；
• 数据路径：pretrain_dataset.yaml与sft_dataset.yaml中两处path都必须正确指向 energon 格式数据目录；
• 环境变量：预训练脚本需WORKSPACE与LOAD_NAME；SFT 脚本额外需要LOAD_ITER；
• 复现性：训练/推理启用了NVTE_ALLOW_NONDETERMINISTIC_ALGO=1，结果存在一定非确定性。

结语

本文完整复现了 Megatron-LM 多模态示例的端到端流程：从 Docker 环境、Mistral/CLIP 权重转换与合并，到 LLaVA-Pretrain 数据管线（webdataset → megatron-energon）、冻结 LM/ViT 的预训练、仅冻结 ViT 的 SFT，再到 COCO Captioning 与 MMMU 的生成与评测。示例的预期表现（COCO CIDEr ≈ 94、MMMU ≈ 38）可作为复现的验收基准。多模态功能仍在积极演进中，若复现遇到问题，可在当前仓库的 issue 系统中反馈；未来版本还将引入训练期零样本评测能力。

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