laoyao_2b_moe/tools/hf_laoyao_moe

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 会导致上下文丢失和异常重复。

Heldout 2.8k 评测

eval_heldout_2p8k.py 复用 tokenizer-swap 实验的评测口径：

• 对所有样本计算 prompt-conditioned completion NLL/PPL。
• 对能从 prompt 中解析出 A/B/C/D 候选项、且 gold completion 以候选 label 开头的样本计算 MCQ accuracy。
• 主 accuracy 是候选 continuation 的平均 token logprob accuracy：mcq_acc_avg_norm。
• mcq_acc_sum 是总 logprob accuracy，容易偏好短答案，只作辅助指标。

示例：

python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \
  --model-dir runs/hf_exports/iter_0107500 \
  --data dataset/val/data/heldout_2p8k_sft_prompt_completion.jsonl \
  --out-dir runs/hf_eval/heldout_2p8k/iter_0107500 \
  --device cuda \
  --max-length 2048

CPU smoke 可以限制样本数：

python3 tools/hf_laoyao_moe/eval_heldout_2p8k.py \
  --model-dir runs/hf_exports/iter_0107500 \
  --max-items 4 \
  --device cpu \
  --dtype float32 \
  --max-length 512

调试

打印逐 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