Update architecture and tokenizer

config.json

@@ -21,7 +21,7 @@
   "rope_mode": "main_sequence",
   "tie_word_embeddings": false,
   "training_sequence_length": 1024,
-  "transformers_version": "5.12.1",
+  "transformers_version": "5.14.1",
   "use_cache": true,
   "use_residual_gate": true,
   "vocab_size": 50277,

huggingface.py

@@ -44,6 +44,7 @@ from torch import nn
 from torch.nn.attention.flex_attention import create_block_mask
 from torch.nn.attention.flex_attention import flex_attention
 import torch.nn.functional as F
+from entmax import entmax_bisect
 
 
 
@@ -671,15 +672,14 @@ class MoSRAHCache(CacheLayerMixin):
             "Use get_heads_lengths() for per-head occupancy."
         )
 
-    def get_max_cache_shape(self) -> int:  # type: ignore[override]
-        """Return the static per-(batch, head) slot capacity of this cache.
-
-        Equal to mosrah_cache_length as supplied at construction, which is derived
-        from config.mosrah_cache_length. Required by the HuggingFace static cache
-        contract; generation machinery uses this to size attention masks.
-        """
+    def get_max_length(self) -> int:
+        """Return the static per-(batch, head) slot capacity of this cache."""
         return self.mosrah_cache_length
 
+    def get_max_cache_shape(self) -> int:  # type: ignore[override]
+        """Compatibility alias for the deprecated cache-shape interface."""
+        return self.get_max_length()
+
     def get_mask_sizes(  # type: ignore[override]
         self,
         cache_position: torch.Tensor,
@@ -987,9 +987,13 @@ class RouterCache(CacheLayerMixin):
         """Not supported — RouterCache tracks block position, not sequence length."""
         raise NotImplementedError("RouterCache does not track sequence length.")
 
+    def get_max_length(self) -> int:
+        """Return an undefined maximum because no token sequence is stored."""
+        return -1
+
     def get_max_cache_shape(self) -> int:
-        """Not supported — RouterCache does not hold KV pairs."""
-        raise NotImplementedError("RouterCache does not have a KV cache shape.")
+        """Compatibility alias for the deprecated cache-shape interface."""
+        return self.get_max_length()
 
     def get_mask_sizes(  # type: ignore[override]
         self,
@@ -1251,9 +1255,14 @@ class LocalSlidingWindowLayerCache(CacheLayerMixin):
         """
         return int(self._total_processed)
 
-    def get_max_cache_shape(self) -> int:
+    def get_max_length(self) -> int:
+        """Return the retained local sliding-window capacity."""
         return self.sliding_window
 
+    def get_max_cache_shape(self) -> int:
+        """Compatibility alias for the deprecated cache-shape interface."""
+        return self.get_max_length()
+
     def get_mask_sizes(  # type: ignore[override]
         self,
         cache_position: torch.Tensor,
@@ -1515,16 +1524,18 @@ class ShramLayerCache(CacheLayerMixin):
             "Update sliding_window_cache or mosrah_cache directly."
         )
 
-    def get_max_cache_shape(self) -> int:  # type: ignore[override]
+    def get_max_length(self) -> int:
         """Return the maximum sequence length this layer cache can serve.
 
         The authoritative upper bound is ``config.inference_sequence_length``, which
         governs the full accumulated token history the model is configured to handle.
-        HuggingFace's static-cache machinery reads this value to determine whether the
-        cache is compileable and to size generation loops.
         """
         return self._inference_sequence_length
 
+    def get_max_cache_shape(self) -> int:  # type: ignore[override]
+        """Compatibility alias for the deprecated cache-shape interface."""
+        return self.get_max_length()
+
     def get_mask_sizes(  # type: ignore[override]
         self,
         cache_position: torch.Tensor,
@@ -2378,6 +2389,54 @@ accumulated cached state returned by that cache.
 
 
 
+# -----------
+# Inlined from: initialization.py
+# -----------
+"""Initialization policy for SHRAM raw projection parameters.
+
+Hugging Face initializes ordinary ``nn.Linear`` and ``nn.Embedding`` modules.
+SHRAM also stores two custom raw-parameter forms that Hugging Face cannot infer:
+
+- a three-dimensional bank of independent expert linear maps
+- a two-dimensional router projection
+
+The expert-bank leading dimension is ownership/storage, not fan geometry. Each
+expert matrix must therefore receive Xavier initialization independently. The
+router follows the proven BALANCE initialization scale instead. Tensor rank is
+used here as the explicit representation contract for these two owners, not as
+a general initialization heuristic. Adding another raw projection shape must
+therefore extend this contract deliberately rather than inheriting a fallback.
+"""
+
+
+
+
+
+ROUTER_INIT_STD = 0.02
+
+
+@torch.no_grad()
+def initialize_projection_parameter(parameter: nn.Parameter) -> None:
+    """Initialize one of SHRAM's supported raw projection representations.
+
+    A rank-three tensor is an independent bank of linear matrices and is
+    initialized matrix-by-matrix with Xavier uniform. A rank-two tensor is the
+    router projection and receives the BALANCE-normal initialization.
+    """
+    if parameter.ndim == 3:
+        for matrix in parameter.unbind(dim=0):
+            nn.init.xavier_uniform_(matrix)
+        return
+
+    if parameter.ndim == 2:
+        nn.init.normal_(parameter, mean=0.0, std=ROUTER_INIT_STD)
+        return
+
+    raise ValueError(
+        "SHRAM raw projection parameters must be rank two (router) or rank "
+        f"three (independent expert bank), got shape {tuple(parameter.shape)}."
+    )
+
 
 
 
@@ -2527,9 +2586,13 @@ class BottleneckedEnsembleAttention(nn.Module):
         return torch.einsum("bltu,lud->bltd", attended_states, self.o_proj)
 
     def _reset_parameters(self) -> None:
-        """Initialize per-head projection weights."""
+        """Initialize projection banks independently of storage rank.
+
+        These tensors store independent per-head linear maps. Their leading
+        head dimension is storage organization rather than fan geometry.
+        """
         for weight in (self.q_proj, self.k_proj, self.v_proj, self.o_proj):
-            nn.init.xavier_uniform_(weight)
+            initialize_projection_parameter(weight)
 
     def _validate_tensor_shape(self, packed_embeddings: torch.Tensor) -> None:
         """Validate the local packed-embedding shape contract required by BEA."""
@@ -2979,8 +3042,10 @@ input hidden state x, the router produces two outputs used downstream:
   - selected_heads (I): which K of the L available expert heads each token
     routes to, determined by a block-balanced causal solver.
   - routing_probs (P): the weights used for the weighted output reduction,
-    gathered from the softmax routing scores at the selected indices and
-    renormalized to sum to 1 per token.
+    gathered from the Entmax routing distribution at the selected indices.
+    Tokens with positive selected mass are renormalized to sum to one; tokens
+    whose mechanically assigned experts all have zero Entmax mass receive an
+    all-zero mixture and therefore no sparse-path contribution.
 
 Routing uses a single learnable projection:
 
@@ -3009,6 +3074,11 @@ Paper ref: Appendix A.Routing.
 
 
 
+
+
+ENTMAX_ALPHA = 1.2
+
+
 class MoSRAHRouter(nn.Module):
     """Token-choice router for MoSRAH sparse attention.
 
@@ -3017,8 +3087,8 @@ class MoSRAHRouter(nn.Module):
     consecutive tokens every expert is used exactly once, giving perfect load
     balance by construction.
 
-    routing_weight is nn.Parameter rather than nn.Linear so that HuggingFace
-    _init_weights does not override its kaiming initialization at construction.
+    routing_weight is an ``nn.Parameter`` rather than ``nn.Linear`` because the
+    router owns a custom raw projection boundary and its initialization.
 
     Attributes:
         routing_weight: Shape (L, embedding_width). Maps input hidden states to
@@ -3034,16 +3104,16 @@ class MoSRAHRouter(nn.Module):
 
     def __init__(self, config: ShramConfig) -> None:
         super().__init__()
-        self.num_mosrah_heads   = config.num_mosrah_heads
+        self.num_mosrah_heads = config.num_mosrah_heads
         self.num_selected_heads = config.num_selected_heads
-        self.block_length       = config.block_length
+        self.block_length = config.block_length
+        self.entmax_alpha = ENTMAX_ALPHA
 
         # Routing projection: maps input (B, N, d) to per-head routing scores (B, N, L).
-        # nn.Parameter ensures HuggingFace _init_weights does not override kaiming init.
         self.routing_weight = nn.Parameter(
             torch.empty(config.num_mosrah_heads, config.embedding_width)
         )
-        nn.init.kaiming_normal_(self.routing_weight)
+        initialize_projection_parameter(self.routing_weight)
 
     def forward(
         self,
@@ -3062,16 +3132,17 @@ class MoSRAHRouter(nn.Module):
         Returns:
             selected_heads: Head indices I of shape (batch, seq_len, num_selected_heads).
                 Each token's K selected head indices from the block-balanced solver.
-            routing_probs: Routing probabilities P of shape (batch, seq_len,
-                num_selected_heads). Gathered from the pre-balance softmax at
-                selected_heads and renormalized to sum to 1 per token.
+            routing_probs: Routing weights P of shape (batch, seq_len,
+                num_selected_heads). Selected positive Entmax mass is normalized
+                to one. An all-zero selected support remains all-zero, allowing
+                the token to make no sparse-path contribution.
             router_diagnostics: Dict of routing scalars:
                 - ``regret_loss``: gradient-carrying mean regret, mean of
                   max(p_max_active − p_chosen, 0) over live (B, num_blocks, L)
                   entries. In [0, 1]. Zero when every expert is assigned at its
                   peak-preference token within the block.
                 - ``logit_regret``: detached logit-space regret; same formula
-                  applied to routing logits rather than softmax probabilities.
+                  applied to routing logits rather than Entmax probabilities.
                   In [0, ∞). Monitoring only.
                 - ``logit_std``: detached mean per-token std of routing logits.
         """
@@ -3101,13 +3172,18 @@ class MoSRAHRouter(nn.Module):
 
         # ── Phase: pre-balance scoring ─────────────────────────────────────────
         #
-        # Establishes the clean routing distribution before any -inf masking.
-        # logit_std is captured here because the block solver's masking would
-        # corrupt the standard deviation. routing_scores is used both for
-        # regret_loss and for the final routing_probs.
+        # Establish the clean routing distribution before any -inf masking.
+        # Entmax and selected-weight normalization are computed in fp32 by
+        # numeric policy. Only the final route weights return to model dtype.
+        # The regret objective also retains the fp32 distribution.
         routing_logits = self._compute_routing_logits(x)                       # (B, N, L)
-        logit_std      = routing_logits.std(dim=-1).mean().detach()
-        routing_scores = F.softmax(routing_logits, dim=-1)                     # (B, N, L)
+        routing_logits_fp32 = routing_logits.float()
+        logit_std = routing_logits_fp32.std(dim=-1).mean().detach()
+        routing_scores_fp32 = entmax_bisect(
+            routing_logits_fp32,
+            alpha=self.entmax_alpha,
+            dim=-1,
+        )                                                                       # (B, N, L), fp32
 
         # ── Phase: block-balanced causal selection ─────────────────────────────
         #
@@ -3136,9 +3212,9 @@ class MoSRAHRouter(nn.Module):
             # token per "block" there is no alternative assignment with higher
             # preference.
             used_in_block = router_cache.get_used_in_block()                   # (B, L)
-            step_logits   = routing_logits[:, 0, :]                            # (B, L)
-            available     = step_logits.masked_fill(used_in_block, float('-inf'))
-            step_heads    = available.topk(K, dim=-1).indices                  # (B, K)
+            step_logits = routing_logits[:, 0, :]                               # (B, L)
+            available = step_logits.masked_fill(used_in_block, float("-inf"))
+            step_heads = available.topk(K, dim=-1).indices                      # (B, K)
 
             router_cache.update_decode(step_heads)
 
@@ -3152,8 +3228,8 @@ class MoSRAHRouter(nn.Module):
             # preceding them have already had their pick each step. The pad is
             # discarded after the solver.
             num_blocks = (N + W - 1) // W
-            N_pad      = num_blocks * W
-            pad_len    = N_pad - N
+            N_pad = num_blocks * W
+            pad_len = N_pad - N
 
             if pad_len > 0:
                 padded_logits = torch.cat(
@@ -3167,14 +3243,20 @@ class MoSRAHRouter(nn.Module):
             # used_in_block tracks which experts have been claimed within each block.
             # No gradient here — expert availability is a hard structural constraint,
             # not a differentiable quantity. Gradient flows through routing_probs.
-            used_in_block   = torch.zeros(B, num_blocks, L, dtype=torch.bool, device=x.device)
+            used_in_block = torch.zeros(
+                B,
+                num_blocks,
+                L,
+                dtype=torch.bool,
+                device=x.device,
+            )
             step_heads_list = []
 
             for step in range(W):
                 step_logits = blocked_logits[:, :, step, :]                    # (B, blk, L)
 
                 # Claimed experts receive -inf so top-K never selects them.
-                available  = step_logits.masked_fill(used_in_block, float('-inf'))
+                available = step_logits.masked_fill(used_in_block, float("-inf"))
                 step_heads = available.topk(K, dim=-1).indices                 # (B, blk, K)
                 step_heads_list.append(step_heads)
 
@@ -3183,7 +3265,7 @@ class MoSRAHRouter(nn.Module):
 
             # Stack W steps and reshape to (B, N_pad, K), then unpad.
             selected_heads_blocked = torch.stack(step_heads_list, dim=2)       # (B, blk, W, K)
-            selected_heads         = selected_heads_blocked.view(B, N_pad, K)[:, :N, :]  # (B, N, K)
+            selected_heads = selected_heads_blocked.view(B, N_pad, K)[:, :N, :]  # (B, N, K)
 
             if router_cache is not None:
                 # Prefill: persist the partial last-block state so decode steps
@@ -3203,27 +3285,38 @@ class MoSRAHRouter(nn.Module):
         # blocks, and a single decode step is not a complete block. Backward is
         # never called during inference so the zero is a correct no-op.
         if router_cache is not None and N == 1:
-            regret_loss  = routing_logits.new_zeros(())
-            logit_regret = routing_logits.new_zeros(()).detach()
+            regret_loss = routing_logits_fp32.new_zeros(())
+            logit_regret = routing_logits_fp32.new_zeros(()).detach()
         else:
             regret_loss, logit_regret = self._compute_regret(
-                routing_scores,
-                routing_logits,
+                routing_scores_fp32,
+                routing_logits_fp32,
                 selected_heads_blocked,
                 active_mask,
             )
 
-        # ── Phase: routing probabilities ────────────────────────────────────────
+        # ── Phase: routing weights ──────────────────────────────────────────────
         #
-        # Gathered from the pre-balance routing_scores to reflect genuine routing
-        # preference; renormalized so they sum to 1 per token.
-        gathered      = routing_scores.gather(dim=-1, index=selected_heads)    # (B, N, K)
-        routing_probs = gathered / gathered.sum(dim=-1, keepdim=True)          # (B, N, K)
+        # Gather from the clean fp32 Entmax distribution. If mechanical balance
+        # assigns only experts outside a token's sparse Entmax support, the token
+        # intentionally receives an all-zero mixture and contributes nothing on
+        # the sparse path. Positive selected mass is normalized in fp32.
+        gathered_fp32 = routing_scores_fp32.gather(                             # (B, N, K)
+            dim=-1,
+            index=selected_heads,
+        )
+        gathered_sum = gathered_fp32.sum(dim=-1, keepdim=True)
+        normalization_denominator = gathered_sum.clamp_min(
+            torch.finfo(gathered_fp32.dtype).tiny
+        )
+        routing_probs = (
+            gathered_fp32 / normalization_denominator
+        ).to(dtype=routing_logits.dtype)                                        # (B, N, K)
 
         router_diagnostics = {
-            "regret_loss":  regret_loss,
+            "regret_loss": regret_loss,
             "logit_regret": logit_regret,
-            "logit_std":    logit_std,
+            "logit_std": logit_std,
         }
         return selected_heads, routing_probs, router_diagnostics
 
@@ -3260,9 +3353,9 @@ class MoSRAHRouter(nn.Module):
         detached — it is a monitoring scalar only, in [0, ∞).
 
         Args:
-            routing_scores:         Softmax routing probabilities, shape (B, N, L).
+            routing_scores:         Entmax routing probabilities, shape (B, N, L).
                                     Gradient flows through this tensor into regret_loss.
-            routing_logits:         Pre-softmax routing logits, shape (B, N, L).
+            routing_logits:         Router logits in fp32, shape (B, N, L).
                                     Used only for the detached logit_regret.
             selected_heads_blocked: Expert assignments from the block solver,
                                     shape (B, num_blocks, W, K). Block geometry
@@ -3274,8 +3367,8 @@ class MoSRAHRouter(nn.Module):
             logit_regret:  Detached scalar in [0, ∞).
         """
         B, num_blocks, W, _K = selected_heads_blocked.shape
-        L   = routing_scores.shape[-1]
-        N   = routing_scores.shape[1]
+        L = routing_scores.shape[-1]
+        N = routing_scores.shape[1]
         N_pad = num_blocks * W
 
         # ── Reshape into block form ─────────────────────────────────────────
@@ -3311,23 +3404,25 @@ class MoSRAHRouter(nn.Module):
         assignment_mask.scatter_(dim=-1, index=selected_heads_blocked, value=1.0)
                                                                                # (B, nb, W, L)
 
-        # ── Prob regret (gradient flows through routing_scores) ─────────────
+        # ── Probability regret (gradient flows through routing_scores) ───────
         #
         # p_chosen: routing score at the assigned token, gated by active_float
         # so dead assignments contribute 0 — the expert accrues full regret
         # against the active maximum rather than no penalty.
         # p_max: peak routing score over active tokens; dead tokens zeroed before
-        # max (safe because softmax outputs are non-negative).
-        p_chosen = (assignment_mask * active_float.unsqueeze(-1) * scores_blocked).sum(dim=2)
-                                                                               # (B, nb, L)
-        p_max    = (active_float.unsqueeze(-1) * scores_blocked).max(dim=2).values
-                                                                               # (B, nb, L)
+        # max (safe because Entmax outputs are non-negative).
+        p_chosen = (
+            assignment_mask * active_float.unsqueeze(-1) * scores_blocked
+        ).sum(dim=2)                                                           # (B, nb, L)
+        p_max = (
+            active_float.unsqueeze(-1) * scores_blocked
+        ).max(dim=2).values                                                     # (B, nb, L)
 
-        regret = (p_max - p_chosen).clamp(min=0.0)                            # (B, nb, L)
+        regret = (p_max - p_chosen).clamp(min=0.0)                             # (B, nb, L)
 
         # Mean over live (B, num_blocks, L) entries. Clamped to 1 for the
         # all-dead edge case where the numerator is already 0.
-        num_live    = block_active.float().sum()                               # scalar
+        num_live = block_active.float().sum()                                  # scalar
         regret_loss = (
             block_active.float().unsqueeze(-1) * regret
         ).sum() / num_live.mul(L).clamp(min=1.0)
@@ -3343,12 +3438,13 @@ class MoSRAHRouter(nn.Module):
         ).sum(dim=2)                                                           # (B, nb, L)
 
         logit_max = logits_blocked.masked_fill(
-            ~active_blocked.unsqueeze(-1), float('-inf')
-        ).max(dim=2).values                                                    # (B, nb, L)
+            ~active_blocked.unsqueeze(-1), float("-inf")
+        ).max(dim=2).values                                                     # (B, nb, L)
         logit_max = logit_max.masked_fill(~block_active.unsqueeze(-1), 0.0)
 
         logit_regret = (
-            block_active.float().unsqueeze(-1) * (logit_max - logit_chosen).clamp(min=0.0)
+            block_active.float().unsqueeze(-1)
+            * (logit_max - logit_chosen).clamp(min=0.0)
         ).sum() / num_live.mul(L).clamp(min=1.0)
         logit_regret = logit_regret.detach()