@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cb6d572e85
Fork 0

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 15, 07:15 UTC

[Model] NemotronH support (#18863)

Signed-off-by: Luis Vega <2478335+vegaluisjose@users.noreply.github.com> Co-authored-by: Luis Vega <2478335+vegaluisjose@users.noreply.github.com>

Luis Vega committed Jun 5, 2025 at 14:29 UTC cb6d572e85a34aec7b4409833bff12af28b0d28b
6 files changed +829
docs/models/supported_models.md
+1
@@ -346,6 +346,7 @@ Specified using `--task generate`.
346 | `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ |
347 | `MPTForCausalLM` | MPT, MPT-Instruct, MPT-Chat, MPT-StoryWriter | `mosaicml/mpt-7b`, `mosaicml/mpt-7b-storywriter`, `mosaicml/mpt-30b`, etc. | | ✅︎ |
348 | `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ |
349 +| `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ |
350 | `OLMoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | | ✅︎ |
351 | `OLMo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | | ✅︎ |
352 | `OLMoEForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ |
tests/models/registry.py
+2
@@ -212,6 +212,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
212 "MptForCausalLM": _HfExamplesInfo("mpt", is_available_online=False),
213 "MPTForCausalLM": _HfExamplesInfo("mosaicml/mpt-7b"),
214 "NemotronForCausalLM": _HfExamplesInfo("nvidia/Minitron-8B-Base"),
215 + "NemotronHForCausalLM": _HfExamplesInfo("nvidia/Nemotron-H-8B-Base-8K",
216 + trust_remote_code=True),
217 "OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),
218 "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),
219 "OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),
vllm/model_executor/models/nemotron_h.py new
+565
@@ -0,0 +1,565 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from https://github.com/vllm-project/vllm/blob/94d8ec8d2bcb4ec55e33022b313c7e978edf05e1/vllm/model_executor/models/bamba.py
4 +# Copyright 2024 HuggingFace Inc. team. All rights reserved.
5 +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +"""Inference-only NemotronH model."""
19 +from collections.abc import Iterable
20 +from typing import Optional
21 +
22 +import torch
23 +from torch import nn
24 +
25 +from vllm.attention.layer import Attention
26 +from vllm.config import CacheConfig, VllmConfig
27 +from vllm.distributed import divide, get_tensor_model_parallel_world_size
28 +from vllm.distributed.parallel_state import get_pp_group
29 +from vllm.forward_context import get_forward_context
30 +from vllm.model_executor.layers.activation import ReLUSquaredActivation
31 +from vllm.model_executor.layers.layernorm import RMSNorm
32 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
33 + QKVParallelLinear,
34 + RowParallelLinear)
35 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
36 +from vllm.model_executor.layers.mamba.mamba2_metadata import (
37 + Mamba2Metadata, prepare_mamba2_metadata)
38 +from vllm.model_executor.layers.mamba.mamba_mixer2 import (
39 + MambaMixer2, extra_groups_for_head_shards)
40 +from vllm.model_executor.layers.quantization import QuantizationConfig
41 +from vllm.model_executor.layers.vocab_parallel_embedding import (
42 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
43 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
44 +from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,
45 + SupportsLoRA, SupportsPP,
46 + SupportsQuant,
47 + SupportsV0Only)
48 +from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
49 + MambaCacheParams)
50 +from vllm.model_executor.models.utils import (
51 + AutoWeightsLoader, make_empty_intermediate_tensors_factory, make_layers,
52 + maybe_prefix)
53 +from vllm.model_executor.sampling_metadata import SamplingMetadata
54 +from vllm.sequence import IntermediateTensors
55 +from vllm.transformers_utils.configs import NemotronHConfig
56 +from vllm.utils import LayerBlockType
57 +
58 +
59 +class NemotronHMLP(nn.Module):
60 +
61 + def __init__(
62 + self,
63 + config: NemotronHConfig,
64 + quant_config: Optional[QuantizationConfig] = None,
65 + bias: bool = False,
66 + ) -> None:
67 + super().__init__()
68 + self.up_proj = MergedColumnParallelLinear(
69 + input_size=config.hidden_size,
70 + output_sizes=[config.intermediate_size],
71 + bias=bias,
72 + quant_config=quant_config,
73 + )
74 + self.down_proj = RowParallelLinear(
75 + input_size=config.intermediate_size,
76 + output_size=config.hidden_size,
77 + bias=bias,
78 + quant_config=quant_config,
79 + )
80 + self.act_fn = ReLUSquaredActivation()
81 +
82 + def forward(self, x: torch.Tensor):
83 + x, _ = self.up_proj(x)
84 + x = self.act_fn(x)
85 + x, _ = self.down_proj(x)
86 + return x
87 +
88 +
89 +class NemotronHMLPDecoderLayer(nn.Module):
90 +
91 + def __init__(
92 + self,
93 + config: NemotronHConfig,
94 + layer_idx: int,
95 + cache_config: Optional[CacheConfig] = None,
96 + quant_config: Optional[QuantizationConfig] = None,
97 + prefix: str = "",
98 + ) -> None:
99 + super().__init__()
100 + self.config = config
101 +
102 + self.mixer = NemotronHMLP(config,
103 + quant_config=quant_config,
104 + bias=config.mlp_bias)
105 +
106 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
107 +
108 + def forward(
109 + self,
110 + hidden_states: torch.Tensor,
111 + residual: Optional[torch.Tensor],
112 + **kwargs,
113 + ):
114 + if residual is None:
115 + residual = hidden_states
116 + hidden_states = self.norm(hidden_states)
117 + else:
118 + hidden_states, residual = self.norm(hidden_states, residual)
119 +
120 + hidden_states = self.mixer(hidden_states)
121 + return hidden_states, residual
122 +
123 +
124 +class NemotronHMambaDecoderLayer(nn.Module):
125 +
126 + def __init__(
127 + self,
128 + config: NemotronHConfig,
129 + layer_idx: int,
130 + cache_config: Optional[CacheConfig] = None,
131 + quant_config: Optional[QuantizationConfig] = None,
132 + prefix: str = "",
133 + ) -> None:
134 + super().__init__()
135 + self.config = config
136 + self.mixer = MambaMixer2(
137 + hidden_size=config.hidden_size,
138 + ssm_state_size=config.ssm_state_size,
139 + conv_kernel_size=config.conv_kernel,
140 + intermediate_size=config.expand * config.hidden_size,
141 + use_conv_bias=config.use_conv_bias,
142 + use_bias=config.use_bias,
143 + n_groups=config.n_groups,
144 + num_heads=config.mamba_num_heads,
145 + head_dim=config.mamba_head_dim,
146 + rms_norm_eps=config.rms_norm_eps,
147 + activation=config.mamba_hidden_act,
148 + quant_config=quant_config,
149 + )
150 +
151 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
152 +
153 + def forward(
154 + self,
155 + hidden_states: torch.Tensor,
156 + residual: Optional[torch.Tensor],
157 + mamba_cache_params: MambaCacheParams,
158 + mamba2_metadata: Mamba2Metadata,
159 + **kwargs,
160 + ):
161 + if residual is None:
162 + residual = hidden_states
163 + hidden_states = self.norm(hidden_states)
164 + else:
165 + hidden_states, residual = self.norm(hidden_states, residual)
166 +
167 + hidden_states = self.mixer(hidden_states, mamba_cache_params,
168 + mamba2_metadata)
169 + return hidden_states, residual
170 +
171 +
172 +class NemotronHAttention(nn.Module):
173 +
174 + def __init__(
175 + self,
176 + config: NemotronHConfig,
177 + layer_idx: int,
178 + cache_config: Optional[CacheConfig] = None,
179 + quant_config: Optional[QuantizationConfig] = None,
180 + prefix: str = "",
181 + ) -> None:
182 + super().__init__()
183 + self.hidden_size = config.hidden_size
184 + tp_size = get_tensor_model_parallel_world_size()
185 + self.total_num_heads = config.num_attention_heads
186 + assert self.total_num_heads % tp_size == 0
187 + self.num_heads = self.total_num_heads // tp_size
188 + self.total_num_kv_heads = config.num_key_value_heads
189 + if self.total_num_kv_heads >= tp_size:
190 + # Number of KV heads is greater than TP size, so we partition
191 + # the KV heads across multiple tensor parallel GPUs.
192 + assert self.total_num_kv_heads % tp_size == 0
193 + else:
194 + # Number of KV heads is less than TP size, so we replicate
195 + # the KV heads across multiple tensor parallel GPUs.
196 + assert tp_size % self.total_num_kv_heads == 0
197 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
198 + self.head_dim = config.hidden_size // self.total_num_heads
199 + self.q_size = self.num_heads * self.head_dim
200 + self.kv_size = self.num_kv_heads * self.head_dim
201 + self.scaling = self.head_dim**-0.5
202 +
203 + self.qkv_proj = QKVParallelLinear(
204 + config.hidden_size,
205 + self.head_dim,
206 + self.total_num_heads,
207 + self.total_num_kv_heads,
208 + bias=False,
209 + quant_config=quant_config,
210 + )
211 + self.o_proj = RowParallelLinear(
212 + self.total_num_heads * self.head_dim,
213 + config.hidden_size,
214 + bias=False,
215 + quant_config=quant_config,
216 + )
217 +
218 + self.attn = Attention(
219 + self.num_heads,
220 + self.head_dim,
221 + self.scaling,
222 + num_kv_heads=self.num_kv_heads,
223 + cache_config=cache_config,
224 + prefix=f"{prefix}.attn",
225 + )
226 +
227 + def forward(
228 + self,
229 + hidden_states: torch.Tensor,
230 + **kwargs,
231 + ) -> torch.Tensor:
232 + qkv, _ = self.qkv_proj(hidden_states)
233 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
234 + attn_output = self.attn(q, k, v)
235 + output, _ = self.o_proj(attn_output)
236 + return output
237 +
238 +
239 +class NemotronHAttentionDecoderLayer(nn.Module):
240 +
241 + def __init__(
242 + self,
243 + config: NemotronHConfig,
244 + layer_idx: int,
245 + cache_config: Optional[CacheConfig] = None,
246 + quant_config: Optional[QuantizationConfig] = None,
247 + prefix: str = "",
248 + ) -> None:
249 + super().__init__()
250 +
251 + self.mixer = NemotronHAttention(
252 + config,
253 + layer_idx,
254 + cache_config,
255 + quant_config,
256 + prefix,
257 + )
258 +
259 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
260 +
261 + def forward(
262 + self,
263 + positions: torch.Tensor,
264 + hidden_states: torch.Tensor,
265 + residual: Optional[torch.Tensor],
266 + **kwargs,
267 + ):
268 + if residual is None:
269 + residual = hidden_states
270 + hidden_states = self.norm(hidden_states)
271 + else:
272 + hidden_states, residual = self.norm(hidden_states, residual)
273 +
274 + hidden_states = self.mixer(hidden_states=hidden_states)
275 + return hidden_states, residual
276 +
277 +
278 +ALL_DECODER_LAYER_TYPES = {
279 + "M": NemotronHMambaDecoderLayer,
280 + "-": NemotronHMLPDecoderLayer,
281 + "*": NemotronHAttentionDecoderLayer,
282 +}
283 +
284 +
285 +class NemotronHModel(nn.Module):
286 +
287 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
288 + super().__init__()
289 +
290 + config: NemotronHConfig = vllm_config.model_config.hf_config
291 + cache_config = vllm_config.cache_config
292 + quant_config = vllm_config.quant_config
293 + lora_config = vllm_config.lora_config
294 +
295 + self.config = config
296 + lora_vocab = ((lora_config.lora_extra_vocab_size *
297 + (lora_config.max_loras or 1)) if lora_config else 0)
298 + self.vocab_size = config.vocab_size + lora_vocab
299 + self.org_vocab_size = config.vocab_size
300 +
301 + self.embed_tokens = VocabParallelEmbedding(
302 + self.vocab_size,
303 + config.hidden_size,
304 + org_num_embeddings=config.vocab_size,
305 + )
306 +
307 + def get_layer(prefix: str):
308 + layer_idx = int(prefix.rsplit(".", 1)[1])
309 + layer_class = ALL_DECODER_LAYER_TYPES[
310 + config.hybrid_override_pattern[layer_idx]]
311 + return layer_class(
312 + config,
313 + layer_idx,
314 + cache_config,
315 + quant_config=quant_config,
316 + prefix=prefix,
317 + )
318 +
319 + self.start_layer, self.end_layer, self.layers = make_layers(
320 + len(config.hybrid_override_pattern),
321 + get_layer,
322 + prefix=f"{prefix}.layers")
323 + self.make_empty_intmd_tensors = make_empty_intermediate_tensors_factory(
324 + ["hidden_states", "residual"], config.hidden_size)
325 +
326 + self.norm_f = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
327 +
328 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
329 + return self.embed_tokens(input_ids)
330 +
331 + def forward(
332 + self,
333 + input_ids: torch.Tensor,
334 + positions: torch.Tensor,
335 + mamba_cache_params: MambaCacheParams,
336 + intermediate_tensors: Optional[IntermediateTensors] = None,
337 + inputs_embeds: Optional[torch.Tensor] = None,
338 + ) -> torch.Tensor:
339 +
340 + attn_metadata = get_forward_context().attn_metadata
341 +
342 + mamba2_metadata = prepare_mamba2_metadata(
343 + chunk_size=self.config.chunk_size,
344 + attn_metadata=attn_metadata,
345 + )
346 +
347 + if get_pp_group().is_first_rank:
348 + if inputs_embeds is not None:
349 + hidden_states = inputs_embeds
350 + else:
351 + hidden_states = self.get_input_embeddings(input_ids)
352 + residual = None
353 + else:
354 + assert intermediate_tensors is not None
355 + hidden_states = intermediate_tensors["hidden_states"]
356 + residual = intermediate_tensors["residual"]
357 +
358 + residual = None
359 + num_non_mamba_layers = 0
360 + for i in range(len(self.layers)):
361 + layer = self.layers[i]
362 + layer_mamba_cache_params = None
363 + if isinstance(layer, NemotronHMambaDecoderLayer):
364 + layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
365 + i - num_non_mamba_layers)
366 + else:
367 + num_non_mamba_layers += 1
368 +
369 + hidden_states, residual = layer(
370 + positions=positions,
371 + hidden_states=hidden_states,
372 + residual=residual,
373 + mamba_cache_params=layer_mamba_cache_params,
374 + mamba2_metadata=mamba2_metadata,
375 + )
376 +
377 + if not get_pp_group().is_last_rank:
378 + return IntermediateTensors({
379 + "hidden_states": hidden_states,
380 + "residual": residual
381 + })
382 + hidden_states, _ = self.norm_f(hidden_states, residual)
383 + return hidden_states
384 +
385 + def load_weights(self, weights: Iterable[tuple[str,
386 + torch.Tensor]]) -> set[str]:
387 + attb_params_mapping = {
388 + "q_proj": "q",
389 + "k_proj": "k",
390 + "v_proj": "v",
391 + }
392 +
393 + params_dict = dict(self.named_parameters())
394 + loaded_params: set[str] = set()
395 + for name, loaded_weight in weights:
396 + if "embeddings" in name:
397 + name = name.replace("embeddings", "embed_tokens")
398 +
399 + if "A_log" in name:
400 + name = name.replace("A_log", "A")
401 + loaded_weight = loaded_weight.to(torch.float32)
402 +
403 + if "D" in name:
404 + loaded_weight = loaded_weight.to(torch.float32)
405 +
406 + if "dt_bias" in name:
407 + loaded_weight = loaded_weight.to(torch.float32)
408 +
409 + # load attn params
410 + if any(proj in name for proj in ["q_proj", "k_proj", "v_proj"]):
411 + weight_name = next(proj
412 + for proj in ["q_proj", "k_proj", "v_proj"]
413 + if proj in name)
414 + name = name.replace(weight_name, "qkv_proj")
415 + param = params_dict[name]
416 + weight_loader = param.weight_loader
417 + weight_loader(param, loaded_weight,
418 + attb_params_mapping[weight_name])
419 + # load other params
420 + else:
421 + param = params_dict[name]
422 + weight_loader = getattr(param, "weight_loader",
423 + default_weight_loader)
424 + weight_loader(param, loaded_weight)
425 +
426 + loaded_params.add(name)
427 + return loaded_params
428 +
429 +
430 +class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
431 + IsHybrid, SupportsV0Only, SupportsQuant):
432 + packed_modules_mapping = {
433 + "qkv_proj": [
434 + "q_proj",
435 + "k_proj",
436 + "v_proj",
437 + ],
438 + "gate_up_proj": ["up_proj", "down_proj"]
439 + }
440 +
441 + # LoRA specific attributes
442 + embedding_modules = {
443 + "embed_tokens": "input_embeddings",
444 + "lm_head": "output_embeddings",
445 + }
446 + embedding_padding_modules = ["lm_head"]
447 +
448 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
449 + config = vllm_config.model_config.hf_config
450 + self.vllm_config = vllm_config
451 + self.model_config = vllm_config.model_config
452 + cache_config = vllm_config.cache_config
453 + lora_config = vllm_config.lora_config
454 + scheduler_config = vllm_config.scheduler_config
455 + assert not cache_config.enable_prefix_caching, \
456 + "NemotronH currently does not support prefix caching"
457 +
458 + self.quant_config = vllm_config.quant_config
459 +
460 + super().__init__()
461 + self.config = config
462 + self.scheduler_config = scheduler_config
463 + self.model = NemotronHModel(vllm_config=vllm_config,
464 + prefix=maybe_prefix(prefix, "model"))
465 + self.unpadded_vocab_size = config.vocab_size
466 + if lora_config:
467 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
468 + self.lm_head = ParallelLMHead(
469 + self.unpadded_vocab_size,
470 + config.hidden_size,
471 + org_num_embeddings=config.vocab_size,
472 + padding_size=DEFAULT_VOCAB_PADDING_SIZE
473 + # We need bigger padding if using lora for kernel
474 + # compatibility
475 + if not lora_config else lora_config.lora_vocab_padding_size,
476 + )
477 + # Used to track and store by the Mamba cache between steps.
478 + self.mamba_cache: Optional[MambaCacheManager] = None
479 +
480 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
481 + config.vocab_size)
482 +
483 + self.make_empty_intmd_tensors = (self.model.make_empty_intmd_tensors)
484 +
485 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
486 + return self.model.get_input_embeddings(input_ids)
487 +
488 + def forward(self,
489 + input_ids: torch.Tensor,
490 + positions: torch.Tensor,
491 + intermediate_tensors: Optional[IntermediateTensors] = None,
492 + inputs_embeds: Optional[torch.Tensor] = None,
493 + **kwargs):
494 + if self.mamba_cache is None:
495 +
496 + num_mamba_layers = self.model_config.get_num_layers_by_block_type(
497 + self.vllm_config.parallel_config, LayerBlockType.mamba)
498 +
499 + self.mamba_cache = MambaCacheManager(
500 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,
501 + *self._get_mamba_cache_shape())
502 + mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
503 + hidden_states = self.model(input_ids, positions, mamba_cache_params,
504 + intermediate_tensors, inputs_embeds)
505 +
506 + return hidden_states
507 +
508 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
509 + return self.mamba_cache.copy_inputs_before_cuda_graphs(
510 + input_buffers, **kwargs)
511 +
512 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
513 + return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
514 +
515 + def _get_mamba_cache_shape(
516 + self) -> tuple[tuple[int, int], tuple[int, int]]:
517 + world_size = get_tensor_model_parallel_world_size()
518 + hidden_size = self.config.hidden_size
519 +
520 + conv_state_shape, temporal_state_shape = None, None
521 +
522 + intermediate_size = self.config.expand * hidden_size
523 +
524 + # if n_groups is not divisible by world_size, need to extend the shards
525 + # to ensure all groups needed by a head is sharded along with it
526 + n_groups = (
527 + self.config.n_groups +
528 + extra_groups_for_head_shards(self.config.n_groups, world_size))
529 +
530 + # - heads and n_groups are TP-ed
531 + conv_dim = (intermediate_size +
532 + 2 * n_groups * self.config.ssm_state_size)
533 + conv_state_shape = (
534 + divide(conv_dim, world_size),
535 + self.config.conv_kernel - 1,
536 + )
537 +
538 + # These are not TP-ed as they depend on A, dt_bias, D
539 + # - they are typically small
540 + # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
541 + temporal_state_shape = (
542 + divide(self.config.mamba_num_heads, world_size),
543 + self.config.mamba_head_dim,
544 + self.config.ssm_state_size,
545 + )
546 + return conv_state_shape, temporal_state_shape
547 +
548 + def compute_logits(
549 + self,
550 + hidden_states: torch.Tensor,
551 + sampling_metadata: SamplingMetadata,
552 + ) -> Optional[torch.Tensor]:
553 + logits = self.logits_processor(self.lm_head, hidden_states,
554 + sampling_metadata)
555 + return logits
556 +
557 + def load_weights(self, weights: Iterable[tuple[str,
558 + torch.Tensor]]) -> set[str]:
559 + # update name in weights before passing to loader
560 + updated_weights = []
561 + for name, loaded_weight in weights:
562 + name = name.replace("backbone", "model")
563 + updated_weights.append((name, loaded_weight))
564 + loader = AutoWeightsLoader(self)
565 + return loader.load_weights(updated_weights)
vllm/model_executor/models/registry.py
+1
@@ -92,6 +92,7 @@ _TEXT_GENERATION_MODELS = {
92 "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
93 "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
94 "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
95 + "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
96 "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
97 "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
98 "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
vllm/transformers_utils/configs/__init__.py
+2
@@ -23,6 +23,7 @@ from vllm.transformers_utils.configs.mlp_speculator import MLPSpeculatorConfig
23 from vllm.transformers_utils.configs.moonvit import MoonViTConfig
24 from vllm.transformers_utils.configs.mpt import MPTConfig
25 from vllm.transformers_utils.configs.nemotron import NemotronConfig
26 +from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig
27 from vllm.transformers_utils.configs.nvlm_d import NVLM_D_Config
28 from vllm.transformers_utils.configs.ovis import OvisConfig
29 from vllm.transformers_utils.configs.skyworkr1v import SkyworkR1VChatConfig
@@ -50,6 +51,7 @@ __all__ = [
51 "MoonViTConfig",
52 "KimiVLConfig",
53 "NemotronConfig",
54 + "NemotronHConfig",
55 "NVLM_D_Config",
56 "OvisConfig",
57 "SkyworkR1VChatConfig",
vllm/transformers_utils/configs/nemotron_h.py new
+258
@@ -0,0 +1,258 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Copyright 2024 HuggingFace Inc. team. All rights reserved.
4 +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
5 +#
6 +# Licensed under the Apache License, Version 2.0 (the "License");
7 +# you may not use this file except in compliance with the License.
8 +# You may obtain a copy of the License at
9 +#
10 +# http://www.apache.org/licenses/LICENSE-2.0
11 +#
12 +# Unless required by applicable law or agreed to in writing, software
13 +# distributed under the License is distributed on an "AS IS" BASIS,
14 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 +# See the License for the specific language governing permissions and
16 +# limitations under the License.
17 +"""NemotronH model configuration"""
18 +
19 +import regex as re
20 +from transformers.configuration_utils import PretrainedConfig
21 +from transformers.utils import logging
22 +
23 +logger = logging.get_logger(__name__)
24 +
25 +
26 +class NemotronHConfig(PretrainedConfig):
27 + r"""
28 + This is the configuration class to store the configuration of a
29 + [`NemotronHModel`]. It is used to instantiate a NemotronH model according
30 + to the specified arguments, defining the model architecture. Instantiating
31 + a configuration with the defaults will yield a similar configuration to
32 + that of the NemotronH-v0.1 model.
33 + Args:
34 + vocab_size (`int`, *optional*, defaults to 131072):
35 + Vocabulary size of the NemotronH model. Defines the number of
36 + different tokens that can be represented by the `inputs_ids`
37 + passed when calling [`NemotronHModel`]
38 + tie_word_embeddings (`bool`, *optional*, defaults to `False`):
39 + Whether the model's input and output word embeddings should be
40 + tied. Note that this is only relevant if the model has a output
41 + word embedding layer.
42 + hidden_size (`int`, *optional*, defaults to 4096):
43 + Dimension of the hidden representations.
44 + intermediate_size (`int`, *optional*, defaults to 21504):
45 + Dimension of the MLP representations.
46 + num_hidden_layers (`int`, *optional*, defaults to 52):
47 + Number of hidden layers in the Transformer encoder.
48 + hybrid_override_pattern (`str`, *optional*, defaults to
49 + `"M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"`):
50 + The pattern of the hybrid model. The pattern is a string of
51 + characters where each character represents
52 + M: Mamba2, *: Attention, -: MLP
53 + num_attention_heads (`int`, *optional*, defaults to 32):
54 + Number of attention heads for each attention layer in the
55 + Transformer encoder.
56 + attention_head_dim (`int`, *optional*, defaults to 128):
57 + Dimension of each attention head.
58 + num_key_value_heads (`int`, *optional*, defaults to 8):
59 + This is the number of key_value heads that should be used to
60 + implement Grouped Query Attention. If
61 + `num_key_value_heads=num_attention_heads`, the model will use
62 + Multi Head Attention (MHA), if `num_key_value_heads=1` the model
63 + will use Multi Query Attention (MQA) otherwise GQA is used.
64 + mlp_hidden_act (`str`, *optional*, defaults to "relu2"):
65 + The non-linear activation function in the MLP layers.
66 + attention_bias (`bool`, *optional*, defaults to `False`):
67 + Whether to use bias in attention layers.
68 + mlp_bias (`bool`, *optional*, defaults to `False`):
69 + Whether to use bias in MLP layers.
70 + use_bias (`bool`, *optional*, defaults to `False`):
71 + Whether to use bias in the model.
72 + initializer_range (`float`, *optional*, defaults to 0.02):
73 + The standard deviation of the truncated_normal_initializer for
74 + initializing all weight matrices.
75 + layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
76 + The epsilon used by the layer normalization layers.
77 + residual_in_fp32 (`bool`, *optional*, defaults to `False`):
78 + Whether or not residuals should be in `float32`. If set to `False`
79 + residuals will keep the same `dtype` as the rest of the model.
80 + use_cache (`bool`, *optional*, defaults to `True`):
81 + Whether or not the model should return the last key/values
82 + attentions (not used by all models). Only relevant if
83 + `config.is_decoder=True`.
84 + num_logits_to_keep (`int` or `None`, *optional*, defaults to 1):
85 + Number of prompt logits to calculate during generation. If `None`,
86 + all logits will be calculated. If an integer value, only last
87 + `num_logits_to_keep` logits will be calculated.
88 + pad_token_id (`int`, *optional*, defaults to 0):
89 + The id of the padding token.
90 + bos_token_id (`int`, *optional*, defaults to 1):
91 + The id of the "beginning-of-sequence" token.
92 + eos_token_id (`int`, *optional*, defaults to 2):
93 + The id of the "end-of-sequence" token.
94 + sliding_window (`int`, *optional*, defaults to None):
95 + Sliding window attention window size.
96 + max_position_embeddings (`int`, *optional*, defaults to 4096):
97 + The maximum sequence length that this model might ever be used
98 + with.
99 + attention_dropout (`float`, *optional*, defaults to 0.0):
100 + The dropout ratio for the attention probabilities.
101 + hidden_dropout (`float`, *optional*, defaults to 0.0):
102 + The dropout ratio for the hidden states.
103 + use_mamba_kernels (`bool`, *optional*, defaults to `True`):
104 + Flag indicating whether or not to use the fast mamba kernels.
105 + These are available only if `mamba-ssm` and `causal-conv1d`
106 + are installed, and the mamba modules are running on a CUDA device.
107 + ssm_state_size (`int`, *optional*, defaults to 128):
108 + The dimension of the mamba state space latents.
109 + mamba_num_heads (`int`, *optional*, defaults to 128):
110 + Number of heads in Mamba layers.
111 + mamba_n_groups (`int`, *optional*, defaults to 8):
112 + Number of groups in Mamba layers.
113 + mamba_head_dim (`int`, *optional*, defaults to 64):
114 + Dimension of each Mamba head.
115 + mamba_d_conv (`int`, *optional*, defaults to 4):
116 + The size of the mamba convolution kernel.
117 + mamba_expand (`int`, *optional*, defaults to 2):
118 + Expanding factor used to determine the mamba intermediate size.
119 + mamba_hidden_act (`str`, *optional*, defaults to "silu"):
120 + The non-linear activation function in the Mamba layers.
121 + mamba_dt_min (`float`, *optional*, defaults to 0.001):
122 + Minimum value for the time step in Mamba.
123 + mamba_dt_max (`float`, *optional*, defaults to 0.1):
124 + Maximum value for the time step in Mamba.
125 + mamba_dt_limit (`tuple`, *optional*, defaults to (0.0, float("inf"))):
126 + Limits for the time step in Mamba.
127 + mamba_dt_init_floor (`float`, *optional*, defaults to 1e-4):
128 + Floor value for time step initialization in Mamba.
129 + mamba_conv_bias (`bool`, *optional*, defaults to `True`):
130 + Whether to use bias in the convolution layer of the mamba mixer
131 + block.
132 + mamba_proj_bias (`bool`, *optional*, defaults to `False`):
133 + Whether to use bias in the input and output projections of the
134 + mamba mixer block.
135 + mamba_chunk_size (`int`, *optional*, defaults to 256):
136 + Size of chunks for Mamba processing.
137 + rescale_prenorm_residual (`bool`, *optional*, defaults to `True`):
138 + Whether to rescale the pre-normalization residual connections.
139 + """
140 +
141 + model_type = "nemotron_h"
142 + keys_to_ignore_at_inference = ["past_key_values"]
143 +
144 + def __init__(
145 + self,
146 + vocab_size=131072,
147 + tie_word_embeddings=False,
148 + hidden_size=4096,
149 + intermediate_size=21504,
150 + num_hidden_layers=52,
151 + 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-",
152 + num_attention_heads=32,
153 + attention_head_dim=128,
154 + num_key_value_heads=8, # nemo: num_query_groups
155 + mlp_hidden_act="relu2",
156 + attention_bias=False,
157 + mlp_bias=False,
158 + use_bias=False,
159 + initializer_range=0.02, # nemo: init_method_std
160 + layer_norm_epsilon=1e-5, # nemo: layernorm_epsilon
161 + residual_in_fp32=False, # Megatron Core default value
162 + use_cache=True,
163 + num_logits_to_keep=1,
164 + pad_token_id=0,
165 + bos_token_id=1,
166 + eos_token_id=2,
167 + sliding_window=None,
168 + max_position_embeddings=4096,
169 + attention_dropout=0.0,
170 + hidden_dropout=0.0, # * ADDED
171 + use_mamba_kernels=True,
172 + ssm_state_size=128, # mamba_state_size
173 + mamba_num_heads=128,
174 + mamba_n_groups=8, # nemo: mamba_ssm_ngroups = num_heads
175 + mamba_head_dim=64,
176 + mamba_d_conv=4,
177 + mamba_expand=2,
178 + mamba_hidden_act="silu",
179 + mamba_dt_min=0.001,
180 + mamba_dt_max=0.1,
181 + mamba_dt_limit=(0.0, float("inf")),
182 + mamba_dt_init_floor=1e-4,
183 + mamba_conv_bias=True,
184 + mamba_proj_bias=False,
185 + mamba_chunk_size=256,
186 + rescale_prenorm_residual=True,
187 + **kwargs,
188 + ):
189 + self.vocab_size = vocab_size
190 + self.tie_word_embeddings = tie_word_embeddings
191 + self.hidden_size = hidden_size
192 + self.intermediate_size = intermediate_size
193 + self.num_hidden_layers = num_hidden_layers
194 + self.hybrid_override_pattern = hybrid_override_pattern
195 + self.num_attention_heads = num_attention_heads
196 + self.attention_head_dim = attention_head_dim
197 + self.sliding_window = sliding_window
198 + self.max_position_embeddings = max_position_embeddings
199 + self.attention_dropout = attention_dropout
200 + self.hidden_dropout = hidden_dropout
201 +
202 + # Validate hybrid_override_pattern
203 + # M: Mamba2, *: Attention, -: MLP
204 + assert len(self.hybrid_override_pattern) == self.num_hidden_layers, (
205 + "hybrid_override_pattern must have same length as "
206 + "num_hidden_layers")
207 + assert re.match(r"^[*-M]+$", self.hybrid_override_pattern), (
208 + "hybrid_override_pattern must only contain characters "
209 + "'M', '*', or '-'")
210 +
211 + # for backward compatibility
212 + if num_key_value_heads is None:
213 + num_key_value_heads = num_attention_heads
214 +
215 + self.num_key_value_heads = num_key_value_heads
216 + self.mlp_hidden_act = mlp_hidden_act
217 + self.attention_bias = attention_bias
218 + self.mlp_bias = mlp_bias
219 + self.use_bias = use_bias
220 + self.initializer_range = initializer_range
221 + self.layer_norm_epsilon = layer_norm_epsilon
222 + self.residual_in_fp32 = residual_in_fp32
223 +
224 + self.use_cache = use_cache
225 + self.num_logits_to_keep = num_logits_to_keep
226 +
227 + self.use_mamba_kernels = use_mamba_kernels
228 + self.n_groups = mamba_n_groups
229 + self.mamba_head_dim = mamba_head_dim
230 + self.ssm_state_size = ssm_state_size
231 + self.mamba_num_heads = mamba_num_heads
232 + self.conv_kernel = mamba_d_conv
233 + self.expand = mamba_expand
234 + self.mamba_hidden_act = mamba_hidden_act
235 + self.time_step_min = mamba_dt_min
236 + self.time_step_max = mamba_dt_max
237 + self.time_step_limit = mamba_dt_limit
238 + self.time_step_floor = mamba_dt_init_floor
239 + self.use_conv_bias = mamba_conv_bias
240 + self.mamba_proj_bias = mamba_proj_bias
241 + self.chunk_size = mamba_chunk_size
242 + self.rescale_prenorm_residual = rescale_prenorm_residual
243 +
244 + super().__init__(
245 + pad_token_id=pad_token_id,
246 + bos_token_id=bos_token_id,
247 + eos_token_id=eos_token_id,
248 + tie_word_embeddings=tie_word_embeddings,
249 + **kwargs,
250 + )
251 +
252 + @property
253 + def layers_block_type(self):
254 + return [
255 + "mamba" if self.hybrid_override_pattern[i] == "M" else
256 + "attention" if self.hybrid_override_pattern[i] == "*" else "mlp"
257 + for i in range(self.num_hidden_layers)
258 + ]