Merge pull request #34 from MotifTechnologies/fix/qk-clip-rotary-full-gamma

torch-ext/optimizer/muon.py

@@ -479,6 +479,18 @@ class Muon(torch.optim.Optimizer):
                             p, qk_clip_state, n
                         ) if qk_clip_state is not None else None
                         if scales_full is not None:
+                            # scales_full is raw per-head γ; apply the content
+                            # exponent per head (MHA √γ, wq_b γ^(1-r), wkv_b
+                            # γ^r). NOTE: this per-head form cannot resolve the
+                            # MLA rotary/value sub-regions (q_pe → full γ, v
+                            # untouched) — see qk_clip; the sharded path scales
+                            # q_pe/v by the content exponent instead.
+                            if not qk_clip_state.is_mla:
+                                e = 0.5
+                            elif qk_clip_state.kind == 'wq_b':
+                                e = 1.0 - qk_clip_state.k_ratio
+                            else:
+                                e = qk_clip_state.k_ratio
                             ratio = p.shape[0] // scales_full.shape[0]
                             idx0 = local_indices[0]
                             if isinstance(idx0, slice):
@@ -486,7 +498,7 @@ class Muon(torch.optim.Optimizer):
                                 idx0 = torch.arange(start,
                                                     idx0.stop,
                                                     device=scales_full.device)
-                            row_scales = scales_full[idx0 // ratio]
+                            row_scales = (scales_full**e)[idx0 // ratio]
                             p._local_tensor.mul_(row_scales.view(-1, 1))
 
                 if not batch_params:
@@ -558,6 +570,18 @@ class Muon(torch.optim.Optimizer):
                             p, qk_clip_state, n
                         ) if qk_clip_state is not None else None
                         if scales_full is not None:
+                            # scales_full is raw per-head γ; apply the content
+                            # exponent per head (MHA √γ, wq_b γ^(1-r), wkv_b
+                            # γ^r). NOTE: this per-head form cannot resolve the
+                            # MLA rotary/value sub-regions (q_pe → full γ, v
+                            # untouched) — see qk_clip; the sharded path scales
+                            # q_pe/v by the content exponent instead.
+                            if not qk_clip_state.is_mla:
+                                e = 0.5
+                            elif qk_clip_state.kind == 'wq_b':
+                                e = 1.0 - qk_clip_state.k_ratio
+                            else:
+                                e = qk_clip_state.k_ratio
                             ratio = p.shape[0] // scales_full.shape[0]
                             idx0 = local_indices[0]
                             if isinstance(idx0, slice):
@@ -565,7 +589,7 @@ class Muon(torch.optim.Optimizer):
                                 idx0 = torch.arange(start,
                                                     idx0.stop,
                                                     device=scales_full.device)
-                            row_scales = scales_full[idx0 // ratio]
+                            row_scales = (scales_full**e)[idx0 // ratio]
                             p._local_tensor.mul_(row_scales.view(-1, 1))
 
     def _setup_parallel(self, names, params, group, qk_logits):

torch-ext/optimizer/qk_clip.py

@@ -1,5 +1,4 @@
 import logging
-import math
 from dataclasses import dataclass
 
 import torch
@@ -127,17 +126,16 @@ def get_qk_clip_info(clip_config, n, qk_logits):
 
 
 def compute_scales(p, qk_clip_state, name: str = ''):
-    """Compute per-head scaling factors for QK clipping.
+    """Compute the per-head raw clip factor γ = threshold / max_logit.
 
-    Returns scales tensor per head if any head exceeds threshold, else None.
+    Returns a per-head tensor of γ (1.0 for heads under the threshold) if any
+    head requires clipping, else None. The per-region exponent on γ is NOT
+    applied here — it is applied at multiply time (``qk_clip`` for the whole
+    tensor, or the sharded path in muon.py) so that one per-head γ can drive
+    different sub-regions (q_nope, q_pe, k_nope, v) with different exponents.
 
-    Scale convention (let γ = threshold/value):
-      MHA/GQA wq/wk : scale = √γ on both Q and K; product = γ.
-      MLA wq_b      : scale = γ^(1 - k_ratio).  Q carries (1-k_ratio) of the
-                       log-reduction; k_ratio=0 → full γ on Q only (K unclipped).
-      MLA wkv_b     : scale = γ^k_ratio applied to k_nope rows only (v rows
-                       untouched). Returns None when k_ratio=0 (K unclipped).
-    Together wq_b × wkv_b deliver γ^(1-k_ratio) × γ^k_ratio = γ end-to-end.
+    wkv_b with k_ratio=0 returns None: the whole reduction is placed on Q and
+    the key projection is left unclipped.
     """
     kind = qk_clip_state.kind
 
@@ -150,32 +148,18 @@ def compute_scales(p, qk_clip_state, name: str = ''):
     threshold = qk_clip_state.threshold
     logit = qk_clip_state.logit
 
-    # Exponent for (threshold/v_ele)^exponent:
-    #   MLA wq_b  → 1 - k_ratio   (Q carries this fraction of log-reduction)
-    #   MLA wkv_b → k_ratio        (K carries the rest; applied to k_nope rows only)
-    #   MHA/GQA   → 0.5            (equal split; equivalent to k_ratio=0.5)
-    #
-    # Principled k_ratio for Q:K = G:1 GQA: 1 / (1 + sqrt(G))
-    #   K's gradient is a sum of G Q-head terms that partially cancel, so its
-    #   effective recovery signal scales as 1/sqrt(G) (random-walk argument).
-    #   Setting clip/recovery equal for both sides:
-    #     (1-r)/1 = r/(1/sqrt(G))  =>  r = 1/(1+sqrt(G))
-    #   G=1 → 0.5 (recovers MHA), G=4 → ≈0.33, G=8 → ≈0.26.
-    if qk_clip_state.is_mla:
-        exponent = (1.0 - qk_clip_state.k_ratio) if kind == 'wq_b' else qk_clip_state.k_ratio
-    else:
-        exponent = 0.5
-
     head_scales = {}
     for logit_idx, head_idx in enumerate(indices):
         v_ele = float(logit[logit_idx])
         if v_ele > threshold:
-            new_scale = (threshold / v_ele) ** exponent
-            if head_idx not in head_scales or new_scale < head_scales[head_idx]:
-                head_scales[head_idx] = new_scale
+            gamma = threshold / v_ele  # raw clip factor; exponent applied later
+            # Within a shared (GQA) group, keep the smallest γ, i.e. the
+            # strongest required clipping.
+            if head_idx not in head_scales or gamma < head_scales[head_idx]:
+                head_scales[head_idx] = gamma
                 logger.info(
                     f"[{kind}] {name} Head {head_idx} exceeded threshold "
-                    f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}"
+                    f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> γ={gamma:.4f}"
                 )
 
     if not head_scales:
@@ -195,36 +179,40 @@ def compute_scales(p, qk_clip_state, name: str = ''):
 
 
 def qk_clip(p, scales, info):
-    """Apply per-head scaling to a Q/K projection weight matrix.
+    """Apply QK-clip to a Q/K projection weight; each sub-region is raised to
+    its own exponent so every logit component is reduced by exactly γ.
 
     Args:
         p: Parameter (nn.Parameter or raw tensor).
-        scales: [n_heads] tensor of per-head scale factors (γ^exponent).
-            Exponent semantics depend on kind — see compute_scales docstring.
-        info: QKClipInfo with kind, head_dim, and MLA sub-head dimensions.
-
-    MLA wq_b: γ^(1-k_ratio) applied uniformly to q_nope and q_pe rows.
-    MLA wkv_b: γ^k_ratio applied to k_nope rows only; v rows untouched.
-    MHA/GQA: √γ applied to all rows in each head.
+        scales: [n_heads] tensor of per-head raw γ (= threshold/value), 1.0 for
+            unclipped heads, as returned by ``compute_scales``.
+        info: QKClipInfo with kind, head_dim, k_ratio and MLA sub-head dims.
+
+    Exponents:
+        MHA        : √γ on all rows.
+        MLA q_nope : γ^(1-k_ratio),  k_nope : γ^k_ratio   (content product = γ).
+        MLA q_pe   : γ (full) — the decoupled rotary key is shared and left
+                     unchanged, so the rotary query carries the whole γ for
+                     every k_ratio; γ^(1-k_ratio) here would under-clip rotary.
+        MLA v rows : untouched.
     """
     W = p.data if isinstance(p, torch.nn.Parameter) else p
+    r = info.k_ratio
 
     if not info.is_mla:
-        # MHA/GQA: uniform √γ applied to all rows in each head
-        W.view(-1, info.head_dim, W.shape[1]).mul_(scales.view(-1, 1, 1))
+        W.view(-1, info.head_dim, W.shape[1]).mul_((scales**0.5).view(-1, 1, 1))
         return
 
-    # MLA: vectorized sub-region scaling within each head
     if info.kind == 'wq_b':
         qk_nope = info.qk_nope_head_dim
         qk_head_dim = qk_nope + info.qk_rope_head_dim
         W_3d = W.view(-1, qk_head_dim, W.shape[1])  # [H, qk_head_dim, in_dim]
-        W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1))  # q_nope → γ^(1-k_ratio)
-        W_3d[:, qk_nope:, :].mul_(scales.view(-1, 1, 1))  # q_pe   → γ^(1-k_ratio)
+        W_3d[:, :qk_nope, :].mul_((scales**(1.0 - r)).view(-1, 1, 1))  # q_nope → γ^(1-r)
+        W_3d[:, qk_nope:, :].mul_(scales.view(-1, 1, 1))  # q_pe → γ (full)
 
     elif info.kind == 'wkv_b':
         qk_nope = info.qk_nope_head_dim
         kv_stride = qk_nope + info.v_head_dim
         W_3d = W.view(-1, kv_stride, W.shape[1])  # [H, kv_stride, in_dim]
-        W_3d[:, :qk_nope, :].mul_(scales.view(-1, 1, 1))  # k_nope → γ^k_ratio
-        # v rows: not touched (k_R shared rotary unchanged)
+        W_3d[:, :qk_nope, :].mul_((scales**r).view(-1, 1, 1))  # k_nope → γ^r
+        # v rows: untouched (shared rotary key k_rope is unchanged)