@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 169 lines 6.03 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import copy
5 from collections.abc import Iterable
6 from functools import partial
7 from typing import TYPE_CHECKING
8
9 import regex
10 import torch
11 import torch.nn as nn
12
13 from vllm.compilation.decorators import support_torch_compile
14 from vllm.config import VllmConfig
15 from vllm.distributed.parallel_state import get_pp_group
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.layernorm import RMSNorm
18 from vllm.model_executor.layers.linear import RowParallelLinear
19 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
20 from vllm.model_executor.models.deepseek_v2 import (
21 DeepseekV2DecoderLayer,
22 DeepseekV2Model,
23 )
24 from vllm.model_executor.models.mistral_large_3 import MistralLarge3ForCausalLM
25
26 from .interfaces import SupportsMultiModal, SupportsMultiModalEmbeddings
27 from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix
28
29 logger = init_logger(__name__)
30
31
32 if TYPE_CHECKING:
33
34 class _EagleMistralLarge3ForCausalLMBase(nn.Module):
35 pass
36
37 else:
38 _EagleMistralLarge3ForCausalLMBase = MistralLarge3ForCausalLM
39
40
41 @support_torch_compile
42 class EagleMistralLarge3Model(DeepseekV2Model):
43 def __init__(
44 self, *, vllm_config: VllmConfig, prefix: str = "", start_layer_id: int = 0
45 ):
46 nn.Module.__init__(self)
47
48 config = copy.deepcopy(vllm_config.model_config.hf_config)
49 config.first_k_dense_replace += start_layer_id
50
51 quant_config = vllm_config.quant_config
52 self.config = config
53 self.vllm_config = vllm_config
54
55 self.vocab_size = config.vocab_size
56
57 assert get_pp_group().world_size == 1
58 self.embed_tokens = VocabParallelEmbedding(
59 config.vocab_size,
60 config.hidden_size,
61 quant_config=quant_config,
62 prefix=f"{prefix}.embed_tokens",
63 )
64
65 self.layers = nn.ModuleList(
66 [
67 DeepseekV2DecoderLayer(
68 vllm_config=vllm_config,
69 prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
70 config=config,
71 )
72 for i in range(self.config.num_hidden_layers)
73 ]
74 )
75 self.start_layer = 0
76 self.end_layer = self.config.num_hidden_layers
77
78 self.fc = RowParallelLinear(
79 self.config.hidden_size * 2,
80 self.config.hidden_size,
81 bias=False,
82 input_is_parallel=False,
83 quant_config=quant_config,
84 return_bias=False,
85 prefix=maybe_prefix(prefix, "fc"),
86 )
87 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
88 self.aux_hidden_state_layers: tuple[int, ...] = ()
89
90 # Needed by load_weights
91 qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
92 qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
93 self.use_mha = config.model_type == "deepseek" or all(
94 dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
95 )
96 self.num_redundant_experts = (
97 vllm_config.parallel_config.eplb_config.num_redundant_experts
98 )
99 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
100 ["hidden_states", "residual"], config.hidden_size
101 )
102
103 def forward(
104 self,
105 input_ids: torch.Tensor,
106 positions: torch.Tensor,
107 hidden_states: torch.Tensor,
108 inputs_embeds: torch.Tensor | None = None,
109 ) -> torch.Tensor:
110 if inputs_embeds is None:
111 inputs_embeds = self.embed_input_ids(input_ids)
112 inputs_embeds = self.fc(torch.cat((inputs_embeds, hidden_states), dim=-1))
113 output = super().forward(
114 input_ids, positions, intermediate_tensors=None, inputs_embeds=inputs_embeds
115 )
116 assert isinstance(output, torch.Tensor)
117 return output
118
119
120 class EagleMistralLarge3ForCausalLM(
121 _EagleMistralLarge3ForCausalLMBase, SupportsMultiModalEmbeddings
122 ):
123 hf_to_vllm_mapper = MistralLarge3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(
124 orig_to_new_regex={
125 regex.compile(r"\Aeagle_linear\.weight\Z"): r"model.fc.weight",
126 regex.compile(r"\Aeagle_linear\.qscale_act\Z"): r"model.fc.input_scale",
127 regex.compile(r"\Aeagle_linear\.qscale_weight\Z"): r"model.fc.weight_scale",
128 },
129 )
130
131 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
132 speculative_config = vllm_config.speculative_config
133 assert speculative_config is not None
134 target_layer_num = vllm_config.model_config.get_num_layers(
135 vllm_config.parallel_config
136 )
137 vllm_config.model_config = speculative_config.draft_model_config
138 # draft model quantization config may differ from target model
139 self.quant_config = VllmConfig.get_quantization_config(
140 speculative_config.draft_model_config, vllm_config.load_config
141 )
142 vllm_config.quant_config = self.quant_config
143 self.model_cls = partial(
144 EagleMistralLarge3Model, start_layer_id=target_layer_num
145 )
146 super().__init__(vllm_config=vllm_config, prefix=prefix)
147
148 def get_language_model(self) -> torch.nn.Module:
149 return self.model
150
151 embed_input_ids = SupportsMultiModal.embed_input_ids # type: ignore
152
153 def forward(
154 self,
155 input_ids: torch.Tensor,
156 positions: torch.Tensor,
157 hidden_states: torch.Tensor,
158 inputs_embeds: torch.Tensor | None = None,
159 ) -> tuple[torch.Tensor, torch.Tensor]:
160 hidden_states = self.model(input_ids, positions, hidden_states, inputs_embeds)
161 return hidden_states, hidden_states
162
163 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
164 # Pretend we've loaded the embedding and lm_head weights
165 # (later copied from target model)
166 return super().load_weights(weights) | {
167 "model.embed_tokens.weight",
168 "lm_head.weight",
169 }