[speechlm2] Temporarily support native Nemotron-V3 with Automodel 0.4

nemo/collections/speechlm2/parts/automodel_compat.py

@@ -0,0 +1,121 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from functools import wraps
+
+from torch import nn
+
+logger = logging.getLogger(__name__)
+
+_COMPATIBILITY_MARKER = "_nemo_speech_nemotron_h_layer_compatibility"
+
+
+class _ModuleDictLayersAdapter:
+    """Expose a ModuleDict through the list operations used by Automodel 0.4."""
+
+    def __init__(self, layers: nn.ModuleDict):
+        self._layers = layers
+        # ModuleDict preserves insertion order. Native Nemotron-V3 registers
+        # its numeric string keys in decoder order; sorting would put "10"
+        # before "2" and produce the wrong layer sequence.
+        self._keys = tuple(layers)
+
+    def __iter__(self):
+        return iter(self._layers.values())
+
+    def __len__(self):
+        return len(self._layers)
+
+    def __getitem__(self, index: int):
+        return self._layers[self._key(index)]
+
+    def __setitem__(self, index: int, layer: nn.Module):
+        self._layers[self._key(index)] = layer
+
+    def _key(self, index: int) -> str:
+        if not isinstance(index, int):
+            raise TypeError(f"ModuleDict layer indices must be integers, got {type(index).__name__}")
+        return self._keys[index]
+
+
+class _BackboneLayersAdapter:
+    """Minimal adapter for Automodel 0.4's exact ``model.backbone.layers`` access.
+
+    This deliberately is not an ``nn.Module`` and does not emulate a complete
+    Hugging Face backbone. The pinned parallelizer only reads ``.layers``; the
+    real Automodel 0.4 smoke test guards that narrow contract.
+    """
+
+    def __init__(self, layers):
+        self.layers = layers
+
+
+def install_nemotron_h_layer_compatibility() -> bool:
+    """Make Automodel 0.4's Nemotron-H parallelizer accept the native Nemotron-V3 layout.
+
+    Automodel 0.4 assumes every ``NemotronHForCausalLM`` stores decoder blocks
+    in ``model.backbone.layers``. The native Nemotron-V3 implementation stores
+    them in ``model.model.layers`` as a ``ModuleDict``. Temporarily expose that
+    container through the legacy interface while the old parallelizer runs.
+
+    Returns:
+        ``True`` when the compatibility wrapper was installed; ``False`` when
+        Automodel already contains the upstream fix or the wrapper was installed
+        previously.
+    """
+    try:
+        from nemo_automodel.components.distributed import parallelizer
+    except ImportError as error:
+        logger.warning("Nemotron-V3 layer compatibility not installed: %s", error)
+        return False
+
+    if hasattr(parallelizer, "_nemotronh_decoder_blocks"):
+        return False
+
+    try:
+        strategy_cls = parallelizer.NemotronHParallelizationStrategy
+    except AttributeError as error:
+        logger.warning("Nemotron-V3 layer compatibility not installed: %s", error)
+        return False
+
+    original_parallelize = strategy_cls.parallelize
+    if getattr(original_parallelize, _COMPATIBILITY_MARKER, False):
+        return False
+
+    @wraps(original_parallelize)
+    def parallelize_with_native_layout(self, model, *args, **kwargs):
+        if hasattr(model, "backbone"):
+            return original_parallelize(self, model, *args, **kwargs)
+
+        inner_model = getattr(model, "model", None)
+        layers = getattr(inner_model, "layers", None)
+        if isinstance(layers, nn.ModuleDict):
+            legacy_layers = _ModuleDictLayersAdapter(layers)
+        elif isinstance(layers, nn.ModuleList):
+            legacy_layers = layers
+        else:
+            return original_parallelize(self, model, *args, **kwargs)
+
+        model.backbone = _BackboneLayersAdapter(legacy_layers)
+        try:
+            return original_parallelize(self, model, *args, **kwargs)
+        finally:
+            del model.backbone
+
+    setattr(parallelize_with_native_layout, _COMPATIBILITY_MARKER, True)
+    strategy_cls.parallelize = parallelize_with_native_layout
+    # TODO(Dongji): Remove this shim after Speech pins an Automodel build containing PR #2638.
+    logger.warning("Installed temporary native Nemotron-V3 compatibility for the Automodel 0.4 parallelizer")
+    return True

nemo/collections/speechlm2/parts/pretrained.py

@@ -88,6 +88,10 @@ def load_pretrained_automodel_llm(
     """
     from nemo_automodel import NeMoAutoModelForCausalLM
 
+    from nemo.collections.speechlm2.parts.automodel_compat import install_nemotron_h_layer_compatibility
+
+    install_nemotron_h_layer_compatibility()
+
     if pretrained_weights:
         return NeMoAutoModelForCausalLM.from_pretrained(
             model_path_or_name,

tests/collections/speechlm2/test_automodel_compat.py

@@ -0,0 +1,130 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import builtins
+import importlib.util
+import logging
+import sys
+from pathlib import Path
+from types import ModuleType
+
+import torch.nn as nn
+
+_COMPAT_PATH = Path(__file__).resolve().parents[3] / "nemo/collections/speechlm2/parts/automodel_compat.py"
+_COMPAT_SPEC = importlib.util.spec_from_file_location("_speechlm2_automodel_compat", _COMPAT_PATH)
+_COMPAT_MODULE = importlib.util.module_from_spec(_COMPAT_SPEC)
+_COMPAT_SPEC.loader.exec_module(_COMPAT_MODULE)
+install_nemotron_h_layer_compatibility = _COMPAT_MODULE.install_nemotron_h_layer_compatibility
+
+
+class _NativeNemotronV3(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Module()
+        self.model.layers = nn.ModuleDict({"0": nn.Linear(2, 2), "1": nn.ReLU()})
+
+
+class _HuggingFaceNemotronH(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = nn.Module()
+        self.backbone.layers = nn.ModuleList([nn.Linear(2, 2), nn.ReLU()])
+
+
+def _install_fake_legacy_automodel(monkeypatch, *, has_upstream_fix=False, has_strategy=True):
+    class LegacyNemotronHParallelizationStrategy:
+        """Minimal reproduction of Automodel 0.4's hard-coded layer access."""
+
+        def parallelize(self, model):
+            layers = model.backbone.layers
+            visited = list(layers)
+            layers[0] = nn.Identity()
+            return visited
+
+    parallelizer = ModuleType("nemo_automodel.components.distributed.parallelizer")
+    if has_strategy:
+        parallelizer.NemotronHParallelizationStrategy = LegacyNemotronHParallelizationStrategy
+    if has_upstream_fix:
+        parallelizer._nemotronh_decoder_blocks = lambda model: model
+
+    distributed = ModuleType("nemo_automodel.components.distributed")
+    distributed.parallelizer = parallelizer
+    components = ModuleType("nemo_automodel.components")
+    components.distributed = distributed
+    automodel = ModuleType("nemo_automodel")
+    automodel.components = components
+
+    monkeypatch.setitem(sys.modules, "nemo_automodel", automodel)
+    monkeypatch.setitem(sys.modules, "nemo_automodel.components", components)
+    monkeypatch.setitem(sys.modules, "nemo_automodel.components.distributed", distributed)
+    monkeypatch.setitem(sys.modules, "nemo_automodel.components.distributed.parallelizer", parallelizer)
+    return parallelizer
+
+
+def test_legacy_parallelizer_supports_native_nemotron_v3(monkeypatch):
+    parallelizer = _install_fake_legacy_automodel(monkeypatch)
+
+    assert install_nemotron_h_layer_compatibility() is True
+
+    model = _NativeNemotronV3()
+    visited = parallelizer.NemotronHParallelizationStrategy().parallelize(model)
+
+    assert len(visited) == 2
+    assert all(isinstance(layer, nn.Module) for layer in visited)
+    assert isinstance(model.model.layers["0"], nn.Identity)
+    assert not hasattr(model, "backbone")
+
+
+def test_legacy_parallelizer_preserves_huggingface_layout(monkeypatch):
+    parallelizer = _install_fake_legacy_automodel(monkeypatch)
+
+    install_nemotron_h_layer_compatibility()
+
+    model = _HuggingFaceNemotronH()
+    original_backbone = model.backbone
+    visited = parallelizer.NemotronHParallelizationStrategy().parallelize(model)
+
+    assert len(visited) == 2
+    assert model.backbone is original_backbone
+    assert isinstance(model.backbone.layers[0], nn.Identity)
+
+
+def test_compatibility_is_not_installed_when_upstream_fix_exists(monkeypatch):
+    parallelizer = _install_fake_legacy_automodel(monkeypatch, has_upstream_fix=True)
+    original_parallelize = parallelizer.NemotronHParallelizationStrategy.parallelize
+
+    assert install_nemotron_h_layer_compatibility() is False
+    assert parallelizer.NemotronHParallelizationStrategy.parallelize is original_parallelize
+
+
+def test_missing_parallelizer_is_best_effort(monkeypatch, caplog):
+    original_import = builtins.__import__
+
+    def import_without_parallelizer(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "nemo_automodel.components.distributed" and "parallelizer" in fromlist:
+            raise ImportError("parallelizer moved")
+        return original_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", import_without_parallelizer)
+    with caplog.at_level(logging.WARNING):
+        assert install_nemotron_h_layer_compatibility() is False
+    assert "Nemotron-V3 layer compatibility not installed: parallelizer moved" in caplog.text
+
+
+def test_missing_strategy_is_best_effort(monkeypatch, caplog):
+    _install_fake_legacy_automodel(monkeypatch, has_strategy=False)
+
+    with caplog.at_level(logging.WARNING):
+        assert install_nemotron_h_layer_compatibility() is False
+    assert "NemotronHParallelizationStrategy" in caplog.text