# 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 仍是更可信的推理基线。

