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>
This commit is contained in:
@@ -2,9 +2,10 @@
|
||||
//!
|
||||
//! A from-scratch decoder built entirely from the [`xtrain_autodiff`] op set:
|
||||
//! token embedding → `n_layers` × {pre-RMSNorm → multi-head causal attention
|
||||
//! (RoPE) → residual; pre-RMSNorm → SwiGLU MLP → residual} → final RMSNorm →
|
||||
//! LM-head matmul. The forward builds an autograd graph; calling `.backward()`
|
||||
//! on the cross-entropy loss fills every parameter's `.grad()`.
|
||||
//! (per-head QK-norm + RoPE) → residual; pre-RMSNorm → SwiGLU MLP → residual} →
|
||||
//! final RMSNorm → LM-head matmul. The forward builds an autograd graph; calling
|
||||
//! `.backward()` on the cross-entropy loss fills every parameter's `.grad()`.
|
||||
//! Per-head QK-norm (Qwen3-style) makes the architecture xserv-compatible (T9).
|
||||
//!
|
||||
//! Conventions (matching the engine, not HuggingFace):
|
||||
//! - Linear weights are `[in, out]` and applied as `x @ W` (no transpose), since
|
||||
|
||||
Reference in New Issue
Block a user