mixtral.py 19.7 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

Pierre Stock's avatar
Pierre Stock committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
# Copyright 2023 The vLLM team.
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# 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 Mixtral model."""
25
from typing import Iterable, Optional, Set, Tuple, Union
Pierre Stock's avatar
Pierre Stock committed
26
27
28

import torch
from torch import nn
29
from transformers import MixtralConfig
Pierre Stock's avatar
Pierre Stock committed
30

31
from vllm.attention import Attention
32
from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, VllmConfig
34
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
35
from vllm.model_executor.layers.fused_moe import FusedMoE
Pierre Stock's avatar
Pierre Stock committed
36
from vllm.model_executor.layers.layernorm import RMSNorm
37
from vllm.model_executor.layers.linear import (QKVParallelLinear,
Philipp Moritz's avatar
Philipp Moritz committed
38
                                               ReplicatedLinear,
Pierre Stock's avatar
Pierre Stock committed
39
                                               RowParallelLinear)
40
from vllm.model_executor.layers.logits_processor import LogitsProcessor
41
from vllm.model_executor.layers.quantization import QuantizationConfig
42
from vllm.model_executor.layers.rotary_embedding import get_rope
Pierre Stock's avatar
Pierre Stock committed
43
from vllm.model_executor.layers.vocab_parallel_embedding import (
44
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
45
46
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
Pierre Stock's avatar
Pierre Stock committed
47
from vllm.model_executor.sampling_metadata import SamplingMetadata
48
from vllm.sequence import IntermediateTensors
Pierre Stock's avatar
Pierre Stock committed
49

50
from .interfaces import SupportsLoRA, SupportsPP
51
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
52
53
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
54

Pierre Stock's avatar
Pierre Stock committed
55

Philipp Moritz's avatar
Philipp Moritz committed
56
57
58
59
60
61
62
63
class MixtralMoE(nn.Module):
    """A tensor-parallel MoE implementation for Mixtral that shards each expert
    across all ranks.

    Each expert's weights are sharded across all ranks and a fused MoE
    kernel is used for the forward pass, and finally we reduce the outputs
    across ranks.
    """
64

65
66
67
68
69
70
71
    def __init__(self,
                 num_experts: int,
                 top_k: int,
                 hidden_size: int,
                 intermediate_size: int,
                 params_dtype: Optional[torch.dtype] = None,
                 quant_config: Optional[QuantizationConfig] = None,
72
                 tp_size: Optional[int] = None,
73
                 dp_size: Optional[int] = None,
74
                 prefix: str = ""):
75
        super().__init__()
Philipp Moritz's avatar
Philipp Moritz committed
76
        self.hidden_size = hidden_size
77

78
        # Gate always runs at half / full precision for now.
79

80
81
        self.gate = ReplicatedLinear(hidden_size,
                                     num_experts,
82
                                     bias=False,
83
                                     params_dtype=params_dtype,
84
85
                                     quant_config=None,
                                     prefix=f"{prefix}.gate")
86

87
88
89
90
91
92
93
94
        self.experts = FusedMoE(num_experts=num_experts,
                                top_k=top_k,
                                hidden_size=hidden_size,
                                intermediate_size=intermediate_size,
                                params_dtype=params_dtype,
                                reduce_results=True,
                                renormalize=True,
                                quant_config=quant_config,
95
                                tp_size=tp_size,
96
                                dp_size=dp_size,
97
                                prefix=f"{prefix}.experts")
98

99
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
100
101
        # NOTE: hidden_states can have either 1D or 2D shape.
        orig_shape = hidden_states.shape
Philipp Moritz's avatar
Philipp Moritz committed
102
        hidden_states = hidden_states.view(-1, self.hidden_size)
103
        # router_logits: (num_tokens, n_experts)
104
        router_logits, _ = self.gate(hidden_states)
105
        final_hidden_states = self.experts(hidden_states, router_logits)
106
        return final_hidden_states.view(orig_shape)
Pierre Stock's avatar
Pierre Stock committed
107
108
109
110


class MixtralAttention(nn.Module):

111
112
    def __init__(
        self,
113
        config: MixtralConfig,
114
115
116
117
118
119
120
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
        rope_theta: float = 10000,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
121
        prefix: str = "",
122
    ) -> None:
Pierre Stock's avatar
Pierre Stock committed
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
        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)
139
140
141
        # MixtralConfig has an optional head_dim argument
        self.head_dim = getattr(config, "head_dim",
                                self.hidden_size // self.total_num_heads)
Pierre Stock's avatar
Pierre Stock committed
142
143
144
145
146
        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

147
        self.qkv_proj = QKVParallelLinear(
Pierre Stock's avatar
Pierre Stock committed
148
149
150
151
152
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
153
            quant_config=quant_config,
154
            prefix=f"{prefix}.qkv_proj",
Pierre Stock's avatar
Pierre Stock committed
155
        )
156
        self.o_proj = RowParallelLinear(
Pierre Stock's avatar
Pierre Stock committed
157
158
159
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
160
            quant_config=quant_config,
161
            prefix=f"{prefix}.o_proj",
Pierre Stock's avatar
Pierre Stock committed
162
163
164
165
166
167
        )
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=int(self.rope_theta),
168
            is_neox_style=True,
Pierre Stock's avatar
Pierre Stock committed
169
        )
170
171
172
173
174
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
                              cache_config=cache_config,
175
176
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
Pierre Stock's avatar
Pierre Stock committed
177
178
179
180
181
182

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
183
        qkv, _ = self.qkv_proj(hidden_states)
Pierre Stock's avatar
Pierre Stock committed
184
185
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
186
        attn_output = self.attn(q, k, v)
187
        output, _ = self.o_proj(attn_output)
Pierre Stock's avatar
Pierre Stock committed
188
189
190
191
192
193
194
        return output


class MixtralDecoderLayer(nn.Module):

    def __init__(
        self,
195
        config: MixtralConfig,
196
        cache_config: Optional[CacheConfig] = None,
197
        quant_config: Optional[QuantizationConfig] = None,
198
        prefix: str = "",
Pierre Stock's avatar
Pierre Stock committed
199
200
201
202
203
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        # Requires transformers > 4.32.0
        rope_theta = getattr(config, "rope_theta", 10000)
204
        self.self_attn = MixtralAttention(
205
            config=config,
Pierre Stock's avatar
Pierre Stock committed
206
207
208
209
210
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            max_position=config.max_position_embeddings,
            num_kv_heads=config.num_key_value_heads,
            rope_theta=rope_theta,
211
            cache_config=cache_config,
212
213
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn")
Philipp Moritz's avatar
Philipp Moritz committed
214
215
216
217
        self.block_sparse_moe = MixtralMoE(
            num_experts=config.num_local_experts,
            top_k=config.num_experts_per_tok,
            hidden_size=config.hidden_size,
218
            intermediate_size=config.intermediate_size,
219
220
            quant_config=quant_config,
            prefix=f"{prefix}.block_sparse_moe")
221
222
223
224
        self.input_layernorm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                                eps=config.rms_norm_eps)
Pierre Stock's avatar
Pierre Stock committed
225
226
227
228

    def forward(
        self,
        positions: torch.Tensor,
229
230
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
Pierre Stock's avatar
Pierre Stock committed
231
    ) -> torch.Tensor:
232
233
234
235
236
237
238
239
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
        hidden_states = self.self_attn(
Pierre Stock's avatar
Pierre Stock committed
240
            positions=positions,
241
            hidden_states=hidden_states,
Pierre Stock's avatar
Pierre Stock committed
242
243
        )

244
245
246
247
248
        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
        hidden_states = self.block_sparse_moe(hidden_states)
        return hidden_states, residual
Pierre Stock's avatar
Pierre Stock committed
249

250

251
@support_torch_compile
252
class MixtralModel(nn.Module):
Pierre Stock's avatar
Pierre Stock committed
253

254
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Pierre Stock's avatar
Pierre Stock committed
255
        super().__init__()
256
257
258
259
260
261

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

262
263
        self.config = config
        self.quant_config = quant_config
264
265
266
267
        lora_vocab = (lora_config.lora_extra_vocab_size *
                      (lora_config.max_loras or 1)) if lora_config else 0
        self.vocab_size = config.vocab_size + lora_vocab
        self.org_vocab_size = config.vocab_size
268
269

        self.embed_tokens = VocabParallelEmbedding(
270
            self.vocab_size,
Pierre Stock's avatar
Pierre Stock committed
271
            config.hidden_size,
272
            org_num_embeddings=config.vocab_size,
Pierre Stock's avatar
Pierre Stock committed
273
        )
274
275

        self.start_layer, self.end_layer, self.layers = make_layers(
276
277
278
279
280
            config.num_hidden_layers,
            lambda prefix: MixtralDecoderLayer(
                config, cache_config, quant_config=quant_config, prefix=prefix
            ),
            prefix=f"{prefix}.layers")
281

282
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
283
284
285
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
Pierre Stock's avatar
Pierre Stock committed
286

287
288
289
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

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

317
318
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
Pierre Stock's avatar
Pierre Stock committed
319
320
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
321
322
323
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
Pierre Stock's avatar
Pierre Stock committed
324
        ]
325

326
327
328
329
330
331
332
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
        expert_params_mapping = FusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="w1",
            ckpt_down_proj_name="w2",
            ckpt_up_proj_name="w3",
            num_experts=self.config.num_local_experts)
Philipp Moritz's avatar
Philipp Moritz committed
333

Pierre Stock's avatar
Pierre Stock committed
334
        params_dict = dict(self.named_parameters())
335
        loaded_params: Set[str] = set()
336
        for name, loaded_weight in weights:
337
338
            if (self.quant_config is not None and
                (scale_name := self.quant_config.get_cache_scale(name))):
339
                # Loading kv cache quantization scales
340
341
342
343
344
345
346
347
348
                param = params_dict[scale_name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
                                 loaded_weight[0])
                weight_loader(param, loaded_weight)
                loaded_params.add(scale_name)
                continue

Pierre Stock's avatar
Pierre Stock committed
349
350
351
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
352
353
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
354
355
                if ((name.endswith(".bias") or name.endswith("_bias"))
                        and name not in params_dict):
CHU Tianxiang's avatar
CHU Tianxiang committed
356
                    continue
357
358
359
                # Skip layers on other devices.
                if is_pp_missing_parameter(name, self):
                    continue
360
361
362
363
364
                if name.endswith("scale"):
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue
CHU Tianxiang's avatar
CHU Tianxiang committed
365
                param = params_dict[name]
Pierre Stock's avatar
Pierre Stock committed
366
367
368
369
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
370
371
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
Philipp Moritz's avatar
Philipp Moritz committed
372
373
374
                    if weight_name not in name:
                        continue
                    name = name.replace(weight_name, param_name)
375
376
377
                    # Skip layers on other devices.
                    if is_pp_missing_parameter(name, self):
                        continue
378
379
380
                    if ((name.endswith(".bias") or name.endswith("_bias"))
                            and name not in params_dict):
                        continue
Philipp Moritz's avatar
Philipp Moritz committed
381
382
383
384
                    param = params_dict[name]
                    weight_loader = param.weight_loader
                    weight_loader(param,
                                  loaded_weight,
385
                                  name,
386
                                  shard_id=shard_id,
Philipp Moritz's avatar
Philipp Moritz committed
387
388
389
390
                                  expert_id=expert_id)
                    break
                else:
                    # Skip loading extra bias for GPTQ models.
391
392
                    if ((name.endswith(".bias") or name.endswith("_bias"))
                            and name not in params_dict):
Philipp Moritz's avatar
Philipp Moritz committed
393
                        continue
394
395
396
                    # Skip layers on other devices.
                    if is_pp_missing_parameter(name, self):
                        continue
397
                    # Remapping the name of FP8 kv-scale.
398
399
400
401
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

Philipp Moritz's avatar
Philipp Moritz committed
402
403
404
405
                    param = params_dict[name]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
406
407
            loaded_params.add(name)
        return loaded_params
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485


class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
    fall_back_to_pt_during_load = False

    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
    }

    # LoRA specific attributes
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }
    embedding_padding_modules = ["lm_head"]

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
        self.config = config
        self.lora_config = lora_config
        self.quant_config = quant_config

        self.model = MixtralModel(vllm_config=vllm_config,
                                  prefix=maybe_prefix(prefix, "model"))
        self.unpadded_vocab_size = config.vocab_size
        if lora_config:
            self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
        self.lm_head = ParallelLMHead(
            self.unpadded_vocab_size,
            config.hidden_size,
            org_num_embeddings=config.vocab_size,
            padding_size=DEFAULT_VOCAB_PADDING_SIZE
            # We need bigger padding if using lora for kernel
            # compatibility
            if not lora_config else lora_config.lora_vocab_padding_size,
            quant_config=quant_config,
        )
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size)
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)

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

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
                                   inputs_embeds)
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        logits = self.logits_processor(self.lm_head, hidden_states,
                                       sampling_metadata)
        return logits

    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
        loader = AutoWeightsLoader(self, skip_prefixes=["rotary_emb.inv_freq"])
        return loader.load_weights(weights)