总目录 › ColossalAI › 第 12 章
Chapter 12

Checkpoint 与 HuggingFace 互转

📌 commit 4f9953be335e 训完模型怎么落盘、怎么转 HF、怎么再次 resume

Checkpoint 是项目生命周期里最容易出问题的环节。本章把 ColossalAI 三类 checkpoint 形态（shard / unshard / HF）讲清楚，并示范在 plugin 之间互转。

12.1三类 checkpoint 形态

形态结构典型大小 (7B)用途
Sharded (ColossalAI) 每张 GPU 一份 pytorch_model.bin，配 model_metadata.json ~14 GB（bf16，每卡 ~1.75GB） 训练 resume；同 plugin 加载
Unshard (single file) 合并成单个 pytorch_model.bin ~14 GB 跨 plugin / 单卡推理
HuggingFace model.safetensors + config.json + tokenizer.json ~14 GB 发布到 HF Hub / vLLM 部署

12.2API：booster.save_model / load_model

看 colossalai/booster/booster.py：

booster.save_model(
    model,
    checkpoint,             # 路径
    shard=True,             # True: 多文件；False: 合并单文件
    gather_dtype=None,
    prefix=None,
    size_per_shard=1024,    # 单文件最大 MB
    use_safetensors=False,  # safetensors 还是 .bin
)

booster.load_model(model, checkpoint, strict=True)

shard 三种策略

策略触发条件产物
shard=True + 内部多卡 plugin 是 ZeRO/Gemini/HybridParallel 每卡一个文件 + index.json
shard=False 把分布式参数 all-gather 后保存 单个大文件（每张卡都跑得动 RAM）
shard=True + HF 格式 plugin 支持 hf_strategy HF safetensors 多分片

12.3训练 resume 完整流程

resume 不只是 load_model，还要 load optimizer / scheduler / RNG。看 train.py 的完整写法：

booster.save_model(model, f"{save_dir}/model")
booster.save_optimizer(optimizer, f"{save_dir}/optim", shard=True)
booster.save_lr_scheduler(lr_scheduler, f"{save_dir}/lr_scheduler")
# RNG 状态自己保
torch.save({
    "epoch": epoch,
    "step": step,
    "torch_rng": torch.get_rng_state(),
    "cuda_rng": torch.cuda.get_rng_state_all(),
}, f"{save_dir}/meta.pt")

恢复：

booster.load_model(model, f"{save_dir}/model")
booster.load_optimizer(optimizer, f"{save_dir}/optim")
booster.load_lr_scheduler(lr_scheduler, f"{save_dir}/lr_scheduler")
meta = torch.load(f"{save_dir}/meta.pt")
torch.set_rng_state(meta["torch_rng"])
torch.cuda.set_rng_state_all(meta["cuda_rng"])
start_epoch = meta["epoch"]
start_step  = meta["step"]

12.4转 HuggingFace

用 ColossalAI plugin 训完后想推到 HF Hub：

方法 A：HybridParallelPlugin / Gemini 直接 dump HF 格式

# 训完后
if dist.get_rank() == 0:
    # 先 all-gather 到 rank 0 的 model
    booster.save_model(
        model,
        "./hf_out",
        shard=True,
        use_safetensors=True,
    )
    tokenizer.save_pretrained("./hf_out")
# 此时 ./hf_out/ 是标准 HF 目录

方法 B：先 unshard 再用 transformers 加载

booster.save_model(model, "./unshard.bin", shard=False)

# 在另一终端
from transformers import LlamaForCausalLM
model = LlamaForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")
state = torch.load("./unshard.bin", map_location="cpu")
model.load_state_dict(state)
model.save_pretrained("./hf_out", safe_serialization=True)

12.5从 HuggingFace 加载到 ColossalAI

反过来，从 HF 模型起步训练：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.bfloat16,
)
# 这一步就是普通 HF 加载，模型是单卡完整版

plugin = HybridParallelPlugin(tp_size=2, pp_size=2, zero_stage=1)
booster = Booster(plugin=plugin)
model, optimizer, ... = booster.boost(model, optimizer, ...)
# booster.boost() 内部会调 ShardFormer 把权重切分到各卡

关键：HF → Colossal 的切分由 ShardFormer 在 booster.boost() 时自动完成，你不用手动 re-shard。

12.6跨 plugin 互转

真实场景：CPT 用 HybridParallel 训完，SFT 想用 Gemini（因为多机不方便配 TP）。流程：

  1. HybridParallel 训完 → booster.save_model(shard=False) 得到 unshard.bin；
  2. 新建 Gemini Booster；
  3. 用 transformers from_pretrained 加载 unshard 模型；
  4. booster.boost(model) Gemini 会自动切 chunk。

原则：plugin 之间的桥梁是"unsharded HF 格式"，所有 plugin 都能产出/读取它。

12.7safetensors 还是 .bin

格式优点缺点
safetensors 更快（mmap） / 更安全（无 pickle） 需要 safetensors 包；老工具链不支持
.bin (pickle) 通用 慢；存在反序列化任意代码执行风险

2025+ 项目默认用 safetensors（HF Hub 也强制）。

12.8大模型 ckpt 的存储优化

12.9常见错误

现象原因处理
save 后 reload loss 不一致 RNG 没保存；或 grad accumulator 未保存 保存 RNG + 不在 step 中间保存
"OSError: Error loading shared model" 不同 rank 看到的 ckpt 路径不同（NFS 延迟） save 后 dist.barrier()；用本地路径
HF 转 Colossal 后 q_proj weight shape 错 GQA 模型的 K/V head 数和 Q 不一致，tp_size 必须整除 K head 数 选合适的 tp_size（如 LLaMA-3 K=8，tp=2 / 4 / 8 都行；tp=3 不行）
ZeRO-3 save 巨慢 每卡都 all-gather 一份，I/O 串行 用 shard=True 让每卡保存自己那份

12.10这章你需要带走的