/Users/seadawn/Documents/MyWork/Megatron-LM/docs/llama_mistral.md

Megatron-LM 对 Llama、Mistral 及其他类 Llama 模型的支持

注意：为简化代码，我们现在仅支持从 Hugging Face 下载的 Llama-3.x 和 Mistral 模型检查点的转换。

Llama-2 和 Llama-3.x 系列模型是一组开源的预训练及微调（用于聊天）模型，在各类基准测试中表现优异。在发布时，Llama-2 和 Llama-3 模型在开源模型中取得了最佳成绩，并可与领先的闭源模型相媲美（参见 https://arxiv.org/pdf/2307.09288.pdf 和 https://ai.meta.com/blog/meta-llama-3/）。

类似地，Mistral-7b 是一个开源模型，提供预训练和微调（用于聊天）版本，在基准测试中同样表现出色。

从架构上看，Llama-2、Llama-3 和 Mistral-7b 非常相似。因此，Megatron 支持加载这三者的检查点用于推理和微调。但每种模型的检查点转换和加载方式略有不同，下文将分别详述。

目录

Llama-2

Llama-2 检查点可加载到 Megatron 中用于推理和微调。加载这些检查点包括三个步骤：

  1. 获取下载检查点的权限。
  2. 将检查点从 Meta/Hugging Face 格式转换为 Megatron 格式。
  3. 设置启动模型所需的参数。

以下各节详细说明这些步骤。最后一节列出了以下两种情况的基准测试结果对比：1）使用 Meta 格式检查点和 Meta 推理代码的 Llama-2；2）使用转换后检查点和 Megatron 推理代码的 Megatron。

下载 Meta 或 Hugging Face 检查点

用户必须首先申请权限，从 Meta 或通过 Hugging Face（HF）下载 Llama-2 检查点。检查点有两种格式：Meta 原生格式（Meta 和 HF 均提供）和 HF 格式（仅 HF 提供）。任一格式均可转换为 Megatron，具体方法见下文。

转换检查点格式

我们建议在训练或微调时传入 --dtype bf16。推理可使用 bfloat16 或 float16。

Meta 格式

Meta 格式的检查点需先转换为 HF 格式作为中间步骤，再转换为 Megatron 格式。需安装 transformers 包，且版本 ≥ 4.31.0（例如 pip install transformers>=4.31.0）。（注意：我们已专门在 4.31.0 和 4.32.0 版本上测试过；使用更新版本时结果可能不同。）假设下载的检查点位于 $CHECKPOINT_DIR（包含 7B、13B、70B 等子目录），以下命令可用于将 Llama-2 格式转换为 bfloat16 的 HF 格式：

python tools/checkpoint/convert.py \
  --model-type GPT \
  --loader llama_mistral \
  --load-dir ${META_FORMAT_DIR} \
  --model-size ${MODEL_SIZE} \
  --checkpoint-type meta \
  --tokenizer-model ${TOKENIZER_MODEL} \
  --saver core \
  --save-dir ${MEGATRON_FORMAT_DIR} \
  --target-tensor-parallel-size ${TP} \
  --target-pipeline-parallel-size ${PP} \
  --bf16

--model-size 的有效值包括：llama2-7B、llama2-13B 和 llama2-70B（仅预训练模型），以及 llama2-7Bf、llama2-13Bf 和 llama2-70Bf（聊天微调模型）。

Hugging Face 格式

HF 检查点可通过 Megatron 自带的 HF 格式 Llama-2 检查点转换器转换为 Megatron 格式（参见脚本 tools/checkpoint/loader_llama_mistral.py）。一个关键参数是每个模型对应的张量并行大小（TP）。下表列出了这些值：

模型大小 张量并行大小 (TP)
7B 1
13B 2
70B 8

使用上表中的 TP 值，以及 Llama-2 分词器模型路径（随原始检查点自动下载；见下文 ${TOKENIZER_MODEL}），从 Megatron 源码根目录运行以下命令，将 HF 格式转换为 Megatron 格式：

python tools/checkpoint/convert.py \
  --model-type GPT \
  --loader llama_mistral \
  --load-dir ${HF_FORMAT_DIR} \
  --model-size ${MODEL_SIZE} \
  --checkpoint-type hf \
  --tokenizer-model ${TOKENIZER_MODEL} \
  --saver core \
  --save-dir ${MEGATRON_FORMAT_DIR} \
  --target-tensor-parallel-size ${TP} \
  --target-pipeline-parallel-size ${PP} \
  --bf16

转换完成后，即可将检查点加载到 Megatron GPT 模型中。

启动模型

启动 Megatron

若用于推理或微调，请使用以下参数：

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

注意：如果转换时使用了旧版模型格式（即 --saver legacy），请参见 使用旧版模型格式。

启动 Meta

Meta 检查点可通过以下方式启动：https://github.com/facebookresearch/llama

启动 Hugging Face

Hugging Face 检查点可通过以下方式启动：https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py

基准测试结果

下表列出了原生 Llama-2（使用 Meta 检查点和 Meta 推理代码）与 Megatron（使用转换后的 HF 检查点和 Megatron 推理代码）之间的基准测试对比。

数值为 Megatron 与 Llama-2 之间的百分比误差，计算公式为：|<llama_score> - <megatron_score>| / <llama_score>，每张表前注明了分数类型。在所有测试中（每个模型大小共 80 项），平均误差为 0.15%。两者基准分数的微小差异源于实现中的细微算术差异，影响因素包括：

  • Megatron 在若干位置（如自注意力和 SwiGLU 中）执行批量矩阵乘法，而 Llama 是分开执行的。
  • Megatron 在自注意力中使用 torch.baddbmm，而 Llama 使用 torch.matmul。
  • Megatron 使用 sin/cos 实现旋转位置编码（RoPE），而 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 / 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%

Llama-3.x

Llama-3.x 检查点可加载到 Megatron 中用于推理和微调。加载过程包括以下步骤：

  1. 获取下载检查点（权重和分词器）的权限。
  2. 将检查点从 Hugging Face 格式转换为 Megatron 格式。
  3. （可选）验证转换后的检查点。
  4. 设置启动模型所需的参数。

以下各节详细说明这些步骤。

下载 Hugging Face 检查点

用户必须首先申请权限，从 Hugging Face 下载 Llama-3.x 检查点。

转换检查点格式

我们建议在训练或微调时传入 --dtype bf16。推理可使用 bfloat16 或 float16。

Hugging Face 格式

HF 检查点可通过 Megatron 自带的 Llama-3.x HF 格式转换器转换为 Megatron 格式（参见脚本 tools/checkpoint/loader_llama_mistral.py）。一个关键参数是每个模型对应的张量并行大小（TP）。下表列出了这些值：

模型大小 张量并行大小 (TP)
1B 1
3B 1
8B 1
70B 8

使用上表中的 TP 值，以及 Llama-3.x 分词器模型路径（随原始检查点自动下载；见下文 ${TOKENIZER_MODEL}），从 Megatron 源码根目录运行以下命令，将 HF 格式转换为 Megatron 格式：

python tools/checkpoint/convert.py \
  --bf16 \
  --model-type GPT \
  --loader llama_mistral \
  --saver core \
  --target-tensor-parallel-size ${TP} \
  --checkpoint-type hf \
  --load-dir ${HF_FORMAT_DIR} \
  --save-dir ${MEGATRON_FORMAT_DIR} \
  --tokenizer-model ${TOKENIZER_MODEL} \
  --model-size llama3

转换完成后，即可将检查点加载到 Megatron GPT 模型中。

（可选）验证检查点

可使用脚本 examples/inference/llama_mistral/run_text_generation_llama3.sh <PATH_TO_CONVERTED_CORE_CHECKPOINT> <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT> 启动 Llama3 的 Megatron-LM 文本生成服务器。对于 Llama3.1，请使用 examples/inference/llama_mistral/run_text_generation_llama3.1.sh。

服务器启动后，可通过以下命令查询：

curl 'http://<TEXT_GENERATION_SERVER_IP>:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8' -d '{"prompts":["<SOME_PROMPT>"], "tokens_to_generate":100, "top_k":1}'

参考生成结果可通过 Hugging Face Transformers 库获取，运行：

python examples/llama_mistral/huggingface_reference.py --model_path <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT> --prompt <SOME_PROMPT>

启动模型

若用于推理或微调，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

Llama3.1 请使用以下参数：

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

注意：如果转换时使用了旧版模型格式（即 --saver legacy），请参见 使用旧版模型格式。

Mistral-7b

Megatron 当前支持加载 Mistral-7b 的 v0.3 版本（该版本不使用滑动窗口注意力，且词表更大，达 32768）用于推理和微调。加载过程包括以下步骤：

  1. 获取下载检查点（权重和分词器）的权限。
  2. 将检查点从 Hugging Face 格式转换为 Megatron 格式。
  3. （可选）验证转换后的检查点。
  4. 设置启动模型所需的参数。

以下各节详细说明这些步骤。

下载 Hugging Face 检查点

用户必须首先申请权限，通过 Hugging Face（HF）下载 Mistral-7b 检查点。

转换检查点格式

HF 检查点可通过 Megatron 自带的 Mistral HF 格式转换器转换为 Megatron 格式（参见脚本 tools/checkpoint/loader_llama_mistral.py）。

使用 Mistral 分词器模型路径（随 HF 检查点一同下载），从 Megatron 源码根目录运行以下命令，将 HF 格式转换为 Megatron core 格式：

python tools/checkpoint/convert.py \
  --bf16 \
  --model-type GPT \
  --loader llama_mistral \
  --saver core \
  --target-tensor-parallel-size ${TP} \
  --checkpoint-type hf \
  --load-dir ${HF_FORMAT_DIR} \
  --save-dir ${MEGATRON_FORMAT_DIR} \
  --tokenizer-model ${TOKENIZER_MODEL} \
  --model-size mistral

转换完成后，即可将检查点加载到 Megatron core GPT 模型中。

（可选）验证检查点

可使用脚本 examples/inference/llama_mistral/run_text_generation_mistral.sh <PATH_TO_CONVERTED_MCORE_CHECKPOINT> <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT> 启动 Mistral-7B 的 Megatron-LM 文本生成服务器。

服务器启动后，可通过以下命令查询：

curl 'http://<TEXT_GENERATION_SERVER_IP>:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8' -d '{"prompts":["<SOME_PROMPT>"], "tokens_to_generate":100, "top_k":1}'

参考生成结果可通过运行以下命令从 Hugging Face Transformers 库获取：

python examples/inference/llama_mistral/huggingface_reference.py --model_path <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT> --prompt <SOME_PROMPT>

启动模型

若用于推理或微调，请使用以下参数：

--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-bias-linear \
--rotary-base 1000000 \
--rotary-percent 1.0 \
--swiglu \
--ffn-hidden-size 14336 \
--num-attention-heads 32

注意：如果转换时使用了旧版模型格式（即 --saver legacy），请参见 使用旧版模型格式。

其他类 Llama 模型支持

注意：实验性功能

许多模型（如 Yi-34B 和 Qwen2.x）使用 Llama 架构，可使用 Llama-3.x 中的命令从 Hugging Face 转换为 Megatron。

已知的数值差异

Megatron 与 Hugging Face 对 Llama3.x 和 Mistral 模型的实现不会产生数值完全一致的结果。以下是部分已知的数值差异来源（非详尽列表）：

  1. TransformerEngine (TE) 在 RMSNorm 中使用模型的 params_dtype，而 Hugging Face 实现使用 fp32。详见：https://github.com/NVIDIA/TransformerEngine/issues/1132
  2. Hugging Face transformers 在自注意力中将 q、k、v 投影分别实现为独立的 GEMM，而 Megatron core 为提高效率将其合并为单个 GEMM。这会导致微小的数值差异。

使用旧版模型格式

本文档中所有检查点转换示例均使用 --saver core，表示将使用更新（且推荐）的 Megatron GPT 模型类，即：

  • 旧类：megatron.legacy.model.gpt_model.GPTModel
  • 新类：megatron.core.models.gpt.gpt_model.GPTModel

推荐使用新格式。但若因特殊需求必须使用旧类（即转换时使用 --saver legacy），则在启动训练或微调时需额外添加以下参数：

  • --use-legacy-models：使用旧版模型类
  • --ckpt-format torch：使用 torch 检查点格式（这是唯一与旧版模型格式兼容的格式）

将 Hugging Face（HF）格式的 Llama 或 Mistral 模型转换为 Megatron 格式时，两者在模型结构组织、权重命名、并行策略等方面存在显著差异，如果不正确处理，会导致加载失败、推理错误，甚至训练崩溃。以下是几个关键差异点及具体示例说明：

1. 权重命名与结构组织不同

Hugging Face：
  • 权重以 state_dict 形式存储，键名如： 
    model.layers.0.self_attn.q_proj.weight
model.layers.0.mlp.gate_proj.weight
lm_head.weight
Megatron（Core 格式）：
  • 权重按 张量并行（TP）和流水线并行（PP） 切分，存储在多个文件中（如 mp_rank_00/model_optim_rng.pt）。
  • 键名采用 Megatron 内部命名规范，例如： 
    transformer.layers.0.self_attention.linear_qkv.weight
transformer.layers.0.mlp.linear_fc1.weight
  • lm_head 与 embedding 默认不共享权重（需显式设置 --untie-embeddings-and-output-weights）。

✅ 问题示例：
若直接用 HF 的 lm_head.weight 覆盖 Megatron 的 output_layer.weight，但未设置 --untie-embeddings-and-output-weights，Megatron 会尝试共享 embedding 和输出层权重，导致维度不匹配（因为 HF 的 Llama 系列通常不共享）。

2. 张量并行（Tensor Parallelism, TP）切分方式不同

  • HF 模型是 单卡完整模型，无并行切分。
  • Megatron 要求在转换时指定 --target-tensor-parallel-size（如 Llama-70B 需 TP=8）。
  • 转换脚本会自动将 QKV、MLP 等权重按列或行切分到不同 TP rank。

✅ 问题示例：
若你用 TP=1 转换 Llama-70B（实际应为 TP=8），然后用 --tensor-model-parallel-size 8 启动 Megatron，会因权重未正确切分而报错：

RuntimeError: shape mismatch in layer ...

3. MLP 结构与激活函数差异

  • HF 的 Llama/Mistral 使用 SwiGLU，由 gate_proj + up_proj + down_proj 三部分组成。
  • Megatron Core 将其合并为两个线性层：
    • linear_fc1 = [gate_proj; up_proj]（拼接）
    • linear_fc2 = down_proj

✅ 问题示例：
若转换脚本未正确拼接 gate_proj 和 up_proj，Megatron 在前向传播时会因输入维度翻倍而崩溃。

4. 位置编码与 RoPE 参数差异

  • HF 的 RoPE 实现在 LlamaRotaryEmbedding 中，base=10000（Llama-2/3）或 1e6（Mistral）。
  • Megatron 需显式传入 --rotary-base 和 --rotary-percent。
  • Llama-3.1 还需 --use-rope-scaling 启用动态缩放。

✅ 问题示例：
若转换 Mistral 时未设置 --rotary-base 1000000，Megatron 会默认用 10000，导致位置编码错误，生成质量严重下降。

5. 分词器（Tokenizer）类型不同

  • HF 使用 AutoTokenizer，自动加载 tokenizer.model 或 tokenizer.json。
  • Megatron 需指定 --tokenizer-type：
    • Llama-2：Llama2Tokenizer
    • Llama-3/Mistral：HuggingFaceTokenizer

✅ 问题示例：
若对 Llama-3 使用 --tokenizer-type Llama2Tokenizer，会因特殊 token（如 <|begin_of_sentence|>）处理错误，导致输入解析失败。

6. 偏置（Bias）与 LayerNorm 实现

  • Llama/Mistral 官方模型 无偏置（bias=False）。
  • Megatron 默认可能启用 bias，需显式关闭： 
    --disable-bias-linear
--apply-layernorm-1p   # Mistral 需要，Llama 不需要

✅ 问题示例：
若未加 --disable-bias-linear，Megatron 会创建不存在的 bias 参数，导致加载失败。

总结：常见转换失败原因

问题类型 表现 解决方案
TP/PP 不匹配 加载时报 shape mismatch 转换时 --target-tensor-parallel-size 必须与启动时一致
RoPE 参数错误 生成乱码或重复 正确设置 --rotary-base、--use-rope-scaling
分词器类型错误 输入解析异常 Llama-3/Mistral 用 HuggingFaceTokenizer
未解绑 embedding/lm_head 维度错误 加 --untie-embeddings-and-output-weights
Bias 未关闭 多出 bias 参数 加 --disable-bias-linear

建议做法

始终使用官方提供的转换命令（如你提供的文档所示），并严格匹配：

  • --model-size
  • --target-tensor-parallel-size
  • --tokenizer-model
  • --bf16 / --fp16

并在启动时使用对应的 Megatron 参数（如 Llama-3 的 --ffn-hidden-size 14336、--group-query-attention 8 等）。

这样才能确保 HF → Megatron 转换后行为一致、数值稳定。

在大语言模型（如 Llama、Mistral）中，“embedding 层权重”和“输出层权重”分别指模型输入和输出阶段使用的两个关键参数矩阵。它们的含义如下：

1. Embedding 层权重（Input Embedding）

  • 作用：将输入的 token ID（整数，如 1234）映射为一个 稠密向量（如 4096 维的浮点向量）。
  • 位置：模型最底层（输入端）。
  • 参数矩阵：通常记作 W_e，形状为 [vocab_size, hidden_size]。
    • vocab_size：词表大小（例如 Llama-2 是 32000）
    • hidden_size：隐藏层维度（例如 Llama-2-7B 是 4096）

✅ 举例：
输入 token ID = 5000 → 查表 W_e[5000] → 得到一个 4096 维的向量，作为模型第一层的输入。

2. 输出层权重（Output / LM Head）

  • 作用：将模型最后一层的 隐藏状态向量（如 4096 维）映射回 词表上的 logits（即每个词的打分），用于预测下一个 token。
  • 位置：模型最顶层（输出端）。
  • 参数矩阵：通常记作 W_o，形状也是 [vocab_size, hidden_size]（或转置形式 [hidden_size, vocab_size]，取决于实现）。

✅ 举例：
最后一层输出向量 h（4096 维） → 计算 logits = h @ W_o.T → 得到长度为 vocab_size 的 logits 向量 → 经 softmax 得到下一个 token 的概率分布。

3. 权重共享（Tying / Untying）

很多模型（如原始 Transformer、BERT）会 共享 embedding 和输出层的权重，即：

W_o = W_e  # 或 W_o = W_e.T

这样做的好处：

  • 减少参数量（节省内存）
  • 理论上有更好的泛化性（输入输出语义空间对齐）

但 Llama 系列和 Mistral 默认不共享！
也就是说：

  • embedding.weight 和 lm_head.weight 是 两个独立的参数矩阵
  • 它们虽然形状相同，但数值不同

因此，在 Megatron-LM 中加载这些模型时，必须显式指定：

--untie-embeddings-and-output-weights

否则 Megatron 会默认尝试共享权重，导致：

  • 维度不匹配（如果 vocab_size 不整除 TP）
  • 或强行共享，造成输出错误（因为 HF 模型并未共享）

总结对比

组件 作用 参数名（HF） 是否共享（Llama/Mistral）
Embedding 层 token ID → 向量 model.embed_tokens.weight ❌ 不共享
输出层（LM Head） 向量 → logits lm_head.weight ❌ 独立参数

💡 关键点：
“解绑 embedding 和输出层权重”（--untie-embeddings-and-output-weights）是正确加载 Llama/Mistral 到 Megatron 的必要条件。

如有需要，我也可以画一个简图说明数据流向。

这是一个典型的错误操作，会导致 Megatron 在加载模型时崩溃。下面从 代码逻辑 和 实际行为 两个层面解释“用 TP=1 转换 Llama-70B，却用 --tensor-model-parallel-size 8 启动”会发生什么，以及 Megatron 内部是如何处理的。

🔧 背景知识回顾

  • TP（Tensor Parallelism）：将模型的权重（如 QKV 投影、MLP）按列或行切分到多个 GPU 上。
  • 转换时的 --target-tensor-parallel-size：决定输出的 Megatron 检查点被切分成多少份。
  • 启动时的 --tensor-model-parallel-size：告诉 Megatron 当前运行时使用多少个 GPU 做张量并行。

✅ 正确做法：转换时的 TP 必须等于启动时的 TP。

❌ 错误操作示例

# 第一步：错误地用 TP=1 转换 Llama-70B（本应 TP=8）
python tools/checkpoint/convert.py \
  --model-type GPT \
  --loader llama_mistral \
  --checkpoint-type hf \
  --load-dir /path/to/llama-70b-hf \
  --save-dir /path/to/megatron-ckpt-tp1 \
  --model-size llama2-70B \
  --tokenizer-model /path/to/tokenizer.model \
  --saver core \
  --target-tensor-parallel-size 1 \    # ❌ 错误：应为 8
  --target-pipeline-parallel-size 1 \
  --bf16

# 第二步：却用 TP=8 启动
python -m torch.distributed.run \
  --nproc_per_node 8 \
  pretrain_gpt.py \
  --tensor-model-parallel-size 8 \     # ❌ 与转换时不一致
  --pipeline-model-parallel-size 1 \
  --load /path/to/megatron-ckpt-tp1 \
  ...

💥 代码层面会发生什么？

1. 检查点结构不匹配

Megatron 启动时，每个 rank（0~7）会尝试加载自己的分片：

  • 它期望在 /path/to/megatron-ckpt-tp1/ 下看到 8 个子目录，如： 
    mp_rank_00/
mp_rank_01/
...
mp_rank_07/
  • 但因为转换时用了 TP=1，实际只有： 
    mp_rank_00/model_optim_rng.pt

➡️ 结果：rank 1~7 找不到自己的 checkpoint，抛出异常。

2. 具体报错（典型错误信息）
RuntimeError: Checkpoint file /path/to/megatron-ckpt-tp1/mp_rank_01/model_optim_rng.pt does not exist.

或（如果启用了 --exit-on-missing-checkpoint）：

AssertionError: Missing checkpoint for model parallel rank 1

这正是文档中强调要加 --exit-on-missing-checkpoint 的原因：尽早暴露错误。

🔍 深入：Megatron 如何加载权重？

在 megatron/core/checkpointing.py 或 load_checkpoint() 函数中，有类似逻辑：

mp_rank = mpu.get_tensor_model_parallel_rank()
checkpoint_name = os.path.join(
    load_dir,
    f"mp_rank_{mp_rank:02d}",
    "model_optim_rng.pt"
)
state_dict = torch.load(checkpoint_name, map_location="cpu")
  • 当 --tensor-model-parallel-size 8 时，mp_rank 会是 0~7。
  • 但转换时只生成了 mp_rank_00/，所以 rank ≥1 的进程会失败。

🛠 正确操作（对比）

✅ 转换时指定 TP=8：

python tools/checkpoint/convert.py \
  ... \
  --target-tensor-parallel-size 8 \   # ✅ 正确
  --save-dir /path/to/megatron-ckpt-tp8

这会生成：

/path/to/megatron-ckpt-tp8/
├── mp_rank_00/
├── mp_rank_01/
├── ...
└── mp_rank_07/

✅ 启动时也用 TP=8：

--tensor-model-parallel-size 8 \
--load /path/to/megatron-ckpt-tp8

每个 rank 加载自己的分片，权重维度匹配，训练/推理正常。

💡 补充：能否“动态重切分”？

目前 Megatron-LM 不支持 在加载时自动将 TP=1 的 checkpoint 重切分为 TP=8。
你必须在转换阶段就指定目标 TP。

有些社区工具（如 NVIDIA NeMo 的 nemo-megatron）支持 checkpoint 重切分（res harding），但标准 Megatron-LM 不支持。

✅ 总结

步骤 错误操作 后果
转换 --target-tensor-parallel-size 1（对 70B） 只生成 1 份权重
启动 --tensor-model-parallel-size 8 8 个 rank 同时启动
结果 rank 1~7 找不到 checkpoint 崩溃

📌 记住：转换时的 --target-tensor-parallel-size 必须等于启动时的 --tensor-model-parallel-size。这是 Megatron 分布式训练的基本前提。

# 学习
Logo
腾讯云开发者社区

腾讯云面向开发者汇聚海量精品云计算使用和开发经验，营造开放的云计算技术生态圈。

更多推荐

终极指南：Flink SQL连接器版本管理从混乱到有序的升级之路

Apache Flink作为流处理领域的佼佼者，其SQL连接器的版本管理一直是开发者面临的核心挑战。本文将系统讲解Flink SQL连接器版本管理的最佳实践，帮助你轻松应对版本兼容性问题，实现从混乱到有序的升级之旅。## 连接器版本管理的常见痛点 😫在Flink应用开发中，连接器版本管理常常让开发者头疼不已。不同版本的连接器可能导致各种兼容性问题，例如API变更、功能差异甚至运行时错误。

avatar 腾讯云开发者社区

Elasticsearch复杂数据类型终极指南：从入门到精通

Elasticsearch作为功能强大的搜索引擎，支持多种复杂数据类型，让开发者能够灵活处理各种结构化和非结构化数据。本文将带你全面了解Elasticsearch中的复杂数据类型，从基础概念到实际应用，助你轻松掌握数据建模的核心技巧。## 内部对象：构建层级化数据结构在Elasticsearch中，对象类型（Object）是最基础的复杂数据类型之一，用于表示具有嵌套关系的数据。例如，我们可

avatar 腾讯云开发者社区

如何快速搭建Neon无服务器PostgreSQL：面向初学者的完整指南

Neon是一款革命性的无服务器PostgreSQL解决方案，它通过分离存储和计算层，实现了自动扩缩容、类代码式数据库分支以及零级扩展能力。本指南将帮助你从零开始搭建Neon开发环境，体验这款创新数据库的强大功能。## 准备工作：环境要求与依赖项在开始搭建Neon环境前，请确保你的系统满足以下要求：- Linux操作系统（推荐Ubuntu 20.04+或Debian 11+）- Git

avatar 腾讯云开发者社区