Add MoE layer support for NemotronH (Nemotron-3 30B/120B)

.agents/skills/adding-a-new-model/SKILL.md

@@ -335,6 +335,27 @@ defaults.
 4. Check weight dtype (numpy arrays must be float32)
 5. Compare layer by layer (moderate diff → norm/residual issue; huge → wrong weights)
 
+### 6. HF `_init_weights` corrupting checkpoint values
+
+**Symptom:** HF reference inference is non-deterministic across model loads —
+different argmax each time, despite identical inputs. Affects golden data
+generation and L4 parity tests.
+
+**Root cause:** Some HF models' `_init_weights` re-initialise parameters
+with random values (e.g. `torch.rand`) AFTER `from_pretrained` loads the
+checkpoint. Known cases:
+- **NemotronH**: `_init_weights` clobbers Mamba2 `dt_bias` with `torch.rand()`
+
+**Fix:** `_fix_nemotron_h_dt_bias()` in `mobius._testing.torch_reference`
+reads correct values from safetensors files and patches them in-place.
+Always call after `from_pretrained` for NemotronH models.
+
+**Diagnosis pattern:** Load the model twice and compare outputs. If argmax
+differs between loads, suspect `_init_weights` corruption. Set
+`torch.manual_seed(42)` before `from_pretrained` — if that makes outputs
+deterministic, `_init_weights` is the culprit. Then compare specific
+parameters between the loaded model and the safetensors checkpoint.
+
 > For additional troubleshooting (gated attention split ordering, DeltaNet
 > scaling, identity node folding, fp32 upcast patterns, multi-token prefill,
 > embedding table off-by-one), read

.agents/skills/moe-models/SKILL.md

@@ -382,3 +382,138 @@ def _dispatch_moe_fallback(self, op, hidden, router_probs):
         ...
     return output
 ```
+
+## NemotronH MoE (sigmoid routing + shared experts + latent projection)
+
+NemotronH uses a non-standard MoE architecture that differs from the
+standard softmax top-k pattern in several ways.
+
+### Architecture
+
+```
+NemotronHMoEBlock
+ ├── NemotronHMoEGate (sigmoid top-k with correction bias)
+ ├── [optional] fc1_latent_proj (hidden → latent_size)
+ ├── Experts[0..N-1] (non-gated FCMLP: up → act → down)
+ ├── [optional] fc2_latent_proj (latent_size → hidden)
+ └── SharedExperts (FCMLP, all tokens, residual add)
+```
+
+### Classes
+
+| Class | File | Purpose |
+|-------|------|---------|
+| `NemotronHMoEGate` | `models/nemotron_h.py` | Sigmoid routing with e_score_correction_bias |
+| `NemotronHMoEBlock` | `models/nemotron_h.py` | MoE dispatch with shared expert + latent proj |
+| `NemotronHMoELayer` | `models/nemotron_h.py` | Pre-norm → MoE block → residual (stateless) |
+
+### Sigmoid gate with correction bias
+
+NemotronH does NOT use softmax routing. Instead:
+
+1. `router_logits = hidden_states @ gate_weight.T`
+2. `probs = sigmoid(router_logits)` — these become final routing weights
+3. `choice_scores = probs + e_score_correction_bias` — bias affects selection only
+4. `selected = topk(choice_scores)` — select top-k using biased scores
+5. `weights = gather(probs, selected)` — gather from UNBIASED sigmoid probs
+6. Normalize + scale by `routed_scaling_factor`
+
+**Key difference from standard MoE**: The correction bias shifts expert
+selection but does NOT affect final routing weights. The `com.microsoft.MoE`
+op's built-in softmax routing is incompatible — you must use the fallback
+loop or pre-compute routing weights and pass them to a modified MoE call.
+
+### Non-gated FCMLP experts
+
+Unlike Mixtral/Qwen (gated: `gate_proj * up_proj → down_proj`), NemotronH
+experts are simple FCMLPs: `up_proj → activation → down_proj`. This maps
+to `fc1_experts_weights` / `fc2_experts_weights` without a gate projection.
+
+### Latent projection (120B only)
+
+The 120B model has `moe_latent_size=1024` (vs `hidden_size=4096`):
+
+```
+hidden → fc1_latent_proj(4096→1024) → experts(1024→inter→1024) → fc2_latent_proj(1024→4096)
+```
+
+Gate routes on original hidden states (NOT latent). Shared expert operates
+on original hidden_size (no latent projection).
+
+### Shared expert
+
+A single FCMLP that runs on ALL tokens (not routed), added as residual:
+
+```python
+output = routed_expert_output + shared_experts(original_hidden)
+```
+
+### HF weight format
+
+HF stores expert weights as 3D stacked tensors:
+```
+experts.up_proj: [num_experts, intermediate_size, input_size]
+experts.down_proj: [num_experts, hidden_size, intermediate_size]
+```
+
+`preprocess_weights()` splits these into per-expert 2D tensors for the
+loop-based dispatch, or keeps them stacked for the fused MoE op path.
+
+### Config fields (NemotronHConfig)
+
+```python
+config.num_local_experts             # Total experts (128 for 30B, 512 for 120B)
+config.num_experts_per_tok           # Top-k (6 for 30B, 22 for 120B)
+config.moe_intermediate_size         # Per-expert hidden dim
+config.moe_latent_size               # Optional latent projection dim (120B: 1024)
+config.shared_expert_intermediate_size  # Shared expert hidden dim
+config.norm_topk_prob                # Whether to normalize routing weights
+config.routed_scaling_factor         # Post-normalization scale
+```
+
+### com.microsoft.MoE compatibility
+
+**Not compatible with NemotronH.** Three blockers:
+
+1. **Squared ReLU activation**: NemotronH experts use `relu2` (squared ReLU:
+   `relu(x)^2`). The fused MoE op only supports `silu`, `gelu`, `relu`,
+   `none` — no squared ReLU. Since the activation is applied between the
+   two matmuls inside the fused op, there is no way to inject a custom
+   activation.
+
+2. **Sigmoid routing with correction bias**: NemotronH gate uses
+   `sigmoid → add_bias → topk` for expert selection, but final routing
+   weights come from the unbiased sigmoid probs. The fused op has no
+   option to bypass its internal softmax/topk routing.
+
+3. **Shared expert + latent projection**: These must run outside the fused
+   op regardless, adding complexity without eliminating the main bottleneck.
+
+NemotronH uses loop-over-experts dispatch. See `NemotronHMoEBlock` docstring.
+
+### Graph size impact
+
+The loop-over-experts fallback creates one subgraph per expert per MoE layer:
+- 30B: 128 experts × 23 MoE layers = 2,944 expert subgraphs → ~40K nodes
+- 120B: 512 experts × 40 MoE layers = 20,480 expert subgraphs → ~270K nodes
+
+The fused MoE op replaces each per-layer loop with a single op, dramatically
+reducing graph size and enabling batched GPU execution.
+
+### HF dt_bias corruption bug
+
+**Critical**: The NemotronH remote-code `_init_weights` re-initialises
+Mamba2 `dt_bias` parameters with `torch.rand()` AFTER `from_pretrained`
+loads checkpoint weights, silently corrupting the model. HF inference
+becomes non-deterministic — different argmax on each model load.
+
+Fix: Use `_fix_nemotron_h_dt_bias()` from `mobius._testing.torch_reference`
+after loading the HF model. This reads correct `dt_bias` values from the
+safetensors files and patches them in-place. Without the fix, golden
+reference data is unreliable.
+
+```python
+from mobius._testing.torch_reference import _fix_nemotron_h_dt_bias
+model = AutoModelForCausalLM.from_pretrained(model_id, ...)
+_fix_nemotron_h_dt_bias(model, model_id)  # Must call before eval()
+```

src/mobius/_configs.py

@@ -2244,6 +2244,7 @@ class NemotronHConfig(ArchitectureConfig):
     mamba_conv_bias: bool = True
     mamba_proj_bias: bool = False
     mamba_time_step_min: float = 0.001
+    moe_latent_size: int | None = None
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> NemotronHConfig:
@@ -2253,15 +2254,20 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig:
         layers_block_type = getattr(config, "layers_block_type", None)
         if layers_block_type is None:
             pattern = getattr(config, "hybrid_override_pattern", "")
-            # Map pattern chars: M=mamba2, *=full_attention, -=mlp
-            char_map = {"M": "mamba2", "*": "full_attention", "-": "mlp"}
+            # Map pattern chars: M=mamba2, *=full_attention, -=mlp, E=moe
+            char_map = {
+                "M": "mamba2",
+                "*": "full_attention",
+                "-": "mlp",
+                "E": "moe",
+            }
             layers_block_type = [char_map.get(c, "mamba2") for c in pattern]
         else:
             # Convert HF names to mobius names
             type_map = {
                 "mamba": "mamba2",
                 "attention": "full_attention",
-                "moe": "mlp",
+                "moe": "moe",
             }
             layers_block_type = [type_map.get(t, t) for t in layers_block_type]
 
@@ -2280,8 +2286,24 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig:
         base_fields = {
             k: v
             for k, v in _shallow_fields(base).items()
-            if k not in ("layer_types", "num_hidden_layers", "hidden_act")
+            if k
+            not in (
+                "layer_types",
+                "num_hidden_layers",
+                "hidden_act",
+                "moe_latent_size",
+                "shared_expert_intermediate_size",
+            )
         }
+
+        # Extract shared expert intermediate size (NemotronH uses a dedicated
+        # field name different from the base ArchitectureConfig default).
+        shared_expert_intermediate_size = getattr(
+            config,
+            "moe_shared_expert_intermediate_size",
+            base.shared_expert_intermediate_size,
+        )
+
         return cls(
             **base_fields,
             num_hidden_layers=n,
@@ -2296,6 +2318,8 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig:
             mamba_conv_bias=getattr(config, "use_conv_bias", True),
             mamba_proj_bias=getattr(config, "mamba_proj_bias", False),
             mamba_time_step_min=getattr(config, "time_step_min", 0.001),
+            moe_latent_size=getattr(config, "moe_latent_size", None),
+            shared_expert_intermediate_size=shared_expert_intermediate_size,
         )
 
 

src/mobius/_testing/torch_reference.py

@@ -5,9 +5,77 @@
 
 from __future__ import annotations
 
+import logging
+
 import numpy as np
 import torch
 
+logger = logging.getLogger(__name__)
+
+
+def _fix_nemotron_h_dt_bias(model: torch.nn.Module, model_id: str) -> None:
+    """Restore Mamba2 ``dt_bias`` from checkpoint after HF clobbers it.
+
+    The NemotronH remote-code ``_init_weights`` re-initialises ``dt_bias``
+    with ``torch.rand`` *after* ``from_pretrained`` loads the checkpoint,
+    silently corrupting the model.  This helper reads the original values
+    back from the safetensors files on disk and patches them in-place.
+    """
+    config = getattr(model, "config", None)
+    model_type = getattr(config, "model_type", None)
+    if model_type != "nemotron_h":
+        return
+
+    try:
+        from safetensors import safe_open
+    except ImportError:
+        return
+
+    import glob
+    import os
+
+    # Locate the cached snapshot directory
+    cache_dir = os.path.join(
+        torch.hub.get_dir().replace("/hub", ""),
+        "huggingface",
+        "hub",
+        f"models--{model_id.replace('/', '--')}",
+    )
+    if not os.path.isdir(cache_dir):
+        # Try the HF_HOME default
+        hf_home = os.environ.get(
+            "HF_HOME",
+            os.path.expanduser("~/.cache/huggingface"),
+        )
+        cache_dir = os.path.join(hf_home, "hub", f"models--{model_id.replace('/', '--')}")
+    snapshot_dirs = sorted(glob.glob(os.path.join(cache_dir, "snapshots", "*")))
+    if not snapshot_dirs:
+        logger.warning("NemotronH dt_bias fix: snapshot dir not found")
+        return
+
+    snapshot = snapshot_dirs[-1]
+    safetensor_files = sorted(glob.glob(os.path.join(snapshot, "*.safetensors")))
+
+    patched = 0
+    state = model.state_dict()
+    for f in safetensor_files:
+        with safe_open(f, framework="pt") as st:
+            for key in st.keys():
+                if "dt_bias" not in key:
+                    continue
+                if key not in state:
+                    continue
+                ckpt_val = st.get_tensor(key)
+                param = state[key]
+                with torch.no_grad():
+                    param.copy_(ckpt_val.to(param.device, dtype=param.dtype))
+                patched += 1
+
+    # Write the fixed values back into the live model
+    if patched:
+        model.load_state_dict(state, strict=False)
+        logger.info("NemotronH: restored %d dt_bias params from checkpoint", patched)
+
 
 def load_torch_model(
     model_id: str,
@@ -33,6 +101,7 @@ def load_torch_model(
         device_map=device,
         trust_remote_code=True,
     )
+    _fix_nemotron_h_dt_bias(model, model_id)
     model.eval()
 
     if tokenizer.pad_token is None: