Sync g0050 training and inference tooling

docs/megatron_to_hf_export_notes.md

@@ -0,0 +1,111 @@
+# Megatron 到 HuggingFace 导出与推理修复记录
+
+本文记录 Laoyao 2B MoE 从 Megatron distributed checkpoint 导出到 HuggingFace 自定义模型时遇到的问题、修复和当前结论。
+
+## 当前结论
+
+Megatron 原生 checkpoint 可以正常通过 Megatron text generation server 推理。导出到 HuggingFace 时，权重本身可以转换为 `safetensors`，但 HF 自定义模型必须显式对齐 Megatron 的模型细节，并且默认禁用 KV cache。
+
+当前最重要的修复是：
+
+- `config.use_cache = false`，因为 HF 自定义模型没有实现 KV cache。
+- `rms_norm_eps = 1e-5`，与 Megatron runtime 参数一致。
+- RoPE 使用 Megatron 默认的 non-interleaved/NeoX half-split 形式。
+- MoE router 模拟 Megatron 的 `moe_expert_capacity_factor=1.25` token dropping 行为。
+- HF dynamic module 使用相对导入，保证 `AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)` 能从导出目录独立加载。
+
+## 已确认的问题
+
+### KV cache
+
+当前 HF 模型没有实现 `past_key_values`、`cache_position` 和 per-layer K/V cache 管理。Transformers `generate()` 默认倾向于使用 cache；如果不禁用 cache，后续 step 可能只输入新 token，而模型会把该 token 当作完整上下文从 position 0 计算，导致上下文丢失和异常重复。
+
+现象：
+
+- `model.generate()` 默认 cache 路径输出 `Paris, and the 1. The...`、中文重复 `在在在...`。
+- 设置 `use_cache=False` 后，输出恢复为连续句子，例如 `Paris. The capital of Belgium...`。
+
+修复：
+
+- `configuration_laoyao_moe.py` 默认 `use_cache=False`。
+- `convert_laoyao_dcp_to_hf.py` 写出的 `config.json` 包含 `"use_cache": false`。
+- `generate_laoyao_hf.py` 调用 `model.generate(..., use_cache=False)`。
+
+### RMSNorm epsilon
+
+Megatron server 日志显示：
+
+```text
+layernorm_epsilon ............................... 1e-05
+```
+
+HF 初始实现使用 `1e-6`，已修正为 `1e-5`。这个差异会造成逐层数值漂移，尤其可能影响 MoE router 的 top-k 边界。
+
+### RoPE layout
+
+Megatron runtime 参数：
+
+```text
+rotary_interleaved .............................. False
+rotary_base ..................................... 10000
+rotary_percent .................................. 1.0
+```
+
+Megatron 源码在 `rotary_interleaved=False` 时使用：
+
+- `emb = torch.cat((freqs, freqs), dim=-1)`
+- `_rotate_half(x) = torch.cat((-x2, x1), dim=-1)`，其中 `x1, x2 = torch.chunk(x, 2, dim=-1)`
+
+因此 HF 侧应使用 NeoX half-split RoPE，而不是 even/odd interleaved RoPE。
+
+### MoE expert capacity
+
+Megatron runtime 参数包含：
+
+```text
+moe_expert_capacity_factor ...................... 1.25
+moe_token_drop_policy ........................... probs
+moe_pad_expert_input_to_capacity ................ False
+```
+
+Megatron `TopKRouter` 在 `moe_expert_capacity_factor` 非空时会调用 `apply_router_token_dropping`，按 expert 维度保留 routing prob 较高的 assignment，超出 capacity 的 token-expert assignment 会被 mask 掉。
+
+HF 初始实现是 dropless top-k routing，已补齐 sparse routing map + capacity mask 逻辑。注意：这不是本次 `1/在` 重复的主因，但它是和 Megatron 对齐必须处理的行为。
+
+### HF dynamic module 导入
+
+HF export 目录中的 `modeling_laoyao_moe.py` 必须使用：
+
+```python
+from .configuration_laoyao_moe import LaoyaoMoeConfig
+```
+
+不能使用裸导入 `from configuration_laoyao_moe import ...`，否则脱离工具目录后 `AutoModelForCausalLM.from_pretrained` 会失败。
+
+## 脚本入口
+
+- `tools/hf_laoyao_moe/convert_laoyao_dcp_to_hf.py`
+  - 从 Megatron DCP checkpoint 读取权重。
+  - 写出 HF `model.safetensors`、`config.json`、tokenizer 文件和 custom code。
+
+- `tools/hf_laoyao_moe/generate_laoyao_hf.py`
+  - 从 HF export 目录加载模型。
+  - 默认 `use_cache=False` 跑生成。
+  - 支持 `--print-token-details` 和 `--top-k-debug` 做逐 token 调试。
+
+- `scripts/serve_laoyao_megatron.sh`
+  - 启动 Megatron 原生 text generation server。
+
+- `scripts/laoyao_megatron_inference_setup.sh`
+  - 容器启动时热补丁 Megatron server：
+    - 修 `inference_max_sequence_length` 参数兼容。
+    - 修 token list concat。
+    - 支持 generated-only logprobs，用于调试。
+
+## 当前限制
+
+- HF export 能正确 full-context 生成，但没有高效 KV cache。
+- HF 模型没有实现 `_reorder_cache()`，不支持 beam search cache 重排。
+- MoE 和 attention 目前是简单 PyTorch 实现，不是高性能推理实现。
+- Megatron 原生 server 仍是更可信的推理基线。
+