# Phase T9: Export to xserv — Design Document

## Goal

闭环 xtrain ↔ xserv：把 xtrain 练出的权重导成 **xserv 的 Qwen3 loader 能直接加载并服务**的格式
（HF 命名的 `config.json` + `model.safetensors` + 复用的 gpt2 `tokenizer.json`），让推理侧 xserv
跑出**与 xtrain 自身一致**的生成结果。

验收信号：**同一份权重 + 同一 prompt，xserv 的贪心生成 token 序列对住 xtrain 的贪心生成**
（logits 在浮点容差内吻合）。这是整条 P0→P6 学习链的收口——训练栈练出来的东西，推理栈真的能用。

> 范围与诚实边界：**不改 xserv**。xserv 是独立项目，T9 只调整 xtrain 的导出去适配 xserv 既有 loader。
> 架构差异中只有一项是「结构性」的（QK-norm，见下），处理方式见 **Key Design Decision 1**。

## 关键第一步：架构 diff（xtrain TinyTransformer vs xserv qwen3.rs）

逐 op 读 `crates/xtrain-model/src/model.rs` 对 `~/projects/xserv/crates/xserv-model/src/qwen3.rs`
（forward 路径 `forward_with_cache`，即 `dump-logits` 走的 prefill）：

| 维度 | xtrain TinyTransformer | xserv Qwen3 | 是否兼容 / 处理 |
|------|------------------------|-------------|----------------|
| RMSNorm 公式 | `x*rsqrt(mean(x²)+eps)*γ`（无均值减） | 同 | ✅ 一致 |
| RMSNorm eps | `cfg.eps`（tiny=1e-5） | `rms_norm_eps`（config.json 提供） | ✅ 导出 eps 到 config |
| **QK-norm** | **T9 前没有** | **强制** per-head RMSNorm(Q,K)（`q_norm`/`k_norm`，shape `[head_dim]`） | **结构性差异** → 见 Decision 1 |
| RoPE 约定 | rotate_half，`freq=θ^(-2i/hd)`，pos=行号 | rotate_half，cos/sin cache `freq=1/θ^(2i/hd)`，pos=token idx | ✅ **逐式一致**（见 rope.cu）|
| RoPE θ | `cfg.rope_theta`（10000） | `rope_theta`（config 提供，默认 1e6） | ✅ 导出 θ 到 config |
| Attention scale | `1/√head_dim` | `1/√head_dim`（attention.rs / flash） | ✅ 一致 |
| Causal mask | 加性 `-1e9` 上三角 | causal flag（online softmax） | ✅ 同义 |
| GQA | MHA（无 kv 分组） | 支持 GQA（`num_kv_heads`） | ✅ 设 `num_key_value_heads = num_attention_heads`（退化为 MHA）|
| SwiGLU | `down(silu(gate)∘up)`，gate/up 独立 proj | 同（融合存 `gate_up_proj`，loader 内部 cat）| ✅ 一致，导出仍分开 gate/up（loader 自己 cat）|
| 偏置 | 无 | `attention_bias=false`，不读 bias | ✅ 一致 |
| final norm | `final_norm` 后接 lm_head | `model.norm` 后接 lm_head | ✅ 一致 |
| embedding tying | 独立 `lm_head` | 独立（`tie_word_embeddings=false`）| ✅ 一致 |
| 2D 权重 layout | `[in,out]`，算 `x@W` | `[out,in]`，算 `x@Wᵀ`（loader `.transpose(0,1)`）| ⚠️ 导出须转置 |
| **dtype** | **f32**（训练 + 自身推理） | **BF16 only**（kernel 全 assert BF16）| ⚠️ 导出转 BF16；数值上不可能 bit-exact，见 Decision 2 |
| vocab | gpt2 BPE，50257 | config 提供 | ✅ 导出 vocab |

**结论**：除 QK-norm 一项外，其余差异都是机械的（命名 / 转置 / dtype / 退化 GQA）。QK-norm 是
xserv 对 Qwen3 的**强制**步骤（`head_rmsnorm(q,q_norm)` / `head_rmsnorm(k,k_norm)` 无条件执行，
且 γ=1 也不是恒等——它仍按每个 head 向量的 RMS 做归一），xtrain 训练时从未施加 → 若不处理，
xserv 的前向数学与 xtrain 不同，闭环不可能成立。

## Key Design Decisions

### Decision 1：给 xtrain 加 per-head QK-norm（而不是伪造 match，也不是停在 blocker）

逃生舱里给了三条路：① 给 xtrain 补 QK-norm 重训以对齐 Qwen3；② 改 xserv（被禁止）；
③ 当作 blocker 报告。选 ①——因为它**真的能闭环**且改动**外科手术级**：

- xserv 的顺序是 `reshape → head_rmsnorm → transpose_for_rope → rope`，即 **QK-norm 在 RoPE 之前**，
  作用在每个 `[head_dim]` 的 head 向量上。
- xtrain 复用既有的 2D `rms_norm` op：把 `[seq,nh,hd]` reshape 成 `[seq*nh, hd]`，用 `[hd]` 的 γ
  做 rms_norm，再 reshape 回去，**插在 reshape 与 rope 之间**——与 xserv 逐步对齐。autograd 全自动
  （rms_norm/reshape 都已有 backward），优化器/checkpoint/DDP 都按 `params()` 泛化迭代，自动兼容。
- 每 block 新增 `q_norm`/`k_norm` 两个 `[head_dim]` leaf；`params()` 顺序在 `wv` 之后、`wo` 之前插入
  `q_norm,k_norm`；`num_params()` 加 `2*head_dim/layer`；PyTorch 对拍参考（parity.py / adamw_parity.py）
  同步加 QK-norm，名单同步——改完仍全套自洽。

这样导出的权重是**真·Qwen3 兼容**的（训练时就带 QK-norm），不是凑出来的假象。

### Decision 2：BF16 是 xserv 的硬约束 → 闭环判据用「贪心 token 一致」而非 bit-exact

xserv 的 Qwen3 前向 **只支持 BF16**（embedding/rmsnorm/gemm/silu/rope kernel 全部
`assert_eq!(dtype, BF16)`，`dump-logits` 也按 bf16 读 logits）。xtrain 是 f32。所以：

- 导出时把所有权重 `bf16::from_f32` 转成 BF16 写入 safetensors。
- **数值上不可能 bit-exact**：BF16 仅 8 位尾数，权重舍入 + 前向全程 BF16 累加，相对 xtrain 的 f32
  会有 ~1e-2 量级的 logits 漂移。因此**闭环判据定为「同 prompt 下贪心 argmax token 序列一致」**
  （+ 报告 logits 的 top-1/分布吻合度），这是 BF16 推理对 f32 训练能给出的最强、且诚实的判据。

### Decision 3：导出复用 train.rs 的 config + checkpoint，零猜测

导出器 `bin/export_safetensors.rs` 用与 `bin/train.rs` **完全相同**的 `Config`（gpt2 vocab、
`n_layers=4`、`Config::tiny()` 其余字段）建空模型 → `checkpoint::load_into` 灌入训练权重 →
按 `params()` 稳定序映射。tokenizer.json 直接 copy 进导出目录，两侧用同一份 BPE。

## 张量名 + layout 映射表

xtrain `params()` 序（T9 后）：
`embed[vocab,dim]` → 每 block `[attn_norm, wq, wk, wv, q_norm, k_norm, wo, ffn_norm, w_gate, w_up, w_down]`
→ `final_norm[dim]` → `lm_head[dim,vocab]`。

| xtrain 参数 | shape | → HF Qwen3 名 | HF shape | 操作 |
|-------------|-------|---------------|----------|------|
| `embed` | `[vocab,dim]` | `model.embed_tokens.weight` | `[vocab,dim]` | keep（行索引两侧同）|
| `attn_norm` | `[dim]` | `model.layers.{i}.input_layernorm.weight` | `[dim]` | keep |
| `wq` | `[dim,dim]` | `…self_attn.q_proj.weight` | `[dim,dim]` | **transpose** |
| `wk` | `[dim,dim]` | `…self_attn.k_proj.weight` | `[dim,dim]` | **transpose** |
| `wv` | `[dim,dim]` | `…self_attn.v_proj.weight` | `[dim,dim]` | **transpose** |
| `q_norm` | `[head_dim]` | `…self_attn.q_norm.weight` | `[head_dim]` | keep |
| `k_norm` | `[head_dim]` | `…self_attn.k_norm.weight` | `[head_dim]` | keep |
| `wo` | `[dim,dim]` | `…self_attn.o_proj.weight` | `[dim,dim]` | **transpose** |
| `ffn_norm` | `[dim]` | `…post_attention_layernorm.weight` | `[dim]` | keep |
| `w_gate` | `[dim,ffn]` | `…mlp.gate_proj.weight` | `[ffn,dim]` | **transpose** |
| `w_up` | `[dim,ffn]` | `…mlp.up_proj.weight` | `[ffn,dim]` | **transpose** |
| `w_down` | `[ffn,dim]` | `…mlp.down_proj.weight` | `[dim,ffn]` | **transpose** |
| `final_norm` | `[dim]` | `model.norm.weight` | `[dim]` | keep |
| `lm_head` | `[dim,vocab]` | `lm_head.weight` | `[vocab,dim]` | **transpose** |

全部以 **BF16** dtype 写入。config.json 字段：`architectures=["Qwen3ForCausalLM"]`、`model_type="qwen3"`、
`vocab_size`、`hidden_size=dim`、`intermediate_size=ffn`、`num_hidden_layers`、`num_attention_heads`、
`num_key_value_heads=num_attention_heads`、`head_dim`、`rms_norm_eps=eps`、`rope_theta`、
`tie_word_embeddings=false`、`attention_bias=false`、`hidden_act="silu"`。

## Module Layout

```
crates/xtrain-model/src/model.rs            # +q_norm/k_norm leaf；attention 插 per-head QK-norm；params() 序更新
crates/xtrain-model/src/config.rs           # num_params() 计入 QK-norm γ
crates/xtrain-train/src/bin/export_safetensors.rs   # 导出器（本 Phase 核心实现）
crates/xtrain-train/Cargo.toml              # +half, +safetensors="0.5"（对齐 xserv）, +bin
crates/xtrain-model/tests/parity{.py,_dump.rs}      # PyTorch 对拍同步加 QK-norm
crates/xtrain-train/tests/adamw_parity{.py,_dump.rs}# 同上
```

## 验证方法

1. **本地**：`cargo check --workspace` + `cargo fmt --all -- --check` 过（导出器 GPU 体 gated 在
   `not(no_cuda)`，host 侧只 check）。
2. **dash5（闭环）**：
   ```bash
   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
   # ① 训练一个小模型 → checkpoint（或复用已训的）
   cargo run -p xtrain-train --release --bin train -- \
       /opt/wjh/models/gpt2/tokenizer.json data/tinystories-valid-3mb.txt \
       <steps> /tmp/xtrain_tinystories.ckpt
   # ② 导出
   cargo run -p xtrain-train --release --bin export_safetensors -- \
       /tmp/xtrain_tinystories.ckpt /opt/wjh/models/gpt2/tokenizer.json /tmp/xtrain_export
   # ③ xserv 加载 + dump logits（同 prompt）
   #   (在 /opt/wjh/projects/xserv) cargo run -p xserv-model --release --bin dump-logits -- /tmp/xtrain_export "<prompt>"
   # ④ 对拍：xserv 贪心 token 序列 vs xtrain 自身贪心（sample.rs generate, temp=0）
   ```
   判据：**贪心 token 序列一致**（BF16 推理 vs f32 训练，logits top-1 吻合；分布在 BF16 容差内）。

## 验证结果（dash5 实跑，capture）

**训练**（CUDA_VISIBLE_DEVICES=0，1200 步，gpt2 vocab，dim 32 / 4 层 / 2 头 / ffn 64，~5M 参，8700 tok/s）：
`loss 10.84 → 3.59`，贪心采样输出连贯英文（QK-norm 加入后训练/采样无回归）。

**导出**：`export: 47 tensors`（embed + 4×11 + final_norm + lm_head），写出 config.json（见上）+
model.safetensors（BF16，6.5 MB）+ tokenizer.json。

**① logits 数值对拍**（同 prompt `"Once upon a time"`，token ids `[7454, 2402, 257, 640]`）：

| rank | xtrain (f32) | xserv (BF16) | token |
|------|--------------|--------------|-------|
| 0 | 11.7711 | 11.7500 | `,` |
| 1 | 10.4724 | 10.5000 | ` there` |
| 2 | 6.6288 | 6.6562 | ` upon` |
| 3 | 6.5125 | 6.5000 | ` to` |
| … | … | … | … |
| 10 | 5.3614 | 5.3438 | ` she` |

top-1 一致（`,`，id 11）；top-11 token 排序完全一致；logit 绝对差 ~1e-2（~0.2–0.9% 相对），
正是 **BF16 推理 vs f32 训练** 的预期舍入漂移，无结构性误差。

**② 贪心生成逐 token 一致**（xserv `xserv-cli --max-tokens 40` vs xtrain `sample.rs generate temp=0`）：

```
prompt "Once upon a time":
  xtrain: Once upon a time, there was a little girl named Lily. Timmy loved to play
          with her mommy. One day, Timmy's mommy's mommy's mommy. "I'm sorry, I
  xserv:  Once upon a time, there was a little girl named Lily. Timmy loved to play
          with her mommy. One day, Timmy's mommy's mommy's mommy. "I'm sorry, I   ← 逐 token 相同

prompt "One day":
  两侧均: One day, Timmy's mommy's mommy's mommy. "I'm sorry, I can't be careful and
          be careful. I'm sorry, I can't have a good time.                         ← 逐 token 相同
```

**结论：闭环成立**。xtrain 练出的权重，经导出后由 xserv 加载并服务，贪心生成与 xtrain 自身**逐 token 一致**，
logits 在 BF16 容差内吻合。整条 P0→P6 学习链收口。

> 注：xtrain 采样每步重跑全量 forward（无 KV cache），xserv 走 KV-cache prefill+decode；两者都是对同一
> logits 的 greedy argmax，故序列一致。BF16 漂移未在 40 步内造成任何 argmax 翻转。