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()
+```

scripts/generate_golden.py

@@ -152,6 +152,34 @@ def _extract_logits_golden(
     }
 
 
+# ---- Compat patches ----
+
+
+def _apply_nemotron_h_generate_patch(model: object) -> None:
+    """Patch NemotronH prepare_inputs_for_generation for transformers 5.x.
+
+    The HF remote code accesses ``cache_position[-1]`` without checking
+    for ``None``, which crashes under transformers >=5.x where
+    ``cache_position`` is no longer passed on the first prefill call.
+    We wrap the method to supply a default ``cache_position`` when missing.
+    """
+    cls_name = type(model).__name__
+    if "NemotronH" not in cls_name:
+        return
+
+    import torch
+
+    original_prepare = model.prepare_inputs_for_generation
+
+    def _patched_prepare(input_ids, **kwargs):
+        if kwargs.get("past_key_values") is not None and kwargs.get("cache_position") is None:
+            seq_len = input_ids.shape[-1]
+            kwargs["cache_position"] = torch.arange(seq_len, device=input_ids.device)
+        return original_prepare(input_ids, **kwargs)
+
+    model.prepare_inputs_for_generation = _patched_prepare
+
+
 # ---- Task-specific generators ----
 # Each generator loads a HF model, runs inference, and calls
 # save_golden_ref() from golden.py.  Heavy imports (torch,
@@ -184,6 +212,8 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
     if "L5" in case.level:
         import torch
 
+        _apply_nemotron_h_generate_patch(model)
+
         with torch.no_grad():
             gen_output = model.generate(
                 torch.from_numpy(input_ids).to(device),

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,93 @@
 
 from __future__ import annotations
 
+import logging
+
 import numpy as np
 import torch
 
+logger = logging.getLogger(__name__)
+
+
+def _fix_nemotron_h_init_weights(model: torch.nn.Module, model_id: str) -> None:
+    """Restore Mamba2 params from checkpoint after HF clobbers them.
+
+    The NemotronH remote-code ``_init_weights`` re-initialises several
+    parameters *after* ``from_pretrained`` loads the checkpoint:
+
+    - ``dt_bias``: overwritten with ``torch.rand(...)``
+    - ``out_proj.weight`` (in Mamba mixer layers): overwritten with
+      ``kaiming_uniform_`` then scaled by ``1/sqrt(n_layers)`` when
+      ``rescale_prenorm_residual`` is True
+
+    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
+
+    from huggingface_hub import snapshot_download
+
+    # Resolve the exact snapshot directory used by HF for this model,
+    # avoiding lexicographic guessing across multiple cached revisions.
+    try:
+        snapshot = snapshot_download(model_id, local_files_only=True)
+    except Exception:
+        logger.warning(
+            "NemotronH init_weights fix: could not resolve snapshot for %s",
+            model_id,
+        )
+        return
+
+    safetensor_files = sorted(glob.glob(os.path.join(snapshot, "*.safetensors")))
+
+    # Collect parameter names that _init_weights corrupts:
+    # 1. All dt_bias params (Mamba2 layers)
+    # 2. mixer.out_proj.weight params (rescale_prenorm_residual)
+    corrupted_suffixes = {"dt_bias"}
+    if getattr(config, "rescale_prenorm_residual", False):
+        corrupted_suffixes.add("mixer.out_proj.weight")
+
+    patched = 0
+    state = model.state_dict()
+    for f in safetensor_files:
+        with safe_open(f, framework="pt") as st:
+            # safe_open objects aren't directly iterable
+            keys = st.keys()
+            for key in keys:
+                if not any(key.endswith(s) for s in corrupted_suffixes):
+                    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 params from checkpoint (dt_bias + out_proj.weight)",
+            patched,
+        )
+    else:
+        logger.warning(
+            "NemotronH init_weights fix: no corrupted params found in checkpoint for %s",
+            model_id,
+        )
+
 
 def load_torch_model(
     model_id: str,
@@ -27,12 +111,22 @@ def load_torch_model(
     import transformers
 
     tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+
+    # NemotronH: disable rescale_prenorm_residual before loading to
+    # prevent _init_weights from corrupting out_proj.weight with
+    # random kaiming_uniform_ initialization after checkpoint loading.
+    config = transformers.AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+    if getattr(config, "model_type", None) == "nemotron_h":
+        config.rescale_prenorm_residual = False
+
     model = transformers.AutoModelForCausalLM.from_pretrained(
         model_id,
+        config=config,
         dtype=dtype,
         device_map=device,
         trust_remote_code=True,
     )
+    _fix_nemotron_h_init_weights(model, model_id)
     model.eval()
 
     if tokenizer.pad_token is None: