Update modeling_neollm.py

modeling_neollm.py

@@ -280,6 +280,8 @@ class MLPAnalysis:
     """
     Internals of a NeoLLMMLP forward pass.
     SwiGLU-like: down_proj(dropout(PolyNorm(gate_proj(fan)) · up_proj(fan)))
+    When use_versatile_ffn=True, the versatile sub-object carries all
+    VersatileFFN internals and fan/gate/up/polynorm fields remain None.
     """
     fan:              Optional[FANAnalysis] = None       # FAN components for MLP
     gate_proj_output: Optional[torch.Tensor] = None     # gate_proj(x_fan) [B,S,I]
@@ -288,6 +290,37 @@ class MLPAnalysis:
     act_times_up:     Optional[torch.Tensor] = None     # PolyNorm(gate)·up [B,S,I]
     output:           Optional[torch.Tensor] = None     # after down_proj [B,S,D]
 
+    # ── VersatileFFN (conditional on use_versatile_ffn) ──────────────────
+    versatile: Optional["VersatileFFNAnalysis"] = None  # None when standard MLP
+
+
+@dataclass
+class VersatileFFNAnalysis:
+    """
+    Internals of a VersatileFFN forward pass.
+    Only populated when use_versatile_ffn=True.
+
+    Reference: Nie et al. (2026). arXiv:2512.14531.
+
+    depth_probs:    Gumbel-Softmax distribution p [B, S, max_depth] — only in
+                    training (hard=True STE). None during inference.
+    expected_loops: E[L] per token [B, S] — differentiable proxy for difficulty.
+                    During inference computed from discrete argmax+1.
+    moe_weight:     λ = (L_max - E[L]) / L_max [B, S] — fusion gate scalar.
+                    Near 1 → width dominates (easy token).
+                    Near 0 → depth dominates (hard token).
+    loop_choice:    argmax(depth_logits) [B, S] — discrete depth selected at
+                    inference. None during training.
+    x_depth:        Output of the depth-versatile path [B, S, D].
+    x_width:        Output of the width-versatile MoE path [B, S, D].
+    """
+    depth_probs:    Optional[torch.Tensor] = None  # [B, S, max_depth] training only
+    expected_loops: Optional[torch.Tensor] = None  # [B, S]
+    moe_weight:     Optional[torch.Tensor] = None  # [B, S]
+    loop_choice:    Optional[torch.Tensor] = None  # [B, S] inference only
+    x_depth:        Optional[torch.Tensor] = None  # [B, S, D]
+    x_width:        Optional[torch.Tensor] = None  # [B, S, D]
+
 
 @dataclass
 class JTokMAnalysis:
@@ -1911,8 +1944,8 @@ def _apply_repo_rope(
     n_groups   = H // H_kv
     rotary_dim = inv_freq.shape[0] * 2   # inv_freq covers half the rotary dim
 
-    # inv_freq arrives from rotary_emb at forward time via repo_rope_args —
-    # already float32 on the correct device, no .to() needed, no DeviceCopy op.
+    # inv_freq is already float32 on the correct device (registered as buffer
+    # via set_repo_inv_freq) — no .to() needed, no DeviceCopy op.
     # No autocast barrier: explicit .float() casts on z_q/z_k are sufficient
     # to maintain float32 precision for the trig ops. Removing the context
     # manager lets Inductor plan all intermediate tensors as part of a single
@@ -2182,9 +2215,37 @@ class NeoLLMAttention(nn.Module):
                 d_p=_d_p,
                 num_heads=config.num_attention_heads,
             )
+            # _repo_inv_freq is registered as a non-persistent buffer by
+            # set_repo_inv_freq(), called from NeoLLMModel.__init__ after
+            # rotary_emb is built. Declaring it here would conflict.
+            self._repo_attn_scaling: float = 1.0
         else:
             self.repo_module = None
 
+    def set_repo_inv_freq(
+        self,
+        inv_freq: torch.Tensor,
+        attention_scaling: float,
+    ) -> None:
+        """
+        Inject the rotary frequency vector from NeoLLMRotaryEmbedding so that
+        REPO can build cos/sin inline from continuous positions.
+
+        Called once by NeoLLMModel.__init__ after rotary_emb is constructed.
+        Only has effect when use_repo=True for this layer.
+
+        Args:
+            inv_freq:          [rotary_dim/2] — frozen inv_freq buffer from
+                               NeoLLMRotaryEmbedding.
+            attention_scaling: float — attention_scaling from the same module.
+        """
+        if self.use_repo:
+            # Register as non-persistent buffer so .to(device) / .cuda() moves
+            # it automatically — eliminates the DeviceCopy op that splits the
+            # CUDAGraph into 2 partitions when _apply_repo_rope runs.
+            self.register_buffer("_repo_inv_freq", inv_freq.float(), persistent=False)
+            self._repo_attn_scaling = attention_scaling
+
     def _apply_momentum_attention(
         self,
         q: torch.Tensor,
@@ -2320,7 +2381,6 @@ class NeoLLMAttention(nn.Module):
         attention_mask: Optional[torch.Tensor] = None,
         first_layer_fan: Optional[torch.Tensor] = None,
         attn_analysis: Optional[AttentionAnalysis] = None,
-        repo_rope_args: Optional[Tuple[torch.Tensor, float]] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
         input_shape = hidden_states.shape[:-1]
@@ -2358,14 +2418,14 @@ class NeoLLMAttention(nn.Module):
             # REPO path: f_ϕ predicts continuous per-head positions from the
             # residual stream, then cos/sin are built inline from those positions
             # so the rotation is differentiable w.r.t. REPOModule parameters.
-            # inv_freq and attention_scaling arrive via repo_rope_args, sourced
-            # directly from rotary_emb at forward time — no buffer on this module,
-            # no meta-tensor issue on lm_eval / to(device) paths.
             # (Li et al., 2026, §3.2 — Eq. 6–7)
             repo_a = attn_analysis.repo if attn_analysis is not None else None
             z = self.repo_module(hidden_states, repo_analysis=repo_a)  # [B, H, S]
-            inv_freq, attn_scaling = repo_rope_args
-            q, k = _apply_repo_rope(q, k, z, inv_freq, attn_scaling)
+            q, k = _apply_repo_rope(
+                q, k, z,
+                self._repo_inv_freq,
+                self._repo_attn_scaling,
+            )
         else:
             # Standard path: integer positions pre-computed by NeoLLMModel.
             q, k = apply_rotary_pos_emb(q, k, cos, sin)
@@ -2502,18 +2562,21 @@ class PolyNorm(nn.Module):
         eps: float = 1e-6,
         proj_eps: float = 1e-6,
         exclusive_init: float = 0.5,
+        exclusive: bool = True,
     ):
         super().__init__()
-        self.weight   = nn.Parameter(torch.ones(3) / 3)
-        self.bias     = nn.Parameter(torch.zeros(1))
-        self.eps      = eps
-        self.proj_eps = proj_eps
-
-        # Dos fuerzas exclusivas aprendibles en (0, 1), una por rama de orden alto.
-        # Se parametrizan con logits para que sigmoid mantenga alpha ∈ (0, 1).
-        exclusive_init = float(min(max(exclusive_init, 1e-4), 1.0 - 1e-4))
-        init = torch.full((2,), exclusive_init, dtype=torch.float32)
-        self.exclusive_logits = nn.Parameter(torch.logit(init))
+        self.weight    = nn.Parameter(torch.ones(3) / 3)
+        self.bias      = nn.Parameter(torch.zeros(1))
+        self.eps       = eps
+        self.exclusive = exclusive
+
+        if exclusive:
+            self.proj_eps = proj_eps
+            # Dos fuerzas exclusivas aprendibles en (0, 1), una por rama de orden alto.
+            # Se parametrizan con logits para que sigmoid mantenga alpha ∈ (0, 1).
+            exclusive_init = float(min(max(exclusive_init, 1e-4), 1.0 - 1e-4))
+            init = torch.full((2,), exclusive_init, dtype=torch.float32)
+            self.exclusive_logits = nn.Parameter(torch.logit(init))
 
     def _norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
@@ -2543,7 +2606,7 @@ class PolyNorm(nn.Module):
         x_sq = x.pow(2)
         x_cu = x * x_sq
 
-        # Tres ramas normalizadas (same as original PolyNorm)
+        # Tres ramas normalizadas
         x1 = x    * x_sq.mean(-1, keepdim=True).add(self.eps).rsqrt()
         x2 = x_sq * (x_sq * x_sq).mean(-1, keepdim=True).add(self.eps).rsqrt()
         x3 = x_cu * (x_cu * x_cu).mean(-1, keepdim=True).add(self.eps).rsqrt()
@@ -2553,26 +2616,31 @@ class PolyNorm(nn.Module):
             analysis.x2_pre_exclusive = x2.detach()
             analysis.x3_pre_exclusive = x3.detach()
 
-        # Fuerzas exclusivas aprendibles
-        alpha2, alpha3 = torch.sigmoid(self.exclusive_logits).unbind()
+        if self.exclusive:
+            # Fuerzas exclusivas aprendibles
+            alpha2, alpha3 = torch.sigmoid(self.exclusive_logits).unbind()
 
-        if analysis is not None:
-            analysis.alpha2  = alpha2.detach()
-            analysis.alpha3  = alpha3.detach()
-            analysis.weights = self.weight.detach()
-            analysis.bias    = self.bias.detach()
+            if analysis is not None:
+                analysis.alpha2  = alpha2.detach()
+                analysis.alpha3  = alpha3.detach()
+                analysis.weights = self.weight.detach()
+                analysis.bias    = self.bias.detach()
 
-        # Precalcular ref (x1) en fp32 y su norma al cuadrado — compartido por x2 y x3
-        x1_f        = x1.float()
-        ref_norm_sq = x1_f.pow(2).sum(-1, keepdim=True).clamp_min(self.proj_eps)
+            # Precalcular ref (x1) en fp32 y su norma al cuadrado — compartido por x2 y x3
+            x1_f        = x1.float()
+            ref_norm_sq = x1_f.pow(2).sum(-1, keepdim=True).clamp_min(self.proj_eps)
 
-        # Ortogonalización parcial de las ramas de orden alto respecto a la lineal
-        x2 = self._exclusive(x2, x1, alpha2, x1_f, ref_norm_sq)
-        x3 = self._exclusive(x3, x1, alpha3, x1_f, ref_norm_sq)
+            # Ortogonalización parcial de las ramas de orden alto respecto a la lineal
+            x2 = self._exclusive(x2, x1, alpha2, x1_f, ref_norm_sq)
+            x3 = self._exclusive(x3, x1, alpha3, x1_f, ref_norm_sq)
 
-        if analysis is not None:
-            analysis.x2_post_exclusive = x2.detach()
-            analysis.x3_post_exclusive = x3.detach()
+            if analysis is not None:
+                analysis.x2_post_exclusive = x2.detach()
+                analysis.x3_post_exclusive = x3.detach()
+        else:
+            if analysis is not None:
+                analysis.weights = self.weight.detach()
+                analysis.bias    = self.bias.detach()
 
         output = (
             self.weight[0] * x3
@@ -2587,6 +2655,379 @@ class PolyNorm(nn.Module):
         return output
 
 
+def compute_versatile_aux_loss(
+    aux_stats_list: list,
+    n_experts: int,
+    weight: float,
+) -> torch.Tensor:
+    """
+    Load-balancing auxiliary loss for VersatileFFN width-path expert routing.
+
+    Identical formula to JTok-M:  L_aux = λ · n_e · Σ_i p_i · f_i
+    averaged over all decoder layers with active VersatileFFN.
+
+    Args:
+        aux_stats_list: list of (p_sum [n_e], f_sum [n_e], N_tokens) per layer.
+        n_experts: total number of virtual experts.
+        weight:    λ coefficient.
+    Returns:
+        Scalar loss tensor.
+    """
+    total_loss = None
+    for p_sum, f_sum, N in aux_stats_list:
+        p_i        = p_sum / N
+        layer_loss = weight * n_experts * (p_i * f_sum).sum()
+        total_loss = layer_loss if total_loss is None else total_loss + layer_loss
+    if total_loss is None:
+        return torch.tensor(0.0)
+    return total_loss / len(aux_stats_list)
+
+
+class VersatileFFN(nn.Module):
+    """
+    VersatileFFN: dual-process feed-forward network with parameter reuse.
+
+    Drop-in replacement for NeoLLMMLP.  Shares the same weight matrices but
+    reuses them across two complementary paths (Nie et al., 2026):
+
+    Width-Versatile path (virtual MoE, Eq. 7–8):
+        Derives N virtual sub-experts by slicing non-overlapping contiguous
+        segments of the intermediate dimension.  A Top-K router selects
+        ``active_experts`` experts per token.  No additional parameters beyond
+        the router ``expert_gate [hidden, N_experts]``.
+
+    Depth-Versatile path (recursive, Eq. 9–11):
+        Applies the full shared MLP (inner SeeDNorm + FANLayer + gate/up/down)
+        recursively up to ``max_depth`` times.  A Gumbel-Softmax loop predictor
+        (``depth_predictor [hidden, max_depth]``) decides per-token depth.
+        During training: always L_max iterations with soft STE weighting.
+        During inference: early exit at argmax(p) to save FLOPs.
+
+    Difficulty-aware fusion (Eq. 12–13):
+        λ = (L_max − E[L]) / L_max ∈ [0, 1)
+        output = λ · Y_width + (1 − λ) · Y_depth
+        Easy tokens (low E[L] → λ → 1) favour the fast width path.
+        Hard tokens (high E[L] → λ → 0) favour the deep recursive path.
+
+    NeoLLM-specific adaptations vs. the OLMo reference implementation:
+        - FANLayer runs once per forward (shared between both paths).
+        - Expert slicing uses contiguous segments (no SwiGLU half-shift)
+          because gate and up are separate projections here.
+        - PolyNorm (shared parameters) is used as the activation in both paths.
+        - Multipliers from LinearWithMultipliers are applied per-slice in the
+          width path to correctly replicate the full MLP forward.
+        - Gumbel temperature stored as a persistent float32 buffer so it
+          survives checkpoint save/load without external tracking.
+        - Width path load-balancing returns (output, aux_stats) for integration
+          with the existing NeoLLMForCausalLM aux-loss accumulation pattern.
+
+    Reference:
+        Nie et al. (2026). "VersatileFFN: Achieving Parameter Efficiency in
+        LLMs via Adaptive Wide-and-Deep Reuse." arXiv:2512.14531.
+    """
+
+    def __init__(self, config: NeoLLMConfig):
+        super().__init__()
+
+        self.total_experts     = getattr(config, "versatile_total_experts",   8)
+        self.active_experts    = getattr(config, "versatile_active_experts",  2)
+        self.max_depth         = getattr(config, "versatile_max_depth",       4)
+        self.hidden_size       = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        # ── Shared MLP weights (identical layout to NeoLLMMLP) ──────────────
+        fan_ratio = getattr(config, "fan_ratio_ffn", 0.0625)
+        fan_dim   = config.hidden_size + int(config.hidden_size * fan_ratio)
+
+        self.fan_layer = FANLayer(hidden_size=config.hidden_size, fan_ratio=fan_ratio)
+        self.gate_proj = LinearWithMultipliers(
+            fan_dim, config.intermediate_size,
+            bias=False, use_row_multiplier=True, use_column_multiplier=False,
+        )
+        self.up_proj   = nn.Linear(fan_dim, config.intermediate_size, bias=False)
+        self.down_proj = LinearWithMultipliers(
+            config.intermediate_size, config.hidden_size,
+            bias=False, use_row_multiplier=True, use_column_multiplier=True,
+        )
+        self.act_fn  = PolyNorm(exclusive=False)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+        # ── Inner normalization for depth-recursive steps ────────��───────────
+        # Applied before each recursive MLP application (mirrors the paper's
+        # ff_norm inside each depth loop).
+        self.ff_norm = SeeDNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # ── Width path: expert router + contiguous segment indices ──────────
+        seg = config.intermediate_size // self.total_experts
+        self.expert_segment = seg
+        self.expert_gate    = nn.Linear(config.hidden_size, self.total_experts, bias=False)
+
+        # Static expert slice indices [total_experts, seg] — non-persistent
+        # buffer so .to(device) moves them automatically.
+        idx_list = [
+            torch.arange(i * seg, (i + 1) * seg)
+            for i in range(self.total_experts)
+        ]
+        self.register_buffer("expert_idx", torch.stack(idx_list), persistent=False)
+
+        # ── Depth path: loop count predictor ────────────────────────────────
+        self.depth_predictor = nn.Linear(config.hidden_size, self.max_depth, bias=False)
+
+        # Gumbel temperature as a persistent scalar buffer.
+        # Decays externally via update_gumbel_temperature() each training step.
+        temp_start = getattr(config, "versatile_gumbel_temp_start", 5.0)
+        self.register_buffer(
+            "gumbel_temp",
+            torch.tensor(temp_start, dtype=torch.float32),
+            persistent=True,
+        )
+        self._gumbel_temp_end   = getattr(config, "versatile_gumbel_temp_end",   0.1)
+        self._gumbel_temp_decay = getattr(config, "versatile_gumbel_temp_decay", 0.99984)
+
+    # ── Public API ────────────────────────────────────────────────────────────
+
+    def update_gumbel_temperature(self) -> None:
+        """Decay Gumbel temperature by one step.  Call once per training step."""
+        new_t = max(
+            self.gumbel_temp.item() * self._gumbel_temp_decay,
+            self._gumbel_temp_end,
+        )
+        self.gumbel_temp.fill_(new_t)
+
+    # ── Private helpers ───────────────────────────────────────────────────────
+
+    def _expert_forward(
+        self,
+        x_fan: torch.Tensor,      # [N, fan_dim]
+        x_in:  torch.Tensor,      # [N, hidden_size] — residual base
+        idx:   torch.Tensor,      # [seg] — which intermediate neurons to use
+    ) -> torch.Tensor:
+        """
+        Virtual expert forward for a single expert (Eq. 6–8 of paper).
+
+        Slices gate_proj, up_proj, down_proj weights to the expert segment.
+        Multipliers from LinearWithMultipliers are applied per-slice to keep
+        the computation exactly equivalent to the full MLP forward.
+
+        Returns x_in + expert_out  (residual included, matching Eq. 8).
+        """
+        # Gate projection: sliced rows + row multiplier
+        g_w   = self.gate_proj.linear.weight[idx]                      # [seg, fan_dim]
+        g_mul = self.gate_proj.row_multiplier.multiplier[idx]          # [seg]
+        gate  = F.linear(x_fan, g_w) * g_mul                          # [N, seg]
+
+        # Up projection: sliced rows, no multiplier on up_proj
+        u_w = self.up_proj.weight[idx]                                 # [seg, fan_dim]
+        up  = F.linear(x_fan, u_w)                                    # [N, seg]
+
+        # Activation — PolyNorm is shape-agnostic (weight [3], bias [1])
+        act = self.act_fn(gate) * up                                   # [N, seg]
+        act = self.dropout(act)
+
+        # Down projection: column multiplier on input, then row on output
+        col_mul = self.down_proj.column_multiplier.multiplier[idx]     # [seg]
+        d_w     = self.down_proj.linear.weight[:, idx]                 # [hidden, seg]
+        row_mul = self.down_proj.row_multiplier.multiplier             # [hidden]
+        out     = F.linear(act * col_mul, d_w) * row_mul              # [N, hidden]
+
+        return x_in + out
+
+    def _full_forward_step(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        One recursive MLP application for the depth path (Eq. 9).
+
+        Includes inner SeeDNorm + FANLayer so each iteration starts from a
+        normalized, periodicity-augmented view of the current hidden state.
+        Residual connection applied inside to match the paper's formulation.
+        """
+        og   = x
+        x    = self.ff_norm(x)
+        x_f  = self.fan_layer(x)
+        gate = self.gate_proj(x_f)
+        up   = self.up_proj(x_f)
+        act  = self.act_fn(gate) * up
+        act  = self.dropout(act)
+        return og + self.down_proj(act)
+
+    # ── Forward ─────────────────────────────────────────────────────────────���─
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        analysis: Optional[MLPAnalysis] = None,
+    ) -> Tuple[torch.Tensor, Optional[tuple]]:
+        """
+        Args:
+            x:        [B, S, hidden_size] — pre-normalized input
+                      (SeeDNorm + LNS already applied by NeoLLMDecoderLayer).
+            analysis: MLPAnalysis with .versatile pre-allocated when analysis
+                      mode is active.  None during training (zero overhead).
+
+        Returns:
+            (output [B, S, hidden_size], aux_stats)
+            aux_stats = (p_sum [N_experts], f_sum [N_experts], N_tokens)
+            for load-balancing loss, or None in inference if width path skipped.
+        """
+        B, S, D = x.shape
+
+        # Depth predictor reads difficulty from the pre-normalized hidden state
+        depth_logits = self.depth_predictor(x)    # [B, S, max_depth]
+
+        # FANLayer runs once — output shared between width and depth paths
+        x_fan = self.fan_layer(x)                 # [B, S, fan_dim]
+
+        # ═════════════════════ TRAINING ══════════════════════════════════════
+        if self.training:
+            # Gumbel-Softmax with hard=True (STE): discrete forward, continuous
+            # backward. Annealed temperature controls sharpness of the selection.
+            depth_probs = F.gumbel_softmax(
+                depth_logits,
+                tau=float(self.gumbel_temp),
+                hard=True,
+                dim=-1,
+            )  # [B, S, max_depth]
+
+            # ── Depth path: always L_max iterations (static graph for compile) ─
+            depth_outputs = []
+            current_x     = x
+            for _ in range(self.max_depth):
+                current_x = self._full_forward_step(current_x)
+                depth_outputs.append(current_x)
+
+            # Soft weighted combination — gradient flows through depth_probs
+            depth_stack = torch.stack(depth_outputs, dim=-1)          # [B,S,D,L]
+            x_depth     = (depth_stack * depth_probs.unsqueeze(2)).sum(dim=-1)  # [B,S,D]
+
+            # ── Width path: Top-K routing over virtual experts ────────────────
+            routing_logits = self.expert_gate(x)                       # [B,S,N]
+            topk_w, topk_i = torch.topk(routing_logits, k=self.active_experts, dim=-1)
+            topk_w         = torch.softmax(topk_w, dim=-1)            # [B,S,k]
+
+            x_flat     = x.reshape(-1, D)                             # [N,D]
+            x_fan_flat = x_fan.reshape(-1, x_fan.shape[-1])          # [N,fan]
+            topk_i_f   = topk_i.reshape(-1, self.active_experts)     # [N,k]
+            topk_w_f   = topk_w.reshape(-1, self.active_experts)     # [N,k]
+            N_tok      = x_flat.shape[0]
+
+            x_moe_flat = torch.zeros_like(x_flat)
+
+            for eid in range(self.total_experts):
+                mask           = (topk_i_f == eid)
+                tok_idx, k_idx = torch.where(mask)
+                if tok_idx.numel() == 0:
+                    continue
+                w_e   = topk_w_f[tok_idx, k_idx].unsqueeze(-1)
+                out_e = self._expert_forward(
+                    x_fan_flat[tok_idx], x_flat[tok_idx], self.expert_idx[eid]
+                )
+                x_moe_flat.index_add_(
+                    0, tok_idx, (out_e * w_e).to(x_moe_flat.dtype)
+                )
+
+            x_moe = x_moe_flat.reshape(B, S, D)
+
+            # Load-balancing aux stats (same pattern as JTok-M)
+            r_probs_flat = torch.softmax(
+                routing_logits.reshape(-1, self.total_experts), dim=-1
+            )  # [N_tok, N_experts]
+            p_sum = r_probs_flat.sum(dim=0)                           # [N_experts]
+            f_counts = torch.zeros(
+                self.total_experts, device=x.device, dtype=x.dtype
+            )
+            for eid in range(self.total_experts):
+                f_counts[eid] = (topk_i_f == eid).float().sum()
+            f_sum     = f_counts / (N_tok * self.active_experts)      # [N_experts]
+            aux_stats = (p_sum, f_sum, N_tok)
+
+            # ── Difficulty-aware fusion (Eq. 12–13) ──────────────────────────
+            loop_idx   = torch.arange(
+                1, self.max_depth + 1, device=x.device, dtype=depth_probs.dtype
+            )
+            expected_L = (depth_probs * loop_idx).sum(dim=-1)         # [B, S]
+            moe_weight = (self.max_depth - expected_L) / self.max_depth  # [B, S]
+            output     = (
+                x_depth * (1.0 - moe_weight.unsqueeze(-1))
+                + x_moe * moe_weight.unsqueeze(-1)
+            )
+            loop_choice = None   # not used during training
+
+        # ═════════════════════ INFERENCE ══════════════════════════════════════
+        else:
+            loop_choice = depth_logits.argmax(dim=-1)                  # [B, S]
+            max_loop    = int(loop_choice.max().item())
+
+            # Depth path: early exit — only compute needed iterations
+            depth_outputs = []
+            current_x     = x
+            for _ in range(max_loop + 1):
+                current_x = self._full_forward_step(current_x)
+                depth_outputs.append(current_x)
+
+            depth_stack = torch.stack(depth_outputs, dim=-1)          # [B,S,D,run]
+            gather_idx  = (
+                loop_choice.unsqueeze(-1).unsqueeze(-1).expand(B, S, D, 1)
+            )
+            x_depth = depth_stack.gather(dim=-1, index=gather_idx).squeeze(-1)
+
+            # Fusion weight from discrete choice
+            expected_L = (loop_choice + 1).float()                    # [B, S]
+            moe_weight = (self.max_depth - expected_L) / self.max_depth  # [B, S]
+
+            # Width path: conditional on λ > 0 (Conditional Parallelism)
+            active_mask = (moe_weight > 1e-6)                         # [B, S]
+            x_moe       = torch.zeros_like(x)
+            aux_stats   = None
+            depth_probs = None
+
+            if active_mask.any():
+                x_flat_all     = x.reshape(-1, D)
+                x_fan_flat_all = x_fan.reshape(-1, x_fan.shape[-1])
+                active_flat    = active_mask.reshape(-1)
+                x_active       = x_flat_all[active_flat]
+                x_fan_active   = x_fan_flat_all[active_flat]
+
+                r_log = self.expert_gate(x_active)                    # [Na, N]
+                tw, ti = torch.topk(r_log, k=self.active_experts, dim=-1)
+                tw     = torch.softmax(tw, dim=-1)
+
+                x_moe_active = torch.zeros_like(x_active)
+                for eid in range(self.total_experts):
+                    mask_e         = (ti == eid)
+                    tok_idx, k_idx = torch.where(mask_e)
+                    if tok_idx.numel() == 0:
+                        continue
+                    w_e   = tw[tok_idx, k_idx].unsqueeze(-1)
+                    out_e = self._expert_forward(
+                        x_fan_active[tok_idx], x_active[tok_idx], self.expert_idx[eid]
+                    )
+                    x_moe_active.index_add_(
+                        0, tok_idx, (out_e * w_e).to(x_moe_active.dtype)
+                    )
+
+                x_moe_flat              = x_moe.reshape(-1, D)
+                x_moe_flat[active_flat] = x_moe_active
+                x_moe                   = x_moe_flat.reshape(B, S, D)
+
+            output = (
+                x_depth * (1.0 - moe_weight.unsqueeze(-1))
+                + x_moe * moe_weight.unsqueeze(-1)
+            )
+
+        # ── Analysis deposits ─────────────────────────────────────────────────
+        if analysis is not None:
+            if analysis.versatile is not None:
+                va              = analysis.versatile
+                va.depth_probs  = depth_probs.detach() if depth_probs is not None else None
+                va.expected_loops = expected_L.detach()
+                va.moe_weight   = moe_weight.detach()
+                va.loop_choice  = loop_choice.detach() if loop_choice is not None else None
+                va.x_depth      = x_depth.detach()
+                va.x_width      = x_moe.detach()
+            analysis.output = output.detach()
+
+        return output, aux_stats
+
+
 class NeoLLMMLP(nn.Module):
     """MLP with FANformer integration and Learnable Multipliers."""
 
@@ -2608,7 +3049,7 @@ class NeoLLMMLP(nn.Module):
             config.intermediate_size, config.hidden_size,
             bias=False, use_row_multiplier=True, use_column_multiplier=True,
         )
-        self.act_fn  = PolyNorm(exclusive_init=0.05)
+        self.act_fn  = PolyNorm(exclusive_init=0.00, exclusive=getattr(config, "polynorm_exclusive", True))
         self.dropout = nn.Dropout(config.dropout_rate)
 
     def forward(
@@ -2664,7 +3105,12 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         self.use_jtokm     = config.use_jtokm
 
         self.self_attn                = NeoLLMAttention(config, layer_idx)
-        self.mlp                      = NeoLLMMLP(config)
+        self.mlp                      = (
+            VersatileFFN(config)
+            if getattr(config, "use_versatile_ffn", False)
+            else NeoLLMMLP(config)
+        )
+        self.use_versatile_ffn        = getattr(config, "use_versatile_ffn", False)
         self.input_layernorm          = SeeDNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = SeeDNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.lns_attn                 = LNS(layer_idx)
@@ -2746,7 +3192,6 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         attn_res_partial: Optional[torch.Tensor] = None,
         layer_analysis: Optional[LayerAnalysis] = None,
         output_attentions: Optional[bool] = False,
-        repo_rope_args: Optional[Tuple[torch.Tensor, float]] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple:
         # ── Snapshot input ────────────────────────────────────────────────
@@ -2784,7 +3229,6 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             position_embeddings=position_embeddings,
             first_layer_fan=first_layer_fan,
             attn_analysis=layer_analysis.attention if layer_analysis is not None else None,
-            repo_rope_args=repo_rope_args,
             **kwargs,
         )
 
@@ -2819,7 +3263,11 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             layer_analysis.lns_mlp_output = h_lns2.detach()
 
         mlp_a   = layer_analysis.mlp if layer_analysis is not None else None
-        delta_m = self.mlp(h_lns2, analysis=mlp_a)
+        if self.use_versatile_ffn:
+            delta_m, versatile_aux = self.mlp(h_lns2, analysis=mlp_a)
+        else:
+            delta_m        = self.mlp(h_lns2, analysis=mlp_a)
+            versatile_aux  = None
 
         if layer_analysis is not None:
             layer_analysis.mlp_contribution = delta_m.detach()
@@ -2850,6 +3298,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             outputs += (attn_weights,)
         if aux_stats is not None:
             outputs += (aux_stats,)
+        if versatile_aux is not None:
+            outputs += (versatile_aux,)
         return outputs
 
 
@@ -3304,6 +3754,17 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
 
         self.post_init()
 
+        # ── REPO: inject inv_freq into every attention layer that uses it ─────
+        # Done after post_init so rotary_emb.inv_freq is already initialized.
+        # Layers below repo_start_layer never call set_repo_inv_freq (their
+        # use_repo flag is False) so the call is harmless for those layers.
+        if getattr(config, "use_repo", False):
+            for layer in self.layers:
+                layer.self_attn.set_repo_inv_freq(
+                    self.rotary_emb.inv_freq,
+                    self.rotary_emb.attention_scaling,
+                )
+
     def get_input_embeddings(self):
         if self.config.use_token_generator:
             return self.token_generator
@@ -3329,6 +3790,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             getattr(cfg, "use_repo", False)
             and layer_idx >= getattr(cfg, "repo_start_layer", cfg.num_hidden_layers // 3)
         )
+        _versatile = getattr(cfg, "use_versatile_ffn", False)
         return LayerAnalysis(
             seednorm_pre_attn  = SeeDNormAnalysis(),
             seednorm_post_attn = SeeDNormAnalysis(),
@@ -3338,8 +3800,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 repo     = REPOAnalysis() if _repo_active else None,
             ),
             mlp                = MLPAnalysis(
-                fan      = FANAnalysis(),
-                polynorm = PolyNormAnalysis(),
+                fan      = FANAnalysis() if not _versatile else None,
+                polynorm = PolyNormAnalysis() if not _versatile else None,
+                versatile = VersatileFFNAnalysis() if _versatile else None,
             ),
             gpas_attn = GPASAnalysis(),
             gpas_mlp  = GPASAnalysis(),
@@ -3435,17 +3898,6 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
         self.first_layer_fan = None
 
-        # ── REPO: pass inv_freq by reference at forward time ──────────────────
-        # rotary_emb.inv_freq is already on the correct device (managed by
-        # NeoLLMRotaryEmbedding as a buffer) — no .to(), no DeviceCopy op.
-        # Computed once here and passed through the decoder layer chain so
-        # NeoLLMAttention never needs to store it as a buffer itself, avoiding
-        # the meta-tensor issue that occurs when lm_eval calls .to(device).
-        repo_rope_args = (
-            (self.rotary_emb.inv_freq, self.rotary_emb.attention_scaling)
-            if getattr(self.config, "use_repo", False) else None
-        )
-
         # ── Attention Residuals state ──────────────────────────────────────
         # Full AttnRes (attn_res_num_blocks=0): sources grows by one entry per
         # decoder layer — all previous outputs are kept, max N=num_layers+1.
@@ -3506,7 +3958,6 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 attn_res_partial=attn_res_partial if use_attn_res else None,
                 layer_analysis=layer_analysis,
                 output_attentions=output_attentions,
-                repo_rope_args=repo_rope_args,
                 **kwargs,
             )
             hidden_states = layer_outputs[0]
@@ -3522,6 +3973,15 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             if self.config.use_jtokm and len(layer_outputs) > (2 if output_attentions else 1):
                 all_aux_stats.append(layer_outputs[-1])
 
+            # Collect VersatileFFN aux stats (second-to-last if jtokm also present,
+            # or last if jtokm is absent). Only non-None during training.
+            if getattr(self.config, "use_versatile_ffn", False):
+                for item in layer_outputs[1:]:
+                    if isinstance(item, tuple) and len(item) == 3:
+                        # (p_sum, f_sum, N_tokens) signature
+                        all_aux_stats.append(("versatile", item))
+                        break
+
             if (self.first_layer_fan is None
                     and hasattr(decoder_layer, "current_layer_fan")):
                 self.first_layer_fan = decoder_layer.current_layer_fan
@@ -3766,12 +4226,30 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             )
             # Add JTok-M load-balancing auxiliary loss
             if self.config.use_jtokm and all_aux_stats:
-                aux_loss = compute_jtokm_aux_loss(
-                    all_aux_stats,
-                    n_e=self.config.jtokm_num_experts,
-                    weight=self.config.jtokm_aux_loss_weight,
-                )
-                loss = loss + aux_loss
+                jtokm_stats = [
+                    s for s in all_aux_stats
+                    if not (isinstance(s, tuple) and len(s) == 2 and s[0] == "versatile")
+                ]
+                if jtokm_stats:
+                    aux_loss = compute_jtokm_aux_loss(
+                        jtokm_stats,
+                        n_e=self.config.jtokm_num_experts,
+                        weight=self.config.jtokm_aux_loss_weight,
+                    )
+                    loss = loss + aux_loss
+            # Add VersatileFFN load-balancing auxiliary loss
+            if getattr(self.config, "use_versatile_ffn", False) and all_aux_stats:
+                versatile_stats = [
+                    s[1] for s in all_aux_stats
+                    if isinstance(s, tuple) and len(s) == 2 and s[0] == "versatile"
+                ]
+                if versatile_stats:
+                    v_loss = compute_versatile_aux_loss(
+                        versatile_stats,
+                        n_experts=self.config.versatile_total_experts,
+                        weight=self.config.versatile_aux_loss_weight,
+                    )
+                    loss = loss + v_loss
             logits = None
         else:
             slice_indices = (
@@ -3812,6 +4290,8 @@ __all__ = [
     "MEAHeadSeeDNorm",
     "HadamardOProj",
     "REPOModule",
+    "VersatileFFN",
+    "compute_versatile_aux_loss",
     # Analysis dataclasses — exported so external tools can type-hint against them
     "AnalysisState",
     "LayerAnalysis",
@@ -3823,6 +4303,7 @@ __all__ = [
     "PolyNormAnalysis",
     "HadamardAnalysis",
     "REPOAnalysis",
+    "VersatileFFNAnalysis",
     "JTokMAnalysis",
     "AttnResAnalysis",
     "GeneratorAnalysis",