diff --git a/src/art/megatron/model_support/handlers/nemotron_h.py b/src/art/megatron/model_support/handlers/nemotron_h.py index 7f9bafcbe..d6d1ce5cd 100644 --- a/src/art/megatron/model_support/handlers/nemotron_h.py +++ b/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" ] diff --git a/src/art/megatron/model_support/registry.py b/src/art/megatron/model_support/registry.py index dd92f2e5e..6edaefff5 100644 --- a/src/art/megatron/model_support/registry.py +++ b/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( diff --git a/tests/integration/megatron/model_support/hf_parity.py b/tests/integration/megatron/model_support/hf_parity.py index e7d944756..6e051da79 100644 --- a/tests/integration/megatron/model_support/hf_parity.py +++ b/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)): diff --git a/tests/integration/megatron/model_support/hf_parity_canonicalization.py b/tests/integration/megatron/model_support/hf_parity_canonicalization.py index ba84b1069..c2dea0271 100644 --- a/tests/integration/megatron/model_support/hf_parity_canonicalization.py +++ b/tests/integration/megatron/model_support/hf_parity_canonicalization.py @@ -7,6 +7,10 @@ _FUSED_MOE_EXPERT_PATTERN = re.compile( r"^(?P.*\.mlp\.experts)\.(?Pgate_up_proj|down_proj)(?:\.weight)?$" ) +_FUSED_NEMOTRON_H_EXPERT_PATTERN = re.compile( + r"^(?Pbackbone\.layers\.\d+\.mixer\.experts)\." + r"(?Pup_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 diff --git a/tests/integration/megatron/model_support/hf_parity_worker.py b/tests/integration/megatron/model_support/hf_parity_worker.py index 4b0c0a73e..92951fd87 100644 --- a/tests/integration/megatron/model_support/hf_parity_worker.py +++ b/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, diff --git a/tests/integration/megatron/model_support/test_workflow.py b/tests/integration/megatron/model_support/test_workflow.py index 115c767bf..69f086438 100644 --- a/tests/integration/megatron/model_support/test_workflow.py +++ b/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: diff --git a/tests/integration/megatron/model_support/workflow_fixtures.py b/tests/integration/megatron/model_support/workflow_fixtures.py index c3ee53fc6..119aec6da 100644 --- a/tests/integration/megatron/model_support/workflow_fixtures.py +++ b/tests/integration/megatron/model_support/workflow_fixtures.py @@ -17,7 +17,7 @@ FIXTURE_CACHE_ENV = "ART_MODEL_SUPPORT_FIXTURE_CACHE" FIXTURE_ROOT_ENV = "ART_MODEL_SUPPORT_FIXTURE_ROOT" FIXTURE_VERSION = 18 -_MODEL_FIXTURE_VERSION_OFFSETS = {"nemotron_h_moe": 7} +_MODEL_FIXTURE_VERSION_OFFSETS = {"nemotron_h_moe": 8} _CANONICAL_CACHE_VERSION = 16 _ROOT = Path("/tmp/art-models/main-merge-oracle") _CACHE_ROOT = Path("/tmp/art-model-support-workflow/hf-cache") @@ -48,7 +48,7 @@ } _TOKENIZER_FIXTURE_VERSION = 3 _FUNCTIONAL_FIXTURE_VERSION = 1 -_FUNCTIONAL_FIXTURE_VERSION_OFFSETS = {"nemotron_h_moe": 5} +_FUNCTIONAL_FIXTURE_VERSION_OFFSETS = {"nemotron_h_moe": 7} _FUNCTIONAL_REMOTE_CODE_FILES = { "nemotron_h_moe": ( "configuration_nemotron_h.py", @@ -70,6 +70,9 @@ "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16": ( "2d59de1cbd51c0adf384eb906b766d1aee0e0517" ), + "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": ( + "b3caaabed0263651a17dc1f2d4ce97e794f76c44" + ), } _MULTIMODAL = {"qwen3_5_dense", "qwen3_5_moe", "gemma4_dense", "gemma4_moe"} @@ -264,6 +267,7 @@ def _common( "moe_shared_expert_intermediate_size": 512, "n_routed_experts": 4, "num_experts_per_tok": 2, + "num_nextn_predict_layers": 0, "tie_word_embeddings": False, }, ), @@ -327,7 +331,18 @@ def _common( } # fmt: on -_FUNCTIONAL_LAYER_FIELDS = ("layer_types", "mlp_layer_types", "indexer_types") +_FUNCTIONAL_LAYER_FIELDS = ( + "layer_types", + "layers_block_type", + "mlp_layer_types", + "indexer_types", +) +_HYBRID_LAYER_SYMBOLS = { + "mamba": "M", + "moe": "E", + "attention": "*", + "mlp": "-", +} _WIDTH_TERMS = ("hidden", "intermediate", "head", "expert", "lora_rank", "topk") @@ -360,7 +375,11 @@ def _plan(depth: int, prefix: str, **values: Any) -> _FunctionalPlan: "dsv4": _plan(6, "layers", auxiliary=("mtp", "num_nextn_predict_layers")), "glm52": _plan(10, "model.layers", auxiliary=(None, "num_nextn_predict_layers")), "gpt_oss_moe": _plan(4, "model.layers"), - "nemotron_h_moe": _plan(6, "backbone.layers"), + "nemotron_h_moe": _plan( + 6, + "backbone.layers", + auxiliary=("mtp.layers", "num_nextn_predict_layers"), + ), } _FUNCTIONAL_PATTERNS = { "qwen3_dense": {"layer_types": ("full_attention",) * 2}, @@ -393,7 +412,20 @@ def _configure( } if model_key in _PLAIN_TEXT: layers, hidden, values = _PLAIN_TEXT[model_key] - text = _set(_common(config, layers=layers, hidden=hidden, **common), **values) + text = _common(config, layers=layers, hidden=hidden, **common) + values = dict(values) + hybrid_pattern = values.pop("hybrid_override_pattern", None) + if hybrid_pattern is not None: + layer_types = getattr(text, "layer_types", None) + if isinstance(layer_types, (list, tuple)): + actual_pattern = "".join( + _HYBRID_LAYER_SYMBOLS[layer_type] for layer_type in layer_types + ) + if actual_pattern != hybrid_pattern: + raise RuntimeError(f"{model_key} reduced hybrid pattern changed") + else: + text.hybrid_override_pattern = hybrid_pattern + text = _set(text, **values) if model_key == "glm52": text.vocab_size = source_vocab_size return config @@ -459,6 +491,19 @@ def _pack_qwen35_experts(path: Path, config: Any) -> None: save_file(tensors, checkpoint, metadata={"format": "pt"}) +def _restore_nemotron_h_embedding_name(path: Path) -> None: + from safetensors.torch import load_file, save_file + + checkpoint = path / "model.safetensors" + tensors = load_file(checkpoint) + source = "backbone.embedding.weight" + target = "backbone.embeddings.weight" + if source not in tensors or target in tensors: + raise RuntimeError("Nemotron-H HF embedding serialization changed") + tensors[target] = tensors.pop(source) + save_file(tensors, checkpoint, metadata={"format": "pt"}) + + def _json_sha256(value: object) -> str: return hashlib.sha256( json.dumps(value, sort_keys=True, separators=(",", ":")).encode() @@ -499,14 +544,21 @@ def _config_shape(config: dict[str, Any], *exclude: str) -> dict[str, object]: } +def _functional_depth(text: dict[str, Any], *, model_key: str) -> int: + depth = text.get("num_hidden_layers") + if depth is None and isinstance(text.get("layers_block_type"), list): + depth = len(text["layers_block_type"]) + if type(depth) is not int or depth < _functional_plan(model_key).depth: + raise RuntimeError(f"{model_key} has invalid production depth") + return depth + + def _functional_config( source: dict[str, Any], *, model_key: str ) -> tuple[dict[str, Any], dict[str, object]]: plan = _functional_plan(model_key) text = _config_text(source, plan) - source_depth = text.get("num_hidden_layers") - if type(source_depth) is not int or source_depth < plan.depth: - raise RuntimeError(f"{model_key} has invalid production depth") + source_depth = _functional_depth(text, model_key=model_key) reduced = json.loads(json.dumps(source)) reduced_text = _config_text(reduced, plan) reduced_text["num_hidden_layers"] = plan.depth @@ -523,12 +575,26 @@ def _functional_config( raise RuntimeError(f"{model_key} production {field} is incomplete") patterns[field] = values[: plan.depth] reduced_text[field] = patterns[field] - hybrid_pattern = text.get("hybrid_override_pattern") + serialized_hybrid_pattern = text.get("hybrid_override_pattern") + hybrid_pattern = serialized_hybrid_pattern + if (layer_types := text.get("layers_block_type")) is not None: + try: + layer_types_pattern = "".join( + _HYBRID_LAYER_SYMBOLS[layer_type] for layer_type in layer_types + ) + except (KeyError, TypeError) as exc: + raise RuntimeError( + f"{model_key} production layers_block_type is invalid" + ) from exc + if hybrid_pattern is not None and hybrid_pattern != layer_types_pattern: + raise RuntimeError(f"{model_key} production hybrid patterns disagree") + hybrid_pattern = layer_types_pattern if hybrid_pattern is not None: if not isinstance(hybrid_pattern, str) or len(hybrid_pattern) != source_depth: raise RuntimeError(f"{model_key} production hybrid pattern is invalid") patterns["hybrid_override_pattern"] = hybrid_pattern[: plan.depth] - reduced_text["hybrid_override_pattern"] = hybrid_pattern[: plan.depth] + if serialized_hybrid_pattern is not None: + reduced_text["hybrid_override_pattern"] = hybrid_pattern[: plan.depth] for field, expected in _FUNCTIONAL_PATTERNS.get(model_key, {}).items(): actual = patterns.get(field) if (tuple(actual) if isinstance(actual, list) else actual) != expected: @@ -630,7 +696,7 @@ def _select_functional_weights( ) -> dict[str, str]: plan = _functional_plan(model_key) text_config = _config_text(config, plan) - source_depth = int(text_config["num_hidden_layers"]) + source_depth = _functional_depth(text_config, model_key=model_key) text_layers: set[int] = set() vision_layers: set[int] = set() auxiliary_layers: set[int] = set() @@ -666,13 +732,19 @@ def _select_functional_weights( count = text_config.get(auxiliary_count) if auxiliary_count else 0 if type(count) is not int or count < 0: raise RuntimeError(f"{model_key} has invalid {auxiliary_count}") + auxiliary_depth = count + if auxiliary_prefix == "mtp.layers" and count: + mtp_layer_types = text_config.get("mtp_layers_block_type") + if not isinstance(mtp_layer_types, list) or not mtp_layer_types: + raise RuntimeError(f"{model_key} MTP layer pattern is invalid") + auxiliary_depth *= len(mtp_layer_types) expected = set( range(source_depth + (count if plan.auxiliary and not auxiliary_prefix else 0)) ) if text_layers != expected: raise RuntimeError(f"{model_key} canonical text-layer coverage changed") if auxiliary_prefix: - if auxiliary_layers != set(range(count)): + if auxiliary_layers != set(range(auxiliary_depth)): raise RuntimeError( f"{model_key} canonical auxiliary-layer coverage changed" ) @@ -890,7 +962,12 @@ def _build( config.save_pretrained(staging) if functional: (staging / "config.json").write_text(json.dumps(reduced, indent=2) + "\n") - for name in _FUNCTIONAL_REMOTE_CODE_FILES.get(model_key, ()): + remote_code_files = ( + _FUNCTIONAL_REMOTE_CODE_FILES.get(model_key, ()) + if source_config.get("auto_map") + else () + ) + for name in remote_code_files: if source_fixture is None: raise RuntimeError( f"{model_key} remote code requires a parent fixture" @@ -944,24 +1021,32 @@ def _build( if model_key == "nemotron_h_moe": config.dtype = configured_dtype fp32 = {} + backbone_prefix = None if model_key == "nemotron_h_moe": + backbone = getattr(model, "backbone", None) + backbone_prefix = "backbone" + if backbone is None: + backbone = getattr(model, "model", None) + backbone_prefix = "model" + if backbone is None: + raise RuntimeError("Nemotron-H HF fixture backbone changed") pattern = str(config.hybrid_override_pattern) for index, symbol in enumerate(pattern): - mixer = model.backbone.layers[index].mixer + mixer = backbone.layers[index].mixer if symbol == "E": torch.nn.init.normal_( mixer.gate.weight, std=float(config.initializer_range) ) if symbol == "M": - fp32[f"backbone.layers.{index}.mixer.A_log"] = ( + fp32[f"{backbone_prefix}.layers.{index}.mixer.A_log"] = ( mixer.A_log.detach().clone() ) - fp32[f"backbone.layers.{index}.mixer.D"] = ( + fp32[f"{backbone_prefix}.layers.{index}.mixer.D"] = ( mixer.D.detach().clone() ) elif symbol == "E": fp32[ - f"backbone.layers.{index}.mixer.gate.e_score_correction_bias" + f"{backbone_prefix}.layers.{index}.mixer.gate.e_score_correction_bias" ] = mixer.gate.e_score_correction_bias.detach().clone() model = model.to(torch.bfloat16) tensors = dict(model.named_parameters()) | dict(model.named_buffers()) @@ -975,10 +1060,13 @@ def _build( layer.post_attention_layernorm.weight.fill_(residual_scale) layer.post_feedforward_layernorm.weight.fill_(residual_scale) if model_key == "nemotron_h_moe": - if model._tied_weights_keys != ["lm_head.weight"]: + if backbone_prefix == "backbone": + if model._tied_weights_keys != ["lm_head.weight"]: + raise RuntimeError("Nemotron-H HF tied-weight metadata changed") + model._tied_weights_keys = {} + model.register_for_auto_class("AutoModelForCausalLM") + elif model._tied_weights_keys != {}: raise RuntimeError("Nemotron-H HF tied-weight metadata changed") - model._tied_weights_keys = {} - model.register_for_auto_class("AutoModelForCausalLM") parameters = sum(parameter.numel() for parameter in model.parameters()) model.save_pretrained( staging, @@ -986,10 +1074,12 @@ def _build( max_shard_size="2GB", **( {"save_original_format": False} - if model_key == "nemotron_h_moe" + if model_key == "nemotron_h_moe" and backbone_prefix == "backbone" else {} ), ) + if model_key == "nemotron_h_moe" and backbone_prefix == "model": + _restore_nemotron_h_embedding_name(staging) del model gc.collect() if model_key == "qwen3_5_moe": diff --git a/tests/integration/megatron/model_support/workflow_resources.py b/tests/integration/megatron/model_support/workflow_resources.py index 71a39661f..d46f8346c 100644 --- a/tests/integration/megatron/model_support/workflow_resources.py +++ b/tests/integration/megatron/model_support/workflow_resources.py @@ -441,6 +441,7 @@ class HandlerWorkflowResources(BaseModel): ), } _B300_THROUGHPUT_FINGERPRINT_OVERRIDES = { + "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": "61c9e114ce17335ccf2a644d49b0fba7f3341d0f8af0b9f7df19f940c33209c7", "Qwen/Qwen3.8-27B": "b07ee7ec6338ec021463a43a90fc96c5c5a036b4a04d90b80e1d22c1eef86774", } _H200_THROUGHPUT_FLOORS = { diff --git a/tests/integration/megatron/model_support/workflow_throughput.py b/tests/integration/megatron/model_support/workflow_throughput.py index 2ea081ea6..d82f1e443 100644 --- a/tests/integration/megatron/model_support/workflow_throughput.py +++ b/tests/integration/megatron/model_support/workflow_throughput.py @@ -36,6 +36,7 @@ _STAGE_DIR_ENV = "ART_MODEL_SUPPORT_WORKFLOW_STAGE_DIR" _LAYER_LIST_FIELDS = ( "layer_types", + "layers_block_type", "mlp_layer_types", "indexer_types", "compress_ratios", @@ -619,8 +620,13 @@ def _sized_config( sized = json.loads(json.dumps(source)) text = _text(sized) source_text = _text(source) - source_layers = int(source_text["num_hidden_layers"]) layer_fields = tuple(field for field in _LAYER_LIST_FIELDS if field in source_text) + source_depth = source_text.get("num_hidden_layers") + if source_depth is None: + if not layer_fields: + raise ValueError(f"{model_key} config is missing its layer count") + source_depth = len(source_text[layer_fields[0]]) + source_layers = int(source_depth) if num_layers > source_layers and layer_fields: raise ValueError( f"cannot expand {model_key} with per-layer fields {layer_fields}" @@ -631,6 +637,12 @@ def _sized_config( if len(values) < num_layers: raise ValueError(f"{model_key} {field} has only {len(values)} entries") text[field] = values[:num_layers] + mtp_count_field = "num_nextn_predict_layers" + if mtp_count_field in source_text: + mtp_count = source_text[mtp_count_field] + if type(mtp_count) is not int or mtp_count < 0: + raise ValueError(f"{model_key} has an invalid {mtp_count_field}") + text[mtp_count_field] = 0 hybrid_pattern = source_text.get("hybrid_override_pattern") if hybrid_pattern is not None: if not isinstance(hybrid_pattern, str) or len(hybrid_pattern) != source_layers: @@ -650,6 +662,7 @@ def _sized_config( "num_hidden_layers", *_LAYER_LIST_FIELDS, "hybrid_override_pattern", + mtp_count_field, ) if field == "num_hidden_layers" or field in source_text ],