Make RoPE config nullable to fix silent NoPE-model corruption (#177)

`_extract_rope_config()` never returned `None`, so NoPE models
(NemotronH, GraniteMoeHybrid, GPT-2/BERT/OPT families) silently received
`rope_type='default'` and had spurious rotary ops injected into their
ONNX graphs. The existing NemotronH fix (#110) is a model-class
workaround; this PR fixes the root cause in the config system.

### Changes

- **`_extract_rope_config()`** returns `None` unless the HF config has a
real RoPE signal: `rope_parameters`, `rope_scaling`, or legacy
`rotary_dim` / `rotary_pct` / `rotary_emb_base`. `rope_theta` alone is
dead data (NemotronH carries it without using RoPE) and is deliberately
not treated as a signal.
- **`ArchitectureConfig.from_transformers()`** propagates `None` to both
the `rope` sub-config and every flat RoPE field when the model is NoPE.
`dataclasses.replace(rope_config, ...)` call sites now guard against
`None`.
- **`ArchitectureConfig.rope_type`** default is `None` (the structural
NoPE signal). `rope_theta` / `partial_rotary_factor` keep inert numeric
defaults so `ArchitectureConfig(rope_type="default", ...)` still works
for direct construction in tests.
- **`initialize_rope()`** returns `None` when `rope_type is None` and
`mrope_section is None`.
- **`TextModel.forward()`** guards `self.rotary_emb(...)` and passes
`position_embeddings=None` down when RoPE is absent.
- **`Attention.__init__`** tolerates `partial_rotary_factor=None`
(treated as the inert 1.0) so NoPE-routed attention doesn't crash on
`math.isclose(None, 1.0)`.

### Tests

- New: `test_from_transformers_nope_model_has_none_rope`,
`test_from_transformers_legacy_rotary_dim_enables_rope`,
`test_rope_theta_alone_is_not_a_rope_signal`, `test_nope_returns_none`
(initialize_rope).
- Existing fakes in `_config_resolver_test.py` / `_configs_test.py` now
include `rope_parameters={"rope_type": "default"}` to match how real HF
`PretrainedConfig.__post_init__` populates the field.
- `make_config()` test helper opts into `rope_type="default"` so
component tests keep the RoPE code path. Direct `BambaConfig` /
`JambaConfig` / `Gemma2Config` test constructors and
`deepseek_ocr2.py`'s internal `ArchitectureConfig` opt in explicitly.

### Intentionally out of scope

- Simplifying the NemotronH / GraniteMoeHybrid text-model workarounds —
they use bespoke layer types (`NemotronHMambaLayer`, etc.) and can't
drop straight onto `TextModel`. They remain correct and now coexist with
a structural defense in the config layer.
- Phase 2 (calling `validate()` in the build path, nullable MoE
sub-config, deprecating flat RoPE fields) and Phase 3 (feature-group
sub-configs, per-model config classes) from the issue.

### Example

```python
# NemotronH config → no RoPE signal → structurally NoPE
class FakeNemotronH:
    model_type = "nemotron_h"
    rope_theta = 10_000.0  # dead data, ignored
    # no rope_parameters, no rope_scaling, no rotary_*
    ...

cfg = ArchitectureConfig.from_transformers(FakeNemotronH())
assert cfg.rope is None
assert cfg.rope_type is None
assert initialize_rope(cfg) is None  # TextModel now skips RoPE automatically
```

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: Copilot

src/mobius/_config_resolver_test.py

@@ -19,7 +19,13 @@
 
 
 def _fake_hf_config(model_type: str, **overrides):
-    """Create a minimal HF-config-like object for testing."""
+    """Create a minimal HF-config-like object for testing.
+
+    ``rope_parameters`` is present by default so the resolver treats the
+    fake config as a RoPE-capable model (matching how real HuggingFace
+    configs populate this field in ``PretrainedConfig.__post_init__``).
+    Pass ``rope_parameters=None`` explicitly to exercise the NoPE path.
+    """
     defaults = {
         "model_type": model_type,
         "vocab_size": 100,
@@ -35,6 +41,7 @@ def _fake_hf_config(model_type: str, **overrides):
         "rms_norm_eps": 1e-6,
         "rope_theta": 10_000.0,
         "rope_scaling": None,
+        "rope_parameters": {"rope_type": "default"},
     }
     defaults.update(overrides)
     return type("FakeHFConfig", (), defaults)()
@@ -141,6 +148,9 @@ def test_gemma3_nested_rope_scaling(self):
         hf = _fake_hf_config(
             "gemma3_text",
             rope_theta=None,  # force fallback to nested lookup
+            # Disable the default rope_parameters so the nested
+            # rope_scaling entries are used for rope_type resolution.
+            rope_parameters=None,
             rope_scaling={
                 "full_attention": {
                     "rope_type": "linear",
@@ -280,6 +290,10 @@ def _deepseek_config(self, **overrides):
             rms_norm_eps=1e-6,
             rope_theta=10000.0,
             rope_scaling=None,
+            # Real HF DeepSeek configs populate rope_parameters in
+            # __post_init__; include it here so _extract_rope_config
+            # treats this as a RoPE-capable model (not NoPE).
+            rope_parameters={"rope_type": "default"},
             # MLA-specific fields
             q_lora_rank=1536,
             kv_lora_rank=512,

src/mobius/_configs.py

@@ -282,14 +282,42 @@ def _first_not_none(*values, default=None):
     return default
 
 
-def _extract_rope_config(config) -> RoPEConfig:
+def _extract_rope_config(config) -> RoPEConfig | None:
     """Extract and normalize RoPE-related config fields.
 
     Reads ``rope_scaling``, ``rope_parameters``, and related attributes
     from a HuggingFace config and returns a :class:`RoPEConfig`.
+
+    Returns ``None`` when the source config has no RoPE signal at all —
+    i.e. it declares neither the modern ``rope_parameters``/``rope_scaling``
+    fields nor the legacy ``rotary_dim``/``rotary_pct``/``rotary_emb_base``
+    fields. This is the "NoPE" case (e.g. NemotronH, GraniteMoeHybrid,
+    GPT-2 family, BERT family, OPT) — the model does not use rotary
+    position embeddings at all, and callers should treat RoPE as absent
+    rather than manufacturing defaults that would silently introduce RoPE
+    operations into the ONNX graph.
+
+    Note: ``rope_theta`` alone is NOT a sufficient RoPE signal because
+    models like NemotronH carry ``rope_theta`` in their config as dead
+    data despite not using RoPE. HuggingFace's ``rope_parameters`` field
+    (populated by ``PretrainedConfig.__post_init__``) is the authoritative
+    modern signal; legacy GPT-J / GPT-NeoX / CodeGen models predate it
+    and use the ``rotary_*`` fields instead.
     """
-    rope_scaling = getattr(config, "rope_scaling", None) or {}
-    rope_parameters = getattr(config, "rope_parameters", None) or {}
+    # Check for RoPE signals BEFORE the `or {}` fallback below —
+    # `or {}` converts None to empty dict, destroying the absence signal.
+    raw_rope_scaling = getattr(config, "rope_scaling", None)
+    raw_rope_parameters = getattr(config, "rope_parameters", None)
+    has_legacy_rope = (
+        getattr(config, "rotary_dim", None) is not None
+        or getattr(config, "rotary_pct", None) is not None
+        or getattr(config, "rotary_emb_base", None) is not None
+    )
+    if raw_rope_scaling is None and raw_rope_parameters is None and not has_legacy_rope:
+        return None
+
+    rope_scaling = raw_rope_scaling or {}
+    rope_parameters = raw_rope_parameters or {}
 
     return RoPEConfig(
         rope_type=_first_not_none(
@@ -734,11 +762,27 @@ class ArchitectureConfig(BaseModelConfig):
 
     rms_norm_eps: float = 1e-6
 
-    # Rotary embedding config
-    rope_type: str = "default"
-    rope_theta: float = 10_000.0
+    # Rotary embedding config.
+    #
+    # ``rope_type`` is the structural signal: ``None`` means "this model
+    # does not use RoPE". ``from_transformers`` populates ``rope_type``
+    # (and the other flat RoPE fields below) from the HuggingFace config
+    # only when RoPE is actually declared — see :func:`_extract_rope_config`.
+    # For NoPE models (NemotronH, GraniteMoeHybrid, GPT-2 family, BERT, OPT,
+    # ...) ``from_transformers`` sets every flat RoPE field to ``None`` so
+    # downstream code (``initialize_rope``, ``TextModel``, ``Attention``) can
+    # structurally detect the absence of RoPE instead of spuriously applying
+    # a "default" rotary encoding.
+    #
+    # The non-``rope_type`` fields keep inert numeric defaults at the
+    # dataclass level so that code that constructs ``ArchitectureConfig``
+    # directly with just ``rope_type="default"`` (e.g. tests, small reproducer
+    # configs) works without having to spell out every RoPE parameter. These
+    # defaults are only consumed when ``rope_type`` is non-``None``.
+    rope_type: str | None = None
+    rope_theta: float | None = 10_000.0
     rope_scaling: dict | None = None
-    partial_rotary_factor: float = 1.0
+    partial_rotary_factor: float | None = 1.0
     rope_local_base_freq: float | None = None
     original_max_position_embeddings: int | None = None
 
@@ -883,14 +927,18 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
             or 0
         )
 
-        # rope_interleave depends on model_type / qk_rope_head_dim
+        # rope_interleave depends on model_type / qk_rope_head_dim.
+        # Only compute it when RoPE is actually in use — for NoPE models
+        # (rope_config is None) we leave the flat ``rope_interleave`` at
+        # its inert ``False`` default.
         rope_interleave = getattr(
             config,
             "rope_interleave",
             (getattr(config, "qk_rope_head_dim", None) or 0) > 0
             or model_type in ("glm", "glm4", "glm4_moe", "chatglm"),
         )
-        rope_config = dataclasses.replace(rope_config, rope_interleave=rope_interleave)
+        if rope_config is not None:
+            rope_config = dataclasses.replace(rope_config, rope_interleave=rope_interleave)
 
         options = dict(
             head_dim=(
@@ -1052,14 +1100,26 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig:
                 )
             ),
             rope=rope_config,
-            # Set flat rope fields for direct access by components
-            rope_type=rope_config.rope_type,
-            rope_theta=rope_config.rope_theta,
-            rope_scaling=rope_config.rope_scaling,
-            partial_rotary_factor=rope_config.partial_rotary_factor,
-            rope_local_base_freq=rope_config.rope_local_base_freq,
-            original_max_position_embeddings=rope_config.original_max_position_embeddings,
-            rope_interleave=rope_config.rope_interleave,
+            # Flat rope field copies: ``None`` for NoPE models so that
+            # ``initialize_rope`` / ``TextModel`` / ``Attention`` can detect
+            # "this model has no RoPE" structurally.
+            rope_type=rope_config.rope_type if rope_config is not None else None,
+            rope_theta=rope_config.rope_theta if rope_config is not None else None,
+            rope_scaling=rope_config.rope_scaling if rope_config is not None else None,
+            partial_rotary_factor=(
+                rope_config.partial_rotary_factor if rope_config is not None else None
+            ),
+            rope_local_base_freq=(
+                rope_config.rope_local_base_freq if rope_config is not None else None
+            ),
+            original_max_position_embeddings=(
+                rope_config.original_max_position_embeddings
+                if rope_config is not None
+                else None
+            ),
+            rope_interleave=(
+                rope_config.rope_interleave if rope_config is not None else False
+            ),
             **mrope_fields,
             max_position_embeddings=getattr(config, "max_position_embeddings", 0),
             tie_word_embeddings=(

src/mobius/_configs_test.py

@@ -33,7 +33,13 @@ def test_default_values(self):
         assert config.hidden_size == DEFAULT_INT
         assert config.num_hidden_layers == DEFAULT_INT
         assert config.rms_norm_eps == pytest.approx(1e-6)
-        assert config.rope_type == "default"
+        # rope_type defaults to None (NoPE) so that directly-constructed
+        # ArchitectureConfig instances without an explicit rope_type are
+        # treated as having no RoPE, matching the signal from_transformers
+        # uses for NoPE models. The other RoPE fields keep inert numeric
+        # defaults so that specifying only rope_type="default" is enough
+        # for test / reproducer configs.
+        assert config.rope_type is None
         assert config.rope_theta == pytest.approx(10_000.0)
         assert config.partial_rotary_factor == pytest.approx(1.0)
         assert config.attn_qkv_bias is False
@@ -125,6 +131,9 @@ class FakeLlamaConfig:
             rms_norm_eps = 1e-5
             rope_theta = 10000.0
             rope_scaling = None
+            # Real HuggingFace LlamaConfig populates rope_parameters in
+            # __post_init__ for any model that declares RoPE support.
+            rope_parameters: ClassVar[dict] = {"rope_type": "default"}
 
         config = ArchitectureConfig.from_transformers(FakeLlamaConfig())
         assert config.vocab_size == 32000
@@ -232,30 +241,128 @@ class FakeConfig:
         assert config.bos_token_id == 2
         assert config.eos_token_id == 1
 
+    def test_from_transformers_nope_model_has_none_rope(self):
+        """NoPE models (e.g. NemotronH) get ``rope=None`` and ``rope_type=None``.
+
+        This is the Phase 1 fix for the silent-RoPE-on-NoPE-models bug:
+        when the HuggingFace config declares neither ``rope_parameters``
+        nor ``rope_scaling`` nor the legacy ``rotary_dim`` / ``rotary_pct``
+        / ``rotary_emb_base`` fields, the resulting ``ArchitectureConfig``
+        must express "no RoPE" structurally so that ``initialize_rope``
+        returns ``None`` and ``TextModel`` skips rotary encoding entirely.
+        """
+
+        class FakeNemotronH:
+            # Minimal NemotronH-like config: carries a stale ``rope_theta``
+            # as dead data but declares NO ``rope_parameters`` / ``rope_scaling``
+            # / ``rotary_*`` fields — so this is a NoPE model.
+            model_type = "nemotron_h"
+            num_attention_heads = 8
+            num_key_value_heads = 2
+            num_hidden_layers = 4
+            vocab_size = 128
+            hidden_size = 64
+            intermediate_size = 128
+            hidden_act = "relu2"
+            max_position_embeddings = 128
+            head_dim = 8
+            pad_token_id = 0
+            rms_norm_eps = 1e-6
+            rope_theta = 10_000.0  # stale — ignored because no rope_parameters
+
+        config = ArchitectureConfig.from_transformers(FakeNemotronH())
+        # Sub-config is None: no RoPE data exists at all.
+        assert config.rope is None
+        # Flat fields are all None: no spurious "default" values.
+        assert config.rope_type is None
+        assert config.rope_theta is None
+        assert config.partial_rotary_factor is None
+        assert config.rope_scaling is None
+        assert config.rope_local_base_freq is None
+        assert config.original_max_position_embeddings is None
+        # rope_interleave stays at its inert False default.
+        assert config.rope_interleave is False
+
+    def test_from_transformers_legacy_rotary_dim_enables_rope(self):
+        """GPT-J / CodeGen-style legacy configs use ``rotary_dim``."""
+
+        class FakeGPTJ:
+            model_type = "gptj"
+            num_attention_heads = 4
+            num_key_value_heads = 4
+            num_hidden_layers = 2
+            vocab_size = 128
+            hidden_size = 64
+            intermediate_size = 128
+            hidden_act = "gelu"
+            max_position_embeddings = 128
+            head_dim = 16
+            pad_token_id = 0
+            rms_norm_eps = 1e-6
+            rotary_dim = 8  # legacy partial-RoPE signal
+
+        config = ArchitectureConfig.from_transformers(FakeGPTJ())
+        # Legacy rotary_dim activates RoPE with partial_rotary_factor = 8/16.
+        assert config.rope is not None
+        assert config.rope_type == "default"
+        assert config.partial_rotary_factor == pytest.approx(0.5)
+
 
 class TestExtractRopeConfig:
     """Unit tests for _extract_rope_config helper."""
 
     def test_defaults_when_no_rope_attrs(self):
-        """Bare config with no rope attrs yields sensible defaults."""
+        """Bare config with no RoPE signal yields ``None`` (NoPE model)."""
 
         class Bare:
             pass
 
         result = _extract_rope_config(Bare())
+        assert result is None
+
+    def test_rope_theta_alone_is_not_a_rope_signal(self):
+        """rope_theta without rope_parameters/rope_scaling is not a RoPE signal.
+
+        For example NemotronH carries ``rope_theta`` as dead data while
+        declaring no actual RoPE support — so the absence of
+        ``rope_parameters`` / ``rope_scaling`` / legacy rotary fields must
+        produce ``None`` (NoPE), not a spurious ``RoPEConfig``.
+        """
+
+        class Cfg:
+            # No rope_scaling, no rope_parameters — just a stale rope_theta.
+            rope_theta = 10_000.0
+
+        assert _extract_rope_config(Cfg()) is None
+
+    def test_rope_parameters_activates_rope(self):
+        """``rope_parameters`` on the HF config is the modern RoPE signal."""
+
+        class Cfg:
+            rope_parameters: ClassVar[dict] = {"rope_type": "default"}
+
+        result = _extract_rope_config(Cfg())
+        assert result is not None
+        assert result.rope_type == "default"
+
+    def test_legacy_rotary_dim_activates_rope(self):
+        """Legacy GPT-J / CodeGen configs use ``rotary_dim``."""
+
+        class Cfg:
+            rotary_dim = 64
+
+        result = _extract_rope_config(Cfg())
+        assert result is not None
         assert result.rope_type == "default"
-        assert result.rope_theta == pytest.approx(10_000.0)
-        assert result.rope_scaling is None
-        assert result.partial_rotary_factor == pytest.approx(1.0)
-        assert result.rope_local_base_freq is None
-        assert result.original_max_position_embeddings is None
 
     def test_rope_theta_from_config_attr(self):
         class Cfg:
             rope_theta = 500_000.0
-            rope_scaling = None
+            # rope_parameters triggers the RoPE path so rope_theta is read.
+            rope_parameters: ClassVar[dict] = {"rope_type": "default"}
 
         result = _extract_rope_config(Cfg())
+        assert result is not None
         assert result.rope_theta == pytest.approx(500_000.0)
 
     def test_rope_type_from_rope_scaling(self):
@@ -268,29 +375,32 @@ class Cfg:
     def test_partial_rotary_factor(self):
         class Cfg:
             partial_rotary_factor = 0.5
-            rope_scaling = None
+            rope_parameters: ClassVar[dict] = {"rope_type": "default"}
 
         result = _extract_rope_config(Cfg())
+        assert result is not None
         assert result.partial_rotary_factor == pytest.approx(0.5)
 
     def test_partial_rotary_factor_zero_is_preserved(self):
         """partial_rotary_factor=0.0 must NOT be replaced by default 1.0."""
 
         class Cfg:
             partial_rotary_factor = 0.0
-            rope_scaling = None
+            rope_parameters: ClassVar[dict] = {"rope_type": "default"}
 
         result = _extract_rope_config(Cfg())
+        assert result is not None
         assert result.partial_rotary_factor == pytest.approx(0.0)
 
     def test_rope_theta_zero_is_preserved(self):
         """rope_theta=0.0 must NOT be replaced by default 10000.0."""
 
         class Cfg:
             rope_theta = 0.0
-            rope_scaling = None
+            rope_parameters: ClassVar[dict] = {"rope_type": "default"}
 
         result = _extract_rope_config(Cfg())
+        assert result is not None
         assert result.rope_theta == pytest.approx(0.0)
 
     def test_mrope_interleaved_from_rope_scaling(self):
@@ -319,9 +429,10 @@ class Cfg:
     def test_original_max_position_embeddings(self):
         class Cfg:
             original_max_position_embeddings = 8192
-            rope_scaling = None
+            rope_parameters: ClassVar[dict] = {"rope_type": "default"}
 
         result = _extract_rope_config(Cfg())
+        assert result is not None
         assert result.original_max_position_embeddings == 8192