laoyao_2b_moe/tools

Tools

本目录保存 Laoyao 2B MoE 训练后的 checkpoint 检查、Megatron DCP 探测、HF 导出和推理调试工具。

Megatron checkpoint 检查

• laoyao_dcp_probe.py: 用 torch.distributed.checkpoint 读取 Megatron distributed checkpoint metadata 和单个 tensor，确认 DCP checkpoint 可以被 CPU/gloo 进程读取。
• inspect_laoyao_ckpt_metadata.py: 打印 checkpoint 的 metadata.json、common.pt 和 run_config.yaml 关键信息。
• dump_laoyao_model_keys.py: dump 模型 key，用于对齐 Megatron 权重名和 HF custom model 权重名。
• laoyao_cpu_infer_probe.py / laoyao_cpu_generate_from_distcp.py: 早期 CPU 推理探测脚本，主要用于定位 checkpoint 读取和模型定义问题。

HuggingFace 导出

hf_laoyao_moe/ 是 Megatron distributed checkpoint 到 HuggingFace custom model 的导出适配层。

核心入口：

python3 tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py \
  --checkpoint-dir runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints/iter_0107500 \
  --tokenizer-dir tokenizer/glm5.2 \
  --output-dir runs/hf_exports/iter_0107500

生成测试：

python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \
  --model-dir runs/hf_exports/iter_0107500 \
  --device cuda \
  --max-new-tokens 32 \
  --prompt "The capital of France is"

注意：HF custom model 当前没有实现 KV cache，因此导出配置和生成脚本默认 use_cache=false。更多细节见 docs/megatron_to_hf_export_notes.md。