gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

model: add per-head QK-norm (Qwen3-compat) for xserv export

Browse Source 
xserv's Qwen3 forward unconditionally applies per-head RMSNorm to Q and K
(q_norm/k_norm, shape [head_dim]) before RoPE — even gamma=1 is a real RMS
divide, not identity. xtrain never had this, so an exact xserv<->xtrain loop
was structurally impossible. Add it (reusing the 2D rms_norm op on the
[seq*nh, hd] head rows, inserted between reshape and rope to mirror
qwen3.rs's order) so the trained model is genuinely Qwen3-compatible.

params() inserts q_norm,k_norm after wv; num_params() counts them; the
PyTorch parity refs (parity.py / adamw_parity.py) + their name lists add the
same step so the dumps stay self-consistent.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-15 17:33:19 +08:00
parent ad82e8bf92
commit 7a4f69e430
7 changed files with 38 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
crates/xtrain-model/src/model.rs
View File

@@ -13,6 +13,8 @@ struct Block {
wq: Var, // [dim, dim]
wk: Var, // [dim, dim]
wv: Var, // [dim, dim]
q_norm: Var, // [head_dim] — per-head QK-norm (Qwen3-style)
k_norm: Var, // [head_dim]
wo: Var, // [dim, dim]
ffn_norm: Var, // [dim]
w_gate: Var, // [dim, ffn_hidden]
@@ -52,6 +54,8 @@ impl TinyTransformer {
wq: mk(&[cfg.dim, cfg.dim]),
wk: mk(&[cfg.dim, cfg.dim]),
wv: mk(&[cfg.dim, cfg.dim]),
q_norm: mk(&[cfg.head_dim]),
k_norm: mk(&[cfg.head_dim]),
wo: mk(&[cfg.dim, cfg.dim]),
ffn_norm: mk(&[cfg.dim]),
w_gate: mk(&[cfg.dim, cfg.ffn_hidden]),
@@ -87,6 +91,8 @@ impl TinyTransformer {
b.wq.clone(),
b.wk.clone(),
b.wv.clone(),
b.q_norm.clone(),
b.k_norm.clone(),
b.wo.clone(),
b.ffn_norm.clone(),
b.w_gate.clone(),
@@ -136,22 +142,29 @@ impl TinyTransformer {
// Project, then lay out as per-head [seq, head_dim] tensors.
// [seq,dim] @ [dim,dim] = [seq,dim]
// reshape [seq, nh, hd]
// qk-norm per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
// rope (kernel expects exactly [tokens, heads, head_dim])
// transpose [nh, seq, hd] → split into nh × [seq, hd]
let to_heads = |proj: Var, rope: bool| -> Vec<Var> {
let to_heads = |proj: Var, norm: Option<&Var>| -> Vec<Var> {
let r = ops::reshape(&proj, &[seq, nh, hd]);
let r = if rope {
ops::rope(&r, self.cfg.rope_theta)
} else {
r
let r = match norm {
// Per-head RMSNorm: flatten the (seq,nh) head rows, norm over hd,
// restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
Some(gamma) => {
let flat = ops::reshape(&r, &[seq * nh, hd]);
let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
let r = ops::reshape(&normed, &[seq, nh, hd]);
ops::rope(&r, self.cfg.rope_theta)
}
None => r,
};
let t = ops::transpose_3d01(&r); // [nh, seq, hd]
ops::split_heads(&t)
};
let q = to_heads(ops::matmul(x, &b.wq), true);
let k = to_heads(ops::matmul(x, &b.wk), true);
let v = to_heads(ops::matmul(x, &b.wv), false);
let q = to_heads(ops::matmul(x, &b.wq), Some(&b.q_norm));
let k = to_heads(ops::matmul(x, &b.wk), Some(&b.k_norm));
let v = to_heads(ops::matmul(x, &b.wv), None);
// Per-head scaled-dot-product attention with causal mask.
let heads_out: Vec<Var> = (0..nh)