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

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>

crates/xtrain-train/tests/adamw_parity.py

@@ -74,7 +74,7 @@ SEQ = len(ids)
 
 NAMES = ["embed"]
 for l in range(NL):
-    for p in ["attn_norm", "wq", "wk", "wv", "wo",
+    for p in ["attn_norm", "wq", "wk", "wv", "q_norm", "k_norm", "wo",
               "ffn_norm", "w_gate", "w_up", "w_down"]:
         NAMES.append(f"l{l}_{p}")
 NAMES += ["final_norm", "lm_head"]
@@ -115,6 +115,9 @@ def forward():
         q = (x @ P[f"l{l}_wq"]).reshape(SEQ, NH, HD)
         k = (x @ P[f"l{l}_wk"]).reshape(SEQ, NH, HD)
         v = (x @ P[f"l{l}_wv"]).reshape(SEQ, NH, HD)
+        # Per-head QK-norm (Qwen3-style), before RoPE.
+        q = rms_norm(q, P[f"l{l}_q_norm"])
+        k = rms_norm(k, P[f"l{l}_k_norm"])
         q = rope(q).transpose(0, 1)
         k = rope(k).transpose(0, 1)
         v = v.transpose(0, 1)

crates/xtrain-train/tests/adamw_parity_dump.rs

@@ -156,6 +156,8 @@ fn param_names(cfg: &Config) -> Vec<String> {
             "wq",
             "wk",
             "wv",
+            "q_norm",
+            "k_norm",
             "wo",
             "ffn_norm",
             "w_gate",