Update Feather training runtime image

overlay/hydra/model.py

@@ -32,19 +32,33 @@ from __future__ import annotations
 
 import os
 
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from mamba_ssm import Mamba3
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from mamba_ssm import Mamba3
+except Exception:
+    Mamba3 = None
 
 from subsystems.hestia_mini import HestiaQAT
 from subsystems.htm import HTMLayer
 from subsystems.mhc_mini import ManifoldHyperConnection
 from subsystems.sdr_semantic import SemanticFoldingSDR
 
-from hydra.engram import GPUEngram
-from hydra.optimizer import MuonAdamW
+from hydra.engram import GPUEngram
+from hydra.optimizer import MuonAdamW
+
+
+class _InertMambaBlock(nn.Module):
+    """Identity fallback used when HYDRA_INERT_MAMBA=1."""
+
+    def __init__(self, d_model: int) -> None:
+        super().__init__()
+        self.d_model = d_model
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
 
 
 def norm(x: torch.Tensor) -> torch.Tensor:
@@ -64,9 +78,10 @@ class PostSemClawModel(nn.Module):
         model(x, y, reduction='mean')           -> scalar loss
     """
 
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self._inert_mamba = os.environ.get("HYDRA_INERT_MAMBA", "0") == "1"
 
         # Token embedding
         self.wte = nn.Embedding(config.vocab_size, config.d_model)
@@ -74,19 +89,29 @@ class PostSemClawModel(nn.Module):
         # Mamba-3 blocks — official mamba-ssm fused CUDA kernel. No fallbacks.
         # RoPE is applied internally by the Mamba3 CUDA kernel via the Angles
         # parameter; external cos/sin buffers are not needed.
-        self.blocks = nn.ModuleList([
-            Mamba3(
-                d_model=config.d_model,
-                d_state=config.d_state,
-                expand=config.expand,
-                headdim=config.headdim,
-                is_mimo=False,          # SISO path uses stable mamba3_siso_combined kernel
-                chunk_size=64,          # upstream-recommended SISO chunk; 16 violated tl.dot M>=16 constraint
-                is_outproj_norm=False,
-                dtype=torch.bfloat16,
-            )
-            for _ in range(config.n_layer)
-        ])
+        if self._inert_mamba or Mamba3 is None:
+            if self._inert_mamba:
+                print("[HYDRA] HYDRA_INERT_MAMBA=1 -> using inert identity blocks", flush=True)
+            else:
+                print("[HYDRA] mamba_ssm unavailable -> using inert identity blocks", flush=True)
+            self.blocks = nn.ModuleList([
+                _InertMambaBlock(config.d_model)
+                for _ in range(config.n_layer)
+            ])
+        else:
+            self.blocks = nn.ModuleList([
+                Mamba3(
+                    d_model=config.d_model,
+                    d_state=config.d_state,
+                    expand=config.expand,
+                    headdim=config.headdim,
+                    is_mimo=False,          # SISO path uses stable mamba3_siso_combined kernel
+                    chunk_size=64,          # upstream-recommended SISO chunk; 16 violated tl.dot M>=16 constraint
+                    is_outproj_norm=False,
+                    dtype=torch.bfloat16,
+                )
+                for _ in range(config.n_layer)
+            ])
 
         # Full-architecture SDR: offline semantic retina + STE (no-bypass).
         self.sdr_semantic = SemanticFoldingSDR(