diff --git a/nemo/collections/speechlm2/parts/automodel_compat.py b/nemo/collections/speechlm2/parts/automodel_compat.py new file mode 100644 index 000000000000..14f8cf40c3ac --- /dev/null +++ b/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 diff --git a/nemo/collections/speechlm2/parts/pretrained.py b/nemo/collections/speechlm2/parts/pretrained.py index e0d485bc0529..397a44cf9e32 100644 --- a/nemo/collections/speechlm2/parts/pretrained.py +++ b/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, diff --git a/tests/collections/speechlm2/test_automodel_compat.py b/tests/collections/speechlm2/test_automodel_compat.py new file mode 100644 index 000000000000..4955c9bd31f1 --- /dev/null +++ b/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