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/14336px作为视觉编码器覆盖从环境搭建、模型权重转换、checkpoint 合并、数据格式转换到预训练、SFT 及 COCO/MMMU 下游评测的全流程。读完本文你将能够独立复现仓库中的多模态训练管线并理解每一步背后的源码实现。需要提前说明的是多模态支持在 Megatron-LM 中仍处于积极开发阶段README 明确标注under active development and is expected to change本示例的目的并非产出 SOTA 模型质量而是演示 megatron 的多模态功能。以下流程已在A100 基座的 DGX 集群上验证在 64 块 GPUtp4四路张量并行上预训练约耗时 1 天指令微调约 11 小时训练速度随 GPU 数量近似线性扩展。一、环境搭建Docker 容器仓库提供了开箱即用的镜像定义 examples/multimodal/Dockerfile。在 Megatron-LM 仓库根目录执行docker build -t megatron-multimodal -f examples/multimodal/Dockerfile .构建完成后即可在该容器内完成本文后续全部步骤。注意容器本身只负责提供运行时环境依赖、TE 等模型权重与数据集仍需按下列小节自行下载与转换。二、模型权重准备多模态模型由语言模型 视觉模型两部分拼装而成二者需要分别下载并转换为 Megatron-Coremcore格式。2.1 语言模型Mistral-7B-Instruct-v0.3按文档 docs/llama_mistral.md 中 Mistral-7B 一节的说明从 HuggingFace 下载 Mistral-7B-Instruct-v0.3 权重并转换为mcore 格式、张量并行度为 4tensor parallel size 4的 checkpoint。转换后请继续使用 HuggingFace 的 tokenizer而非自行训练的词表因为后续训练脚本通过--tokenizer-model mistralai/Mistral-7B-Instruct-v0.3直接引用该 tokenizer 名称。2.2 视觉模型OpenAI CLIP ViT-L/14336px本示例使用 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-te从源码 examples/multimodal/model_converter/clip_converter.py 可以清晰看到转换器做了什么加载 OpenAI 原始 CLIP 权重clip.load(ViT-L/14336px, ...)仅保留visual前缀即文本塔被丢弃跳过visual.proj与ln_post等本模型不使用的位置将 PyTorch 多头注意力的参数布局重排为 Megatron 的linear_qkv布局通过kv_channels64、num_heads16、hidden_dim1024计算 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.weightmlp.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。该脚本同时支持 NVLMTP8与 Mistral-CLIPTP4两种模型类型默认走 TP4 分支。注意若加载 checkpoint 时遇到错误可尝试设置环境变量TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD1。该变量仅在处理可信 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.zip3.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__为图片 idjpg字段存图片字节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加载。该脚本的核心训练配置如下可在脚本中调整配置项预训练取值说明--tensor-model-parallel-size4张量并行度与权重转换、合并时的 TP 保持一致--pipeline-model-parallel-size1流水线并行度--num-layers / --hidden-size / --num-attention-heads32 / 4096 / 32Mistral-7B 主干结构--ffn-hidden-size14336FFN 隐层宽度--num-query-groups8GQA 查询组数--seq-length / --decoder-seq-length576 / 1024图像 patch 序列长度 / 文本解码序列长度--max-position-embeddings4096最大位置编码--micro-batch-size / --global-batch-size1 / 256微批次 1全局批次 256DEBUG 模式为 32--train-iters / --lr-decay-iters20000 / 20000训练步数与学习率衰减步数--lr / --min-lr / --lr-decay-style0.00015 / 1e-5 / cosine学习率策略--lr-warmup-fraction0.01warmup 比例--hidden-dropout / --attention-dropout0.1 / 0.0正则化DEBUG 模式 hidden-dropout0.0--save-interval / --eval-interval / --eval-iters1000 / 1000 / 10保存与评估节奏--tokenizer-type MultimodalTokenizer—多模态专用 tokenizer--tokenizer-model mistralai/Mistral-7B-Instruct-v0.3—复用 HuggingFace tokenizer--tokenizer-prompt-format mistral—提示词格式--prompt-pathmanual_prompts.json手工提示词模板--patch-dim / --img-h / --img-w14 / 336 / 336CLIP ViT-L/14336px 的 patch 与输入分辨率--freeze-LM / --freeze-ViT二者均开启预训练阶段冻结语言模型与视觉塔只训练投影层等--eod-mask-loss—按 EOD 掩码损失--dataloader-type external—使用外部energondataloader--ckpt-format torch—使用 torch 格式 checkpoint其他值得注意的点脚本使用--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 roperotary-base 1000000等 Megatron-Core 标准配置--pretrained-checkpoint ${CHECKPOINT_DIR}指向组合好的多模态 checkpoint配合--load/--save ${FINETUNE_DIR}完成训练过程断点保存--allow-missing-vision-projection-checkpoint允许视觉投影层权重缺失时继续加载环境变量NVTE_ALLOW_NONDETERMINISTIC_ALGO1允许非确定性注意力算法NVTE_APPLY_QK_LAYER_SCALING0关闭 QK 层缩放训练入口为 examples/multimodal/train.py它通过megatron.core.models.multimodal.llava_model.LLaVAModel构建模型并仅在流水线首末级 stage 上运行 dataloadertrain_valid_test_dataloaders_provider。一切正常时你会观察到与下图类似的训练/验证 loss 曲线该曲线在global batch size 256下获得更改该值曲线会相应变化。对 LLaVA 类模型而言loss 曲线并不能可靠预测下游任务表现因此必须通过后续的测试生成与多指标评测来判断模型质量文档表示未来版本会加入训练期零样本评测。断点续训可多次执行预训练脚本恢复时脚本会自动加载最新的 model、optimizer 与 dataloader 状态。四、指令微调SFTSFT 阶段让预训练模型学会遵循指令完成任务。准备一份指令微调数据集同样需要 megatron-energon 格式README 明确说明仓库不提供该数据集的制作指引需自行准备编辑 examples/multimodal/sft_dataset.yaml将train与val两个 split 的path分别指向微调数据集的训练与验证部分运行examples/multimodal/sft_mistral_clip.shSFT 脚本 examples/multimodal/sft_mistral_clip.sh 需要三个环境变量WORKSPACE、LOAD_NAME预训练输出目录名、LOAD_ITER要加载的预训练迭代号用于定位${WORKSPACE}/${LOAD_NAME}/checkpoints下的具体 iter 目录。与预训练相比SFT 配置的关键差异其余架构参数保持一致配置项SFT 取值与预训练对比--decoder-seq-length2048预训练为 1024SFT 文本序列更长--global-batch-size128预训练为 256--lr / --min-lr1e-6 / 1e-7学习率降低约两个数量级--save-interval / --eval-interval500 / 500更频繁保存与评估--clip-grad0.5预训练为 1.0--weight-decay0.1预训练为 1e-2--freeze-LM不设置仅冻结 ViT语言模型参与训练--allow-missing-vision-projection-checkpoint不设置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取值是评测基准名称例如captioningCOCO 标题生成或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下载 COCO 2014 测试图片集wget http://images.cocodataset.org/zips/test2014.zip下载 COCO 测试图片标注Karpathy 测试标注https://storage.googleapis.com/sfr-vision-language-research/datasets/coco_karpathy_test.json先用--task captioning运行文本生成计算 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 模型评估MMMUMMMU 官方仓库目前无法通过 pip 安装需在examples/multimodal目录下手动克隆cd examples/multimodal git clone https://github.com/MMMU-Benchmark/MMMU.gitMMMU 数据集由代码自动从 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_LOAD1重试仅限可信文件存在任意代码执行风险TP 一致性CLIP 转换--tensor-parallel-size 4、Mistral 转换、checkpoint 合并脚本TP4 分支与训练脚本--tensor-model-parallel-size 4必须保持张量并行度一致磁盘空间LLaVA-Pretrain 解压后约 79GB请提前规划存储数据路径pretrain_dataset.yaml与sft_dataset.yaml中两处path都必须正确指向 energon 格式数据目录环境变量预训练脚本需WORKSPACE与LOAD_NAMESFT 脚本额外需要LOAD_ITER复现性训练/推理启用了NVTE_ALLOW_NONDETERMINISTIC_ALGO1结果存在一定非确定性。结语本文完整复现了 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考