A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Model] NemotronH Support (#22349)
Signed-off-by: Daniel Afrimi <danielafrimi8@gmail.com>
danielafrimi committed
Aug 11, 2025 at 14:09 UTC
14a5d903ab826b723a24a2d89631006394de76a1
2 files changed
+23
-7
vllm/model_executor/models/nemotron_h.py
+21
-5
@@ -64,20 +64,32 @@ class NemotronHMLP(nn.Module):
64
def __init__(
65
self,
66
config: NemotronHConfig,
67
+ layer_idx: int,
68
quant_config: Optional[QuantizationConfig] = None,
69
bias: bool = False,
70
prefix: str = "",
71
) -> None:
72
super().__init__()
73
+
74
+ hybrid_override_pattern = config.hybrid_override_pattern
75
+ mlp_index = hybrid_override_pattern[:layer_idx + 1].count("-") - 1
76
+ if isinstance(config.intermediate_size, list):
77
+ if len(config.intermediate_size) == 1:
78
+ intermediate_size = config.intermediate_size[0]
79
+ else:
80
+ intermediate_size = config.intermediate_size[mlp_index]
81
+ else:
82
+ intermediate_size = config.intermediate_size
83
+
84
self.up_proj = ColumnParallelLinear(
85
input_size=config.hidden_size,
74
- output_size=config.intermediate_size,
86
+ output_size=intermediate_size,
87
bias=bias,
88
quant_config=quant_config,
89
prefix=f"{prefix}.up_proj",
90
)
91
self.down_proj = RowParallelLinear(
80
- input_size=config.intermediate_size,
92
+ input_size=intermediate_size,
93
output_size=config.hidden_size,
94
bias=bias,
95
quant_config=quant_config,
@@ -110,6 +122,7 @@ class NemotronHMLPDecoderLayer(nn.Module):
122
quant_config=quant_config,
123
bias=config.mlp_bias,
124
prefix=f"{prefix}.mixer",
125
+ layer_idx=layer_idx,
126
)
127
128
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -146,7 +159,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
159
hidden_size=config.hidden_size,
160
ssm_state_size=config.ssm_state_size,
161
conv_kernel_size=config.conv_kernel,
149
- intermediate_size=config.expand * config.hidden_size,
162
+ intermediate_size=config.mamba_num_heads * config.mamba_head_dim,
163
use_conv_bias=config.use_conv_bias,
164
use_bias=config.use_bias,
165
n_groups=config.n_groups,
@@ -205,7 +218,10 @@ class NemotronHAttention(nn.Module):
218
# the KV heads across multiple tensor parallel GPUs.
219
assert tp_size % self.total_num_kv_heads == 0
220
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
208
- self.head_dim = config.hidden_size // self.total_num_heads
221
+ if hasattr(config, "head_dim") and config.head_dim is not None:
222
+ self.head_dim = config.head_dim
223
+ else:
224
+ self.head_dim = config.hidden_size // self.total_num_heads
225
self.q_size = self.num_heads * self.head_dim
226
self.kv_size = self.num_kv_heads * self.head_dim
227
self.scaling = self.head_dim**-0.5
@@ -481,7 +497,7 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
497
"""
498
parallel_config = vllm_config.parallel_config
499
hf_config = vllm_config.model_config.hf_config
484
- intermediate_size = hf_config.expand * hf_config.hidden_size
500
+ intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim
501
502
return MambaStateShapeCalculator.mamba2_state_shape(
503
intermediate_size=intermediate_size,
vllm/transformers_utils/configs/nemotron_h.py
+2
-2
@@ -151,7 +151,7 @@ class NemotronHConfig(PretrainedConfig):
151
num_hidden_layers=52,
152
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-",
153
num_attention_heads=32,
154
- attention_head_dim=128,
154
+ head_dim=128,
155
num_key_value_heads=8, # nemo: num_query_groups
156
mlp_hidden_act="relu2",
157
attention_bias=False,
@@ -194,7 +194,7 @@ class NemotronHConfig(PretrainedConfig):
194
self.num_hidden_layers = num_hidden_layers
195
self.hybrid_override_pattern = hybrid_override_pattern
196
self.num_attention_heads = num_attention_heads
197
- self.attention_head_dim = attention_head_dim
197
+ self.head_dim = head_dim
198
self.sliding_window = sliding_window
199
self.max_position_embeddings = max_position_embeddings
200
self.attention_dropout = attention_dropout