Checkpoint：Megatron ↔ HuggingFace 互转

12.1 Megatron checkpoint 长什么样

Megatron 的 checkpoint 不是单一文件，而是按并行布局切分的目录树。一个 TP=4 × PP=2 训练的 checkpoint 目录大致：

iter_0001000/
├── mp_rank_00_000/                  # TP=0, PP=0
│   └── model_optim_rng.pt
├── mp_rank_01_000/                  # TP=1, PP=0
│   └── model_optim_rng.pt
├── mp_rank_02_000/                  # TP=2, PP=0
│   └── model_optim_rng.pt
├── mp_rank_03_000/                  # TP=3, PP=0
│   └── model_optim_rng.pt
├── mp_rank_00_001/                  # TP=0, PP=1
│   └── model_optim_rng.pt
├── ... 共 4 × 2 = 8 个目录
└── latest_checkpointed_iteration.txt

每个 model_optim_rng.pt 含：

• 当前 rank 持有的参数（按 TP / PP 切片后的）；
• 对应的 optimizer state；
• RNG 状态；
• 训练步数、ENV、args 等元信息。

12.2 三种 checkpoint 格式

torch_dist 的好处是"训练用 TP=8 PP=16，恢复用 TP=4 PP=8"完全没问题。Megatron 0.10+ 默认推荐这个。

12.3 互转的两个方向

核心工具是 tools/checkpoint/convert.py，它接受一对 loader（读什么格式）+ saver（写什么格式）。

12.4 HF → Megatron（用 HF 模型起步训练）

常见诉求：拿 HuggingFace 上的 Llama-3-70B 当起点做继续预训练。

python tools/checkpoint/convert.py \
    --model-type GPT \
    --loader llama_mistral \
    --saver core \
    --load-dir /path/to/hf/Meta-Llama-3-70B \
    --save-dir /path/to/megatron/llama3_70b_tp8_pp4 \
    --target-tensor-parallel-size 8 \
    --target-pipeline-parallel-size 4 \
    --tokenizer-model /path/to/tokenizer.model \
    --model-size llama3-70B \
    --bf16

关键参数：

12.5 Megatron → HF（训练完导出给生态用）

python tools/checkpoint/convert.py \
    --model-type GPT \
    --loader core \
    --saver llama_mistral \
    --load-dir /path/to/megatron/output/iter_0010000 \
    --save-dir /path/to/hf/my_llama3_70b_continued \
    --target-tensor-parallel-size 1 \
    --target-pipeline-parallel-size 1 \
    --tokenizer-model /path/to/tokenizer.model

注意：

• output TP/PP 必须设 1，HF 格式不分布式切；
• 转完会得到标准 pytorch_model-XX.bin 或 model.safetensors + config.json；
• tokenizer 文件单独 copy，loader 不会自动复制。

12.6 改变并行布局（Megatron → Megatron）

训练时用 TP=8 PP=16（128 卡），后续用 8 卡微调，需要 TP=4 PP=2：

python tools/checkpoint/convert.py \
    --model-type GPT \
    --loader core \
    --saver core \
    --load-dir /old/megatron/ckpt \
    --save-dir /new/megatron/ckpt \
    --target-tensor-parallel-size 4 \
    --target-pipeline-parallel-size 2

新格式 torch_dist 训练时直接换并行配置就行，不需要离线转换：

# 训练时
--ckpt-format torch_dist

# 恢复时换并行配置直接 --load 即可

12.7 dist_checkpointing：新版机制详解

megatron/core/dist_checkpointing/ 下的代码实现了"布局无关"的 checkpoint。原理：

1. 保存时按"逻辑张量"组织（每个参数有一个全局 key），每张卡只存自己持有的切片；
2. 所有 rank 协同写入同一 checkpoint 目录；
3. 恢复时按当前进程的 TP/PP/DP 重新切分加载——和保存时不一致也没问题。

这种"分片但可重组"的设计让checkpoint 不再绑定具体并行布局。

12.8 异步保存（H100 神器）

--ckpt-format torch_dist \
--async-save

原理：把"序列化 + 落盘" 放到后台线程，训练继续进行。能省下大型 checkpoint 的同步时间（70B 模型一次 save 同步会卡 30-60 秒）。

12.9 loader / saver 列表（按模型）

没列出的模型（如 Qwen3、DeepSeek-V3）通常用 NVIDIA-NeMo/Megatron-Bridge，那是 NV 官方做的更通用工具。

12.10 Megatron-Bridge：推荐的统一互转工具

仓库 README 顶部就推荐：

Megatron Bridge provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production-ready recipes.

用法：

pip install megatron-bridge

# HF → Megatron
python -m megatron_bridge.tools.convert_hf_to_megatron \
    --hf_model meta-llama/Llama-3-70B \
    --output ckpt/llama3_70b_mcore \
    --tp 8 --pp 4 \
    --dtype bf16

# Megatron → HF
python -m megatron_bridge.tools.convert_megatron_to_hf \
    --megatron_ckpt ckpt/llama3_70b_mcore/iter_0010000 \
    --output hf/my_llama3_70b \
    --hf_model_id meta-llama/Llama-3-70B          # 用原模型的 config 模板

Bridge 比 tools/checkpoint/convert.py 的优势：

• 支持更多模型；
• 对 MoE / multimodal / 新结构（GroupedQueryAttention / SwiGLU / RoPE 变体）覆盖更全；
• 有 production-ready recipes（如 Llama-3.1-70B 完整训练 + 导出）。

12.11 转完之后的验证

检查 HF 格式模型转换正确性：

# 1. config.json 存在
ls path/to/hf/model/config.json

# 2. 能正常加载 + 输出和原模型相近
python -c "
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tok = AutoTokenizer.from_pretrained('path/to/hf/model')
model = AutoModelForCausalLM.from_pretrained('path/to/hf/model', torch_dtype=torch.bfloat16).cuda()
inputs = tok('Hello world', return_tensors='pt').to('cuda')
print(model.generate(**inputs, max_new_tokens=20))
"

# 3. 数值差异检查（与原 HF 模型对比）
python tools/checkpoint/checkpoint_inspector.py \
    --check-dir path/to/megatron/ckpt

12.12 常见坑

12.13 Async dist-checkpointing（2025-Q2 新）

megatron/core/dist_checkpointing/strategies/ 下两个新文件：

启用方式：--ckpt-format torch_dist --async-save。训练循环在 save_step 立即返回，不阻塞 forward/backward；后台线程完成后再回调释放 buffer。1B+ 模型每 100 步保存一次也基本无感。

12.14 torch_dist 格式的优势

2026 新工程建议直接用 torch_dist，旧的 zarr / per-rank pickle 都进入维护状态。

12.15 这章你需要带走什么

• Megatron 用切片目录存 checkpoint；
• 新版 torch_dist 让 checkpoint 不再绑定并行布局；
• tools/checkpoint/convert.py 是内置工具，但模型覆盖有限；
• NVIDIA-NeMo/Megatron-Bridge 是 production-ready 推荐工具；
• 转换前后必须做数值验证，否则训完模型可能乱码。