Chapter 12
把 nanotron checkpoint 转 HF
nanotron 训完的 ckpt 不能直接给 HF 用。本章过转换流程。
12.1nanotron ckpt 结构
每 (pp, tp) 一份 safetensors，weight 名按 nanotron 命名：
./checkpoints/iteration_50000/
├── config.yaml
├── pp-0-of-2/
│ ├── tp-0-of-4/
│ │ └── model_weights.safetensors
│ └── ...
└── metadata.json12.2转换流程
nanotron 仓库 tools/llama3/convert_nanotron_to_hf.py 提供 Llama 系转换：
python tools/llama3/convert_nanotron_to_hf.py \
--checkpoint_path ./checkpoints/iteration_50000 \
--tokenizer_name meta-llama/Meta-Llama-3-8B \
--save_path ./hf_out核心步骤：
- 把 pp / tp 各自 shard 加载，合并成单进程；
- 把 nanotron 风格命名重命名成 HF 命名；
- 把 fused QKV / gate_up 拆回 q_proj/k_proj/v_proj 和 gate_proj/up_proj；
- 保存为 HF safetensors 分片。
12.2命名映射表（节选）
| nanotron 名 | HF 名 |
|---|---|
| model.token_embeddings.pp_block.token_embedding.weight | model.embed_tokens.weight |
| model.decoder.{N}.pp_block.input_layernorm.weight | model.layers.{N}.input_layernorm.weight |
| model.decoder.{N}.pp_block.attn.qkv_proj.weight | ★ 拆成 q_proj/k_proj/v_proj.weight |
| model.decoder.{N}.pp_block.attn.o_proj.weight | model.layers.{N}.self_attn.o_proj.weight |
| model.decoder.{N}.pp_block.mlp.gate_up_proj.weight | ★ 拆成 gate_proj/up_proj.weight |
| model.decoder.{N}.pp_block.mlp.down_proj.weight | model.layers.{N}.mlp.down_proj.weight |
| model.lm_head.pp_block.weight | lm_head.weight |
转换的关键难点：QKV 是 fused 的，要按 num_heads / num_kv_heads / head_dim 切对。
12.3HF → nanotron
反方向用 tools/llama3/convert_hf_to_nanotron.py：
python tools/llama3/convert_hf_to_nanotron.py \
--hf_model meta-llama/Meta-Llama-3-8B \
--save_path ./nanotron_init \
--tp 4 --pp 2注意 --tp / --pp 决定目标 shard 数，要跟之后训练 config 一致。
12.4转完用 HF 推理
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("./hf_out", torch_dtype=torch.bfloat16)
tok = AutoTokenizer.from_pretrained("./hf_out")
out = model.generate(**tok("Hello", return_tensors="pt"), max_new_tokens=50)
print(tok.decode(out[0]))转完可以直接进 vLLM / TGI / lmdeploy 生态。
12.5常见错误
| 现象 | 处理 |
|---|---|
| HF 推理后乱码 | QKV 拆错；查 num_key_value_heads 设对 |
| shape mismatch | nanotron config 跟 HF 不一致（hidden_size / vocab_size 等） |
| tokenizer 不一致 | 转时显式 --tokenizer_name |
| RoPE base 错 | nanotron config 里 rope_theta 跟 HF 一致 |
12.6这章你需要带走的
- nanotron → HF 用
tools/llama3/convert_nanotron_to_hf.py； - 关键 = 命名映射 + QKV 拆分；
- 反方向 HF → nanotron 用
convert_hf_to_nanotron.py； - 转完进 HF 生态：vLLM / TGI / lmdeploy 都能用。