xtrain/docs/08-export-xserv.md

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（闭环）：

   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 翻转。