Support Nemotron 3.5 Lightning

src/art/megatron/model_support/handlers/nemotron_h.py

@@ -38,6 +38,8 @@
 
 _MOE_FFN_ALIGNMENT = 128
 _LOGICAL_MOE_FFN_ATTR = "art_nemotron_h_logical_moe_ffn_hidden_size"
+_ART_LORA_LAYER_PREFIX = "base_model.model.backbone.layers."
+_TRANSFORMERS_LORA_LAYER_PREFIX = "base_model.model.model.layers."
 _EXPERT_MAPPING_AXES = {
     "decoder.layers.*.mlp.experts.linear_fc1.weight*": 0,
     "decoder.layers.*.mlp.experts.linear_fc2.weight*": 1,
@@ -114,19 +116,45 @@ def _padding_sizes_from_hf_config(config: Any) -> tuple[int, int]:
 
 
 @lru_cache(maxsize=8)
-def _expert_sizes_from_adapter_base(base_model: str) -> tuple[int, int, int]:
+def _adapter_base_metadata(base_model: str) -> tuple[int, int, int, str]:
     config_path = Path(base_model) / "config.json"
     if not config_path.exists():
         from huggingface_hub import hf_hub_download
 
         config_path = Path(hf_hub_download(base_model, "config.json"))
-    config = json.loads(config_path.read_text(encoding="utf-8"))
-    config = config.get("text_config") or config
+    source = json.loads(config_path.read_text(encoding="utf-8"))
+    config = source.get("text_config") or source
     logical, internal = _padding_sizes_from_hf_config(config)
     experts = int(config.get("n_routed_experts", 0) or 0)
     if experts <= 0:
         raise RuntimeError("Nemotron-H config is missing n_routed_experts")
-    return logical, internal, experts
+    prefix = (
+        _ART_LORA_LAYER_PREFIX
+        if source.get("auto_map")
+        else _TRANSFORMERS_LORA_LAYER_PREFIX
+    )
+    return logical, internal, experts, prefix
+
+
+def _external_lora_layer_prefix(adapter_config: dict[str, Any]) -> str:
+    base_model = adapter_config.get("base_model_name_or_path")
+    if not isinstance(base_model, str) or not base_model:
+        raise RuntimeError(
+            "Nemotron-H LoRA conversion requires base_model_name_or_path"
+        )
+    return _adapter_base_metadata(base_model)[3]
+
+
+def _replace_lora_layer_prefix(
+    tensors: dict[str, torch.Tensor],
+    *,
+    source: str,
+    target: str,
+) -> dict[str, torch.Tensor]:
+    return {
+        target + key.removeprefix(source) if key.startswith(source) else key: tensor
+        for key, tensor in tensors.items()
+    }
 
 
 _PACKED_EXPERT_LORA_SLOTS = tuple(
@@ -194,7 +222,7 @@ def _convert_lora_padding(
         raise RuntimeError(
             "Nemotron-H LoRA conversion requires base_model_name_or_path"
         )
-    logical, internal, experts = _expert_sizes_from_adapter_base(base_model)
+    logical, internal, experts, _ = _adapter_base_metadata(base_model)
     if (
         not pad
         and (
@@ -542,6 +570,9 @@ def configure_provider_for_runtime(self, provider: Any) -> None:
             )
         if getattr(provider, "virtual_pipeline_model_parallel_size", None) is not None:
             raise ValueError("Nemotron-H does not support virtual pipeline parallelism")
+        provider.mtp_num_layers = None
+        provider.mtp_hybrid_override_pattern = None
+        provider.mtp_loss_scaling_factor = None
         _configure_moe_padding(provider)
         provider.use_mamba_mem_eff_path = True
 
@@ -660,11 +691,16 @@ def to_vllm_lora_tensors(
         *,
         adapter_config: dict[str, Any],
     ) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
+        tensors = _convert_lora_padding(
+            tensors,
+            adapter_config=adapter_config,
+            pad=False,
+        )
         return (
-            _convert_lora_padding(
+            _replace_lora_layer_prefix(
                 tensors,
-                adapter_config=adapter_config,
-                pad=False,
+                source=_ART_LORA_LAYER_PREFIX,
+                target=_external_lora_layer_prefix(adapter_config),
             ),
             adapter_config,
         )
@@ -675,6 +711,11 @@ def from_vllm_lora_tensors(
         *,
         adapter_config: dict[str, Any],
     ) -> dict[str, torch.Tensor]:
+        tensors = _replace_lora_layer_prefix(
+            tensors,
+            source=_external_lora_layer_prefix(adapter_config),
+            target=_ART_LORA_LAYER_PREFIX,
+        )
         return _convert_lora_padding(
             tensors,
             adapter_config=adapter_config,
@@ -720,6 +761,12 @@ def get_forward_kwargs(self, model: Any, **kwargs: Any) -> dict[str, Any]:
             "packed_seq_params": None,
         }
 
+    def prepare_hf_reference_config(self, config: Any) -> None:
+        if int(getattr(config, "num_nextn_predict_layers", 0) or 0):
+            config.num_nextn_predict_layers = 0
+            if hasattr(config, "mtp_layers_block_type"):
+                config.mtp_layers_block_type = []
+
     def prepare_hf_reference_model_class(self, model_class: type[Any]) -> type[Any]:
         if model_class.__name__ != "NemotronHForCausalLM":
             raise TypeError("Nemotron-H HF reference model class changed")
@@ -745,23 +792,30 @@ def dtype_plan(model: Any, dtype: torch.dtype) -> dict[str, torch.dtype]:
         )
 
     def prepare_hf_reference_model(self, model: Any) -> Any:
+        backbone = getattr(model, "backbone", None)
+        prefix = "backbone"
+        if backbone is None:
+            backbone = getattr(model, "model", None)
+            prefix = "model"
+        if backbone is None:
+            raise RuntimeError("Nemotron-H HF backbone changed")
         pattern = str(model.config.hybrid_override_pattern)
         fp32_names = {
             name
             for index, symbol in enumerate(pattern)
             for name in (
                 (
-                    f"backbone.layers.{index}.mixer.A_log",
-                    f"backbone.layers.{index}.mixer.D",
+                    f"{prefix}.layers.{index}.mixer.A_log",
+                    f"{prefix}.layers.{index}.mixer.D",
                 )
                 if symbol == "M"
-                else (f"backbone.layers.{index}.mixer.gate.e_score_correction_bias",)
+                else (f"{prefix}.layers.{index}.mixer.gate.e_score_correction_bias",)
                 if symbol == "E"
                 else ()
             )
         }
         state = model.state_dict()
-        params_dtype = model.backbone.embeddings.weight.dtype
+        params_dtype = backbone.embeddings.weight.dtype
         invalid = {
             name: tensor.dtype
             for name, tensor in state.items()
@@ -773,7 +827,7 @@ def prepare_hf_reference_model(self, model: Any) -> Any:
         if invalid:
             raise RuntimeError(f"Nemotron-H HF reference precision changed: {invalid}")
         expected_names = [
-            f"backbone.layers.{index}.mixer"
+            f"{prefix}.layers.{index}.mixer"
             for index, symbol in enumerate(pattern)
             if symbol == "M"
         ]

src/art/megatron/model_support/registry.py

@@ -277,7 +277,10 @@
     key="nemotron_h_moe",
     handler_key=_NEMOTRON_H_HANDLER_KEY,
     is_moe=True,
-    model_names=("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",),
+    model_names=(
+        "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
+        "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16",
+    ),
     default_target_modules=_NEMOTRON_H_TARGET_MODULES,
     native_vllm_lora_status=_VALIDATED_NATIVE_VLLM_LORA_STATUS,
     dependency_floor=DependencyFloor(

tests/integration/megatron/model_support/hf_parity.py

@@ -307,7 +307,9 @@ def set_hf_config_num_layers(config: Any, num_layers: int) -> str:
             if isinstance(layer_types, (list, tuple)):
                 setattr(config_view, "layer_types", list(layer_types[:num_layers]))
             hybrid_pattern = getattr(config_view, "hybrid_override_pattern", None)
-            if isinstance(hybrid_pattern, str):
+            if isinstance(hybrid_pattern, str) and not isinstance(
+                layer_types, (list, tuple)
+            ):
                 config_view.hybrid_override_pattern = hybrid_pattern[:num_layers]
             mlp_only_layers = getattr(config_view, "mlp_only_layers", None)
             if isinstance(mlp_only_layers, (list, tuple)):

tests/integration/megatron/model_support/hf_parity_canonicalization.py

@@ -7,6 +7,10 @@
 _FUSED_MOE_EXPERT_PATTERN = re.compile(
     r"^(?P<prefix>.*\.mlp\.experts)\.(?P<param>gate_up_proj|down_proj)(?:\.weight)?$"
 )
+_FUSED_NEMOTRON_H_EXPERT_PATTERN = re.compile(
+    r"^(?P<prefix>backbone\.layers\.\d+\.mixer\.experts)\."
+    r"(?P<param>up_proj|down_proj)(?:\.weight)?$"
+)
 
 
 def _strip_language_model_prefix(key: str) -> str:
@@ -40,8 +44,21 @@ def hf_tensor_map_to_art_canonical(
     expected_keys: set[str],
 ) -> dict[str, torch.Tensor]:
     canonical: dict[str, torch.Tensor] = {}
+    uses_backbone_prefix = any(
+        expected.startswith("backbone.") for expected in expected_keys
+    )
     for key, value in hf_tensor_map.items():
+        if key.startswith("model.") and uses_backbone_prefix:
+            key = f"backbone.{key.removeprefix('model.')}"
         match = _FUSED_MOE_EXPERT_PATTERN.match(key)
+        nemotron_h_match = _FUSED_NEMOTRON_H_EXPERT_PATTERN.match(key)
+        if match is None and nemotron_h_match is not None:
+            prefix = nemotron_h_match.group("prefix")
+            param = nemotron_h_match.group("param")
+            if value.ndim == 3 and f"{prefix}.0.{param}.weight" in expected_keys:
+                for expert in range(int(value.shape[0])):
+                    canonical[f"{prefix}.{expert}.{param}.weight"] = value[expert]
+                continue
         if match is None:
             canonical[key] = value
             continue

tests/integration/megatron/model_support/hf_parity_worker.py

@@ -514,15 +514,17 @@ def _load_hf_model(
             if isinstance(auto_map, dict)
             else None
         )
-        if not isinstance(class_reference, str) or not class_reference:
-            raise RuntimeError("HF reference model class is unavailable")
-        model_class = prepare_model_class(
-            get_class_from_dynamic_module(
+        if isinstance(class_reference, str) and class_reference:
+            reference_model_class = get_class_from_dynamic_module(
                 class_reference,
                 base_model,
                 revision=getattr(config, "_commit_hash", None),
             )
-        )
+        else:
+            reference_model_class = cast(Any, AutoModelForCausalLM)._model_mapping[
+                type(config)
+            ]
+        model_class = prepare_model_class(reference_model_class)
     model = model_class.from_pretrained(
         base_model,
         config=config,

tests/integration/megatron/model_support/test_workflow.py

@@ -886,7 +886,7 @@ def test_validated_architecture_representative_models_are_fixed() -> None:
     ]
 
 
-def test_qwen38_uses_its_measured_throughput_fingerprint() -> None:
+def test_aliases_use_their_measured_throughput_fingerprints() -> None:
     qwen35 = handler_workflow_resources_for_base_model("Qwen/Qwen3.5-27B")
     qwen38 = handler_workflow_resources_for_base_model("Qwen/Qwen3.8-27B")
     assert qwen35 is not None and qwen35.e2e_throughput is not None
@@ -900,6 +900,16 @@ def test_qwen38_uses_its_measured_throughput_fingerprint() -> None:
     assert qwen38_config.thresholds["b300"].calibration_fingerprint == (
         "b07ee7ec6338ec021463a43a90fc96c5c5a036b4a04d90b80e1d22c1eef86774"
     )
+    lightning = handler_workflow_resources_for_base_model(
+        "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
+    )
+    assert lightning is not None and lightning.e2e_throughput is not None
+    assert lightning.e2e_throughput.throughput is not None
+    assert lightning.e2e_throughput.throughput.thresholds[
+        "b300"
+    ].calibration_fingerprint == (
+        "61c9e114ce17335ccf2a644d49b0fba7f3341d0f8af0b9f7df19f940c33209c7"
+    )
 
 
 def test_dsv4_runtime_stages_use_full_model_resources() -> None: