@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 196 lines 6.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Iterable
5 from typing import TYPE_CHECKING
6
7 import torch
8 import torch.nn as nn
9 from transformers import LlamaConfig
10
11 from vllm.compilation.decorators import support_torch_compile
12 from vllm.config import VllmConfig
13 from vllm.logger import init_logger
14 from vllm.model_executor.layers.linear import ReplicatedLinear
15 from vllm.model_executor.layers.logits_processor import LogitsProcessor
16 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
17 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
18 from vllm.model_executor.models.llama import (
19 LlamaDecoderLayer as BaseLlamaDecoderLayer,
20 )
21 from vllm.model_executor.models.llama import LlamaForCausalLM
22
23 from .utils import (
24 AutoWeightsLoader,
25 WeightsMapper,
26 get_draft_quant_config,
27 maybe_prefix,
28 process_eagle_weight,
29 )
30
31 logger = init_logger(__name__)
32
33
34 if TYPE_CHECKING:
35
36 class _EagleLlamaForCausalLMBase(nn.Module):
37 pass
38
39 else:
40 _EagleLlamaForCausalLMBase = LlamaForCausalLM
41
42
43 class LlamaDecoderLayer(BaseLlamaDecoderLayer):
44 def __init__(
45 self,
46 vllm_config: VllmConfig,
47 disable_input_layernorm: bool,
48 prefix: str = "",
49 config: LlamaConfig | None = None,
50 ) -> None:
51 super().__init__(vllm_config, prefix=prefix, config=config)
52
53 # Skip the input_layernorm
54 # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
55 if disable_input_layernorm:
56 del self.input_layernorm
57 self.input_layernorm = nn.Identity()
58
59 def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
60 """Use drafter's quantization config instead of verifier's."""
61 return get_draft_quant_config(vllm_config)
62
63
64 @support_torch_compile
65 class LlamaModel(nn.Module):
66 hf_to_vllm_mapper = WeightsMapper(
67 orig_to_new_stacked={
68 # weight_name: (param_name, shard_id)
69 ".q_proj": (".qkv_proj", "q"),
70 ".k_proj": (".qkv_proj", "k"),
71 ".v_proj": (".qkv_proj", "v"),
72 ".gate_proj": (".gate_up_proj", 0),
73 ".up_proj": (".gate_up_proj", 1),
74 }
75 )
76
77 def __init__(
78 self,
79 *,
80 vllm_config: VllmConfig,
81 prefix: str = "",
82 start_layer_id: int = 0,
83 ) -> None:
84 super().__init__()
85 speculative_config = vllm_config.speculative_config
86 assert speculative_config is not None
87 self.config = speculative_config.draft_model_config.hf_config
88 self.vocab_size = self.config.vocab_size
89
90 # Get drafter's quantization config
91 self.quant_config = get_draft_quant_config(vllm_config)
92
93 self.embed_tokens = VocabParallelEmbedding(
94 self.config.vocab_size,
95 self.config.hidden_size,
96 prefix=maybe_prefix(prefix, "embed_tokens"),
97 )
98
99 self.layers = nn.ModuleList(
100 [
101 LlamaDecoderLayer(
102 vllm_config,
103 i == 0,
104 prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
105 config=self.config,
106 )
107 for i in range(self.config.num_hidden_layers)
108 ]
109 )
110 self.fc = ReplicatedLinear(
111 input_size=self.config.hidden_size * 2,
112 output_size=self.config.hidden_size,
113 bias=False,
114 params_dtype=vllm_config.model_config.dtype,
115 quant_config=self.quant_config,
116 prefix=maybe_prefix(prefix, "fc"),
117 return_bias=False,
118 )
119
120 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
121 return self.embed_tokens(input_ids)
122
123 def forward(
124 self,
125 input_ids: torch.Tensor,
126 positions: torch.Tensor,
127 hidden_states: torch.Tensor,
128 ) -> tuple[torch.Tensor, torch.Tensor]:
129 input_embeds = self.embed_tokens(input_ids)
130 hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
131 residual = None
132 for layer in self.layers:
133 hidden_states, residual = layer(
134 positions,
135 hidden_states,
136 residual,
137 )
138 hidden_states = hidden_states + residual
139 return hidden_states, hidden_states
140
141 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
142 loader = AutoWeightsLoader(self)
143 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
144
145
146 class EagleLlamaForCausalLM(_EagleLlamaForCausalLMBase):
147 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
148 nn.Module.__init__(self)
149 speculative_config = vllm_config.speculative_config
150 assert speculative_config is not None
151 self.config = speculative_config.draft_model_config.hf_config
152 # Ensure draft_vocab_size is set
153 # default to the base vocab size when absent
154 if getattr(self.config, "draft_vocab_size", None) is None:
155 base_vocab_size = getattr(self.config, "vocab_size", None)
156 self.config.draft_vocab_size = base_vocab_size
157 target_layer_num = vllm_config.model_config.get_num_layers(
158 vllm_config.parallel_config
159 )
160 self.model = LlamaModel(
161 vllm_config=vllm_config,
162 prefix=maybe_prefix(prefix, "model"),
163 start_layer_id=target_layer_num,
164 )
165
166 logit_scale = getattr(self.config, "logit_scale", 1.0)
167 self.logits_processor = LogitsProcessor(
168 self.config.vocab_size, scale=logit_scale
169 )
170
171 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
172 return self.model.embed_input_ids(input_ids)
173
174 def forward(
175 self,
176 input_ids: torch.Tensor,
177 positions: torch.Tensor,
178 hidden_states: torch.Tensor,
179 inputs_embeds: torch.Tensor | None = None,
180 ) -> tuple[torch.Tensor, torch.Tensor]:
181 if inputs_embeds is not None:
182 raise NotImplementedError(
183 f"{type(self).__name__} does not support multimodal inputs yet."
184 )
185 return self.model(input_ids, positions, hidden_states)
186
187 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
188 def transform(inputs):
189 name, loaded_weight = inputs
190 if "lm_head" not in name:
191 name = "model." + name
192 process_eagle_weight(self, name)
193 return name, loaded_weight
194
195 loader = AutoWeightsLoader(self)
196 loader.load_weights(map(transform, weights))