laoyao_2b_moe/tools/hf_laoyao_moe/README.md

HF Laoyao MoE 工具说明

本目录包含把 Laoyao 2B MoE Megatron 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

导出目录包含：

• model.safetensors
• config.json
• configuration_laoyao_moe.py
• modeling_laoyao_moe.py
• tokenizer 文件

生成测试

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

当前 HF custom model 没有实现 KV cache，因此生成脚本强制 use_cache=False。不要删除这个设置；否则 Transformers 默认 cache path 会导致上下文丢失和异常重复。

调试

打印逐 token 结果：

python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \
  --model-dir runs/hf_exports/iter_0107500 \
  --device cuda \
  --fix-mistral-regex \
  --max-new-tokens 8 \
  --print-token-details \
  --prompt "The capital of France is"

打印每步 top-k：

python3 tools/hf_laoyao_moe/generate_laoyao_hf.py \
  --model-dir runs/hf_exports/iter_0107500 \
  --device cuda \
  --fix-mistral-regex \
  --max-new-tokens 5 \
  --top-k-debug 8 \
  --prompt "The capital of France is"

已对齐的 Megatron 细节

• rms_norm_eps=1e-5
• RoPE rotary_interleaved=False，即 NeoX half-split layout
• GQA QKV group-major layout
• tied embedding/output head
• MoE top-k router with moe_expert_capacity_factor=1.25
• use_cache=false