Megatron-LM 加载 Llama / Mistral 权重：从 Hugging Face 转换到推理与微调完整指南-开发者社区

Megatron-LM 加载 Llama / Mistral 权重：从 Hugging Face 转换到推理与微调完整指南

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

本指南基于 docs/llama_mistral.md，系统讲解 Megatron-LM 如何加载 Llama-2、Llama-3.x 与 Mistral-7B 的 Hugging Face 权重并用于推理和微调。你将从零掌握三步流程：申请并下载 HF 权重 → 用 Megatron-Bridge 转换为 Megatron 格式 → 为每种模型配置正确的启动参数，同时理解架构差异、数值差异来源及同源 Llama 系模型的扩展方法。

背景与适用模型

Llama-2、Llama-3.x 与 Mistral-7B 是开放权重的大语言模型家族，发布时在多种基准上取得领先的开放模型成绩，并与当时的头部闭源模型竞争。三者架构高度相似（RMSNorm、SwiGLU、RoPE、GQA 等），因此 Megatron-LM 可以用同一套 GPT 模型框架加载它们的权重用于推理与微调。

当前仓库的转换流程只简化支持从 Hugging Face 下载的 llama-3.x 和 mistral 检查点（详见 models/index.md 的 Megatron Bridge 章节），其他模型的转换请查阅该文档。

总体流程

无论哪个模型，加载权重的步骤都相同：

申请访问权限并下载检查点（权重 + 分词器）。
用 Megatron-Bridge 的 checkpoint converter 将 Hugging Face 格式转换为 Megatron 格式。
（可选）校验转换后的检查点。
配置模型启动参数。

官方强烈建议训练或微调时使用--dtype bf16，推理则可用 bfloat16 或 float16。

Llama-2 加载

下载 Hugging Face 检查点

用户需要先向 Hugging Face 申请 Llama-2 权重下载权限（参考 transformers 的 Llama2 文档）。HF 格式的权重可以转换为 Megatron 格式，转换方法见下文。

转换为 Megatron 格式

使用 Megatron-Bridge 的转换脚本convert_checkpoints.py，执行 import 子命令：

python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \ --hf-model meta-llama/Llama-2-7B \ --megatron-path ./checkpoints/llama2_7b \ --torch-dtype bfloat16 \ --device-map auto

转换完成后即可将权重加载进 Megatron GPT 模型。

启动参数

无论是推理还是微调，都使用以下参数（注意--tokenizer-model ${TOKENIZER_MODEL}需替换为下载的tokenizer.model路径）：

--tensor-model-parallel-size ${TP} \ --pipeline-model-parallel-size 1 \ --seq-length 4096 \ --max-position-embeddings 4096 \ --tokenizer-type Llama2Tokenizer \ --tokenizer-model ${TOKENIZER_MODEL} \ --load ${CHECKPOINT_DIR} \ --exit-on-missing-checkpoint \ --use-checkpoint-args \ --no-load-optim \ --no-load-rng \ --untie-embeddings-and-output-weights \ --use-rotary-position-embeddings \ --normalization RMSNorm \ --no-position-embedding \ --no-masked-softmax-fusion \ --attention-softmax-in-fp32

这些参数的含义与底层实现对应如下（参考 megatron/training/arguments.py）：

--tokenizer-type Llama2Tokenizer：使用 SentencePiece 系列分词器，Megatron 的 tokenizer 构建逻辑将Llama2Tokenizer归入SP_TOKENIZERS，并通过args.tokenizer_model加载分词模型（见 build_tokenizer.py）。
--untie-embeddings-and-output-weights：词嵌入与输出层权重解耦（Llama 系列模型的输出层使用独立的权重矩阵）。
--use-rotary-position-embeddings与--no-position-embedding：该写法是旧参数风格，当前版本中这两个参数已被标记为 deprecated（--no-position-embedding提示改用--position-embedding-type），新代码建议直接使用--position-embedding-type rope。
--normalization RMSNorm：归一化层使用 RMSNorm。
--no-masked-softmax-fusion：关闭 masked softmax 融合，避免融合 kernel 引入数值差异。
--attention-softmax-in-fp32：注意力 softmax 在 fp32 精度下计算。

基准结果对比

下表对比了原生 Llama-2（Meta 检查点 + Meta 推理代码）与 Megatron（转换后的 HF 检查点 + Megatron 推理代码）的基准成绩。数值为两者间的百分比误差，计算公式为|<llama_score> - <megatron_score>| / <llama_score>。全部测试（每个模型规模共 80 项）平均误差为 0.15%，主要源于实现层面的微小算术差异，包括：

Megatron 在自注意力与 SwiGLU 等若干位置执行批量矩阵乘法，而 Llama 分开执行；
Megatron 在自注意力中使用torch.baddbmm，Llama 使用torch.matmul；
Megatron 的 RoPE 采用sin/cos实现，Llama 采用polar/complex实现；
Llama 初始化时调用torch.set_default_dtype(torch.float16)，Megatron 不会。

Big Bench（得分类型：多选题正确率）

| bigbench / standard | 7b | 13b | 70b | | -- | -- | -- | -- | | date_understanding | 0.29% | 0.13% | 0.12% | | general_knowledge | 0.00% | 0.00% | 0.00% | | human_organs_senses | 0.00% | 0.00% | 0.00% | | intent_recognition | 0.00% | 0.11% | 0.00% | | riddle_sense | 0.00% | 0.00% | 0.00% | | similarities_abstraction | 0.00% | 0.58% | 0.00% | | simple_arithmetic_json_multiple_choice | 0.00% | 0.00% | 0.00% | | undo_permutation | 0.19% | 0.19% | 0.18% |

Multilingual（得分类型：多选题正确率）

| multilingual / xcopa | 7b | 13b | 70b | | -- | -- | -- | -- | | en-template-mGPT-remove-punctuation | 0.08% | 0.00% | 0.00% | | et-template-mGPT-remove-punctuation | 0.00% | 0.13% | 0.25% | | ht-template-mGPT-remove-punctuation | 0.26% | 0.13% | 0.26% | | id-template-mGPT-remove-punctuation | 0.11% | 0.00% | 0.19% | | it-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% | | qu-template-mGPT-remove-punctuation | 0.00% | 0.00% | 0.27% | | sw-template-mGPT-remove-punctuation | 0.14% | 0.13% | 0.13% | | th-template-mGPT-remove-punctuation | 0.25% | 0.13% | 0.13% | | tr-template-mGPT-remove-punctuation | 0.26% | 0.00% | 0.34% | | vi-template-mGPT-remove-punctuation | 0.00% | 0.11% | 0.00% | | zh-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |

LM Evaluation Harness（得分类型：多选题正确率）

| lm-eval | 7b | 13b | 70b | | -- | -- | -- | -- | | boolq | 0.04% | 0.04% | 0.07% | | hellaswag | 0.02% | 0.03% | 0.03% | | piqa | 0.00% | 0.00% | 0.07% | | winogrande | 0.00% | 0.11% | 0.20% |

MMLU（得分类型：多选题正确率）

注：括号内为该超类别的子任务数量。

| mmlu | 7b | 13b | 70b | | -- | -- | -- | -- | | stem [18] | 0.79% | 0.05% | 0.01% | | humanities [13] | 0.19% | 0.01% | 0.02% | | other (business, health, misc.) [14] | 0.08% | 0.06% | 0.12% | | social sciences [12] | 0.37% | 0.21% | 0.01% |

可以看到绝大多数误差在 0.5% 以下，证明转换后的权重在数值上高度接近原生实现。

Llama-3.x 加载

Llama-3.x 检查点同样可加载进 Megatron 用于推理与微调，步骤与 Llama-2 一致（含可选的转换后校验）。

下载 Hugging Face 检查点

先向 Hugging Face 的 meta-llama 组织申请 Llama-3.x 权重下载权限。

转换为 Megatron 格式

python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \ --hf-model meta-llama/Llama-3.2-1B \ --megatron-path ./checkpoints/llama3_2_1b \ --torch-dtype bfloat16 \ --device-map auto

启动参数

Llama 3.0 参数

--tensor-model-parallel-size ${TP} \ --pipeline-model-parallel-size 1 \ --seq-length 8192 \ --max-position-embeddings 8192 \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model ${TOKENIZER_MODEL} \ --load ${CHECKPOINT_DIR} \ --exit-on-missing-checkpoint \ --use-checkpoint-args \ --no-load-optim \ --no-load-rng \ --untie-embeddings-and-output-weights \ --normalization RMSNorm \ --position-embedding-type rope \ --no-masked-softmax-fusion \ --attention-softmax-in-fp32 \ --disable-bias-linear \ --transformer-impl transformer_engine \ --group-query-attention 8 \ --attention-dropout 0.0 \ --hidden-dropout 0.0 \ --rotary-base 500000 \ --rotary-percent 1.0 \ --ffn-hidden-size 14336 \ --num-attention-heads 32 \ --swiglu \ --bf16 \

Llama 3.1 参数

Llama 3.1 支持更长的上下文窗口（131072），且启用了 RoPE 缩放，因此与 3.0 的参数相比，--max-position-embeddings变为 131072，并新增--use-rope-scaling：

--tensor-model-parallel-size ${TP} \ --pipeline-model-parallel-size 1 \ --seq-length 8192 \ --max-position-embeddings 131072 \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model ${TOKENIZER_MODEL} \ --load ${CHECKPOINT_DIR} \ --exit-on-missing-checkpoint \ --use-checkpoint-args \ --no-load-optim \ --no-load-rng \ --untie-embeddings-and-output-weights \ --normalization RMSNorm \ --position-embedding-type rope \ --no-masked-softmax-fusion \ --attention-softmax-in-fp32 \ --disable-bias-linear \ --transformer-impl transformer_engine \ --group-query-attention 8 \ --attention-dropout 0.0 \ --hidden-dropout 0.0 \ --rotary-base 500000 \ --rotary-percent 1.0 \ --use-rope-scaling \ --ffn-hidden-size 14336 \ --num-attention-heads 32 \ --swiglu \ --bf16 \

关键参数说明（源码依据 megatron/training/arguments.py）：

--tokenizer-type HuggingFaceTokenizer：直接加载 HF 的 tokenizer，通过--tokenizer-model指定分词模型文件路径。
--position-embedding-type rope：新版参数写法，等价于旧的--use-rotary-position-embeddings；--position-embedding-type的合法取值包括learned_absolute、rope、yarn、mrope、relative、none。
--rotary-base 500000：RoPE 频率计算的基础数，Llama-3 系列使用 500000（默认值为 10000）；--rotary-percent 1.0表示全部 head 维度参与旋转。
--use-rope-scaling：启用 Llama-3.x 的 RoPE 缩放，配合--rope-scaling-factor（默认 8.0）实现长上下文外推。
--group-query-attention 8：启用 GQA，每组查询头共享 8 个 KV 头（即--num-query-groups为 8）。
--disable-bias-linear：线性层不使用 bias，与 Llama 架构一致。
--transformer-impl transformer_engine：使用 TransformerEngine 实现，配合--bf16以获得最佳性能。
--swiglu与--ffn-hidden-size 14336：SwiGLU 激活 + FFN 中间维度 14336（对应 Llama-3 8B 的 4:1 隐藏比）。
--num-attention-heads 32、--attention-dropout 0.0、--hidden-dropout 0.0：Llama-3 训练时无 dropout。

仓库内配套训练脚本

如需在仓库内从零预训练或继续训练 Llama-3 架构模型，可参考 examples/llama/train_llama3_8b_h100_fp8.sh，其模型参数与上面的一致（32 层、hidden 4096、FFN 14336、32 头、8 query groups、RoPE、RMSNorm、SwiGLU、--untie-embeddings-and-output-weights、--disable-bias-linear），并额外演示了 FP8 混合精度（--fp8-format hybrid）、序列并行、分布式优化器（--use-distributed-optimizer）等训练增强；数据与分词器则通过--mock-data+NullTokenizer或--data-path+HuggingFaceTokenizer切换。完整用法见 examples/llama/README.md。

Mistral-7B 加载

Megatron 当前支持加载 Mistral-7B v0.3 版本（该版本不使用滑动窗口注意力，并提供更大的 32768 词表）用于推理与微调。

下载 Hugging Face 检查点

通过 Hugging Face 申请下载权限，两个变体均可：基础模型 Mistral-7B-v0.3 与指令模型 Mistral-7B-Instruct-v0.3。

转换为 Megatron 格式

python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \ --hf-model mistralai/Mistral-7B-Instruct-v0.3 \ --megatron-path ./checkpoints/mistral_7b \ --torch-dtype bfloat16 \ --device-map auto

启动参数

--tensor-model-parallel-size ${TP} \ --pipeline-model-parallel-size 1 \ --seq-length 4096 \ --max-position-embeddings 4096 \ --tokenizer-type HuggingFaceTokenizer \ --tokenizer-model ${TOKENIZER_MODEL} \ --load ${CHECKPOINT_DIR} \ --exit-on-missing-checkpoint \ --use-checkpoint-args \ --no-load-optim \ --no-load-rng \ --untie-embeddings-and-output-weights \ --normalization RMSNorm \ --position-embedding-type rope \ --no-masked-softmax-fusion \ --attention-softmax-in-fp32 \ --apply-layernorm-1p \ --transformer-impl transformer_engine \ --group-query-attention 8 \ --disable-bia-linear \ --rotary-base 1000000 \ --rotary-percent 1.0 \ --swiglu \ --ffn-hidden-size 14336 \ --num-attention-heads 32 \

注意：原文档中 Mistral 参数段存在两处笔误——--disable-bia-linear（应为--disable-bias-linear）与--num-attention-heads 32后缺少续行符\，实际使用时应按 Llama-3 段落的写法修正。Mistral-7B 与 Llama-3 架构相近，因此参数高度重合，主要区别在于--rotary-base 1000000（Mistral 的 RoPE base）与--apply-layernorm-1p（归一化层权重乘以 1p 缩放系数）。

其他 Llama 系模型支持（实验性）

许多使用 Llama 架构的模型（如 Yi-34B、Qwen2.x）都可以复用 Llama-3.x 章节的转换命令完成 HF → Megatron 转换。此路径为实验特性，建议核对目标模型的词表、层数、hidden size、RoPE base 等超参数是否与转换参数匹配。

已知数值差异

不期望 Megatron 与 Hugging Face 对 llama3.x 和 mistral 的实现产生逐位一致的数值结果，主要差异来源包括（非穷尽列表）：

TransformerEngine（TE）在 RMSNorm 中使用模型的params_dtype，而 Hugging Face 实现使用 fp32（参考 TransformerEngine issue #1132）。
Hugging Facetransformers将自注意力中的 q、k、v 投影拆成三个独立的 GEMM，而 Megatron Core 出于效率将它们合并为单个 GEMM，导致微小数值差异。

这些差异幅度通常远小于基准测试误差，不影响下游任务的可用性。

总结与实操建议

实操要点：

训练/微调统一使用--dtype bf16，推理可按需选择 bfloat16 或 float16。
--exit-on-missing-checkpoint与--use-checkpoint-args确保只加载缺失权重时报错、并以检查点内记录的参数为准，避免参数不匹配导致的静默错误。
转换工具由 Megatron-Bridge 仓库提供（Megatron-Bridge/examples/conversion/convert_checkpoints.py），命令中的--torch-dtype bfloat16与--device-map auto建议保持一致。
若需要继续训练而非只做推理，可在上述参数基础上补充训练相关参数（数据路径、batch size、学习率、优化器等），仓库内的 examples/llama/train_llama3_8b_h100_fp8.sh 与 examples/mixtral/train_mixtral_8x7b_distributed.sh 提供了完整的可运行示例。
加载成功与否受 tokenizer 路径（--tokenizer-model）、检查点目录（--load）与并行切分参数（--tensor-model-parallel-size等）影响，出现错误时优先检查这三处。

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

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

Megatron-LM 加载 Llama / Mistral 权重：从 Hugging Face 转换到推理与微调完整指南