bloom.py 12 KB
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
# coding=utf-8
2
3
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/bloom/modeling_bloom.py
4
# Copyright 2023 The vLLM team.
Woosuk Kwon's avatar
Woosuk Kwon committed
5
6
7
8
9
10
11
12
13
14
15
16
17
# 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
18
"""Inference-only BLOOM model compatible with HuggingFace weights."""
Woosuk Kwon's avatar
Woosuk Kwon committed
19
import math
20
from typing import Iterable, List, Optional, Tuple
Woosuk Kwon's avatar
Woosuk Kwon committed
21
22
23
24
25

import torch
from torch import nn
from transformers import BloomConfig

26
from vllm.attention import Attention, AttentionMetadata
27
from vllm.config import CacheConfig
28
29
from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size)
Woosuk Kwon's avatar
Woosuk Kwon committed
30
from vllm.model_executor.layers.activation import get_act_fn
31
32
33
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
34
from vllm.model_executor.layers.logits_processor import LogitsProcessor
35
36
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
Woosuk Kwon's avatar
Woosuk Kwon committed
37
from vllm.model_executor.layers.sampler import Sampler
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
41
from vllm.model_executor.sampling_metadata import SamplingMetadata
42
from vllm.sequence import SamplerOutput
Woosuk Kwon's avatar
Woosuk Kwon committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71


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):

72
73
74
    def __init__(
        self,
        config: BloomConfig,
75
        cache_config: Optional[CacheConfig] = None,
76
        quant_config: Optional[QuantizationConfig] = None,
77
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
78
79
80
81
82
83
84
85
86
87
        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

88
        self.query_key_value = QKVParallelLinear(
Woosuk Kwon's avatar
Woosuk Kwon committed
89
            self.hidden_size,
90
91
            self.head_dim,
            self.total_num_heads,
Woosuk Kwon's avatar
Woosuk Kwon committed
92
            bias=True,
93
            quant_config=quant_config,
Woosuk Kwon's avatar
Woosuk Kwon committed
94
95
96
97
98
        )
        self.dense = RowParallelLinear(
            self.hidden_size,
            self.hidden_size,
            bias=True,
99
            quant_config=quant_config,
Woosuk Kwon's avatar
Woosuk Kwon committed
100
101
102
103
104
105
106
107
108
109
        )

        # 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
110
111
112
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              scaling,
113
114
                              alibi_slopes=alibi_slopes,
                              cache_config=cache_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
115
116
117
118
119

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
120
121
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
122
123
124
125
    ) -> torch.Tensor:
        del position_ids  # Unused.
        qkv, _ = self.query_key_value(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
126
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
127
128
129
130
131
132
        output, _ = self.dense(attn_output)
        return output


class BloomMLP(nn.Module):

133
134
135
    def __init__(
        self,
        config: BloomConfig,
136
        quant_config: Optional[QuantizationConfig] = None,
137
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
138
139
        super().__init__()
        hidden_size = config.hidden_size
140
141
142
        self.dense_h_to_4h = ColumnParallelLinear(
            hidden_size,
            4 * hidden_size,
143
            quant_config=quant_config,
144
        )
145
        self.gelu_impl = get_act_fn("gelu", quant_config, 4 * hidden_size)
146
147
148
        self.dense_4h_to_h = RowParallelLinear(
            4 * hidden_size,
            hidden_size,
149
            quant_config=quant_config,
150
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
151
152
153

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.dense_h_to_4h(x)
154
        x = self.gelu_impl(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
155
156
157
158
159
160
        x, _ = self.dense_4h_to_h(x)
        return x


class BloomBlock(nn.Module):

161
162
163
    def __init__(
        self,
        config: BloomConfig,
164
        cache_config: Optional[CacheConfig] = None,
165
        quant_config: Optional[QuantizationConfig] = None,
166
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
167
168
169
170
171
        super().__init__()
        hidden_size = config.hidden_size

        self.input_layernorm = nn.LayerNorm(hidden_size,
                                            eps=config.layer_norm_epsilon)
172
173
        self.self_attention = BloomAttention(config, cache_config,
                                             quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
174
175
        self.post_attention_layernorm = nn.LayerNorm(
            hidden_size, eps=config.layer_norm_epsilon)
176
        self.mlp = BloomMLP(config, quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
177
178
179
180
181
182
183
        self.apply_residual_connection_post_layernorm = (
            config.apply_residual_connection_post_layernorm)

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
184
185
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
    ) -> 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,
            kv_cache=kv_cache,
201
            attn_metadata=attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
        )
        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


class BloomModel(nn.Module):

219
220
221
    def __init__(
        self,
        config: BloomConfig,
222
        cache_config: Optional[CacheConfig] = None,
223
        quant_config: Optional[QuantizationConfig] = None,
224
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
225
226
227
228
229
        super().__init__()
        self.embed_dim = config.hidden_size

        # Embedding + LN Embedding
        self.word_embeddings = VocabParallelEmbedding(
230
231
232
            config.vocab_size,
            self.embed_dim,
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
233
234
235
236
        self.word_embeddings_layernorm = nn.LayerNorm(
            self.embed_dim, eps=config.layer_norm_epsilon)

        # Transformer blocks
237
        self.h = nn.ModuleList([
238
            BloomBlock(config, cache_config, quant_config)
239
240
            for _ in range(config.num_hidden_layers)
        ])
Woosuk Kwon's avatar
Woosuk Kwon committed
241
242
243
244
245
246
247
248

        # Final Layer Norm
        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
249
250
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
251
252
253
254
255
256
257
258
259
    ) -> torch.Tensor:
        hidden_states = self.word_embeddings(input_ids)
        hidden_states = self.word_embeddings_layernorm(hidden_states)
        for i in range(len(self.h)):
            layer = self.h[i]
            hidden_states = layer(
                position_ids,
                hidden_states,
                kv_caches[i],
260
                attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
261
262
263
264
265
266
267
            )
        hidden_states = self.ln_f(hidden_states)
        return hidden_states


class BloomForCausalLM(nn.Module):

268
269
270
    def __init__(
        self,
        config: BloomConfig,
271
        cache_config: Optional[CacheConfig] = None,
272
        quant_config: Optional[QuantizationConfig] = None,
273
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
274
275
        super().__init__()
        self.config = config
276
        self.quant_config = quant_config
277
        self.transformer = BloomModel(config, cache_config, quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
278
        self.lm_head_weight = self.transformer.word_embeddings.weight
279
280
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
Woosuk Kwon's avatar
Woosuk Kwon committed
281
282
283
284
285

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
286
287
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
288
    ) -> torch.Tensor:
Woosuk Kwon's avatar
Woosuk Kwon committed
289
        hidden_states = self.transformer(input_ids, positions, kv_caches,
290
                                         attn_metadata)
291
292
        return hidden_states

293
294
295
296
297
298
    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        logits = self.logits_processor(self.lm_head_weight, hidden_states,
                                       sampling_metadata)
        return logits

299
300
    def sample(
        self,
301
        logits: torch.Tensor,
302
        sampling_metadata: SamplingMetadata,
303
    ) -> Optional[SamplerOutput]:
304
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
305
306
        return next_tokens

307
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
308
        params_dict = dict(self.named_parameters(remove_duplicate=False))
309
        for name, loaded_weight in weights:
310
            if name == "lm_head.weight":
311
312
313
314
                continue
            if not name.startswith("transformer."):
                name = "transformer." + name
            param = params_dict[name]
Woosuk Kwon's avatar
Woosuk Kwon committed
315
316

            if "query_key_value" in name:
317
318
319
                # 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).
Woosuk Kwon's avatar
Woosuk Kwon committed
320
                # Thus, we need weight conversion.
321
                output_dim = getattr(param, "output_dim", None)
Woosuk Kwon's avatar
Woosuk Kwon committed
322
                num_heads = self.config.num_attention_heads
323
324
325
326
327
328
329
330
331
332
333
334
                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)