diff --git a/nemo_automodel/components/distributed/parallelizer.py b/nemo_automodel/components/distributed/parallelizer.py index 817bbb69c6..f52e029358 100644 --- a/nemo_automodel/components/distributed/parallelizer.py +++ b/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"], } diff --git a/tests/unit_tests/distributed/test_parallelization_strategies.py b/tests/unit_tests/distributed/test_parallelization_strategies.py index 37ff12c401..218939b34c 100644 --- a/tests/unit_tests/distributed/test_parallelization_strategies.py +++ b/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."""