@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1cf3753b90
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 14, 05:30 UTC

[MODEL] `Apertus` and `XIELU` (#23068)

Signed-off-by: EduardDurech <39579228+EduardDurech@users.noreply.github.com> Co-authored-by: AllenHaoHuang <allenhuangdd@gmail.com>

EduardDurech committed Aug 29, 2025 at 14:29 UTC 1cf3753b901ba874a830c19555bb31fe37f91231
6 files changed +696 -1
tests/models/language/generation/test_common.py
+2 -1
@@ -92,7 +92,8 @@ AITER_MODEL_LIST = [
92 pytest.param(
93 "allenai/OLMoE-1B-7B-0924-Instruct",
94 marks=[pytest.mark.cpu_model],
95 - )
95 + ),
96 + pytest.param("swiss-ai/Apertus-8B"), # apertus
97 ])
98 @pytest.mark.parametrize("max_tokens", [32])
99 @pytest.mark.parametrize("num_logprobs", [5])
tests/models/registry.py
+3
@@ -137,6 +137,9 @@ class _HfExamplesInfo:
137 # yapf: disable
138 _TEXT_GENERATION_EXAMPLE_MODELS = {
139 # [Decoder-only]
140 + "ApertusForCausalLM": _HfExamplesInfo("swiss-ai/Apertus-8B",
141 + min_transformers_version="4.56.0",
142 + trust_remote_code=True),
143 "AquilaModel": _HfExamplesInfo("BAAI/AquilaChat-7B",
144 trust_remote_code=True),
145 "AquilaForCausalLM": _HfExamplesInfo("BAAI/AquilaChat2-7B",
tests/models/test_registry.py
+3
@@ -24,6 +24,9 @@ from .registry import HF_EXAMPLE_MODELS
24
25 @pytest.mark.parametrize("model_arch", ModelRegistry.get_supported_archs())
26 def test_registry_imports(model_arch):
27 + # Skip if transformers version is incompatible
28 + model_info = HF_EXAMPLE_MODELS.get_hf_info(model_arch)
29 + model_info.check_transformers_version(on_fail="skip")
30 # Ensure all model classes can be imported successfully
31 model_cls = ModelRegistry._try_load_model_cls(model_arch)
32 assert model_cls is not None
vllm/model_executor/layers/activation.py
+111
@@ -10,11 +10,14 @@ import torch.nn.functional as F
10
11 from vllm.distributed import (divide, get_tensor_model_parallel_rank,
12 get_tensor_model_parallel_world_size)
13 +from vllm.logger import init_logger
14 from vllm.model_executor.custom_op import CustomOp
15 from vllm.model_executor.utils import set_weight_attrs
16 from vllm.platforms import current_platform
17 from vllm.utils import LazyDict
18
19 +logger = init_logger(__name__)
20 +
21
22 @CustomOp.register("fatrelu_and_mul")
23 class FatreluAndMul(CustomOp):
@@ -363,6 +366,112 @@ class ReLUSquaredActivation(CustomOp):
366 return self.forward_native(x)
367
368
369 +@CustomOp.register("xielu")
370 +class XIELU(CustomOp):
371 + """
372 + Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010
373 + If the user has installed the nickjbrowning/XIELU, we import xIELU CUDA
374 + Otherwise, we emit a single warning and use xIELU Python
375 + """
376 +
377 + def __init__(
378 + self,
379 + alpha_p_init: float = 0.8,
380 + alpha_n_init: float = 0.8,
381 + beta: float = 0.5,
382 + eps: float = -1e-6,
383 + dtype: torch.dtype = torch.bfloat16,
384 + with_vector_loads: bool = False,
385 + ):
386 + super().__init__()
387 + self.alpha_p = nn.Parameter(
388 + torch.log(torch.exp(torch.tensor(alpha_p_init, dtype=dtype)) -
389 + 1).unsqueeze(0))
390 + self.alpha_n = nn.Parameter(
391 + torch.log(
392 + torch.exp(torch.tensor(alpha_n_init - beta, dtype=dtype)) -
393 + 1).unsqueeze(0))
394 + self.register_buffer("beta", torch.tensor(beta, dtype=dtype))
395 + self.register_buffer("eps", torch.tensor(eps, dtype=dtype))
396 + self.with_vector_loads = with_vector_loads
397 + # Temporary until xIELU CUDA fully implemented
398 + self._beta_scalar = float(self.beta.detach().cpu().float().item())
399 + self._eps_scalar = float(self.eps.detach().cpu().float().item())
400 +
401 + self._xielu_cuda_obj = None
402 + try:
403 + import xielu.ops # noqa: F401
404 +
405 + self._xielu_cuda_obj = torch.classes.xielu.XIELU()
406 + msg = "Using experimental xIELU CUDA."
407 + try:
408 + from torch._dynamo import allow_in_graph
409 +
410 + self._xielu_cuda_fn = allow_in_graph(self._xielu_cuda)
411 + msg += " Enabled torch._dynamo for xIELU CUDA."
412 + except Exception as err:
413 + msg += (f" Could not enable torch._dynamo for xIELU ({err}) - "
414 + "this may result in slower performance.")
415 + self._xielu_cuda_fn = self._xielu_cuda
416 + logger.warning_once(msg)
417 + except Exception as err:
418 + logger.warning_once(
419 + "CUDA-fused xIELU not available (%s) –"
420 + " falling back to a Python version.\n"
421 + "For CUDA xIELU (experimental), `pip install git+https://github.com/nickjbrowning/XIELU`",
422 + str(err),
423 + )
424 +
425 + def _xielu_python(self, x: torch.Tensor) -> torch.Tensor:
426 + alpha_p = nn.functional.softplus(self.alpha_p)
427 + alpha_n = self.beta + nn.functional.softplus(self.alpha_n)
428 + return torch.where(
429 + x > 0,
430 + alpha_p * x * x + self.beta * x,
431 + (torch.expm1(torch.min(x, self.eps)) - x) * alpha_n +
432 + self.beta * x,
433 + )
434 +
435 + def _xielu_cuda(self, x: torch.Tensor) -> torch.Tensor:
436 + """Firewall function to prevent torch.compile from seeing .item()"""
437 + assert self._xielu_cuda_obj is not None, (
438 + "XIELU CUDA object must not be None")
439 + original_shape = x.shape
440 + # CUDA kernel expects 3D tensors, reshape if needed
441 + while x.dim() < 3:
442 + x = x.unsqueeze(0)
443 + if x.dim() > 3:
444 + x = x.view(-1, 1, x.size(-1))
445 + if original_shape != x.shape:
446 + logger.warning_once(
447 + "Warning: xIELU input tensor expects 3 dimensions"
448 + " but got (shape: %s). Reshaping to (shape: %s).",
449 + original_shape,
450 + x.shape,
451 + )
452 + result = self._xielu_cuda_obj.forward(
453 + x,
454 + self.alpha_p,
455 + self.alpha_n,
456 + # Temporary until xIELU CUDA fully implemented ->
457 + # self.{beta,eps}.item()
458 + self._beta_scalar,
459 + self._eps_scalar,
460 + self.with_vector_loads,
461 + )
462 + return result.view(original_shape)
463 +
464 + def forward(self, input: torch.Tensor) -> torch.Tensor:
465 + if self._xielu_cuda_obj is not None and input.is_cuda:
466 + if not torch._dynamo.is_compiling():
467 + return self._xielu_cuda_fn(input)
468 + else:
469 + logger.warning_once(
470 + "torch._dynamo is compiling, using Python version of xIELU."
471 + )
472 + return self._xielu_python(input)
473 +
474 +
475 class ScaledActivation(nn.Module):
476 """An activation function with post-scale parameters.
477
@@ -426,6 +535,8 @@ _ACTIVATION_REGISTRY = LazyDict({
535 lambda: nn.Tanh(),
536 "sigmoid":
537 lambda: nn.Sigmoid(),
538 + "xielu":
539 + lambda: XIELU(),
540 })
541
542
vllm/model_executor/models/apertus.py new
+576
@@ -0,0 +1,576 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Adapted from
5 +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6 +# Copyright 2025 The Swiss AI Initiative.
7 +# Copyright 2023 The vLLM team.
8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
9 +#
10 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
11 +# and OPT implementations in this library. It has been modified from its
12 +# original forms to accommodate the architectural differences made by
13 +# the Swiss AI Initiative that trained the model.
14 +#
15 +# Licensed under the Apache License, Version 2.0 (the "License");
16 +# you may not use this file except in compliance with the License.
17 +# You may obtain a copy of the License at
18 +#
19 +# http://www.apache.org/licenses/LICENSE-2.0
20 +#
21 +# Unless required by applicable law or agreed to in writing, software
22 +# distributed under the License is distributed on an "AS IS" BASIS,
23 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
24 +# See the License for the specific language governing permissions and
25 +# limitations under the License.
26 +"""Inference-only Apertus model compatible with HuggingFace weights."""
27 +from collections.abc import Iterable
28 +from typing import Any, Optional, Union
29 +
30 +import torch
31 +from torch import nn
32 +from transformers import ApertusConfig
33 +
34 +from vllm.attention import Attention, AttentionType
35 +from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
36 +from vllm.compilation.decorators import support_torch_compile
37 +from vllm.config import CacheConfig, VllmConfig
38 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
39 +from vllm.model_executor.layers.activation import XIELU
40 +from vllm.model_executor.layers.layernorm import RMSNorm
41 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
42 + QKVParallelLinear,
43 + RowParallelLinear)
44 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
45 +from vllm.model_executor.layers.quantization import QuantizationConfig
46 +from vllm.model_executor.layers.rotary_embedding import get_rope
47 +from vllm.model_executor.layers.vocab_parallel_embedding import (
48 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
49 +from vllm.model_executor.model_loader.weight_utils import (
50 + default_weight_loader, maybe_remap_kv_scale_name)
51 +from vllm.model_executor.sampling_metadata import SamplingMetadata
52 +from vllm.sequence import IntermediateTensors
53 +
54 +from .interfaces import SupportsLoRA, SupportsPP
55 +from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,
56 + is_pp_missing_parameter,
57 + make_empty_intermediate_tensors_factory, make_layers,
58 + maybe_prefix)
59 +
60 +
61 +class ApertusMLP(nn.Module):
62 +
63 + def __init__(
64 + self,
65 + hidden_size: int,
66 + intermediate_size: int,
67 + hidden_act: str,
68 + quant_config: Optional[QuantizationConfig] = None,
69 + bias: bool = False,
70 + prefix: str = "",
71 + reduce_results: bool = True,
72 + ) -> None:
73 + super().__init__()
74 + self.up_proj = ColumnParallelLinear(
75 + input_size=hidden_size,
76 + output_size=intermediate_size,
77 + bias=bias,
78 + quant_config=quant_config,
79 + prefix=f"{prefix}.up_proj",
80 + )
81 + self.down_proj = RowParallelLinear(
82 + input_size=intermediate_size,
83 + output_size=hidden_size,
84 + bias=bias,
85 + quant_config=quant_config,
86 + reduce_results=reduce_results,
87 + prefix=f"{prefix}.down_proj",
88 + )
89 + if hidden_act != "xielu":
90 + raise ValueError(f"Unsupported activation: {hidden_act}. "
91 + "Only xIELU is supported for now.")
92 + self.act_fn = XIELU()
93 +
94 + def forward(self, x):
95 + x, _ = self.up_proj(x)
96 + x = self.act_fn(x)
97 + x, _ = self.down_proj(x)
98 + return x
99 +
100 +
101 +class ApertusAttention(nn.Module):
102 +
103 + def __init__(
104 + self,
105 + config: ApertusConfig,
106 + hidden_size: int,
107 + num_heads: int,
108 + num_kv_heads: int,
109 + rope_theta: float = 10000,
110 + rope_scaling: Optional[dict[str, Any]] = None,
111 + max_position_embeddings: int = 8192,
112 + quant_config: Optional[QuantizationConfig] = None,
113 + bias: bool = False,
114 + bias_o_proj: bool = False,
115 + cache_config: Optional[CacheConfig] = None,
116 + prefix: str = "",
117 + attn_type: str = AttentionType.DECODER,
118 + ) -> None:
119 + super().__init__()
120 + layer_idx = extract_layer_index(prefix)
121 + self.hidden_size = hidden_size
122 + tp_size = get_tensor_model_parallel_world_size()
123 + self.total_num_heads = num_heads
124 + assert self.total_num_heads % tp_size == 0
125 + self.num_heads = self.total_num_heads // tp_size
126 + self.total_num_kv_heads = num_kv_heads
127 + if self.total_num_kv_heads >= tp_size:
128 + # Number of KV heads is greater than TP size, so we partition
129 + # the KV heads across multiple tensor parallel GPUs.
130 + assert self.total_num_kv_heads % tp_size == 0
131 + else:
132 + # Number of KV heads is less than TP size, so we replicate
133 + # the KV heads across multiple tensor parallel GPUs.
134 + assert tp_size % self.total_num_kv_heads == 0
135 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
136 + # MistralConfig has an optional head_dim introduced by Mistral-Nemo
137 + head_dim = getattr(config, "head_dim", None)
138 + if head_dim is None:
139 + head_dim = self.hidden_size // self.total_num_heads
140 + self.head_dim = head_dim
141 + # Phi models introduced a partial_rotary_factor parameter in the config
142 + self.partial_rotary_factor = getattr(config, "partial_rotary_factor",
143 + 1)
144 + self.q_size = self.num_heads * self.head_dim
145 + self.kv_size = self.num_kv_heads * self.head_dim
146 + self.scaling = self.head_dim**-0.5
147 + self.rope_theta = rope_theta
148 + self.max_position_embeddings = max_position_embeddings
149 +
150 + self.qkv_proj = QKVParallelLinear(
151 + hidden_size=hidden_size,
152 + head_size=self.head_dim,
153 + total_num_heads=self.total_num_heads,
154 + total_num_kv_heads=self.total_num_kv_heads,
155 + bias=bias,
156 + quant_config=quant_config,
157 + prefix=f"{prefix}.qkv_proj",
158 + )
159 +
160 + self.o_proj = RowParallelLinear(
161 + input_size=self.total_num_heads * self.head_dim,
162 + output_size=hidden_size,
163 + bias=bias_o_proj,
164 + quant_config=quant_config,
165 + prefix=f"{prefix}.o_proj",
166 + )
167 +
168 + self._init_rotary_emb(config,
169 + rope_scaling=rope_scaling,
170 + quant_config=quant_config)
171 +
172 + sliding_window = None
173 + if layer_types := getattr(config, "layer_types", None):
174 + is_sliding = layer_types[layer_idx] == "sliding_attention"
175 + if is_sliding:
176 + sliding_window = config.sliding_window
177 +
178 + attn_cls = (EncoderOnlyAttention
179 + if attn_type == AttentionType.ENCODER_ONLY else Attention)
180 +
181 + self.attn = attn_cls(
182 + self.num_heads,
183 + self.head_dim,
184 + self.scaling,
185 + num_kv_heads=self.num_kv_heads,
186 + cache_config=cache_config,
187 + quant_config=quant_config,
188 + per_layer_sliding_window=sliding_window,
189 + attn_type=attn_type,
190 + prefix=f"{prefix}.attn",
191 + )
192 +
193 + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
194 + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
195 +
196 + def forward(
197 + self,
198 + positions: torch.Tensor,
199 + hidden_states: torch.Tensor,
200 + ) -> torch.Tensor:
201 + qkv, _ = self.qkv_proj(hidden_states)
202 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
203 + q = self.q_norm(q.contiguous().view(-1, self.head_dim)).view_as(q)
204 + k = self.k_norm(k.contiguous().view(-1, self.head_dim)).view_as(k)
205 + q, k = self.rotary_emb(positions, q, k)
206 + attn_output = self.attn(q, k, v)
207 + output, _ = self.o_proj(attn_output)
208 + return output
209 +
210 + def _init_rotary_emb(self, config: ApertusConfig,
211 + rope_scaling: Optional[dict[str, Any]],
212 + quant_config: Optional[QuantizationConfig]) -> None:
213 + is_neox_style = True
214 + is_gguf = quant_config and quant_config.get_name() == "gguf"
215 + if is_gguf and config.model_type == "apertus":
216 + is_neox_style = False
217 +
218 + self.rotary_emb = get_rope(
219 + self.head_dim,
220 + rotary_dim=int(self.partial_rotary_factor * self.head_dim),
221 + max_position=self.max_position_embeddings,
222 + base=self.rope_theta,
223 + rope_scaling=rope_scaling,
224 + is_neox_style=is_neox_style,
225 + partial_rotary_factor=self.partial_rotary_factor,
226 + )
227 +
228 +
229 +class ApertusDecoderLayer(nn.Module):
230 +
231 + def __init__(
232 + self,
233 + config: ApertusConfig,
234 + cache_config: Optional[CacheConfig] = None,
235 + quant_config: Optional[QuantizationConfig] = None,
236 + prefix: str = "",
237 + ) -> None:
238 + super().__init__()
239 + self.hidden_size = config.hidden_size
240 + rope_theta = getattr(config, "rope_theta", 10000)
241 + rope_scaling = getattr(config, "rope_scaling", None)
242 + if rope_scaling is not None and getattr(
243 + config, "original_max_position_embeddings", None):
244 + rope_scaling["original_max_position_embeddings"] = (
245 + config.original_max_position_embeddings)
246 + max_position_embeddings = getattr(config, "max_position_embeddings",
247 + 8192)
248 + # Support abacusai/Smaug-72B-v0.1 with attention_bias
249 + # Support internlm/internlm-7b with bias
250 + attention_bias = getattr(config, "attention_bias", False) or getattr(
251 + config, "bias", False)
252 + bias_o_proj = attention_bias
253 + # support internlm/internlm3-8b with qkv_bias
254 + if hasattr(config, 'qkv_bias'):
255 + attention_bias = config.qkv_bias
256 +
257 + # Apertus defaults to causal attention as it is a decoder-only model.
258 + # You can override the HF config with `is_causal=False` to enable
259 + # bidirectional attention, which is used in some embedding models
260 + # (e.g. parasail-ai/GritLM-7B-vllm)
261 + if getattr(config, "is_causal", True):
262 + attn_type = AttentionType.DECODER
263 + else:
264 + attn_type = AttentionType.ENCODER_ONLY
265 +
266 + self.self_attn = ApertusAttention(
267 + config=config,
268 + hidden_size=self.hidden_size,
269 + num_heads=config.num_attention_heads,
270 + num_kv_heads=getattr(config, "num_key_value_heads",
271 + config.num_attention_heads),
272 + rope_theta=rope_theta,
273 + rope_scaling=rope_scaling,
274 + max_position_embeddings=max_position_embeddings,
275 + quant_config=quant_config,
276 + bias=attention_bias,
277 + bias_o_proj=bias_o_proj,
278 + cache_config=cache_config,
279 + prefix=f"{prefix}.self_attn",
280 + attn_type=attn_type,
281 + )
282 + self.mlp = ApertusMLP(
283 + hidden_size=self.hidden_size,
284 + intermediate_size=config.intermediate_size,
285 + hidden_act=config.hidden_act,
286 + quant_config=quant_config,
287 + bias=getattr(config, "mlp_bias", False),
288 + prefix=f"{prefix}.mlp",
289 + )
290 + self.attention_layernorm = RMSNorm(config.hidden_size,
291 + eps=config.rms_norm_eps)
292 + self.feedforward_layernorm = RMSNorm(config.hidden_size,
293 + eps=config.rms_norm_eps)
294 +
295 + def forward(
296 + self,
297 + positions: torch.Tensor,
298 + hidden_states: torch.Tensor,
299 + residual: Optional[torch.Tensor],
300 + ) -> tuple[torch.Tensor, torch.Tensor]:
301 + # Self Attention
302 + if residual is None:
303 + residual = hidden_states
304 + hidden_states = self.attention_layernorm(hidden_states)
305 + else:
306 + hidden_states, residual = self.attention_layernorm(
307 + hidden_states, residual)
308 + hidden_states = self.self_attn(positions=positions,
309 + hidden_states=hidden_states)
310 +
311 + # Fully Connected
312 + hidden_states, residual = self.feedforward_layernorm(
313 + hidden_states, residual)
314 + hidden_states = self.mlp(hidden_states)
315 + return hidden_states, residual
316 +
317 +
318 +@support_torch_compile
319 +class ApertusModel(nn.Module):
320 +
321 + def __init__(self,
322 + *,
323 + vllm_config: VllmConfig,
324 + prefix: str = "",
325 + layer_type: type[nn.Module] = ApertusDecoderLayer):
326 + super().__init__()
327 +
328 + config = vllm_config.model_config.hf_config
329 + cache_config = vllm_config.cache_config
330 + quant_config = vllm_config.quant_config
331 + lora_config = vllm_config.lora_config
332 +
333 + self.config = config
334 + self.quant_config = quant_config
335 + lora_vocab = (lora_config.lora_extra_vocab_size *
336 + (lora_config.max_loras or 1)) if lora_config else 0
337 + self.vocab_size = config.vocab_size + lora_vocab
338 + self.org_vocab_size = config.vocab_size
339 + if get_pp_group().is_first_rank or (config.tie_word_embeddings
340 + and get_pp_group().is_last_rank):
341 + self.embed_tokens = VocabParallelEmbedding(
342 + self.vocab_size,
343 + config.hidden_size,
344 + org_num_embeddings=config.vocab_size,
345 + quant_config=quant_config,
346 + )
347 + else:
348 + self.embed_tokens = PPMissingLayer()
349 + self.start_layer, self.end_layer, self.layers = make_layers(
350 + config.num_hidden_layers,
351 + lambda prefix: layer_type(config=config,
352 + cache_config=cache_config,
353 + quant_config=quant_config,
354 + prefix=prefix),
355 + prefix=f"{prefix}.layers",
356 + )
357 + if get_pp_group().is_last_rank:
358 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
359 + else:
360 + self.norm = PPMissingLayer()
361 +
362 + self.aux_hidden_state_layers = tuple[int, ...]()
363 +
364 + self.make_empty_intermediate_tensors = (
365 + make_empty_intermediate_tensors_factory(
366 + ["hidden_states", "residual"], config.hidden_size))
367 +
368 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
369 + return self.embed_tokens(input_ids)
370 +
371 + def forward(
372 + self,
373 + input_ids: Optional[torch.Tensor],
374 + positions: torch.Tensor,
375 + intermediate_tensors: Optional[IntermediateTensors],
376 + inputs_embeds: Optional[torch.Tensor] = None,
377 + ) -> Union[torch.Tensor, IntermediateTensors, tuple[torch.Tensor,
378 + list[torch.Tensor]]]:
379 + if get_pp_group().is_first_rank:
380 + if inputs_embeds is not None:
381 + hidden_states = inputs_embeds
382 + else:
383 + hidden_states = self.get_input_embeddings(input_ids)
384 + residual = None
385 + else:
386 + assert intermediate_tensors is not None
387 + hidden_states = intermediate_tensors["hidden_states"]
388 + residual = intermediate_tensors["residual"]
389 +
390 + aux_hidden_states = []
391 + for idx, layer in enumerate(
392 + self.layers[self.start_layer:self.end_layer]):
393 + if idx in self.aux_hidden_state_layers:
394 + aux_hidden_states.append(hidden_states + residual)
395 + hidden_states, residual = layer(positions, hidden_states, residual)
396 +
397 + if not get_pp_group().is_last_rank:
398 + return IntermediateTensors({
399 + "hidden_states": hidden_states,
400 + "residual": residual
401 + })
402 +
403 + hidden_states, _ = self.norm(hidden_states, residual)
404 +
405 + if len(aux_hidden_states) > 0:
406 + return hidden_states, aux_hidden_states
407 + return hidden_states
408 +
409 + def load_weights(self, weights: Iterable[tuple[str,
410 + torch.Tensor]]) -> set[str]:
411 + stacked_params_mapping = [
412 + # (param_name, shard_name, shard_id)
413 + (".qkv_proj", ".q_proj", "q"),
414 + (".qkv_proj", ".k_proj", "k"),
415 + (".qkv_proj", ".v_proj", "v"),
416 + ]
417 + params_dict = dict(self.named_parameters())
418 + loaded_params: set[str] = set()
419 + for name, loaded_weight in weights:
420 + if "rotary_emb.inv_freq" in name:
421 + continue
422 + if ("rotary_emb.cos_cached" in name
423 + or "rotary_emb.sin_cached" in name):
424 + # Models trained using ColossalAI may include these tensors in
425 + # the checkpoint. Skip them.
426 + continue
427 + if (self.quant_config is not None and
428 + (scale_name := self.quant_config.get_cache_scale(name))):
429 + # Loading kv cache quantization scales
430 + param = params_dict[scale_name]
431 + weight_loader = getattr(param, "weight_loader",
432 + default_weight_loader)
433 + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
434 + loaded_weight[0])
435 + weight_loader(param, loaded_weight)
436 + loaded_params.add(scale_name)
437 + continue
438 + if "scale" in name:
439 + # Remapping the name of FP8 kv-scale.
440 + name = maybe_remap_kv_scale_name(name, params_dict)
441 + if name is None:
442 + continue
443 + for param_name, weight_name, shard_id in stacked_params_mapping:
444 + if weight_name not in name:
445 + continue
446 + name = name.replace(weight_name, param_name)
447 + # Skip loading extra bias for GPTQ models.
448 + if name.endswith(".bias") and name not in params_dict:
449 + continue
450 +
451 + if is_pp_missing_parameter(name, self):
452 + continue
453 +
454 + param = params_dict[name]
455 + weight_loader = param.weight_loader
456 + weight_loader(param, loaded_weight, shard_id)
457 + break
458 + else:
459 + # Skip loading extra bias for GPTQ models.
460 + if name.endswith(".bias") and name not in params_dict:
461 + continue
462 +
463 + if is_pp_missing_parameter(name, self):
464 + continue
465 +
466 + param = params_dict[name]
467 + weight_loader = getattr(param, "weight_loader",
468 + default_weight_loader)
469 + weight_loader(param, loaded_weight)
470 + loaded_params.add(name)
471 + return loaded_params
472 +
473 +
474 +class ApertusForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
475 + packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
476 +
477 + # LoRA specific attributes
478 + embedding_modules = {
479 + "embed_tokens": "input_embeddings",
480 + "lm_head": "output_embeddings"
481 + }
482 + embedding_padding_modules = ["lm_head"]
483 +
484 + def __init__(self,
485 + *,
486 + vllm_config: VllmConfig,
487 + prefix: str = "",
488 + layer_type: type[nn.Module] = ApertusDecoderLayer):
489 + super().__init__()
490 + config = vllm_config.model_config.hf_config
491 + quant_config = vllm_config.quant_config
492 + lora_config = vllm_config.lora_config
493 + self.config = config
494 + self.lora_config = lora_config
495 +
496 + self.model = self._init_model(vllm_config=vllm_config,
497 + prefix=maybe_prefix(prefix, "model"),
498 + layer_type=layer_type)
499 +
500 + if get_pp_group().is_last_rank:
501 + self.unpadded_vocab_size = config.vocab_size
502 + if lora_config:
503 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
504 + self.lm_head = ParallelLMHead(
505 + self.unpadded_vocab_size,
506 + config.hidden_size,
507 + org_num_embeddings=config.vocab_size,
508 + padding_size=(
509 + DEFAULT_VOCAB_PADDING_SIZE
510 + # We need bigger padding if using lora for kernel
511 + # compatibility
512 + if not lora_config else
513 + lora_config.lora_vocab_padding_size),
514 + quant_config=quant_config,
515 + prefix=maybe_prefix(prefix, "lm_head"),
516 + )
517 + if config.tie_word_embeddings:
518 + self.lm_head = self.lm_head.tie_weights(
519 + self.model.embed_tokens)
520 +
521 + logit_scale = getattr(config, "logit_scale", 1.0)
522 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
523 + config.vocab_size,
524 + logit_scale)
525 + else:
526 + self.lm_head = PPMissingLayer()
527 +
528 + self.make_empty_intermediate_tensors = (
529 + self.model.make_empty_intermediate_tensors)
530 +
531 + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
532 + self.model.aux_hidden_state_layers = layers
533 +
534 + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
535 + num_layers = len(self.model.layers)
536 + return (2, num_layers // 2, num_layers - 3)
537 +
538 + def _init_model(self,
539 + vllm_config: VllmConfig,
540 + prefix: str = "",
541 + layer_type: type[nn.Module] = ApertusDecoderLayer):
542 + return ApertusModel(vllm_config=vllm_config,
543 + prefix=prefix,
544 + layer_type=layer_type)
545 +
546 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
547 + return self.model.get_input_embeddings(input_ids)
548 +
549 + def forward(
550 + self,
551 + input_ids: torch.Tensor,
552 + positions: torch.Tensor,
553 + intermediate_tensors: Optional[IntermediateTensors] = None,
554 + inputs_embeds: Optional[torch.Tensor] = None,
555 + ) -> Union[torch.Tensor, IntermediateTensors]:
556 + model_output = self.model(input_ids, positions, intermediate_tensors,
557 + inputs_embeds)
558 + return model_output
559 +
560 + def compute_logits(
561 + self,
562 + hidden_states: torch.Tensor,
563 + sampling_metadata: SamplingMetadata,
564 + ) -> Optional[torch.Tensor]:
565 + logits = self.logits_processor(self.lm_head, hidden_states,
566 + sampling_metadata)
567 + return logits
568 +
569 + def load_weights(self, weights: Iterable[tuple[str,
570 + torch.Tensor]]) -> set[str]:
571 + loader = AutoWeightsLoader(
572 + self,
573 + skip_prefixes=(["lm_head."]
574 + if self.config.tie_word_embeddings else None),
575 + )
576 + return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -39,6 +39,7 @@ logger = init_logger(__name__)
39 # yapf: disable
40 _TEXT_GENERATION_MODELS = {
41 # [Decoder-only]
42 + "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"),
43 "AquilaModel": ("llama", "LlamaForCausalLM"),
44 "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2
45 "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"),