bloom.py 14.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/bloom/modeling_bloom.py
6
# Copyright 2023 The vLLM team.
Woosuk Kwon's avatar
Woosuk Kwon committed
7
8
9
10
11
12
13
14
15
16
17
18
19
# Copyright 2022 HuggingFace Inc. team and BigScience workshop.
#
# 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.
Woosuk Kwon's avatar
Woosuk Kwon committed
20
"""Inference-only BLOOM model compatible with HuggingFace weights."""
Woosuk Kwon's avatar
Woosuk Kwon committed
21
import math
22
from collections.abc import Iterable
23
from itertools import islice
24
from typing import Optional, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
25
26
27
28
29

import torch
from torch import nn
from transformers import BloomConfig

30
from vllm.attention import Attention
31
from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
34
                              get_tensor_model_parallel_world_size)
Woosuk Kwon's avatar
Woosuk Kwon committed
35
from vllm.model_executor.layers.activation import get_act_fn
36
37
38
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
39
from vllm.model_executor.layers.logits_processor import LogitsProcessor
40
from vllm.model_executor.layers.quantization import QuantizationConfig
41
from vllm.model_executor.layers.vocab_parallel_embedding import (
42
    ParallelLMHead, VocabParallelEmbedding)
43
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
44
from vllm.sequence import IntermediateTensors
Woosuk Kwon's avatar
Woosuk Kwon committed
45

46
from .interfaces import SupportsPP, SupportsQuant
47
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
48
49
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
50

Woosuk Kwon's avatar
Woosuk Kwon committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
    closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
    base = torch.tensor(
        2**(-(2**-(math.log2(closest_power_of_2) - 3))),
        dtype=torch.float32,
    )
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != total_num_heads:
        extra_base = torch.tensor(
            2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
            dtype=torch.float32,
        )
        num_remaining_heads = min(closest_power_of_2,
                                  total_num_heads - closest_power_of_2)
        extra_powers = torch.arange(start=1,
                                    end=1 + 2 * num_remaining_heads,
                                    step=2,
                                    dtype=torch.int32)
        slopes = torch.cat(
            [slopes, torch.pow(extra_base, extra_powers)], dim=0)
    return slopes


class BloomAttention(nn.Module):

79
80
81
    def __init__(
        self,
        config: BloomConfig,
82
        cache_config: Optional[CacheConfig] = None,
83
        quant_config: Optional[QuantizationConfig] = None,
84
        prefix: str = "",
85
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
86
87
88
89
90
91
92
93
94
95
        super().__init__()
        self.hidden_size = config.hidden_size
        self.total_num_heads = config.n_head
        self.head_dim = self.hidden_size // self.total_num_heads
        assert self.head_dim * self.total_num_heads == self.hidden_size

        tp_world_size = get_tensor_model_parallel_world_size()
        assert self.total_num_heads % tp_world_size == 0
        self.num_heads = self.total_num_heads // tp_world_size

96
        self.query_key_value = QKVParallelLinear(
Woosuk Kwon's avatar
Woosuk Kwon committed
97
            self.hidden_size,
98
99
            self.head_dim,
            self.total_num_heads,
Woosuk Kwon's avatar
Woosuk Kwon committed
100
            bias=True,
101
            quant_config=quant_config,
Woosuk Kwon's avatar
Woosuk Kwon committed
102
103
104
105
106
        )
        self.dense = RowParallelLinear(
            self.hidden_size,
            self.hidden_size,
            bias=True,
107
            quant_config=quant_config,
Woosuk Kwon's avatar
Woosuk Kwon committed
108
109
110
111
112
113
114
115
116
117
        )

        # Create the alibi slopes and slice them.
        tp_rank = get_tensor_model_parallel_rank()
        head_start = tp_rank * self.num_heads
        head_end = (tp_rank + 1) * self.num_heads
        alibi_slopes = _get_alibi_slopes(self.total_num_heads)
        alibi_slopes = alibi_slopes[head_start:head_end].tolist()

        scaling = self.head_dim**-0.5
118
119
120
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              scaling,
121
                              alibi_slopes=alibi_slopes,
122
                              cache_config=cache_config,
123
124
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
Woosuk Kwon's avatar
Woosuk Kwon committed
125
126
127
128
129
130
131
132
133

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        del position_ids  # Unused.
        qkv, _ = self.query_key_value(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
134
        attn_output = self.attn(q, k, v)
Woosuk Kwon's avatar
Woosuk Kwon committed
135
136
137
138
139
140
        output, _ = self.dense(attn_output)
        return output


class BloomMLP(nn.Module):

141
142
143
    def __init__(
        self,
        config: BloomConfig,
144
        quant_config: Optional[QuantizationConfig] = None,
145
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
146
147
        super().__init__()
        hidden_size = config.hidden_size
148
149
150
        self.dense_h_to_4h = ColumnParallelLinear(
            hidden_size,
            4 * hidden_size,
151
            quant_config=quant_config,
152
        )
153
        self.gelu_impl = get_act_fn("gelu")
154
155
156
        self.dense_4h_to_h = RowParallelLinear(
            4 * hidden_size,
            hidden_size,
157
            quant_config=quant_config,
158
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
159
160
161

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.dense_h_to_4h(x)
162
        x = self.gelu_impl(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
163
164
165
166
167
168
        x, _ = self.dense_4h_to_h(x)
        return x


class BloomBlock(nn.Module):

169
170
171
    def __init__(
        self,
        config: BloomConfig,
172
        cache_config: Optional[CacheConfig] = None,
173
        quant_config: Optional[QuantizationConfig] = None,
174
        prefix: str = "",
175
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
176
177
178
179
180
        super().__init__()
        hidden_size = config.hidden_size

        self.input_layernorm = nn.LayerNorm(hidden_size,
                                            eps=config.layer_norm_epsilon)
181
182
183
184
        self.self_attention = BloomAttention(config,
                                             cache_config,
                                             quant_config,
                                             prefix=f"{prefix}.self_attention")
Woosuk Kwon's avatar
Woosuk Kwon committed
185
186
        self.post_attention_layernorm = nn.LayerNorm(
            hidden_size, eps=config.layer_norm_epsilon)
187
        self.mlp = BloomMLP(config, quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
        self.apply_residual_connection_post_layernorm = (
            config.apply_residual_connection_post_layernorm)

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        # Layer norm at the beginning of the transformer layer.
        layernorm_output = self.input_layernorm(hidden_states)

        # Layer norm post the self attention.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = hidden_states

        # Self attention.
        attention_output = self.self_attention(
            position_ids=position_ids,
            hidden_states=layernorm_output,
        )
        attention_output = attention_output + residual
        layernorm_output = self.post_attention_layernorm(attention_output)

        # Get residual
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = attention_output

        # MLP.
        output = self.mlp(layernorm_output) + residual
        return output


224
@support_torch_compile
Woosuk Kwon's avatar
Woosuk Kwon committed
225
226
class BloomModel(nn.Module):

227
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Woosuk Kwon's avatar
Woosuk Kwon committed
228
        super().__init__()
229
230
231
232

        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
233
        self.config = config
234

Woosuk Kwon's avatar
Woosuk Kwon committed
235
236
237
238
        self.embed_dim = config.hidden_size

        # Embedding + LN Embedding
        self.word_embeddings = VocabParallelEmbedding(
239
240
241
            config.vocab_size,
            self.embed_dim,
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
242
243
244
245
        self.word_embeddings_layernorm = nn.LayerNorm(
            self.embed_dim, eps=config.layer_norm_epsilon)

        # Transformer blocks
246
247
        self.start_layer, self.end_layer, self.h = make_layers(
            config.num_hidden_layers,
248
249
            lambda prefix: BloomBlock(
                config, cache_config, quant_config, prefix=prefix),
250
            prefix=f"{prefix}.h")
Woosuk Kwon's avatar
Woosuk Kwon committed
251
252
253

        # Final Layer Norm
        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
254
255
256
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))
Woosuk Kwon's avatar
Woosuk Kwon committed
257

258
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
259
        return self.word_embeddings(input_ids)
260

Woosuk Kwon's avatar
Woosuk Kwon committed
261
262
263
264
    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
265
        intermediate_tensors: Optional[IntermediateTensors],
266
        inputs_embeds: Optional[torch.Tensor] = None,
267
268
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
269
270
271
272
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
273
            hidden_states = self.word_embeddings_layernorm(hidden_states)
274
275
276
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
277
        for layer in islice(self.h, self.start_layer, self.end_layer):
278
            hidden_states = layer(position_ids, hidden_states)
279
280
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
Woosuk Kwon's avatar
Woosuk Kwon committed
281
282
283
        hidden_states = self.ln_f(hidden_states)
        return hidden_states

284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            if is_pp_missing_parameter(name, self):
                continue
            param = params_dict[name]

            if "query_key_value" in name:
                # NOTE: BLOOM's fused QKV's output_dim has the shape of
                # (num_heads * 3 * head_size), while the
                # required shape is (3 * num_heads * head_size).
                # Thus, we need weight conversion.
                output_dim = getattr(param, "output_dim", None)
                num_heads = self.config.num_attention_heads
                if output_dim is not None:
                    loaded_weight_shape = loaded_weight.shape
                    loaded_weight = loaded_weight.view(
                        loaded_weight_shape[:output_dim] + (num_heads, 3, -1) +
                        loaded_weight_shape[output_dim + 1:])
                    loaded_weight = loaded_weight.transpose(
                        output_dim, output_dim + 1)
                    loaded_weight = loaded_weight.reshape(loaded_weight_shape)

            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)
            loaded_params.add(name)

        return loaded_params

Woosuk Kwon's avatar
Woosuk Kwon committed
316

317
class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant):
Woosuk Kwon's avatar
Woosuk Kwon committed
318

319
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Woosuk Kwon's avatar
Woosuk Kwon committed
320
        super().__init__()
321
322
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
Woosuk Kwon's avatar
Woosuk Kwon committed
323
        self.config = config
324
        self.quant_config = quant_config
325
326
327
        self.transformer = BloomModel(vllm_config=vllm_config,
                                      prefix=maybe_prefix(
                                          prefix, "transformer"))
328
329
330
331
        if self.config.tie_word_embeddings:
            self.lm_head = self.transformer.word_embeddings
        else:
            self.lm_head = ParallelLMHead(self.config.vocab_size,
332
333
334
                                          self.config.hidden_size,
                                          prefix=maybe_prefix(
                                              prefix, "lm_head"))
335

336
        self.logits_processor = LogitsProcessor(config.vocab_size)
337
338
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
Woosuk Kwon's avatar
Woosuk Kwon committed
339

340
341
342
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.transformer.get_input_embeddings(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
343
344
345
346
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
347
        intermediate_tensors: Optional[IntermediateTensors] = None,
348
        inputs_embeds: Optional[torch.Tensor] = None,
349
    ) -> Union[torch.Tensor, IntermediateTensors]:
350
351
        hidden_states = self.transformer(input_ids, positions,
                                         intermediate_tensors, inputs_embeds)
352
353
        return hidden_states

354
355
356
357
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> Optional[torch.Tensor]:
358
        logits = self.logits_processor(self.lm_head, hidden_states)
359
360
        return logits

361
362
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
363
364
365
366
367
368
369
370
371
372
373
374
        loader = AutoWeightsLoader(self, skip_prefixes=["lm_head.weight"])
        weights = _add_transformer_prefix(weights)
        return loader.load_weights(weights)


def _add_transformer_prefix(
    weights: Iterable[tuple[str, torch.Tensor]]
) -> Iterable[tuple[str, torch.Tensor]]:
    for name, tensor in weights:
        if not name.startswith('transformer.'):
            name = 'transformer.' + name
        yield name, tensor