@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3aa2b6a637
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, 17:15 UTC

[Model] Update support for NemotronNAS models (#15008)

Signed-off-by: Nave Assaf <nassaf@nvidia.com>

Naveassaf committed Mar 31, 2025 at 15:35 UTC 3aa2b6a637140554f60669027279f5301e5d6e05
8 files changed +524 -133
docs/source/models/supported_models.md
+1 -1
@@ -224,7 +224,7 @@ See [this page](#generative-models) for more information on how to use generativ
224 * ✅︎
225 - * `DeciLMForCausalLM`
226 * DeciLM
227 - * `Deci/DeciLM-7B`, `Deci/DeciLM-7B-instruct`, etc.
227 + * `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc.
228 *
229 * ✅︎
230 - * `DeepseekForCausalLM`
tests/models/registry.py
+1 -1
@@ -112,7 +112,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
112 "Cohere2ForCausalLM": _HfExamplesInfo("CohereForAI/c4ai-command-r7b-12-2024", # noqa: E501
113 trust_remote_code=True),
114 "DbrxForCausalLM": _HfExamplesInfo("databricks/dbrx-instruct"),
115 - "DeciLMForCausalLM": _HfExamplesInfo("Deci/DeciLM-7B-instruct",
115 + "DeciLMForCausalLM": _HfExamplesInfo("nvidia/Llama-3_3-Nemotron-Super-49B-v1", # noqa: E501
116 trust_remote_code=True),
117 "DeepseekForCausalLM": _HfExamplesInfo("deepseek-ai/deepseek-llm-7b-chat"),
118 "DeepseekV2ForCausalLM": _HfExamplesInfo("deepseek-ai/DeepSeek-V2-Lite-Chat", # noqa: E501
vllm/config.py
+20 -1
@@ -411,6 +411,7 @@ class ModelConfig:
411
412 self.is_attention_free = self._init_attention_free()
413 self.is_hybrid = self._init_is_hybrid()
414 + self.has_noops = self._init_has_noops()
415 self.has_inner_state = self._init_has_inner_state()
416
417 if current_platform.is_neuron():
@@ -510,6 +511,10 @@ class ModelConfig:
511 def _init_is_hybrid(self) -> bool:
512 return self.registry.is_hybrid_model(self.architectures)
513
514 + def _init_has_noops(self) -> bool:
515 + architectures = getattr(self.hf_config, "architectures", [])
516 + return self.registry.is_noops_model(architectures)
517 +
518 def _init_has_inner_state(self) -> bool:
519 return self.registry.model_has_inner_state(self.architectures)
520
@@ -872,6 +877,14 @@ class ModelConfig:
877 return getattr(self.hf_config.attn_config, "kv_n_heads",
878 self.hf_config.num_attention_heads)
879
880 + if self.hf_config.model_type == "nemotron-nas":
881 + for block in self.hf_config.block_configs:
882 + if not block.attention.no_op:
883 + return self.hf_config.num_attention_heads \
884 + // block.attention.n_heads_in_group
885 +
886 + raise RuntimeError("Couldn't determine number of kv heads")
887 +
888 if self.is_attention_free:
889 return 0
890
@@ -940,7 +953,9 @@ class ModelConfig:
953 # This function relies on 'layers_block_type' in hf_config,
954 # for w/o this attribute, we will need to have workarounds like so
955 attn_block_type = block_type == LayerBlockType.attention
943 - is_transformer = not self.is_hybrid and not self.is_attention_free
956 + is_transformer = not self.is_hybrid and \
957 + not self.has_noops and \
958 + not self.is_attention_free
959 start, end = self.get_layers_start_end_indices(parallel_config)
960
961 if is_transformer:
@@ -951,6 +966,10 @@ class ModelConfig:
966 # Note that this code assumes there
967 # is only one type of attention-free block type.
968 return 0 if attn_block_type else end - start
969 + elif self.has_noops:
970 + block_configs = self.hf_config.block_configs
971 + return sum(not bc.attention.no_op
972 + for bc in block_configs[start:end])
973 else:
974 # Hybrid model
975 layers_block_type_value = getattr(self.hf_config,
vllm/model_executor/models/decilm.py deleted
-124
@@ -1,124 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -# Adapted from
4 -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
5 -# Copyright 2023 DeciAI Research Team. All rights reserved.
6 -# Copyright 2023 The vLLM team.
7 -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8 -#
9 -# This code is based on MistralAI GPT-NeoX library and the GPT-NeoX
10 -# and OPT implementations in this library. It has been modified from its
11 -# original forms to accommodate minor architectural differences compared
12 -# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13 -#
14 -# Licensed under the Apache License, Version 2.0 (the "License");
15 -# you may not use this file except in compliance with the License.
16 -# You may obtain a copy of the License at
17 -#
18 -# http://www.apache.org/licenses/LICENSE-2.0
19 -#
20 -# Unless required by applicable law or agreed to in writing, software
21 -# distributed under the License is distributed on an "AS IS" BASIS,
22 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 -# See the License for the specific language governing permissions and
24 -# limitations under the License.
25 -"""Inference-only DeciLM model compatible with HuggingFace weights."""
26 -
27 -from typing import Iterable, Set, Tuple
28 -
29 -import torch
30 -
31 -from vllm.config import VllmConfig
32 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader
33 -from vllm.model_executor.models.llama import LlamaForCausalLM
34 -
35 -from .utils import is_pp_missing_parameter
36 -
37 -
38 -class DeciLMForCausalLM(LlamaForCausalLM):
39 - """
40 - Implementation for https://huggingface.co/Deci/DeciLM-7b-instruct.
41 - Based on the llama executor.
42 -
43 - The main difference is that DeciLM uses Variable Grouped Query Attention.
44 - The constant number of GQA heads in the decoder is overridden with a value
45 - per layer.
46 -
47 - Usually, in the HuggingFace implementation, instead of
48 - "config.num_key_value_heads", we use
49 - "config.num_key_value_heads_per_layer[i]" which varies.
50 -
51 - Currently, PagedAttention does not work well with variable GQA, so we
52 - normalize the weights upon loading, and use uniform GQA with the max value
53 - instead.
54 - """
55 -
56 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
57 - config = vllm_config.model_config.hf_config
58 - config.num_key_value_heads = max(config.num_key_value_heads_per_layer)
59 - delattr(config, "num_key_value_heads_per_layer")
60 - super().__init__(vllm_config=vllm_config)
61 -
62 - def load_weights(self, weights: Iterable[Tuple[str,
63 - torch.Tensor]]) -> Set[str]:
64 - stacked_params_mapping = [
65 - # (param_name, shard_name, shard_id)
66 - ("qkv_proj", "q_proj", "q"),
67 - ("qkv_proj", "k_proj", "k"),
68 - ("qkv_proj", "v_proj", "v"),
69 - ("gate_up_proj", "gate_proj", 0),
70 - ("gate_up_proj", "up_proj", 1),
71 - ]
72 - params_dict = dict(self.named_parameters())
73 - loaded_params: Set[str] = set()
74 - for name, loaded_weight in weights:
75 - if "rotary_emb.inv_freq" in name:
76 - continue
77 -
78 - if "k_proj" in name or "v_proj" in name:
79 - loaded_weight = self._degroup_weight(loaded_weight)
80 -
81 - for (param_name, weight_name, shard_id) in stacked_params_mapping:
82 - if weight_name not in name:
83 - continue
84 - name = name.replace(weight_name, param_name)
85 - # Skip loading extra bias for GPTQ models.
86 - if name.endswith(".bias") and name not in params_dict:
87 - continue
88 - if is_pp_missing_parameter(name, self):
89 - continue
90 - param = params_dict[name]
91 - weight_loader = param.weight_loader
92 - weight_loader(param, loaded_weight, shard_id)
93 - break
94 - else:
95 - # Skip loading extra bias for GPTQ models.
96 - if name.endswith(".bias") and name not in params_dict:
97 - continue
98 - if is_pp_missing_parameter(name, self):
99 - continue
100 - param = params_dict[name]
101 - weight_loader = getattr(param, "weight_loader",
102 - default_weight_loader)
103 - weight_loader(param, loaded_weight)
104 - loaded_params.add(name)
105 - return loaded_params
106 -
107 - def _degroup_weight(self, loaded_weight: torch.Tensor) -> torch.Tensor:
108 - hidden_size = self.config.hidden_size
109 - head_size = self.config.hidden_size // self.config.num_attention_heads
110 - target_num_kv_heads = self.config.num_key_value_heads
111 - num_kv_heads = loaded_weight.shape[0] // head_size
112 - n_repeats = target_num_kv_heads / num_kv_heads
113 - assert n_repeats == int(n_repeats)
114 -
115 - n_repeats = int(n_repeats)
116 - loaded_weight = loaded_weight.view(num_kv_heads, head_size,
117 - hidden_size)
118 - loaded_weight = torch.repeat_interleave(loaded_weight,
119 - repeats=n_repeats,
120 - dim=0)
121 - loaded_weight = loaded_weight.reshape(target_num_kv_heads * head_size,
122 - hidden_size)
123 -
124 - return loaded_weight
vllm/model_executor/models/interfaces.py
+29
@@ -411,6 +411,35 @@ def is_hybrid(
411 return isinstance(model, IsHybrid)
412
413
414 +@runtime_checkable
415 +class HasNoOps(Protocol):
416 + has_noops: ClassVar[Literal[True]] = True
417 +
418 +
419 +@runtime_checkable
420 +class _HasNoOpsType(Protocol):
421 + has_noops: ClassVar[Literal[True]]
422 +
423 +
424 +@overload
425 +def has_noops(model: object) -> TypeIs[HasNoOps]:
426 + ...
427 +
428 +
429 +@overload
430 +def has_noops(model: Type[object]) -> TypeIs[Type[HasNoOps]]:
431 + ...
432 +
433 +
434 +def has_noops(
435 + model: Union[Type[object], object]
436 +) -> Union[TypeIs[Type[HasNoOps]], TypeIs[HasNoOps]]:
437 + if isinstance(model, type):
438 + return isinstance(model, _HasNoOpsType)
439 +
440 + return isinstance(model, HasNoOps)
441 +
442 +
443 @runtime_checkable
444 class SupportsCrossEncoding(Protocol):
445 """The interface required for all models that support cross encoding."""
vllm/model_executor/models/nemotron_nas.py new
+454
@@ -0,0 +1,454 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from
4 +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
5 +# Copyright 2023 The vLLM team.
6 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
7 +#
8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9 +# and OPT implementations in this library. It has been modified from its
10 +# original forms to accommodate minor architectural differences compared
11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12 +#
13 +# Licensed under the Apache License, Version 2.0 (the "License");
14 +# you may not use this file except in compliance with the License.
15 +# You may obtain a copy of the License at
16 +#
17 +# http://www.apache.org/licenses/LICENSE-2.0
18 +#
19 +# Unless required by applicable law or agreed to in writing, software
20 +# distributed under the License is distributed on an "AS IS" BASIS,
21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 +# See the License for the specific language governing permissions and
23 +# limitations under the License.
24 +"""Inference-only deci model compatible with HuggingFace weights."""
25 +from typing import Iterable, Optional, Set, Tuple, Type, Union
26 +
27 +import torch
28 +from torch import nn
29 +from transformers import LlamaConfig
30 +
31 +from vllm.compilation.decorators import support_torch_compile
32 +from vllm.config import CacheConfig, VllmConfig
33 +from vllm.distributed import get_pp_group
34 +from vllm.model_executor.layers.layernorm import RMSNorm
35 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
36 +from vllm.model_executor.layers.quantization import QuantizationConfig
37 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
38 +from vllm.model_executor.layers.vocab_parallel_embedding import (
39 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
40 +from vllm.model_executor.model_loader.weight_utils import (
41 + default_weight_loader, maybe_remap_kv_scale_name)
42 +from vllm.model_executor.models.llama import LlamaAttention, LlamaMLP
43 +from vllm.model_executor.sampling_metadata import SamplingMetadata
44 +from vllm.sequence import IntermediateTensors
45 +
46 +from .interfaces import HasNoOps, SupportsLoRA, SupportsPP
47 +from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
48 + make_empty_intermediate_tensors_factory, make_layers,
49 + maybe_prefix)
50 +
51 +
52 +def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int:
53 + # DeciLM-specific code
54 + intermediate_size = int(2 * ffn_mult * n_embd / 3)
55 + return _find_multiple(intermediate_size, 256)
56 +
57 +
58 +def _find_multiple(n: int, k: int) -> int:
59 + # DeciLM-specific code
60 + if n % k == 0:
61 + return n
62 + return n + k - (n % k)
63 +
64 +
65 +class DeciLMDecoderLayer(nn.Module):
66 +
67 + def __init__(
68 + self,
69 + config: LlamaConfig,
70 + layer_idx: int,
71 + cache_config: Optional[CacheConfig] = None,
72 + quant_config: Optional[QuantizationConfig] = None,
73 + prefix: str = "",
74 + ) -> None:
75 + super().__init__()
76 + block_config = config.block_configs[layer_idx]
77 + self._is_no_op_attention = block_config.attention.no_op
78 + self._is_no_op_ffn = block_config.ffn.no_op
79 +
80 + self.hidden_size = config.hidden_size
81 + rope_theta = getattr(config, "rope_theta", 10000)
82 + rope_scaling = getattr(config, "rope_scaling", None)
83 + if rope_scaling is not None and getattr(
84 + config, "original_max_position_embeddings", None):
85 + rope_scaling["original_max_position_embeddings"] = (
86 + config.original_max_position_embeddings)
87 + max_position_embeddings = getattr(config, "max_position_embeddings",
88 + 8192)
89 + # Support abacusai/Smaug-72B-v0.1 with attention_bias
90 + # Support internlm/internlm-7b with bias
91 + attention_bias = getattr(config, "attention_bias", False) or getattr(
92 + config, "bias", False)
93 + bias_o_proj = attention_bias
94 + # support internlm/internlm3-8b with qkv_bias
95 + if hasattr(config, "qkv_bias"):
96 + attention_bias = config.qkv_bias
97 +
98 + if not self._is_no_op_attention:
99 + num_kv_heads = (config.num_attention_heads //
100 + block_config.attention.n_heads_in_group)
101 + self.self_attn = LlamaAttention(
102 + config=config,
103 + hidden_size=self.hidden_size,
104 + num_heads=config.num_attention_heads,
105 + num_kv_heads=num_kv_heads,
106 + rope_theta=rope_theta,
107 + rope_scaling=rope_scaling,
108 + max_position_embeddings=max_position_embeddings,
109 + quant_config=quant_config,
110 + bias=attention_bias,
111 + bias_o_proj=bias_o_proj,
112 + cache_config=cache_config,
113 + prefix=f"{prefix}.self_attn",
114 + )
115 + self.input_layernorm = RMSNorm(config.hidden_size,
116 + eps=config.rms_norm_eps)
117 +
118 + if not self._is_no_op_ffn:
119 + ffn_mult = block_config.ffn.ffn_mult
120 + intermediate_size = _ffn_mult_to_intermediate_size(
121 + ffn_mult, config.hidden_size)
122 +
123 + self.mlp = LlamaMLP(
124 + hidden_size=self.hidden_size,
125 + intermediate_size=intermediate_size,
126 + hidden_act=config.hidden_act,
127 + quant_config=quant_config,
128 + bias=getattr(config, "mlp_bias", False),
129 + prefix=f"{prefix}.mlp",
130 + )
131 + self.post_attention_layernorm = RMSNorm(config.hidden_size,
132 + eps=config.rms_norm_eps)
133 +
134 + def forward(
135 + self,
136 + positions: torch.Tensor,
137 + hidden_states: torch.Tensor,
138 + residual: Optional[torch.Tensor],
139 + ) -> Tuple[torch.Tensor, torch.Tensor]:
140 + # Self Attention
141 +
142 + if self._is_no_op_attention:
143 + pass
144 + else:
145 + if (residual is None):
146 + residual = hidden_states
147 + hidden_states = self.input_layernorm(hidden_states)
148 + else:
149 + hidden_states, residual = self.input_layernorm(
150 + hidden_states, residual)
151 + hidden_states = self.self_attn(
152 + positions=positions,
153 + hidden_states=hidden_states,
154 + )
155 +
156 + # Fully Connected
157 + if not self._is_no_op_ffn:
158 + hidden_states, residual = self.post_attention_layernorm(
159 + hidden_states, residual)
160 + hidden_states = self.mlp(hidden_states)
161 + return hidden_states, residual
162 +
163 +
164 +@support_torch_compile
165 +class DeciModel(nn.Module):
166 +
167 + def __init__(
168 + self,
169 + *,
170 + vllm_config: VllmConfig,
171 + prefix: str = "",
172 + layer_type: Type[DeciLMDecoderLayer] = DeciLMDecoderLayer,
173 + ):
174 + super().__init__()
175 +
176 + config = vllm_config.model_config.hf_config
177 + cache_config = vllm_config.cache_config
178 + quant_config = vllm_config.quant_config
179 + lora_config = vllm_config.lora_config
180 +
181 + self.config = config
182 + self.quant_config = quant_config
183 + self.padding_idx = config.pad_token_id
184 + lora_vocab = ((lora_config.lora_extra_vocab_size *
185 + (lora_config.max_loras or 1)) if lora_config else 0)
186 + self.vocab_size = config.vocab_size + lora_vocab
187 + self.org_vocab_size = config.vocab_size
188 + if get_pp_group().is_first_rank or (config.tie_word_embeddings
189 + and get_pp_group().is_last_rank):
190 + self.embed_tokens = VocabParallelEmbedding(
191 + self.vocab_size,
192 + config.hidden_size,
193 + org_num_embeddings=config.vocab_size,
194 + quant_config=quant_config,
195 + )
196 + else:
197 + self.embed_tokens = PPMissingLayer()
198 +
199 + def get_layer(prefix: str):
200 + layer_idx = int(prefix.rsplit(".", 1)[1])
201 + return layer_type(
202 + config,
203 + layer_idx,
204 + cache_config,
205 + quant_config=quant_config,
206 + prefix=prefix,
207 + )
208 +
209 + self.start_layer, self.end_layer, self.layers = make_layers(
210 + config.num_hidden_layers,
211 + get_layer,
212 + prefix=f"{prefix}.layers",
213 + )
214 + if get_pp_group().is_last_rank:
215 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
216 + else:
217 + self.norm = PPMissingLayer()
218 +
219 + self.make_empty_intermediate_tensors = (
220 + make_empty_intermediate_tensors_factory(
221 + ["hidden_states", "residual"], config.hidden_size))
222 +
223 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
224 + return self.embed_tokens(input_ids)
225 +
226 + def forward(
227 + self,
228 + input_ids: Optional[torch.Tensor],
229 + positions: torch.Tensor,
230 + intermediate_tensors: Optional[IntermediateTensors],
231 + inputs_embeds: Optional[torch.Tensor] = None,
232 + ) -> Union[torch.Tensor, IntermediateTensors]:
233 + if get_pp_group().is_first_rank:
234 + if inputs_embeds is not None:
235 + hidden_states = inputs_embeds
236 + else:
237 + hidden_states = self.get_input_embeddings(input_ids)
238 + residual = None
239 + else:
240 + assert intermediate_tensors is not None
241 + hidden_states = intermediate_tensors["hidden_states"]
242 + residual = intermediate_tensors["residual"]
243 +
244 + kv_cache_index = 0
245 + for i in range(self.start_layer, self.end_layer):
246 + layer = self.layers[i]
247 + if not layer._is_no_op_attention:
248 + hidden_states, residual = layer(positions, hidden_states,
249 + residual)
250 + kv_cache_index += 1
251 + else:
252 + hidden_states, residual = layer(positions, hidden_states,
253 + residual)
254 +
255 + if not get_pp_group().is_last_rank:
256 + return IntermediateTensors({
257 + "hidden_states": hidden_states,
258 + "residual": residual
259 + })
260 +
261 + hidden_states, _ = self.norm(hidden_states, residual)
262 + return hidden_states
263 +
264 + def load_weights(self, weights: Iterable[Tuple[str,
265 + torch.Tensor]]) -> Set[str]:
266 + stacked_params_mapping = [
267 + # (param_name, shard_name, shard_id)
268 + (".qkv_proj", ".q_proj", "q"),
269 + (".qkv_proj", ".k_proj", "k"),
270 + (".qkv_proj", ".v_proj", "v"),
271 + (".gate_up_proj", ".gate_proj", 0),
272 + (".gate_up_proj", ".up_proj", 1),
273 + ]
274 + params_dict = dict(self.named_parameters())
275 + loaded_params: Set[str] = set()
276 + for name, loaded_weight in weights:
277 + if "rotary_emb.inv_freq" in name:
278 + continue
279 + if ("rotary_emb.cos_cached" in name
280 + or "rotary_emb.sin_cached" in name):
281 + # Models trained using ColossalAI may include these tensors in
282 + # the checkpoint. Skip them.
283 + continue
284 + if self.quant_config is not None and (
285 + scale_name := self.quant_config.get_cache_scale(name)):
286 + # Loading kv cache quantization scales
287 + param = params_dict[scale_name]
288 + weight_loader = getattr(param, "weight_loader",
289 + default_weight_loader)
290 + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
291 + loaded_weight[0])
292 + weight_loader(param, loaded_weight)
293 + loaded_params.add(scale_name)
294 + continue
295 + if "scale" in name:
296 + # Remapping the name of FP8 kv-scale.
297 + name = maybe_remap_kv_scale_name(name, params_dict)
298 + if name is None:
299 + continue
300 + for param_name, weight_name, shard_id in stacked_params_mapping:
301 + if weight_name not in name:
302 + continue
303 + name = name.replace(weight_name, param_name)
304 + # Skip loading extra bias for GPTQ models.
305 + if name.endswith(".bias") and name not in params_dict:
306 + continue
307 +
308 + if is_pp_missing_parameter(name, self):
309 + continue
310 +
311 + param = params_dict[name]
312 + weight_loader = param.weight_loader
313 + weight_loader(param, loaded_weight, shard_id)
314 + break
315 + else:
316 + # Skip loading extra bias for GPTQ models.
317 + if name.endswith(".bias") and name not in params_dict:
318 + continue
319 +
320 + if is_pp_missing_parameter(name, self):
321 + continue
322 +
323 + param = params_dict[name]
324 + weight_loader = getattr(param, "weight_loader",
325 + default_weight_loader)
326 + weight_loader(param, loaded_weight)
327 + loaded_params.add(name)
328 + return loaded_params
329 +
330 +
331 +class DeciLMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, HasNoOps):
332 + packed_modules_mapping = {
333 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
334 + "gate_up_proj": ["gate_proj", "up_proj"],
335 + }
336 +
337 + # LoRA specific attributes
338 + supported_lora_modules = [
339 + "qkv_proj",
340 + "o_proj",
341 + "gate_up_proj",
342 + "down_proj",
343 + "embed_tokens",
344 + "lm_head",
345 + ]
346 + embedding_modules = {
347 + "embed_tokens": "input_embeddings",
348 + "lm_head": "output_embeddings",
349 + }
350 + embedding_padding_modules = ["lm_head"]
351 +
352 + # Mistral/Llama models can also be loaded with --load-format mistral
353 + # from consolidated.safetensors checkpoints
354 + mistral_mapping = {
355 + "layers": "model.layers",
356 + "attention": "self_attn",
357 + "wq": "q_proj",
358 + "wk": "k_proj",
359 + "wv": "v_proj",
360 + "wo": "o_proj",
361 + "attention_norm": "input_layernorm",
362 + "feed_forward": "mlp",
363 + "w1": "gate_proj",
364 + "w2": "down_proj",
365 + "w3": "up_proj",
366 + "ffn_norm": "post_attention_layernorm",
367 + "tok_embeddings": "model.embed_tokens",
368 + "output": "lm_head",
369 + "norm": "model.norm",
370 + }
371 +
372 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
373 + super().__init__()
374 + config = vllm_config.model_config.hf_config
375 + quant_config = vllm_config.quant_config
376 + lora_config = vllm_config.lora_config
377 + self.config = config
378 + self.lora_config = lora_config
379 +
380 + self.model = self._init_model(vllm_config=vllm_config,
381 + prefix=maybe_prefix(prefix, "model"))
382 +
383 + if get_pp_group().is_last_rank:
384 + self.unpadded_vocab_size = config.vocab_size
385 + if lora_config:
386 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
387 + self.lm_head = ParallelLMHead(
388 + self.unpadded_vocab_size,
389 + config.hidden_size,
390 + org_num_embeddings=config.vocab_size,
391 + padding_size=(
392 + DEFAULT_VOCAB_PADDING_SIZE
393 + # We need bigger padding if using lora for kernel
394 + # compatibility
395 + if not lora_config else
396 + lora_config.lora_vocab_padding_size),
397 + quant_config=quant_config,
398 + prefix=maybe_prefix(prefix, "lm_head"),
399 + )
400 + if config.tie_word_embeddings:
401 + self.lm_head = self.lm_head.tie_weights(
402 + self.model.embed_tokens)
403 +
404 + logit_scale = getattr(config, "logit_scale", 1.0)
405 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
406 + config.vocab_size,
407 + logit_scale)
408 + else:
409 + self.lm_head = PPMissingLayer()
410 +
411 + self.sampler = get_sampler()
412 +
413 + self.make_empty_intermediate_tensors = (
414 + self.model.make_empty_intermediate_tensors)
415 +
416 + def _init_model(self, vllm_config: VllmConfig, prefix: str = ""):
417 + return DeciModel(vllm_config=vllm_config, prefix=prefix)
418 +
419 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
420 + return self.model.get_input_embeddings(input_ids)
421 +
422 + def forward(
423 + self,
424 + input_ids: torch.Tensor,
425 + positions: torch.Tensor,
426 + intermediate_tensors: Optional[IntermediateTensors] = None,
427 + inputs_embeds: Optional[torch.Tensor] = None,
428 + ) -> Union[torch.Tensor, IntermediateTensors]:
429 + model_output = self.model(input_ids, positions, intermediate_tensors,
430 + inputs_embeds)
431 + return model_output
432 +
433 + def compute_logits(
434 + self,
435 + hidden_states: torch.Tensor,
436 + sampling_metadata: SamplingMetadata,
437 + ) -> Optional[torch.Tensor]:
438 + logits = self.logits_processor(self.lm_head, hidden_states,
439 + sampling_metadata)
440 + return logits
441 +
442 + def sample(self, logits: torch.Tensor,
443 + sampling_metadata: SamplingMetadata) -> Optional[SamplerOutput]:
444 + next_tokens = self.sampler(logits, sampling_metadata)
445 + return next_tokens
446 +
447 + def load_weights(self, weights: Iterable[Tuple[str,
448 + torch.Tensor]]) -> Set[str]:
449 + loader = AutoWeightsLoader(
450 + self,
451 + skip_prefixes=(["lm_head."]
452 + if self.config.tie_word_embeddings else None),
453 + )
454 + return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+15 -5
@@ -21,9 +21,10 @@ import torch.nn as nn
21 from vllm.logger import init_logger
22 from vllm.utils import is_in_doc_build
23
24 -from .interfaces import (has_inner_state, is_attention_free, is_hybrid,
25 - supports_cross_encoding, supports_multimodal,
26 - supports_pp, supports_transcription, supports_v0_only)
24 +from .interfaces import (has_inner_state, has_noops, is_attention_free,
25 + is_hybrid, supports_cross_encoding,
26 + supports_multimodal, supports_pp,
27 + supports_transcription, supports_v0_only)
28 from .interfaces_base import is_text_generation_model
29
30 logger = init_logger(__name__)
@@ -44,7 +45,7 @@ _TEXT_GENERATION_MODELS = {
45 "CohereForCausalLM": ("commandr", "CohereForCausalLM"),
46 "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"),
47 "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
47 - "DeciLMForCausalLM": ("decilm", "DeciLMForCausalLM"),
48 + "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
49 "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
50 "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
51 "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
@@ -118,7 +119,7 @@ _EMBEDDING_MODELS = {
119 "RobertaModel": ("roberta", "RobertaEmbeddingModel"),
120 "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
121 "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
121 - "DeciLMForCausalLM": ("decilm", "DeciLMForCausalLM"),
122 + "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
123 "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
124 "GlmForCausalLM": ("glm", "GlmForCausalLM"),
125 "GritLM": ("gritlm", "GritLM"),
@@ -235,6 +236,7 @@ class _ModelInfo:
236 has_inner_state: bool
237 is_attention_free: bool
238 is_hybrid: bool
239 + has_noops: bool
240 supports_transcription: bool
241 supports_v0_only: bool
242
@@ -252,6 +254,7 @@ class _ModelInfo:
254 is_hybrid=is_hybrid(model),
255 supports_transcription=supports_transcription(model),
256 supports_v0_only=supports_v0_only(model),
257 + has_noops=has_noops(model),
258 )
259
260
@@ -511,6 +514,13 @@ class _ModelRegistry:
514 model_cls, _ = self.inspect_model_cls(architectures)
515 return model_cls.is_hybrid
516
517 + def is_noops_model(
518 + self,
519 + architectures: Union[str, List[str]],
520 + ) -> bool:
521 + model_cls, _ = self.inspect_model_cls(architectures)
522 + return model_cls.has_noops
523 +
524 def is_transcription_model(
525 self,
526 architectures: Union[str, List[str]],
vllm/model_executor/models/utils.py
+4 -1
@@ -497,7 +497,10 @@ def set_cpu_offload_max_bytes(max_bytes: int) -> None:
497
498
499 def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
500 - device = next(module.parameters()).device
500 + if (params := next(module.parameters(), None)) is None:
501 + return module
502 +
503 + device = params.device
504
505 if device == torch.device("cpu"):
506 return module