From f8bbf56de7b2bbbea959334e2a24f4bef06d0034 Mon Sep 17 00:00:00 2001 From: danielafrimi <45691845+danielafrimi@users.noreply.github.com> Date: Sat, 7 Mar 2026 09:22:03 +0200 Subject: [PATCH] Refactor NemotronHConfig to canonical layers_block_type and add MTP block-type support (#19950) Signed-off-by: dafrimi --- python/sglang/srt/configs/nemotron_h.py | 199 ++++++++++++++++++++++-- 1 file changed, 182 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/configs/nemotron_h.py b/python/sglang/srt/configs/nemotron_h.py index 833e97d87..d24eec766 100644 --- a/python/sglang/srt/configs/nemotron_h.py +++ b/python/sglang/srt/configs/nemotron_h.py @@ -15,7 +15,6 @@ """NemotronH model configuration""" -import regex as re from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging @@ -31,6 +30,8 @@ MAMBA = "M" ATTENTION = "*" MLP = "-" MOE = "E" +DEFAULT_LAYERS_BLOCK_TYPE = ["mamba", "moe", "attention", "moe"] +DEFAULT_MTP_LAYERS_BLOCK_TYPE = ["attention", "moe"] class NemotronHConfig(PretrainedConfig): @@ -53,13 +54,17 @@ class NemotronHConfig(PretrainedConfig): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 21504): Dimension of the MLP representations. - num_hidden_layers (`int`, *optional*, defaults to 52): - Number of hidden layers in the Transformer encoder. + num_hidden_layers (`int`, *optional*): + Deprecated. Kept only for backward compatibility. The effective + layer count is derived from `layers_block_type`. hybrid_override_pattern (`str`, *optional*, defaults to `"M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"`): - The pattern of the hybrid model. The pattern is a string of - characters where each character represents - M: Mamba2, *: Attention, -: MLP + Deprecated compatibility field. Pattern string where each + character represents Mamba2 (`M`), Attention (`*`), MLP (`-`), + or MoE (`E`). + layers_block_type (`list[str]`, *optional*): + Canonical layer layout. Each entry is one of: + `"mamba"`, `"attention"`, `"mlp"`, `"moe"`. num_attention_heads (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer in the Transformer encoder. @@ -151,14 +156,72 @@ class NemotronHConfig(PretrainedConfig): model_type = "nemotron_h" keys_to_ignore_at_inference = ["past_key_values"] + @staticmethod + def _validate_layers_block_type( + layers_block_type, expected_length=None, param_name="layers_block_type" + ): + """ + Validate layers_block_type list. + Args: + layers_block_type: List of layer types to validate. + expected_length: If provided, validate the list has this length. + param_name: Parameter name for error messages. + Raises: + ValueError: If validation fails. + """ + if not isinstance(layers_block_type, list): + raise ValueError( + f"{param_name} must be a list of strings. Got type: {type(layers_block_type)}" + ) + if expected_length is not None and len(layers_block_type) != expected_length: + raise ValueError( + f"{param_name} must have length {expected_length}. Got length {len(layers_block_type)}." + ) + valid_types = {"mamba", "attention", "mlp", "moe"} + if not all(block_type in valid_types for block_type in layers_block_type): + invalid = set(layers_block_type) - valid_types + raise ValueError( + f"{param_name} contains invalid types: {invalid}. Must be one of: {valid_types}" + ) + + @staticmethod + def _resolve_layers_block_type( + layers_block_type, hybrid_override_pattern, kwargs + ) -> list[str]: + """Resolve canonical layers_block_type from new and legacy config fields.""" + # Prefer explicit kwargs override first (legacy HF path), otherwise use + # the function argument value from config fields. + pattern = kwargs.pop("hybrid_override_pattern", hybrid_override_pattern) + if layers_block_type is None: + if pattern is not None: + layers_block_type = NemotronHConfig._pattern_to_list(pattern) + else: + # Last-resort fallback to preserve compatibility when neither + # canonical nor legacy pattern fields are provided. + layers_block_type = DEFAULT_LAYERS_BLOCK_TYPE + return layers_block_type + + @staticmethod + def _resolve_mtp_layers_block_type(mtp_layers_block_type, kwargs) -> list[str]: + """Resolve canonical mtp_layers_block_type from new and legacy config fields.""" + if "mtp_hybrid_override_pattern" in kwargs: + pattern = kwargs.pop("mtp_hybrid_override_pattern") + if mtp_layers_block_type is None or mtp_layers_block_type == [ + "attention", + "moe", + ]: + mtp_layers_block_type = NemotronHConfig._pattern_to_list(pattern) + return mtp_layers_block_type + def __init__( self, vocab_size=131072, tie_word_embeddings=False, hidden_size=4096, intermediate_size=21504, - num_hidden_layers=52, + num_hidden_layers=None, # Deprecated, only for backward compatibility hybrid_override_pattern="M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-", + layers_block_type=None, num_attention_heads=32, head_dim=128, num_key_value_heads=8, # nemo: num_query_groups @@ -204,14 +267,33 @@ class NemotronHConfig(PretrainedConfig): n_group=1, topk_group=1, norm_topk_prob=True, + num_nextn_predict_layers=0, + mtp_layers_block_type=DEFAULT_MTP_LAYERS_BLOCK_TYPE, **kwargs, ): + # Compatibility parsing: normalize legacy pattern fields into canonical list fields. + layers_block_type = self._resolve_layers_block_type( + layers_block_type, hybrid_override_pattern, kwargs + ) + mtp_layers_block_type = self._resolve_mtp_layers_block_type( + mtp_layers_block_type, kwargs + ) + + # num_hidden_layers is deprecated and ignored as a source of truth. + if ( + num_hidden_layers is not None + and len(layers_block_type) != num_hidden_layers + ): + logger.warning( + f"num_hidden_layers ({num_hidden_layers}) is deprecated and doesn't match " + f"layers_block_type length ({len(layers_block_type)}). Using layers_block_type length." + ) + + # Core model attributes. self.vocab_size = vocab_size self.tie_word_embeddings = tie_word_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.hybrid_override_pattern = hybrid_override_pattern self.num_attention_heads = num_attention_heads self.head_dim = head_dim self.sliding_window = sliding_window @@ -219,14 +301,10 @@ class NemotronHConfig(PretrainedConfig): self.attention_dropout = attention_dropout self.hidden_dropout = hidden_dropout - # Validate hybrid_override_pattern - # M: Mamba2, *: Attention, -: MLP - assert ( - len(self.hybrid_override_pattern) == self.num_hidden_layers - ), "hybrid_override_pattern must have same length as num_hidden_layers" - assert re.match( - r"^[*\-ME]+$", self.hybrid_override_pattern - ), "hybrid_override_pattern must only contain characters 'M', '*', '-' or 'E'" + self._validate_layers_block_type( + layers_block_type, expected_length=None, param_name="layers_block_type" + ) + self.layers_block_type = layers_block_type # for backward compatibility if num_key_value_heads is None: @@ -244,6 +322,7 @@ class NemotronHConfig(PretrainedConfig): self.use_cache = use_cache self.num_logits_to_keep = num_logits_to_keep + # Mamba attributes. self.use_mamba_kernels = use_mamba_kernels self.mamba_n_groups = mamba_n_groups self.mamba_head_dim = mamba_head_dim @@ -260,6 +339,7 @@ class NemotronHConfig(PretrainedConfig): self.mamba_proj_bias = mamba_proj_bias self.mamba_chunk_size = mamba_chunk_size self.rescale_prenorm_residual = rescale_prenorm_residual + # MoE attributes. self.n_routed_experts = n_routed_experts self.n_shared_experts = n_shared_experts self.moe_intermediate_size = moe_intermediate_size @@ -270,6 +350,20 @@ class NemotronHConfig(PretrainedConfig): self.n_group = n_group self.topk_group = topk_group self.norm_topk_prob = norm_topk_prob + # MTP attributes. + self.num_nextn_predict_layers = num_nextn_predict_layers + + if self.num_nextn_predict_layers > 0: + if mtp_layers_block_type is None: + raise ValueError( + "mtp_layers_block_type is required when num_nextn_predict_layers > 0. " + "Please provide an explicit list of layer types for MTP layers. " + "Example: mtp_layers_block_type=['attention', 'moe']" + ) + self._validate_layers_block_type( + mtp_layers_block_type, None, "mtp_layers_block_type" + ) + self.mtp_layers_block_type = mtp_layers_block_type super().__init__( pad_token_id=pad_token_id, @@ -312,3 +406,74 @@ class NemotronHConfig(PretrainedConfig): return Mamba2CacheParams( shape=shape, layers=self.mamba_layer_ids, dtype=mamba2_state_dtype(self) ) + + @property + def num_hidden_layers(self) -> int: + """ + Number of hidden layers derived from the length of layers_block_type. + This property replaces the deprecated num_hidden_layers parameter. + """ + return len(self.layers_block_type) + + @num_hidden_layers.setter + def num_hidden_layers(self, value): + """ + Setter for backward compatibility when loading configs. + The value is ignored since num_hidden_layers is computed from layers_block_type. + """ + pass + + @property + def hybrid_override_pattern(self) -> str: + """ + Backward compatibility property. + Returns the pattern string representation of layers_block_type. + """ + return self._list_to_pattern(self.layers_block_type) + + @hybrid_override_pattern.setter + def hybrid_override_pattern(self, value): + """ + Setter for backward compatibility when loading configs. + """ + self.layers_block_type = self._pattern_to_list(value) + + @property + def mtp_hybrid_override_pattern(self) -> str: + """ + Backward compatibility property. + Returns the pattern string representation of mtp_layers_block_type. + """ + return self._list_to_pattern(self.mtp_layers_block_type) + + @mtp_hybrid_override_pattern.setter + def mtp_hybrid_override_pattern(self, value): + """Setter for backward compatibility when loading configs.""" + self.mtp_layers_block_type = self._pattern_to_list(value) + + @staticmethod + def _list_to_pattern(layers_list: list[str]) -> str: + """Convert list of layer types back to pattern string (for backward compatibility).""" + reverse_mapping = { + "mamba": MAMBA, + "moe": MOE, + "attention": ATTENTION, + "mlp": MLP, + } + return "".join(reverse_mapping[layer_type] for layer_type in layers_list) + + @staticmethod + def _pattern_to_list(pattern: str) -> list[str]: + """Convert pattern string to list of layer types (for backward compatibility).""" + if any(char not in {MAMBA, MOE, ATTENTION, MLP} for char in pattern): + raise ValueError( + "Pattern must only contain characters 'M', '*', '-' or 'E'. " + f"Got: {pattern}" + ) + pattern_mapping = { + MAMBA: "mamba", + MOE: "moe", + ATTENTION: "attention", + MLP: "mlp", + } + return [pattern_mapping[char] for char in pattern]