36 lines
1.5 KiB
Markdown
36 lines
1.5 KiB
Markdown
# 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 的导出适配层。
|
|
|
|
核心入口：
|
|
|
|
```bash
|
|
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
|
|
```
|
|
|
|
生成测试：
|
|
|
|
```bash
|
|
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`。
|