chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Refactor NemotronHConfig to canonical layers_block_type and add MTP block-type support (#19950)

Browse Source 
Signed-off-by: dafrimi <dafrimi@nvidia.com>
This commit is contained in:
danielafrimi
2026-03-07 09:22:03 +02:00
committed by GitHub
parent b91fb8393e
commit f8bbf56de7
1 changed files with 182 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

199
python/sglang/srt/configs/nemotron_h.py
View File

@@ -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]