Update modeling_neollm.py

modeling_neollm.py

@@ -1878,77 +1878,16 @@ class SeeDNorm(nn.Module):
         return f"dim={self.dim}, eps={self.eps}"
 
 
-def _canonical_norm_type(value: str, *, allow_siamese: bool = False) -> str:
-    """Canonicalise NeoLLM normalisation names."""
-    raw = str(value).lower().replace("_", "").replace("-", "")
-    aliases = {
-        "rms": "rmsnorm",
-        "rmsnorm": "rmsnorm",
-        "seed": "seednorm",
-        "seednorm": "seednorm",
-        "seeednorm": "seednorm",
-        "siamese": "siamesenorm",
-        "siamesenorm": "siamesenorm",
-    }
-    out = aliases.get(raw)
-    valid = {"seednorm", "rmsnorm"} | ({"siamesenorm"} if allow_siamese else set())
-    if out not in valid:
-        raise ValueError(
-            f"Invalid normalization type {value!r}. Valid values are {sorted(valid)}."
-        )
-    return out
-
-
-def _uses_siamesenorm(config: NeoLLMConfig) -> bool:
-    norm_type = getattr(config, "norm_type", None)
-    if norm_type is None:
-        return bool(getattr(config, "use_siamesenorm", False))
-    return _canonical_norm_type(norm_type, allow_siamese=True) == "siamesenorm"
-
-
-def _active_base_norm_type(config: NeoLLMConfig) -> str:
-    """
-    Return the local normalizer used inside modules.
-
-    ``norm_type="siamesenorm"`` is a two-stream topology, not a scalar
-    normalization layer.  Its internal local normalizers are selected by
-    ``siamese_norm_base_type`` and default to RMSNorm, matching the released
-    SiameseNorm/OLMo configuration style.
+def _make_norm(dim: int, eps: float, use_seednorm: bool = True) -> nn.Module:
     """
-    norm_type = getattr(config, "norm_type", None)
-    if norm_type is None:
-        norm_type = (
-            "siamesenorm"
-            if bool(getattr(config, "use_siamesenorm", False))
-            else ("seednorm" if bool(getattr(config, "use_seednorm", True)) else "rmsnorm")
-        )
-    norm_type = _canonical_norm_type(norm_type, allow_siamese=True)
-    if norm_type == "siamesenorm":
-        return _canonical_norm_type(
-            getattr(config, "siamese_norm_base_type", "rmsnorm"),
-            allow_siamese=False,
-        )
-    return norm_type
-
+    Build the active Transformer normalisation module.
 
-def _make_norm(
-    dim: int,
-    eps: float,
-    use_seednorm: bool = True,
-    norm_type: Optional[str] = None,
-) -> nn.Module:
+    ``use_seednorm=True`` preserves the previous SeeDNorm path.
+    ``use_seednorm=False`` uses standard RMSNorm at the same sites, so the
+    module keeps the same pre-norm topology without SeeDNorm's dynamic
+    alpha/beta/gamma re-scaling parameters.
     """
-    Build a local normalization module.
-
-    ``norm_type="seednorm"`` instantiates SeeDNorm.
-    ``norm_type="rmsnorm"`` instantiates ``nn.RMSNorm``.
-    SiameseNorm is handled at the decoder-layer topology level and must not be
-    requested as a single local layer.
-    """
-    if norm_type is None:
-        norm_type = "seednorm" if use_seednorm else "rmsnorm"
-    norm_type = _canonical_norm_type(norm_type, allow_siamese=False)
-    if norm_type == "seednorm":
+    if use_seednorm:
         return SeeDNorm(dim, eps=eps)
     return nn.RMSNorm(dim, eps=eps)
 
@@ -3457,11 +3396,10 @@ class NeoLLMAttention(nn.Module):
                 enable_multipliers=self.use_learnable_multipliers,
             )
 
-        self.norm_type = _active_base_norm_type(config)
-        self.use_seednorm = self.norm_type == "seednorm"
+        self.use_seednorm = bool(getattr(config, "use_seednorm", True))
         self.mea_norm_type = str(getattr(config, "mea_norm_type", "seednorm")).lower()
-        self.q_norm = _make_norm(self.head_dim, eps=config.rms_norm_eps, norm_type=self.norm_type)
-        self.k_norm = _make_norm(self.head_dim, eps=config.rms_norm_eps, norm_type=self.norm_type)
+        self.q_norm = _make_norm(self.head_dim, eps=config.rms_norm_eps, use_seednorm=self.use_seednorm)
+        self.k_norm = _make_norm(self.head_dim, eps=config.rms_norm_eps, use_seednorm=self.use_seednorm)
 
         if self.use_mea_attention:
             self.mea_key_mix = nn.Parameter(
@@ -4993,7 +4931,7 @@ class VersatileFFN(nn.Module):
         self.ff_norm = _make_norm(
             config.hidden_size,
             eps=config.rms_norm_eps,
-            norm_type=_active_base_norm_type(config),
+            use_seednorm=bool(getattr(config, "use_seednorm", True)),
         )
 
         # ── Width path: expert router + contiguous segment indices ──────────
@@ -5520,10 +5458,7 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         self.hidden_size   = config.hidden_size
         self.layer_idx     = layer_idx
         self.use_jtokm     = config.use_jtokm
-        self.use_siamesenorm = _uses_siamesenorm(config)
-        self.norm_type     = _active_base_norm_type(config)
-        self.use_seednorm  = self.norm_type == "seednorm"
-        self.siamese_depth_scaling = bool(getattr(config, "siamese_depth_scaling", True))
+        self.use_seednorm  = bool(getattr(config, "use_seednorm", True))
         # Controls only the first pre-attention normalisation applied directly
         # to the embedding stream. Defaults to True for checkpoint/config
         # backward compatibility.  When False, layer 0 does not instantiate
@@ -5547,7 +5482,7 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             _make_norm(
                 config.hidden_size,
                 eps=config.rms_norm_eps,
-                norm_type=self.norm_type,
+                use_seednorm=self.use_seednorm,
             )
             if self.has_input_layernorm
             else None
@@ -5555,27 +5490,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         self.post_attention_layernorm = _make_norm(
             config.hidden_size,
             eps=config.rms_norm_eps,
-            norm_type=self.norm_type,
+            use_seednorm=self.use_seednorm,
         )
-        if self.use_siamesenorm:
-            # SiameseNorm best variant from the released source: the X stream
-            # uses the ordinary Hybrid/Post-like normalizers, while the Y
-            # stream has its own Pre-Norm normalizers.  The fused MLP input is
-            # normalized explicitly because the paper's ablation finds that
-            # normalized fused inputs are necessary.
-            self.siamese_attn_y_norm = _make_norm(
-                config.hidden_size, eps=config.rms_norm_eps, norm_type=self.norm_type
-            )
-            self.siamese_mlp_y_norm = _make_norm(
-                config.hidden_size, eps=config.rms_norm_eps, norm_type=self.norm_type
-            )
-            self.siamese_mlp_input_norm = _make_norm(
-                config.hidden_size, eps=config.rms_norm_eps, norm_type=self.norm_type
-            )
-        else:
-            self.siamese_attn_y_norm = None
-            self.siamese_mlp_y_norm = None
-            self.siamese_mlp_input_norm = None
         self.lns_attn                 = LNS(layer_idx)
         self.lns_mlp                  = LNS(layer_idx)
         self.gpas_attn                = GPAS(config.hidden_size)
@@ -5834,136 +5750,10 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             # LAUREL-LR (paper eq. 3): f(x) + BAx + x
             return delta + lr_delta + residual
 
-    def _siamese_scale(self) -> float:
-        """Depth-wise residual scale for the SiameseNorm X stream."""
-        if not self.siamese_depth_scaling:
-            return 1.0
-        return math.sqrt(2.0 * float(self.layer_idx + 1))
-
-    def _forward_siamesenorm(
-        self,
-        hidden_states: torch.Tensor,
-        siamese_states: Optional[torch.Tensor],
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        first_layer_fan: Optional[torch.Tensor] = None,
-        z_tilde: Optional[torch.Tensor] = None,
-        B_vals: Optional[torch.Tensor] = None,
-        layer_analysis: Optional[LayerAnalysis] = None,
-        output_attentions: Optional[bool] = False,
-        repo_rope_args: Optional[Tuple[torch.Tensor, float]] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> Tuple:
-        """
-        SiameseNorm two-stream decoder block.
-
-        X stream receives the shared residual updates divided by
-        c_l = sqrt(2(l+1)) when depth scaling is enabled.  Y stream receives
-        the same shared updates without depth scaling, preserving the Pre-Norm
-        identity-gradient path.  Attention and MLP parameters are shared.
-        """
-        if siamese_states is None:
-            siamese_states = hidden_states
-
-        if layer_analysis is not None:
-            layer_analysis.hidden_states_input = hidden_states.detach()
-
-        # ── Attention fused input: Norm_X(x) + Norm_Y(y) ─────────────────
-        x_residual_attn = hidden_states
-        y_residual_attn = siamese_states
-        sn_pre = layer_analysis.seednorm_pre_attn if layer_analysis is not None else None
-
-        if self.input_layernorm is None:
-            x_normed = x_residual_attn
-            if sn_pre is not None:
-                sn_pre.output = x_normed.detach()
-        else:
-            x_normed = _apply_norm(self.input_layernorm, x_residual_attn, analysis=sn_pre)
-        y_normed = _apply_norm(self.siamese_attn_y_norm, y_residual_attn, analysis=None)
-        h_lns = self.lns_attn(x_normed + y_normed)
-        if layer_analysis is not None:
-            layer_analysis.lns_attn_output = h_lns.detach()
-
-        attn_delta, attn_weights, self.current_layer_fan = self.self_attn(
-            hidden_states=h_lns,
-            key_value_states=None,
-            attention_mask=attention_mask,
-            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,
-            position_ids=position_ids,
-            **kwargs,
-        )
-        if layer_analysis is not None:
-            layer_analysis.attn_contribution = attn_delta.detach()
-
-        scale = self._siamese_scale()
-        gpas_attn_a = layer_analysis.gpas_attn if layer_analysis is not None else None
-        x_after_attn = self.gpas_attn(x_residual_attn + attn_delta / scale, analysis=gpas_attn_a)
-        y_after_attn = y_residual_attn + attn_delta
-
-        if layer_analysis is not None:
-            layer_analysis.h_tilde = x_after_attn.detach()
-
-        # ── MLP fused input: NormInput(Norm_X(x) + Norm_Y(y)) ─────────────
-        x_residual_mlp = x_after_attn
-        sn_post = layer_analysis.seednorm_post_attn if layer_analysis is not None else None
-        x_mlp_normed = _apply_norm(self.post_attention_layernorm, x_after_attn, analysis=sn_post)
-        y_mlp_normed = _apply_norm(self.siamese_mlp_y_norm, y_after_attn, analysis=None)
-        mlp_fused = _apply_norm(self.siamese_mlp_input_norm, x_mlp_normed + y_mlp_normed, analysis=None)
-        h_lns2 = self.lns_mlp(mlp_fused)
-        if layer_analysis is not None:
-            layer_analysis.lns_mlp_output = h_lns2.detach()
-
-        mlp_a = layer_analysis.mlp if layer_analysis is not None else None
-        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()
-
-        aux_stats = None
-        shared_mlp_update = delta_m
-        if self.use_jtokm and z_tilde is not None and B_vals is not None:
-            orig_shape = x_after_attn.shape
-            h_flat = x_after_attn.reshape(-1, self.hidden_size)
-            z_flat = z_tilde.reshape(-1, z_tilde.shape[-1])
-            B_flat = B_vals.reshape(-1, B_vals.shape[-2], B_vals.shape[-1])
-
-            jtokm_a = layer_analysis.jtokm if layer_analysis is not None else None
-            delta_r, aux_stats = self.jtokm(h_flat, z_flat, B_flat, analysis=jtokm_a)
-            delta_r = delta_r.reshape(orig_shape)
-            shared_mlp_update = shared_mlp_update + delta_r
-
-        gpas_mlp_a = layer_analysis.gpas_mlp if layer_analysis is not None else None
-        hidden_states = self.gpas_mlp(
-            x_residual_mlp + shared_mlp_update / scale,
-            analysis=gpas_mlp_a,
-        )
-        siamese_states = y_after_attn + shared_mlp_update
-
-        if layer_analysis is not None:
-            layer_analysis.hidden_states_output = hidden_states.detach()
-
-        outputs = (hidden_states, siamese_states)
-        if output_attentions:
-            outputs += (attn_weights,)
-        if aux_stats is not None:
-            outputs += (aux_stats,)
-        if versatile_aux is not None:
-            outputs += (versatile_aux,)
-        return outputs
-
     def forward(
         self,
         hidden_states: torch.Tensor,
-        siamese_states: Optional[torch.Tensor] = None,
-        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: Optional[torch.Tensor] = None,
         first_layer_fan: Optional[torch.Tensor] = None,
         z_tilde: Optional[torch.Tensor] = None,
@@ -5977,22 +5767,6 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         position_ids: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple:
-        if self.use_siamesenorm:
-            return self._forward_siamesenorm(
-                hidden_states=hidden_states,
-                siamese_states=siamese_states,
-                position_embeddings=position_embeddings,
-                attention_mask=attention_mask,
-                first_layer_fan=first_layer_fan,
-                z_tilde=z_tilde,
-                B_vals=B_vals,
-                layer_analysis=layer_analysis,
-                output_attentions=output_attentions,
-                repo_rope_args=repo_rope_args,
-                position_ids=position_ids,
-                **kwargs,
-            )
-
         # ── Snapshot input ────────────────────────────────────────────────
         if layer_analysis is not None:
             layer_analysis.hidden_states_input = hidden_states.detach()
@@ -6694,27 +6468,11 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             [NeoLLMDecoderLayer(config, layer_idx)
              for layer_idx in range(config.num_hidden_layers)]
         )
-        self.use_siamesenorm = _uses_siamesenorm(config)
-        self.norm_type       = _active_base_norm_type(config)
         self.norm        = _make_norm(
             config.hidden_size,
             eps=config.rms_norm_eps,
-            norm_type=self.norm_type,
+            use_seednorm=bool(getattr(config, "use_seednorm", True)),
         )
-        if self.use_siamesenorm:
-            self.siamese_final_y_norm = _make_norm(
-                config.hidden_size,
-                eps=config.rms_norm_eps,
-                norm_type=self.norm_type,
-            )
-            self.siamese_final_norm = _make_norm(
-                config.hidden_size,
-                eps=config.rms_norm_eps,
-                norm_type=self.norm_type,
-            )
-        else:
-            self.siamese_final_y_norm = None
-            self.siamese_final_norm = None
         self.rotary_emb  = NeoLLMRotaryEmbedding(config=config)
         self.gradient_checkpointing = False
         self.first_layer_fan        = None if getattr(config, "use_fan_residual", True) else False
@@ -6900,7 +6658,6 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         )
 
         hidden_states     = inputs_embeds
-        siamese_states    = hidden_states if self.use_siamesenorm else None
         all_hidden_states = () if output_hidden_states else None
         all_attentions    = () if output_attentions    else None
         all_aux_stats     = []
@@ -7015,7 +6772,6 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
 
             layer_outputs = decoder_layer(
                 hidden_states,
-                siamese_states=siamese_states,
                 position_embeddings=position_embeddings,
                 attention_mask=causal_mask,
                 first_layer_fan=self.first_layer_fan,
@@ -7031,10 +6787,6 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 **kwargs,
             )
             hidden_states = layer_outputs[0]
-            output_offset = 1
-            if self.use_siamesenorm:
-                siamese_states = layer_outputs[1]
-                output_offset = 2
 
             # Update AttnRes partial sum — the new partial is the layer output
             if use_attn_res:
@@ -7045,34 +6797,27 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 dca_sources = dca_sources + [hidden_states]
 
             if output_attentions:
-                all_attentions = all_attentions + (layer_outputs[output_offset],)
-                output_offset += 1
-
-            # Collect auxiliary losses without relying on positional offsets,
-            # because SiameseNorm adds the Y stream as an additional output.
-            # JTok-M returns a raw (p_sum, f_sum, N_tokens) tuple; VersatileFFN
-            # uses the same tuple shape, so keep the first such tuple for JTok-M
-            # when active and tag later tuples as VersatileFFN stats.
-            jtokm_aux_taken = False
-            for item in layer_outputs[output_offset:]:
-                if isinstance(item, tuple) and len(item) == 3:
-                    if self.config.use_jtokm and not jtokm_aux_taken:
-                        all_aux_stats.append(item)
-                        jtokm_aux_taken = True
-                    elif getattr(self.config, "use_versatile_ffn", False):
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+            # Collect JTok-M aux stats (last element if present)
+            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 (getattr(self.config, "use_fan_residual", True)
                     and self.first_layer_fan is None
                     and hasattr(decoder_layer, "current_layer_fan")):
                 self.first_layer_fan = decoder_layer.current_layer_fan
 
-        if self.use_siamesenorm:
-            hidden_states = self.norm(hidden_states)
-            siamese_states = self.siamese_final_y_norm(siamese_states)
-            hidden_states = self.siamese_final_norm(hidden_states + siamese_states)
-        else:
-            hidden_states = self.norm(hidden_states)
+        hidden_states = self.norm(hidden_states)
 
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)