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:
@@ -42,6 +42,7 @@ impl Config {
|
||||
/// Total learnable parameter count (for logging / sanity).
|
||||
pub fn num_params(&self) -> usize {
|
||||
let per_layer = 2 * self.dim // 2 rmsnorm gammas
|
||||
+ 2 * self.head_dim // q/k per-head norm gammas
|
||||
+ 3 * self.dim * self.dim // q/k/v proj
|
||||
+ self.dim * self.dim // out proj
|
||||
+ 2 * self.dim * self.ffn_hidden // gate/up proj
|
||||
|
||||
Reference in New Issue
Block a user