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

import torch
from torch import nn
from transformers import BloomConfig
24
25
import os
import re
Woosuk Kwon's avatar
Woosuk Kwon committed
26

27
from vllm.attention import Attention, AttentionMetadata
28
from vllm.compilation.decorators import support_torch_compile
29
from vllm.config import CacheConfig, VllmConfig
30
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
31
                              get_tensor_model_parallel_world_size)
Woosuk Kwon's avatar
Woosuk Kwon committed
32
from vllm.model_executor.layers.activation import get_act_fn
33
34
35
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
from vllm.model_executor.layers.quantization import QuantizationConfig
Joe Runde's avatar
Joe Runde committed
38
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
39
from vllm.model_executor.layers.vocab_parallel_embedding import (
40
    ParallelLMHead, VocabParallelEmbedding)
41
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
42
from vllm.model_executor.sampling_metadata import SamplingMetadata
43
from vllm.sequence import IntermediateTensors
44
45
from vllm import _custom_ops as ops
from vllm.model_executor.utils import pad_weight, gemm_bank_conf
Woosuk Kwon's avatar
Woosuk Kwon committed
46

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

Woosuk Kwon's avatar
Woosuk Kwon committed
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
79

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

80
81
82
    def __init__(
        self,
        config: BloomConfig,
83
        cache_config: Optional[CacheConfig] = None,
84
        quant_config: Optional[QuantizationConfig] = None,
85
        prefix: str = "",
86
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
87
88
89
90
91
92
93
94
95
96
        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

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

        # 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
119
120
121
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              scaling,
122
                              alibi_slopes=alibi_slopes,
123
                              cache_config=cache_config,
124
125
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
zhuwenwen's avatar
zhuwenwen committed
126
        
127
128
129
130
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
Woosuk Kwon's avatar
Woosuk Kwon committed
131
132
133
134
135

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
136
137
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
138
139
140
    ) -> torch.Tensor:
        del position_ids  # Unused.
        qkv, _ = self.query_key_value(hidden_states)
zhuwenwen's avatar
zhuwenwen committed
141
142
        # if os.environ.get('FA_PAD') == '1' and self.quant_method is None:
        #     qkv = qkv[...,:-32]
Woosuk Kwon's avatar
Woosuk Kwon committed
143
        q, k, v = qkv.chunk(chunks=3, dim=-1)
144
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
145
146
147
148
149
150
        output, _ = self.dense(attn_output)
        return output


class BloomMLP(nn.Module):

151
152
153
    def __init__(
        self,
        config: BloomConfig,
154
        quant_config: Optional[QuantizationConfig] = None,
155
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
156
157
        super().__init__()
        hidden_size = config.hidden_size
158
159
160
        self.dense_h_to_4h = ColumnParallelLinear(
            hidden_size,
            4 * hidden_size,
161
            quant_config=quant_config,
162
        )
163
        self.gelu_impl = get_act_fn("gelu")
164
165
166
        self.dense_4h_to_h = RowParallelLinear(
            4 * hidden_size,
            hidden_size,
167
            quant_config=quant_config,
168
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
169
170
171

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.dense_h_to_4h(x)
172
        x = self.gelu_impl(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
173
174
175
176
177
178
        x, _ = self.dense_4h_to_h(x)
        return x


class BloomBlock(nn.Module):

179
180
181
    def __init__(
        self,
        config: BloomConfig,
182
        cache_config: Optional[CacheConfig] = None,
183
        quant_config: Optional[QuantizationConfig] = None,
184
        prefix: str = "",
185
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
186
187
188
189
190
        super().__init__()
        hidden_size = config.hidden_size

        self.input_layernorm = nn.LayerNorm(hidden_size,
                                            eps=config.layer_norm_epsilon)
191
192
193
194
        self.self_attention = BloomAttention(config,
                                             cache_config,
                                             quant_config,
                                             prefix=f"{prefix}.self_attention")
Woosuk Kwon's avatar
Woosuk Kwon committed
195
196
        self.post_attention_layernorm = nn.LayerNorm(
            hidden_size, eps=config.layer_norm_epsilon)
197
        self.mlp = BloomMLP(config, quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
198
199
200
201
202
203
204
        self.apply_residual_connection_post_layernorm = (
            config.apply_residual_connection_post_layernorm)

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
205
206
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
    ) -> 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,
222
            attn_metadata=attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
        )
        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


238
@support_torch_compile
Woosuk Kwon's avatar
Woosuk Kwon committed
239
240
class BloomModel(nn.Module):

241
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Woosuk Kwon's avatar
Woosuk Kwon committed
242
        super().__init__()
243
244
245
246
247

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

Woosuk Kwon's avatar
Woosuk Kwon committed
248
249
250
251
        self.embed_dim = config.hidden_size

        # Embedding + LN Embedding
        self.word_embeddings = VocabParallelEmbedding(
252
253
254
            config.vocab_size,
            self.embed_dim,
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
255
256
257
258
        self.word_embeddings_layernorm = nn.LayerNorm(
            self.embed_dim, eps=config.layer_norm_epsilon)

        # Transformer blocks
259
260
        self.start_layer, self.end_layer, self.h = make_layers(
            config.num_hidden_layers,
261
262
            lambda prefix: BloomBlock(
                config, cache_config, quant_config, prefix=prefix),
263
            prefix=f"{prefix}.h")
Woosuk Kwon's avatar
Woosuk Kwon committed
264
265
266

        # Final Layer Norm
        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
267
268
269
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))
Woosuk Kwon's avatar
Woosuk Kwon committed
270

271
272
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.word_embeddings_layernorm(self.word_embeddings(input_ids))
Woosuk Kwon's avatar
Woosuk Kwon committed
273
274
275
276
277

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
278
279
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
280
        intermediate_tensors: Optional[IntermediateTensors],
281
        inputs_embeds: Optional[torch.Tensor] = None,
282
283
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
284
285
286
287
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
288
289
290
291
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
        for i in range(self.start_layer, self.end_layer):
Woosuk Kwon's avatar
Woosuk Kwon committed
292
293
294
295
            layer = self.h[i]
            hidden_states = layer(
                position_ids,
                hidden_states,
296
                kv_caches[i - self.start_layer],
297
                attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
298
            )
299
300
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
Woosuk Kwon's avatar
Woosuk Kwon committed
301
302
303
304
        hidden_states = self.ln_f(hidden_states)
        return hidden_states


305
class BloomForCausalLM(nn.Module, SupportsPP):
Woosuk Kwon's avatar
Woosuk Kwon committed
306

307
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Woosuk Kwon's avatar
Woosuk Kwon committed
308
        super().__init__()
309
310
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
Woosuk Kwon's avatar
Woosuk Kwon committed
311
        self.config = config
312
        self.quant_config = quant_config
313
314
315
        self.transformer = BloomModel(vllm_config=vllm_config,
                                      prefix=maybe_prefix(
                                          prefix, "transformer"))
316
317
318
319
320
321
        if self.config.tie_word_embeddings:
            self.lm_head = self.transformer.word_embeddings
        else:
            self.lm_head = ParallelLMHead(self.config.vocab_size,
                                          self.config.hidden_size)

322
        self.logits_processor = LogitsProcessor(config.vocab_size)
zhuwenwen's avatar
zhuwenwen committed
323

Joe Runde's avatar
Joe Runde committed
324
        self.sampler = get_sampler()
325
326
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
327
328
329
330
331
332
333
334
335
        
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
              
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
        self.use_gemm_pad = os.environ.get('GEMM_PAD') == '1'
        self.use_fa_pad = os.environ.get('FA_PAD') == '1'
Woosuk Kwon's avatar
Woosuk Kwon committed
336

337
338
339
    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
340
341
342
343
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
344
345
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
346
        intermediate_tensors: Optional[IntermediateTensors] = None,
347
        inputs_embeds: Optional[torch.Tensor] = None,
348
    ) -> Union[torch.Tensor, IntermediateTensors]:
Woosuk Kwon's avatar
Woosuk Kwon committed
349
        hidden_states = self.transformer(input_ids, positions, kv_caches,
350
351
                                         attn_metadata, intermediate_tensors,
                                         inputs_embeds)
352
353
        return hidden_states

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

363
364
    def sample(
        self,
365
        logits: torch.Tensor,
366
        sampling_metadata: SamplingMetadata,
367
    ) -> Optional[SamplerOutput]:
368
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
369
370
        return next_tokens

371
372
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
373
        params_dict = dict(self.named_parameters(remove_duplicate=False))
374
        loaded_params: Set[str] = set()
375
        for name, loaded_weight in weights:
376
            if name == "lm_head.weight":
377
378
379
                continue
            if not name.startswith("transformer."):
                name = "transformer." + name
380
381
            if is_pp_missing_parameter(name, self):
                continue
382
            param = params_dict[name]
Woosuk Kwon's avatar
Woosuk Kwon committed
383
384

            if "query_key_value" in name:
385
386
387
                # 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
388
                # Thus, we need weight conversion.
389
                output_dim = getattr(param, "output_dim", None)
Woosuk Kwon's avatar
Woosuk Kwon committed
390
                num_heads = self.config.num_attention_heads
391
392
393
394
395
396
397
398
399
400
401
402
                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)
403
            loaded_params.add(name)
404
405
406
407
408
409
410
411
412
413
        
        if self.use_llama_nn and self.quant_method is None:
            lay_key_words = [
                "self_attention.query_key_value.weight",
                "self_attention.dense.weight",
                "mlp.dense_h_to_4h.weight",
                "mlp.dense_4h_to_h.weight"
            ]
            combined_words = "|".join(lay_key_words)
            
zhuwenwen's avatar
zhuwenwen committed
414
415
            # lay_qkv_words = ["self_attention.query_key_value.weight"]   
            # qkv_words = "|".join(lay_qkv_words)  
416
            
zhuwenwen's avatar
zhuwenwen committed
417
418
            # lay_qkv_bias_words = ["self_attention.query_key_value.bias"]   
            # qkv_bias_words = "|".join(lay_qkv_bias_words) 
419
420
            
            for layername, weight in params_dict.items():
zhuwenwen's avatar
zhuwenwen committed
421
422
                # if self.use_fa_pad and (re.findall(qkv_bias_words, layername)):
                #     weight.data = pad_weight(weight.data, 32)
423
424
425
                    
                matches = re.findall(combined_words, layername)
                if matches:   
zhuwenwen's avatar
zhuwenwen committed
426
427
                    # if self.use_gemm_pad and gemm_bank_conf(weight.data.shape[0]):
                    #     weight.data = pad_weight(weight.data, 32)  
428
                    
zhuwenwen's avatar
zhuwenwen committed
429
430
431
                    # if self.use_fa_pad and (re.findall(qkv_words, layername)):
                    #     if not gemm_bank_conf(weight.data.shape[0]):
                    #         weight.data = pad_weight(weight.data, 32)
432
433
434
435
436
437
438
439
                        
                    _weight = torch.zeros_like(weight.data)
                    ori_shape =_weight.shape
                    
                    ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
                    weight.data.copy_(_weight)
                    
                    weight.data=weight.data.reshape(ori_shape[1],-1)
zhuwenwen's avatar
zhuwenwen committed
440
                    
441
        return loaded_params
442