企业官网建设流程全解析

Megatron-LM 实战指南：Llama-3.1-Nemotron-Nano-VL-8B-V1 多模态模型的权重转换、预训练、SFT 与推理全流程

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

本指南以 Megatron-LM 仓库中examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/目录下官方示例（README.md）为骨架，完整讲解基于 Llama-3.1-Nemotron-Nano-VL-8B-V1 架构的多模态视觉语言模型（VLM）在 Megatron-Core 中的落地流程。读完本文，你将掌握：从 Hugging Face 权重出发完成语言模型与 RADIO 视觉模型的格式转换与合并、使用 Megatron Energon 准备多模态数据、运行预训练与指令微调（SFT），以及通过统一推理脚本执行生成与 MMMU/TextVQA 等基准评测。

模型与检查点概览

Llama-3.1-Nemotron-Nano-VL-8B-V1 是一款以 Llama-3.1-8B-Instruct 为语言主干、RADIO 系列视觉编码器为视觉主干的多模态模型，属于典型的 LLaVA 式架构（视觉塔 + 投影层 + 语言模型）。

官方提供两种格式的检查点：

如果希望完全从零走一遍流程，可以跳过直接下载 mcore 检查点，而按照后文「模型转换」与「训练」章节，从 HF 权重自行构建模型并完成预训练与 SFT。

环境准备：Docker 镜像与依赖

多模态训练涉及 RADIO 视觉塔、Megatron Energon 数据加载等较多依赖，官方在 examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/Dockerfile 中给出了完整的容器构建方案：

• 基础镜像为nvcr.io/nvidia/pytorch:25.04-py3（NGCT 官方 PyTorch 容器，内含 NVIDIA 驱动栈与 PyTorch 生态）；
• 安装系统工具（git、curl、wget、htop、tmux 等）与通用 Python 依赖（einops、webdataset、braceexpand、transformers、datasets、accelerate、timm、tiktoken、sentencepiece 等）；
• 编译安装causal-conv1d与mamba（Mamba/SSM 相关仓库的卷积依赖）；
• 安装mmf（--no-deps避免依赖冲突）、open_clip_torch、open-flamingo[eval]、tensorstore==0.1.45；
• 安装git+https://github.com/NVIDIA/Megatron-Energon.git#egg=megatron-energon[av_decode]，这是后续多模态数据加载的核心组件。

构建命令需在 Megatron-LM 仓库根目录执行：

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

SLURM 批量模式训练时（脚本中的srun分支），需将构建好的容器镜像路径填入训练脚本的--container-image与--container-mounts参数。

数据准备：Megatron Energon 多模态数据集

本示例的多模态数据加载完全依赖 Megatron Energon（仓库内通过pip install git+https://github.com/NVIDIA/Megatron-Energon.git#egg=megatron-energon[av_decode]引入），对应训练脚本中的--dataloader-type external开关。

数据准备的基本思路是把图像与文本问答对组织成 Energon 可识别的格式：

1. 将图文样本打包为 WebDataset（wds）格式，通常一个样本包含jpg图像字段与json元数据字段（context、answers 等）；
2. 使用 Energon 的energon prepare ./交互式命令生成dataset.yaml（选择训练/验证/测试划分比例、VQASample样本类型、字段映射等）；
3. 将dataset.yaml中的path指向实际的 wds 目录。

训练脚本默认引用examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/pretrain_blend.yaml（预训练）与sft_blend.yaml（SFT）作为--data-path传入，这两个数据清单文件需要根据你自己的数据集与划分按 Energon 格式准备（例如参考 examples/multimodal/README.md 中 LLaVA-Pretrain 数据集从下载、解压到convert_llava_pretrain_to_wds.py转换、再energon prepare的完整示例）。预训练阶段官方使用--split 100,0,0（全部数据用于训练），SFT 阶段同样为100,0,0。

从 Hugging Face 权重构建 Megatron-Core 模型

该模型由「语言模型 + 视觉模型」两部分拼接而成，转换也分三步进行。

语言模型转换

以 Hugging Face 上的meta-llama/Llama-3.1-8B-Instruct为起点，下载后使用仓库统一的转换工具 tools/checkpoint/convert.py 转成 Megatron-Core 格式：

export LLAMA_DOWNLOAD_DIR=<downloaded hf model directory> CUDA_DEVICE_MAX_CONNECTIONS=1 python tools/checkpoint/convert.py --bf16 --model-type GPT --loader llama_mistral --saver core \ --target-tensor-parallel-size 4 --checkpoint-type hf \ --load-dir $LLAMA_DOWNLOAD_DIR --save-dir llama3p1 --tokenizer-model $LLAMA_DOWNLOAD_DIR \ --saver-transformer-impl transformer_engine --model-size llama3

关键参数说明：

转换得到的llama3p1目录即 mcore 格式的语言模型检查点。

视觉模型转换：RADIO

视觉主干使用 NVIDIA RADIO 系列编码器，通过 examples/multimodal/model_converter/radio_converter.py 将其转换为 mcore 兼容格式：

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

从 radio_converter.py 源码可以看到转换的内部逻辑：

• 通过torch.hub.load('NVlabs/RADIO', 'radio_model', version=...)加载 RADIO 官方权重，version默认为radio_v2.5-h；
• 按kv_channels=80、hidden_dim=1280、num_heads=16等超参把 PyTorch 多头注意力权重重排为 Megatron 的 QKV 拼接布局，并按目标 TP 大小对embedder.weight、linear_qkv等张量做维度切分（chunk_dim）生成tensor_parallel_size份分片权重；
• 名称映射覆盖patch_generator（embedder / class_token / position_embeddings）、blocks.*（decoder.layers.*.self_attention 等）等 RADIO 特有结构；
• 权重由 fp16 转 fp32 后保存，供训练脚本以--bf16加载。

--model-type radio_v2.5-h与--version c-radio_v2-vlm-h指定了具体的 RADIO 变体，训练时对应--vision-model-type radio与 megatron/core/models/multimodal/llava_model.py 中RADIOViTModel的实现。

合并语言与视觉检查点

使用官方脚本 examples/multimodal/combine_lm_vision_checkpoints.sh 将两部分检查点合并为单个多模态检查点目录：

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

该脚本底层调用 examples/multimodal/combine_state_dicts.py：按mp_rank_XX/model_optim_rng.pt逐 rank 读取语言模型与视觉模型权重，以language_model、vision_model前缀交替拼接成新的model_optim_rng.pt写入输出目录，并在末尾生成latest_checkpointed_iteration.txt（内容为1）供训练脚本识别。脚本内置了 TP=4（Mistral-CLIP 示例）与 TP=8（NVLM 示例）两套 rank 展开模式，本模型使用 TP=4 分支。

提示：若合并或加载时遇到权重加载错误，可尝试设置TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1（仅在信任检查点文件来源时使用）。

预训练（Pretraining）

官方预训练脚本为 examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/pretraining_llama_3p1_nemotron_nano_vl_8b_v1.sh，要求从 Megatron-LM 仓库根目录启动。

运行模式：交互与 SLURM 批处理自动切换

脚本通过which srun的返回码自动判断环境：存在srun则进入批处理模式（MODEL_NAME固定为pretraining_llama_3p1_nemotron_nano_vl_8b_v1，特殊 token 需要转义，8 GPU 并行，全局 batch 1024）；否则进入交互模式（MODEL_NAME带时间戳，特殊 token 不加转义，4 GPU，全局 batch 1，用于快速验证流程）。交互模式等价于：

torchrun --nproc_per_node 4 examples/multimodal/train.py <OPTIONS>

批处理模式下通过srun --container-image <镜像> --container-mounts <挂载> --output=${LOGS_DIR}/%x_%j_$DATETIME.log在容器内执行python -u examples/multimodal/train.py。注意特殊 token 在批处理模式下必须写成\<image\>这类转义形式，以避免被 shell 解释为重定向。

关键训练参数

脚本OPTIONS变量中完整定义了模型与训练配置，核心参数如下：

多模态专项参数

--tokenizer-type MultimodalTokenizer \ --tokenizer-model meta-llama/Llama-3.1-8B-Instruct \ --tokenizer-prompt-format llama3p1 \ --vision-model-type radio \ --patch-dim 16 \ --img-h 512 --img-w 512 \ --pixel-shuffle --use-tiling --max-num-tiles 12 --use-thumbnail \ --image-tag-type internvl \ --disable-vision-class-token \ --force-system-message \ --use-area-weighted-aspect-ratio \ --freeze-ViT --freeze-LM \ --special-tokens <image> <img> </img> <quad> </quad> <ref> </ref> <box> </box> \ --allow-missing-vision-projection-checkpoint \ --eod-mask-loss --use-loss-scaling \ --prompt-path examples/multimodal/manual_prompts.json \ --inference-max-seq-length 32768

这些参数的语义如下：

• tokenizer 相关：MultimodalTokenizer负责在文本序列中插入图像 token 占位符；--tokenizer-prompt-format llama3p1指定 Llama-3.1 风格的系统提示词组装；--force-system-message强制注入系统消息；--special-tokens注册图像与多模态标记（图像、引用框<ref></ref>、边界框<box></box>、区域<quad></quad>等）；
• 视觉输入：--patch-dim 16（patch 尺寸）、--img-h/--img-w 512（输入图像分辨率）、--vision-model-type radio（RADIO 视觉塔）、--disable-vision-class-token（丢弃视觉 class token）；
• 动态分辨率（tiling）：--use-tiling --max-num-tiles 12 --use-thumbnail --pixel-shuffle将高分辨率输入切分为最多 12 个 tile 并附加缩略图，--use-area-weighted-aspect-ratio依据面积加权选择宽高比。开启 tiling 后--seq-length被压到 256（图像 token 由 tiling 机制产生，语言侧仅需少量序列长度），而--decoder-seq-length保持 4096 容纳展开后的视觉 token；
• 冻结策略：预训练阶段--freeze-ViT、--freeze-LM冻结视觉塔与语言模型，仅训练视觉投影层与相关连接模块——这正是多模态「对齐」预训练的典型做法（参见下文源码分析）；
• 加载容错：--allow-missing-vision-projection-checkpoint允许投影层权重缺失（如从仅含 LM+Vision 主干权重的合并检查点启动时，投影层使用随机初始化）；
• 损失与日志：--eod-mask-loss屏蔽结束符后的损失，--use-loss-scaling对多模态损失做缩放，--log-params-norm、--log-num-zeros-in-grad输出参数范数与梯度零值统计，--tensorboard-dir输出 TensorBoard 日志。

环境变量

脚本额外设置两个 TE 相关环境变量：

export NVTE_APPLY_QK_LAYER_SCALING=0 export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1

前者关闭 QK 层缩放（与--attention-softmax-in-fp32配合保证 fp32 softmax 数值行为），后者允许 Transformer Engine 使用非确定性算法以获得更高性能；如需严格可复现训练，可将NONDETERMINISTIC_ATTN置 0。

指令微调（SFT）

SFT 脚本为 examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/sft_llama_3p1_nemotron_nano_vl_8b_v1.sh，入口同样是examples/multimodal/train.py，数据来自sft_blend.yaml。其与预训练的主要差异：

SFT 阶段移除了冻结参数，对语言模型、视觉塔、投影层进行全量微调；为缓解 16K 解码序列与高分辨率视觉 token 带来的显存压力，启用了完整的逐层重计算（recompute 全部 32 层且包含视觉部分）。SFT 以预训练检查点为起点，通过--pretrained-checkpoint指定（脚本中默认路径为${WORKSPACE}/output/${LOAD_NAME}/checkpoints/pretraining_llama_3p1_nemotron_nano_vl_8b_v1），并用--load/--dataloader-save支持中途断点续训。

推理与评测

统一推理脚本为 examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/text_generation.sh，核心命令：

export LLAMA_NEMOTRON_NANO_VL_PATH=<path to the megatron tp=4 checkpoint> examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/text_generation.sh --model-path $LLAMA_NEMOTRON_NANO_VL_PATH \ --task inference --output-path inference-example --tensor-model-parallel-size 4

脚本支持的命令行参数：

脚本内部以torchrun --nproc_per_node ${TP}启动 examples/multimodal/run_text_generation.py，传入与训练一致的模型结构参数（--language-model-type=llama3.1_8b、--vision-model-type radio、RoPE、GQA 等，并依赖--use-checkpoint-args从检查点元数据恢复超参），解码配置为--temperature 1.0 --top_k 1（贪心解码）、--seed 153，同时--no-load-rng --no-load-optim跳过优化器/RNG 状态加载，--exit-on-missing-checkpoint在检查点缺失时立即报错。注意脚本中--tokenizer-model示例路径为内部 Lustre 路径，使用前必须替换为本地 Llama-3.1-8B-Instruct tokenizer 目录。

评测时只需切换--task为基准名：如MMMU或TextVQA，并配合--gt-path提供 groundtruth，生成结果后由 examples/multimodal/evaluation/evaluate_mmmu.py、examples/multimodal/evaluation/evaluate_textvqa.py 等评测脚本计算得分（仓库examples/multimodal/evaluation/下还提供 AI2D、ChartQA、COCO、MathVista 等十余个基准的评测实现，可作为扩展参考）。

源码级原理：LLaVAModel 与冻结/动态分辨率机制

理解上述参数在底层如何生效，可以阅读多模态模型核心实现 megatron/core/models/multimodal/llava_model.py（LLaVAModel类，MegatronModule 子类）：

• 三模块结构：模型由language_model（Llama-3.1 主干）、vision_model（RADIOViTModel，通过--vision-model-type选择radio等变体）、vision_projection（视觉-语言投影层）组成；pixel_shuffle开启时投影输入尺寸放大 4 倍（对应--pixel-shuffle），max_num_tiles控制动态分辨率 tile 数；
• 冻结实现：源码第 557-579 行附近的freeze_language_model、freeze_vision_model、freeze_vision_projection三个开关分别对对应子模块执行requires_grad_(False)——训练脚本中的--freeze-ViT、--freeze-LM正是预训练阶段只冻结视觉塔与语言模型、保留投影层可训的实现入口；
• token 处理：IMAGE_TOKEN = "<image>"等常量与image-tag-type（如internvl）共同决定图像 token 在序列中的组织方式；drop_vision_class_token对应--disable-vision-class-token，在送入语言模型前丢弃视觉 class token。

此外，tools/checkpoint/convert.py中--saver-transformer-impl transformer_engine与脚本--transformer-impl transformer_engine --use-te一一对应，保证转换、训练、推理三阶段使用同一 Transformer 实现，避免权重布局不兼容。

注意事项与排错提示

1. 并行度一致性：语言模型转换（--target-tensor-parallel-size 4）、RADIO 转换（--tensor-parallel-size 4）、合并脚本、训练脚本TP=4以及推理脚本--tensor-model-parallel-size 4必须保持一致，否则权重分片无法对齐加载；
2. tokenizer 路径：训练脚本使用--tokenizer-model meta-llama/Llama-3.1-8B-Instruct（HF 模型名，需本地已缓存），推理脚本则是显式目录路径，务必替换为实际下载目录；
3. 数据清单文件：pretrain_blend.yaml/sft_blend.yaml属于按 Energon 规范准备的数据描述文件，需结合你的数据集自行生成，参考 examples/multimodal/README.md 中 LLaVA-Pretrain 的准备流程（下载 → 解压 →convert_llava_pretrain_to_wds.py转 wds →energon prepare生成 dataset.yaml）；
4. 批处理模式转义：SLURM 批处理模式下特殊 token 必须转义为\<image\>形式，交互模式下则无需转义；
5. 显存与重计算：SFT 长解码序列（16K）必须配合--recompute-granularity full --recompute-vision，预训练阶段显存压力主要来自视觉 token，可优先调整--max-num-tiles；
6. 续训：预训练与 SFT 脚本均可重复执行续训，自动加载最新的模型、优化器与 dataloader 状态（--load+--dataloader-save）；
7. 权重加载失败：可尝试TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1（仅限可信检查点）。

至此，从 HF 权重转换、RADIO 视觉塔转换、检查点合并，到预训练、SFT 与推理评测的完整闭环即全部打通；如需深入多模态通用框架（Mistral-CLIP 示例、NVLM、InternViT 等），可继续阅读 examples/multimodal/README.md 与仓库中其余多模态示例。

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