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

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-2Meta 检查点 Meta 推理代码与 Megatron转换后的 HF 检查点 Megatron 推理代码的基准成绩。数值为两者间的百分比误差计算公式为|llama_score - megatron_score| / llama_score。全部测试每个模型规模共 80 项平均误差为 0.15%主要源于实现层面的微小算术差异包括Megatron 在自注意力与 SwiGLU 等若干位置执行批量矩阵乘法而 Llama 分开执行Megatron 在自注意力中使用torch.baddbmmLlama 使用torch.matmulMegatron 的 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 500000RoPE 频率计算的基础数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 14336SwiGLU 激活 FFN 中间维度 14336对应 Llama-3 8B 的 4:1 隐藏比。--num-attention-heads 32、--attention-dropout 0.0、--hidden-dropout 0.0Llama-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-dataNullTokenizer或--data-pathHuggingFaceTokenizer切换。完整用法见 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 1000000Mistral 的 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 的实现产生逐位一致的数值结果主要差异来源包括非穷尽列表TransformerEngineTE在 RMSNorm 中使用模型的params_dtype而 Hugging Face 实现使用 fp32参考 TransformerEngine issue #1132。Hugging Facetransformers将自注意力中的 q、k、v 投影拆成三个独立的 GEMM而 Megatron Core 出于效率将它们合并为单个 GEMM导致微小数值差异。这些差异幅度通常远小于基准测试误差不影响下游任务的可用性。总结与实操建议| 模型 | 转换命令 | tokenizer 类型 | 关键差异参数 | | -- | -- | -- | -- | | Llama-2 |convert_checkpoints.py import --hf-model meta-llama/Llama-2-7B ...|Llama2Tokenizer|--max-position-embeddings 4096、--rotary-base默认 10000 | | Llama-3.0 |... --hf-model meta-llama/Llama-3.2-1B ...|HuggingFaceTokenizer|--rotary-base 500000、--group-query-attention 8、--ffn-hidden-size 14336| | Llama-3.1 | 同上 |HuggingFaceTokenizer| 额外--use-rope-scaling、--max-position-embeddings 131072| | Mistral-7B v0.3 |... --hf-model mistralai/Mistral-7B-Instruct-v0.3 ...|HuggingFaceTokenizer|--rotary-base 1000000、--apply-layernorm-1p|实操要点训练/微调统一使用--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),仅供参考