gemma.py 15.8 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

Xiang Xu's avatar
Xiang Xu committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Copyright 2023 The vLLM team.
# Copyright (c) Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Inference-only Gemma model compatible with HuggingFace weights."""
18
from functools import cache
19
from typing import Iterable, Optional, Set, Tuple, Union
Xiang Xu's avatar
Xiang Xu committed
20
21
22
23
24

import torch
from torch import nn
from transformers import GemmaConfig

25
from vllm.attention import Attention
26
from vllm.compilation.decorators import support_torch_compile
27
from vllm.config import CacheConfig, VllmConfig
28
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
29
from vllm.logger import init_logger
30
from vllm.model_executor.layers.activation import GeluAndMul
31
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
32
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
Xiang Xu's avatar
Xiang Xu committed
33
34
                                               QKVParallelLinear,
                                               RowParallelLinear)
35
from vllm.model_executor.layers.logits_processor import LogitsProcessor
36
from vllm.model_executor.layers.quantization import QuantizationConfig
Woosuk Kwon's avatar
Woosuk Kwon committed
37
from vllm.model_executor.layers.rotary_embedding import get_rope
Xiang Xu's avatar
Xiang Xu committed
38
39
from vllm.model_executor.layers.vocab_parallel_embedding import (
    VocabParallelEmbedding)
40
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
Xiang Xu's avatar
Xiang Xu committed
41
from vllm.model_executor.sampling_metadata import SamplingMetadata
42
from vllm.sequence import IntermediateTensors
Xiang Xu's avatar
Xiang Xu committed
43

44
from .interfaces import SupportsLoRA, SupportsPP
45
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
46
47
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
48

49
50
51
logger = init_logger(__name__)


52
@cache
53
54
55
56
57
58
59
60
61
62
63
def _get_gemma_act_fn(
    hidden_act: Optional[str],
    hidden_activation: Optional[str],
) -> nn.Module:
    if hidden_activation is None:
        if hidden_act is not None:
            logger.warning(
                "Gemma's activation function was incorrectly set to exact GeLU "
                "in the config JSON file when it was initially released. "
                "Changing the activation function to approximate GeLU "
                "(`gelu_pytorch_tanh`). If you want to use the legacy "
64
65
                "`%s`, edit the config JSON to set "
                "`hidden_activation=%s` instead of `hidden_act`. "
66
                "See https://github.com/huggingface/transformers/pull/29402 "
67
                "for more details.", hidden_act, hidden_act)
68
69
70
71
72
73
74
75
76
        return GeluAndMul(approximate="tanh")
    elif hidden_activation == "gelu_pytorch_tanh":
        return GeluAndMul(approximate="tanh")
    elif hidden_activation == "gelu":
        return GeluAndMul(approximate="none")
    else:
        raise ValueError(f"Activation function {hidden_act} is not "
                         "supported for Gemma models.")

Xiang Xu's avatar
Xiang Xu committed
77
78
79
80
81
82
83

class GemmaMLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
84
85
        hidden_act: Optional[str] = None,
        hidden_activation: Optional[str] = None,
86
        quant_config: Optional[QuantizationConfig] = None,
87
        prefix: str = "",
Xiang Xu's avatar
Xiang Xu committed
88
89
    ) -> None:
        super().__init__()
90
        self.gate_up_proj = MergedColumnParallelLinear(
91
92
            hidden_size,
            [intermediate_size] * 2,
93
            bias=False,
94
95
96
97
98
99
100
101
102
103
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
104
        self.act_fn = _get_gemma_act_fn(hidden_act, hidden_activation)
Xiang Xu's avatar
Xiang Xu committed
105
106

    def forward(self, x):
107
108
109
110
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x
Xiang Xu's avatar
Xiang Xu committed
111
112
113
114


class GemmaAttention(nn.Module):

115
116
117
118
119
120
121
122
123
124
125
126
    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        max_position_embeddings: int = 8192,
        rope_theta: float = 10000,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
Xiang Xu's avatar
Xiang Xu committed
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        super().__init__()
        self.hidden_size = hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = head_dim
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.rope_theta = rope_theta

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
155
            quant_config=quant_config,
156
            prefix=f"{prefix}.qkv_proj",
Xiang Xu's avatar
Xiang Xu committed
157
158
159
160
161
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
162
            quant_config=quant_config,
163
            prefix=f"{prefix}.o_proj",
Xiang Xu's avatar
Xiang Xu committed
164
165
        )

Woosuk Kwon's avatar
Woosuk Kwon committed
166
        self.rotary_emb = get_rope(
Xiang Xu's avatar
Xiang Xu committed
167
168
            self.head_dim,
            rotary_dim=self.head_dim,
Woosuk Kwon's avatar
Woosuk Kwon committed
169
            max_position=max_position_embeddings,
Xiang Xu's avatar
Xiang Xu committed
170
171
172
            base=self.rope_theta,
            is_neox_style=True,
        )
173
174
175
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
176
                              num_kv_heads=self.num_kv_heads,
177
                              cache_config=cache_config,
178
179
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
Xiang Xu's avatar
Xiang Xu committed
180
181
182
183
184
185
186
187
188

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
189
        attn_output = self.attn(q, k, v)
Xiang Xu's avatar
Xiang Xu committed
190
191
192
193
194
195
196
197
198
        output, _ = self.o_proj(attn_output)
        return output


class GemmaDecoderLayer(nn.Module):

    def __init__(
        self,
        config: GemmaConfig,
199
        cache_config: Optional[CacheConfig] = None,
200
        quant_config: Optional[QuantizationConfig] = None,
201
        prefix: str = "",
Xiang Xu's avatar
Xiang Xu committed
202
203
204
205
206
207
208
209
210
211
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = GemmaAttention(
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            head_dim=config.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            rope_theta=config.rope_theta,
212
            cache_config=cache_config,
213
            quant_config=quant_config,
214
            prefix=f"{prefix}.self_attn",
Xiang Xu's avatar
Xiang Xu committed
215
216
217
218
        )
        self.mlp = GemmaMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
219
220
            hidden_act=config.hidden_act,
            hidden_activation=getattr(config, "hidden_activation", None),
221
            quant_config=quant_config,
222
            prefix=f"{prefix}.mlp",
Xiang Xu's avatar
Xiang Xu committed
223
        )
224
225
226
227
        self.input_layernorm = GemmaRMSNorm(config.hidden_size,
                                            eps=config.rms_norm_eps)
        self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size,
                                                     eps=config.rms_norm_eps)
Xiang Xu's avatar
Xiang Xu committed
228
229
230
231
232

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
233
        residual: Optional[torch.Tensor],
Xiang Xu's avatar
Xiang Xu committed
234
235
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
236
237
238
239
240
241
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
Xiang Xu's avatar
Xiang Xu committed
242
243
244
245
246
247
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        # Fully Connected
248
249
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
Xiang Xu's avatar
Xiang Xu committed
250
        hidden_states = self.mlp(hidden_states)
251
        return hidden_states, residual
Xiang Xu's avatar
Xiang Xu committed
252
253


254
@support_torch_compile
Xiang Xu's avatar
Xiang Xu committed
255
256
class GemmaModel(nn.Module):

257
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Xiang Xu's avatar
Xiang Xu committed
258
        super().__init__()
259
260
261
262
263

        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

Xiang Xu's avatar
Xiang Xu committed
264
265
266
267
268
269
        self.config = config

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
270
271
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
272
273
            lambda prefix: GemmaDecoderLayer(
                config, cache_config, quant_config, prefix=prefix),
274
            prefix=f"{prefix}.layers")
275
        self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Xiang Xu's avatar
Xiang Xu committed
276

277
278
279
280
281
282
        # Normalize the embedding by sqrt(hidden_size)
        # The normalizer's data type should be downcasted to the model's
        # data type such as bfloat16, not float32.
        # See https://github.com/huggingface/transformers/pull/29402
        normalizer = self.config.hidden_size**0.5
        self.register_buffer("normalizer", torch.tensor(normalizer))
283
284
285
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
286

Roger Wang's avatar
Roger Wang committed
287
288
289
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Xiang Xu's avatar
Xiang Xu committed
290
291
292
293
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
294
        intermediate_tensors: Optional[IntermediateTensors],
Roger Wang's avatar
Roger Wang committed
295
        inputs_embeds: Optional[torch.Tensor] = None,
296
297
298
299
300
301
302
303
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
            hidden_states *= self.normalizer
            residual = None
Roger Wang's avatar
Roger Wang committed
304
        else:
305
306
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
307
        for layer in self.layers[self.start_layer:self.end_layer]:
308
            hidden_states, residual = layer(
Xiang Xu's avatar
Xiang Xu committed
309
310
                positions,
                hidden_states,
311
                residual,
Xiang Xu's avatar
Xiang Xu committed
312
            )
313
314
315
316
317
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })
318
        hidden_states, _ = self.norm(hidden_states, residual)
Xiang Xu's avatar
Xiang Xu committed
319
320
        return hidden_states

321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        for name, loaded_weight in weights:
            for (param_name, shard_name, shard_id) in stacked_params_mapping:
                if shard_name not in name:
                    continue
                name = name.replace(shard_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)

        return loaded_params

Xiang Xu's avatar
Xiang Xu committed
361

362
class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
363
364
365
366
367
368
369
370
371
372
373
374
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

375
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Xiang Xu's avatar
Xiang Xu committed
376
        super().__init__()
377
378
379
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
380

Xiang Xu's avatar
Xiang Xu committed
381
        self.config = config
382
383
        # currently all existing Gemma models have `tie_word_embeddings` enabled
        assert config.tie_word_embeddings
384
385
        self.lora_config = lora_config

386
        self.quant_config = quant_config
387
        self.model = GemmaModel(vllm_config=vllm_config,
388
                                prefix=maybe_prefix(prefix, "model"))
389
        self.logits_processor = LogitsProcessor(config.vocab_size)
390
391
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
Xiang Xu's avatar
Xiang Xu committed
392

393
394
395
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

Xiang Xu's avatar
Xiang Xu committed
396
397
398
399
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
400
        intermediate_tensors: Optional[IntermediateTensors] = None,
401
        inputs_embeds: Optional[torch.Tensor] = None,
402
    ) -> Union[torch.Tensor, IntermediateTensors]:
403
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
404
                                   inputs_embeds)
Xiang Xu's avatar
Xiang Xu committed
405
406
        return hidden_states

407
408
409
410
411
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
412
413
        logits = self.logits_processor(self.model.embed_tokens, hidden_states,
                                       sampling_metadata)
414
415
        return logits

416
417
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
418
419
420
421
422
423
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=(["lm_head."]
                           if self.config.tie_word_embeddings else None),
        )
        return loader.load_weights(weights)