fix(parallelizer): resolve NemotronH decoder blocks for Nemotron-V3

nemo_automodel/components/distributed/parallelizer.py

@@ -380,6 +380,24 @@ def parallelize(
         return model
 
 
+def _nemotronh_decoder_blocks(model: nn.Module) -> tuple[nn.Module, list[nn.Module]]:
+    """Return ``(container, blocks)`` for a NemotronH model's decoder blocks.
+
+    Two distinct classes share the name ``NemotronHForCausalLM``:
+
+    * the HF model keeps its blocks in ``model.backbone.layers`` (an ``nn.ModuleList``), while
+    * the native Nemotron-V3 model (``NemotronV3Model``) keeps them in ``model.model.layers``
+      (an ``nn.ModuleDict`` keyed ``"0".."N-1"``).
+
+    ``container`` is the underlying ``ModuleList``/``ModuleDict`` (so callers can write rewrapped
+    blocks back into the model), and ``blocks`` is the ordered list of block modules.
+    """
+    inner = model.backbone if hasattr(model, "backbone") else model.model
+    container = inner.layers
+    blocks = list(container.values()) if isinstance(container, nn.ModuleDict) else list(container)
+    return container, blocks
+
+
 class NemotronHParallelizationStrategy(ParallelizationStrategy):
     """Specialized parallelization strategy for NemotronH models."""
 
@@ -401,7 +419,7 @@ def parallelize(
         assert not sequence_parallel, "Sequence parallelism is not supported for NemotronHForCausalLM"
         logger.info("Custom parallel plan is not supported for NemotronHForCausalLM. Using NemotronH-specific TP plan.")
 
-        layers: torch.nn.ModuleList = model.backbone.layers
+        block_container, layers = _nemotronh_decoder_blocks(model)
         tp_mesh = device_mesh[tp_mesh_name]
         if tp_mesh.size() > 1:
             model_tp_plan: dict[str, ParallelStyle] = {
@@ -415,7 +433,7 @@ def parallelize(
 
             parallelize_module(model, tp_mesh, model_tp_plan)
 
-            for layer in model.backbone.layers:
+            for layer in layers:
                 if layer.block_type == "mlp":
                     parallelize_module(layer, tp_mesh, mlp_tp_plan)
 
@@ -450,12 +468,16 @@ def parallelize(
                         )
 
         if activation_checkpointing:
-            for i in range(len(layers)):
-                if layers[i].block_type == "mlp":
-                    layers[i] = checkpoint_wrapper(layers[i])
-
-                if layers[i].block_type == "mamba":
-                    layers[i] = checkpoint_wrapper(layers[i])
+            # Write rewrapped blocks back into the real container (ModuleList -> int key,
+            # ModuleDict -> str key) so the model, not just the local handle, is updated.
+            block_items = (
+                block_container.items() if isinstance(block_container, nn.ModuleDict) else enumerate(block_container)
+            )
+            for key, layer in list(block_items):
+                if getattr(layer, "block_type", None) in ("mlp", "mamba"):
+                    block_container[key] = checkpoint_wrapper(layer)
+            # Refresh the local handle so the FSDP wrap below sees the wrapped blocks.
+            _, layers = _nemotronh_decoder_blocks(model)
 
         dp_mesh = get_fsdp_dp_mesh(device_mesh, dp_replicate_mesh_name, dp_shard_cp_mesh_name)
 
@@ -1528,7 +1550,7 @@ def _reduce_attrs(model, fqns: List[str]) -> List[nn.Module]:
         ],
     }
     LLM_MODEL_CLS_TO_LAYERS = {
-        "NemotronHForCausalLM": ["backbone.layers"],
+        "NemotronHForCausalLM": ["backbone.layers", "model.layers"],
         GPT2LMHeadModel: ["transformer.h"],
     }
 

tests/unit_tests/distributed/test_parallelization_strategies.py

@@ -35,6 +35,8 @@
     NemotronHParallelizationStrategy,
     ParallelizationStrategy,
     WanParallelizationStrategy,
+    _extract_model_layers,
+    _nemotronh_decoder_blocks,
     fsdp2_strategy_parallelize,
     get_parallelization_strategy,
 )
@@ -104,6 +106,55 @@ def forward(self, x):
         return x
 
 
+class MockNemotronV3Model(nn.Module):
+    """Mock of the native Nemotron-V3 model: decoder blocks live in ``model.model.layers``
+    as a ``ModuleDict`` (keyed "0".."N-1") and there is no ``backbone`` attribute."""
+
+    def __init__(self, num_layers=4):
+        super().__init__()
+
+        class MockInner(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = nn.ModuleDict()
+                for i in range(num_layers):
+                    layer = nn.Module()
+                    setattr(layer, "block_type", "mlp" if i % 2 == 0 else "attention")
+                    self.layers[str(i)] = layer
+
+        self.model = MockInner()
+        self.__class__.__name__ = "NemotronHForCausalLM"
+
+    def forward(self, x):
+        return x
+
+
+class TestNemotronHLayoutResolution:
+    """Both classes named ``NemotronHForCausalLM`` must resolve their decoder blocks
+    without an AttributeError (AM-448): the HF model exposes ``backbone.layers``
+    (``ModuleList``) and the native Nemotron-V3 model exposes ``model.layers``
+    (``ModuleDict``)."""
+
+    def test_helper_hf_backbone_modulelist(self):
+        container, blocks = _nemotronh_decoder_blocks(MockNemotronHModel())
+        assert isinstance(container, nn.ModuleList)
+        assert len(blocks) == 2
+
+    def test_helper_native_model_moduledict(self):
+        container, blocks = _nemotronh_decoder_blocks(MockNemotronV3Model(num_layers=4))
+        assert isinstance(container, nn.ModuleDict)
+        assert len(blocks) == 4  # ordered values of the ModuleDict
+
+    def test_extract_model_layers_native_has_no_backbone(self):
+        # Regression for AM-448: the registry still lists "backbone.layers", but the native
+        # model has no `backbone`; _reduce_attrs must skip it and resolve "model.layers"
+        # instead of raising.
+        assert len(_extract_model_layers(MockNemotronV3Model(num_layers=4))) == 4
+
+    def test_extract_model_layers_hf_backbone(self):
+        assert len(_extract_model_layers(MockNemotronHModel())) == 2
+
+
 @pytest.fixture
 def mock_device_mesh():
     """Create a mock device mesh for testing."""