Skip SpanEngine/Grammar/DebtLedger during training (inference-only ops on 200K logits)

chimera/model.py

@@ -32,16 +32,10 @@ from .multimodal import VisionEncoder, AudioEncoder
 
 
 class CausalLMOutput(dict):
-    """Light HF-compatible output dict supporting tuple unpacking."""
-
-    def __init__(self, loss: Optional[torch.Tensor] = None,
-                 logits: Optional[torch.Tensor] = None,
-                 hidden_states: Optional[torch.Tensor] = None,
-                 caches: Optional[list] = None,
-                 evolution_metrics: Optional[dict] = None):
-        super().__init__(loss=loss, logits=logits,
-                         hidden_states=hidden_states, caches=caches,
-                         evolution_metrics=evolution_metrics)
+    def __init__(self, loss=None, logits=None, hidden_states=None,
+                 caches=None, evolution_metrics=None):
+        super().__init__(loss=loss, logits=logits, hidden_states=hidden_states,
+                         caches=caches, evolution_metrics=evolution_metrics)
         self.loss = loss
         self.logits = logits
         self.hidden_states = hidden_states
@@ -53,8 +47,7 @@ class CausalLMOutput(dict):
         yield self.logits
 
 
-def expand_layer_pattern(config: dict) -> List[str]:
-    """Expand the layer-pattern shorthand into a list."""
+def expand_layer_pattern(config):
     backbone = config.get("backbone", {})
     pattern_str = backbone.get("layer_pattern", "GD XM GD TM GD XM GD SK")
     aliases = backbone.get("layer_aliases", {
@@ -68,12 +61,9 @@ def expand_layer_pattern(config: dict) -> List[str]:
 
 
 class Chimera51Block(nn.Module):
-    """One block with evolution-aware forward."""
-
     _RECURRENT = {"gated_deltanet", "xlstm_m", "titans_mac", "tsp_span_knot"}
 
-    def __init__(self, config: dict, layer_type: str, layer_idx: int,
-                 use_moe: bool = False):
+    def __init__(self, config, layer_type, layer_idx, use_moe=False):
         super().__init__()
         h = int(config["hidden_size"])
         eps = float(config.get("rms_norm_eps", 1e-6))
@@ -87,22 +77,18 @@ class Chimera51Block(nn.Module):
         self.attn_norm = RMSNorm(h, eps=eps)
 
         if layer_type == "gated_deltanet":
-            self.attn = GatedDeltaNetLayer(h, heads, head_dim, norm_eps=eps,
-                                           chunk_size=chunk_sz, use_ternary=ternary)
+            self.attn = GatedDeltaNetLayer(h, heads, head_dim, norm_eps=eps, chunk_size=chunk_sz, use_ternary=ternary)
         elif layer_type == "xlstm_m":
             mem_h = config.get("xlstm", {}).get("memory_size_per_head", [head_dim, head_dim])
-            self.attn = MLSTMLayer(h, heads, int(mem_h[0]), norm_eps=eps,
-                                   use_ternary=ternary)
+            self.attn = MLSTMLayer(h, heads, int(mem_h[0]), norm_eps=eps, use_ternary=ternary)
         elif layer_type == "titans_mac":
             tc = config.get("titans", {})
-            self.attn = TitansMACLayer(h, heads, head_dim,
-                                       memory_depth=int(tc.get("memory_depth", 2)),
+            self.attn = TitansMACLayer(h, heads, head_dim, memory_depth=int(tc.get("memory_depth", 2)),
                                        persistent_slots=int(tc.get("persistent_memory_slots", 64)),
                                        local_window=int(tc.get("local_window_size", 1024)),
                                        norm_eps=eps, use_ternary=ternary)
         elif layer_type == "tsp_span_knot":
-            self.attn = TSPSpanKnotLayer(h, heads, head_dim, norm_eps=eps,
-                                         chunk_size=chunk_sz, use_ternary=ternary)
+            self.attn = TSPSpanKnotLayer(h, heads, head_dim, norm_eps=eps, chunk_size=chunk_sz, use_ternary=ternary)
         else:
             raise ValueError(f"Unknown layer type: {layer_type}")
 
@@ -110,45 +96,33 @@ class Chimera51Block(nn.Module):
         self.use_moe = bool(use_moe)
         if self.use_moe:
             moe_cfg = config.get("backbone", {}).get("moe", {})
-            self.mlp = MoELayer(
-                hidden_size=h,
+            self.mlp = MoELayer(hidden_size=h,
                 moe_intermediate_size=int(moe_cfg.get("moe_intermediate_size", h * 2)),
                 n_routed_experts=int(moe_cfg.get("n_routed_experts", 16)),
                 n_shared_experts=int(moe_cfg.get("n_shared_experts", 1)),
                 num_experts_per_tok=int(moe_cfg.get("num_experts_per_tok", 2)),
-                use_ternary=ternary,
-            )
+                use_ternary=ternary)
         else:
             inter = int(config.get("intermediate_size", int(h * 8 / 3)))
             inter = 256 * ((inter + 255) // 256)
             self.mlp = SwiGLUMLP(h, inter, use_ternary=ternary)
 
-        # Evolution modulation projection (learnable scale)
         self.evo_gate = nn.Linear(h, h, bias=False)
         nn.init.zeros_(self.evo_gate.weight)
 
-    def forward(self, x: torch.Tensor, cache: Optional[dict] = None,
-                evo_modulation: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, dict]:
-        # Apply attention with pre-norm
+    def forward(self, x, cache=None, evo_modulation=None):
         normed = self.attn_norm(x)
         attn_out, new_cache = self.attn(normed, cache=cache)
         x = x + attn_out
-
-        # Apply MLP with pre-norm
         x = x + self.mlp(self.mlp_norm(x))
-
-        # Apply evolution modulation (gated residual)
         if evo_modulation is not None:
             gate = torch.sigmoid(self.evo_gate(x))
             x = x + gate * evo_modulation
-
         return x, new_cache
 
 
 class Chimera51ForCausalLM(nn.Module):
-    """Chimera 5.x causal language model with functional self-evolution."""
-
-    def __init__(self, config: dict):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         h = int(config["hidden_size"])
@@ -159,19 +133,15 @@ class Chimera51ForCausalLM(nn.Module):
         self.embed = nn.Embedding(vocab, h)
         layer_types = expand_layer_pattern(config)
         moe_layers = set(int(i) for i in config.get("backbone", {}).get("moe", {}).get("layers", []))
-
         self.layers = nn.ModuleList([
             Chimera51Block(config, layer_types[i], i, use_moe=(i in moe_layers))
             for i in range(n_layers)
         ])
-
         self.norm = RMSNorm(h, eps=eps)
         self.lm_head = nn.Linear(h, vocab, bias=False)
-
         if config.get("tie_word_embeddings", True):
             self.lm_head.weight = self.embed.weight
 
-        # Parcae looping controller
         loop_cfg = config.get("looping", {})
         self.looping_enabled = bool(loop_cfg.get("enabled", True)) and n_layers >= 3
         if self.looping_enabled:
@@ -181,24 +151,20 @@ class Chimera51ForCausalLM(nn.Module):
             self.loop_controller = ParcaeLoopController(
                 h, loop_range=tuple(loop_cfg.get("loop_range", [1, 6])),
                 loop_default=int(loop_cfg.get("loop_default", 2)),
-                adaptive_exit_threshold=float(loop_cfg.get("adaptive_exit_threshold", 0.01)),
-            )
+                adaptive_exit_threshold=float(loop_cfg.get("adaptive_exit_threshold", 0.01)))
 
-        # Inference systems
         si_cfg = config.get("span_inference", {})
         self.span_engine = SpanInferenceEngine(h, si_cfg) if si_cfg.get("enabled", True) else None
         self.grammar = GrammarFST(config.get("grammar", {}))
         self.entropy_valve = EntropyValve(config.get("entropy_valve", {}))
         self.debt_ledger = DebtLedger(config.get("debt_ledger", {}))
 
-        # Self-evolution — FUNCTIONAL
         evo_cfg = dict(config.get("self_evolution", {}))
         evo_cfg["_semantic_memory_config"] = config.get("semantic_memory", {})
         self.evolution = SelfEvolutionEngine(evo_cfg, h)
         self.evo_weight = float(config.get("evolution_loss_weight", 0.01))
         self.evo_every_n_layers = int(config.get("evolution_every_n_layers", 4))
 
-        # Multimodal
         mm_cfg = dict(config.get("multimodal", {}))
         mm_cfg["hidden_size"] = h
         if mm_cfg.get("enabled", False):
@@ -212,19 +178,19 @@ class Chimera51ForCausalLM(nn.Module):
         self._init_weights()
         self._wire_semantic_memory()
 
-    def enable_gradient_checkpointing(self) -> None:
+    def enable_gradient_checkpointing(self):
         self.gradient_checkpointing = True
 
-    def disable_gradient_checkpointing(self) -> None:
+    def disable_gradient_checkpointing(self):
         self.gradient_checkpointing = False
 
-    def _wire_semantic_memory(self) -> None:
+    def _wire_semantic_memory(self):
         mem = self.evolution.semantic_memory
         for layer in self.layers:
             if hasattr(layer.attn, "set_semantic_memory"):
                 layer.attn.set_semantic_memory(mem)
 
-    def _init_weights(self) -> None:
+    def _init_weights(self):
         init_range = float(self.config.get("initializer_range", 0.006))
         for module in self.modules():
             if isinstance(module, (nn.Linear, BitLinear)):
@@ -238,11 +204,7 @@ class Chimera51ForCausalLM(nn.Module):
             if isinstance(module, BitLinear):
                 module.invalidate_packed()
 
-    def _run_layers(self, x: torch.Tensor, start: int, end: int,
-                    caches: Optional[list],
-                    compute_logits: bool = False,
-                    labels: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor], list]:
-        """Run layers with evolution hooks. Returns (x, logits_if_computed, caches)."""
+    def _run_layers(self, x, start, end, caches, compute_logits=False, labels=None):
         all_metrics = []
         logits = None
         evolution_loss = torch.tensor(0.0, device=x.device)
@@ -250,21 +212,15 @@ class Chimera51ForCausalLM(nn.Module):
         for i in range(start, min(end + 1, len(self.layers))):
             layer = self.layers[i]
             cache = caches[i] if caches is not None else None
-
-            # Evolution modulation every N layers (lightweight)
             evo_mod = None
             if i % self.evo_every_n_layers == 0 and self.evolution is not None:
                 evo_result = self.evolution(
                     hidden_states=x.detach() if not x.requires_grad else x,
-                    layer_idx=i,
-                    loss=None
-                )
+                    layer_idx=i, loss=None)
                 evo_mod = evo_result['modulation']
                 if evo_result['evolution_loss'] is not None:
                     evolution_loss = evolution_loss + evo_result['evolution_loss']
                 all_metrics.append(evo_result.get('metrics', {}))
-
-                # TTT update for target layers (only in training, no backprop)
                 if self.training and evo_result.get('ttt_delta') is not None:
                     with torch.no_grad():
                         if hasattr(layer.mlp, 'w_down'):
@@ -280,24 +236,16 @@ class Chimera51ForCausalLM(nn.Module):
                 if caches is not None:
                     caches[i] = new_cache
 
-            # Compute probe logits for entropy valve (every few layers)
             if compute_logits and i == end:
                 logits = self.lm_head(self.norm(x[:, -1:, :]))
 
         return x, logits, caches, evolution_loss, all_metrics
 
-    def forward(self, input_ids: torch.Tensor,
-                labels: Optional[torch.Tensor] = None,
-                pixel_values: Optional[torch.Tensor] = None,
-                mel_features: Optional[torch.Tensor] = None,
-                num_loops: Optional[int] = None,
-                caches: Optional[list] = None,
-                use_cache: bool = False,
-                logits_to_keep: int = 0,
-                return_evolution_metrics: bool = False):
+    def forward(self, input_ids, labels=None, pixel_values=None,
+                mel_features=None, num_loops=None, caches=None,
+                use_cache=False, logits_to_keep=0, return_evolution_metrics=False):
         x = self.embed(input_ids)
 
-        # Multimodal prepend
         if pixel_values is not None and self.vision_encoder is not None:
             v = self.vision_encoder(pixel_values)
             if v is not None:
@@ -313,31 +261,23 @@ class Chimera51ForCausalLM(nn.Module):
         total_evo_loss = torch.tensor(0.0, device=x.device)
         all_evo_metrics = []
 
-        # Prelude + Loop + Coda with evolution
         if self.looping_enabled and hasattr(self, "loop_controller"):
-            # Prelude
             x, probe_logits, caches, evo_loss, metrics = self._run_layers(
                 x, self.prelude_start, self.prelude_end, caches,
                 compute_logits=not self.training, labels=labels)
             total_evo_loss = total_evo_loss + evo_loss
             all_evo_metrics.extend(metrics)
 
-            # Determine loop depth
             effective = num_loops
             if effective is None and not self.training and probe_logits is not None:
                 effective = self.entropy_valve.get_loop_count(probe_logits)
             elif effective is None:
-                # FIX: During training, use the loop_controller.loop_default directly
-                # instead of running the loop classifier (which calls .item() and is
-                # expensive). The ProgressiveLoopScheduler already sets loop_default.
                 effective = self.loop_controller.loop_default
 
-            # Loop body
             loop_fn = lambda inp: self._run_layers(
                 inp, self.loop_start, self.loop_end, caches, labels=labels)[0]
             x = self.loop_controller(x, loop_fn, num_loops=effective)
 
-            # Coda
             x, _, caches, evo_loss, metrics = self._run_layers(
                 x, self.coda_start, self.coda_end, caches, labels=labels)
             total_evo_loss = total_evo_loss + evo_loss
@@ -349,30 +289,29 @@ class Chimera51ForCausalLM(nn.Module):
             total_evo_loss = total_evo_loss + evo_loss
             all_evo_metrics.extend(metrics)
 
-        # Final norm and logits
         if logits_to_keep and labels is None:
             keep = int(logits_to_keep)
             tail = x[:, -keep:, :]
             tail = self.norm(tail)
-            if self.span_engine is not None:
+            if self.span_engine is not None and not self.training:
                 tail = self.span_engine(tail)
             logits = self.lm_head(tail)
         else:
             x = self.norm(x)
-            if self.span_engine is not None:
+            if self.span_engine is not None and not self.training:
                 x = self.span_engine(x)
             logits = self.lm_head(x)
 
-        logits = self.grammar(logits)
-        logits = self.debt_ledger(logits)
+        # Inference-only post-processing on 200K-dim logits — skip during training
+        if not self.training:
+            logits = self.grammar(logits)
+            logits = self.debt_ledger(logits)
 
-        # Self-feedback refinement check (inference only)
         if not self.training and self.evolution is not None:
             should_refine = self.evolution.self_feedback.should_refine(logits)
             if should_refine:
                 all_evo_metrics.append({'refinement_triggered': True})
 
-        # Compute loss
         loss = None
         if labels is not None:
             seq_len = min(logits.size(1), labels.size(1))
@@ -380,49 +319,38 @@ class Chimera51ForCausalLM(nn.Module):
             shift_labels = labels[:, :seq_len].contiguous()
             ce_loss = F.cross_entropy(
                 shift_logits.view(-1, shift_logits.size(-1)),
-                shift_labels.view(-1),
-                ignore_index=-100,
-            )
-            # Add evolution loss (contrastive memory evaluation)
+                shift_labels.view(-1), ignore_index=-100)
             loss = ce_loss + self.evo_weight * total_evo_loss
         else:
             ce_loss = None
 
-        # Store episodic case after forward (for inference mode)
-        if not self.training and self.evolution is not None:
-            last_hidden = x[:, -1, :].detach()
-
         return CausalLMOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=x,
+            loss=loss, logits=logits, hidden_states=x,
             caches=caches if use_cache else None,
             evolution_metrics={
                 'ce_loss': ce_loss.item() if ce_loss is not None else None,
                 'evo_loss': total_evo_loss.item(),
                 'layer_metrics': all_evo_metrics,
-            } if return_evolution_metrics else None
-        )
+            } if return_evolution_metrics else None)
 
     @torch.no_grad()
-    def prepare_for_inference(self) -> None:
-        """Pre-pack every BitLinear so the first generation step is fast."""
+    def prepare_for_inference(self):
         for module in self.modules():
             if isinstance(module, BitLinear):
                 module.prepare_for_inference()
 
-    def get_mode_config(self, mode: str = "balanced") -> dict:
+    def get_mode_config(self, mode="balanced"):
         modes = self.config.get("modes", {})
         return modes.get(mode, modes.get("balanced", {}))
 
-    def count_parameters(self) -> dict:
+    def count_parameters(self):
         total = sum(p.numel() for p in self.parameters())
         ternary = sum(p.numel() for _, m in self.named_modules()
                       if isinstance(m, BitLinear) for p in m.parameters())
         return {"total": total, "ternary": ternary, "fp32": total - ternary}
 
     @classmethod
-    def from_config_file(cls, path: str) -> "Chimera51ForCausalLM":
+    def from_config_file(cls, path):
         with open(path, "r", encoding="utf-8") as fh:
             config = json.load(fh)
         return cls(config)