diff --git a/docs/08-export-xserv.md b/docs/08-export-xserv.md new file mode 100644 index 0000000..066a40e --- /dev/null +++ b/docs/08-export-xserv.md @@ -0,0 +1,138 @@ +# 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 \ + /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 "" + # ④ 对拍:xserv 贪心 token 序列 vs xtrain 自身贪心(sample.rs generate, temp=0) + ``` + 判据:**贪心 token 序列一致**(BF16 推理 vs f32 训练,logits top-1 吻合;分布在 BF16 容差内)。 + +## 验证结果 + +> 待 dash5 实跑回填(push origin main → dash5 pull → 跑 ②③④ → capture)。