mixtral.py 24.6 KB
Newer Older
Pierre Stock's avatar
Pierre Stock committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# coding=utf-8
# 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."""
24
from typing import Iterable, List, Optional, Tuple
Pierre Stock's avatar
Pierre Stock committed
25
26
27

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

30
from vllm import _custom_ops as ops
31
from vllm.attention import Attention, AttentionMetadata
32
from vllm.config import CacheConfig, LoRAConfig
33
34
35
from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size,
                              tensor_model_parallel_all_reduce)
Philipp Moritz's avatar
Philipp Moritz committed
36
from vllm.model_executor.layers.fused_moe import fused_moe
Pierre Stock's avatar
Pierre Stock committed
37
from vllm.model_executor.layers.layernorm import RMSNorm
38
from vllm.model_executor.layers.linear import (QKVParallelLinear,
Philipp Moritz's avatar
Philipp Moritz committed
39
                                               ReplicatedLinear,
Pierre Stock's avatar
Pierre Stock committed
40
                                               RowParallelLinear)
41
from vllm.model_executor.layers.logits_processor import LogitsProcessor
42
43
44
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
45
from vllm.model_executor.layers.rotary_embedding import get_rope
Pierre Stock's avatar
Pierre Stock committed
46
47
from vllm.model_executor.layers.sampler import Sampler
from vllm.model_executor.layers.vocab_parallel_embedding import (
48
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
49
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
Pierre Stock's avatar
Pierre Stock committed
50
from vllm.model_executor.sampling_metadata import SamplingMetadata
Philipp Moritz's avatar
Philipp Moritz committed
51
from vllm.model_executor.utils import set_weight_attrs
Pierre Stock's avatar
Pierre Stock committed
52
from vllm.sequence import SamplerOutput
53
from vllm.utils import print_warning_once
Pierre Stock's avatar
Pierre Stock committed
54
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

    def __init__(
        self,
        num_experts: int,
Philipp Moritz's avatar
Philipp Moritz committed
68
        top_k: int,
69
70
        hidden_size: int,
        intermediate_size: int,
Philipp Moritz's avatar
Philipp Moritz committed
71
        params_dtype: Optional[torch.dtype] = None,
72
        tp_size: Optional[int] = None,
73
        quant_config: Optional[QuantizationConfig] = None,
Philipp Moritz's avatar
Philipp Moritz committed
74
    ):
75
        super().__init__()
76
        self.tp_size = tp_size or get_tensor_model_parallel_world_size()
Philipp Moritz's avatar
Philipp Moritz committed
77
78
79
        self.num_total_experts = num_experts
        self.top_k = top_k
        self.hidden_size = hidden_size
80
        self.intermediate_size = intermediate_size // self.tp_size
81
82
        self.quant_config = quant_config

83
84
        # FIXME(pcmoritz): Make this more general to support different
        # quantization schemes
85
        self.use_fp8 = isinstance(quant_config, Fp8Config)
86

Philipp Moritz's avatar
Philipp Moritz committed
87
88
89
        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        self.params_dtype = params_dtype
90

91
        # Gate always runs at half / full precision for now.
Philipp Moritz's avatar
Philipp Moritz committed
92
        self.gate = ReplicatedLinear(self.hidden_size,
93
94
                                     self.num_total_experts,
                                     bias=False,
Philipp Moritz's avatar
Philipp Moritz committed
95
                                     params_dtype=self.params_dtype,
96
                                     quant_config=None)
97

98
        if self.use_fp8 and self.quant_config.is_checkpoint_fp8_serialized:
99
100
101
            params_dtype = torch.float8_e4m3fn

        self.w13_weight = nn.Parameter(
Philipp Moritz's avatar
Philipp Moritz committed
102
103
104
            torch.empty(self.num_total_experts,
                        2 * self.intermediate_size,
                        self.hidden_size,
105
106
                        dtype=params_dtype))
        self.w2_weight = nn.Parameter(
Philipp Moritz's avatar
Philipp Moritz committed
107
108
109
            torch.empty(self.num_total_experts,
                        self.hidden_size,
                        self.intermediate_size,
110
                        dtype=params_dtype))
Philipp Moritz's avatar
Philipp Moritz committed
111

112
        set_weight_attrs(self.w13_weight, {
113
114
            "weight_loader": self.weight_loader,
        })
115
        set_weight_attrs(self.w2_weight, {
116
117
118
            "weight_loader": self.weight_loader,
        })

119
120
121
122
123
124
125
126
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
155
156
157
158
159
160
161
162
163
        # Used for fp8.
        self.w13_scale = None
        self.w2_scale = None
        self.a13_scale = None
        self.a2_scale = None

        if self.use_fp8:
            # WEIGHT_SCALE (for fp8)
            self.w13_scale = nn.Parameter(torch.ones(self.num_total_experts,
                                                     dtype=torch.float32),
                                          requires_grad=False)
            self.w2_scale = nn.Parameter(torch.ones(self.num_total_experts,
                                                    dtype=torch.float32),
                                         requires_grad=False)

            # If loading fp8 checkpoint, pass the weight loaders.
            # If loading an fp16 checkpoint, do not (we will quantize in
            #   process_weights_after_loading()
            if quant_config.is_checkpoint_fp8_serialized:
                set_weight_attrs(self.w13_scale, {
                    "weight_loader": self.weight_loader,
                })
                set_weight_attrs(self.w2_scale, {
                    "weight_loader": self.weight_loader,
                })

            # ACT_SCALE (for fp8)
            if quant_config.activation_scheme == "static":
                if not quant_config.is_checkpoint_fp8_serialized:
                    raise ValueError(
                        "Found static activation scheme for checkpoint that "
                        "was not serialized fp8.")
                self.a13_scale = nn.Parameter(torch.zeros(
                    self.num_total_experts, dtype=torch.float32),
                                              requires_grad=False)
                self.a2_scale = nn.Parameter(torch.zeros(
                    self.num_total_experts, dtype=torch.float32),
                                             requires_grad=False)

                set_weight_attrs(self.a13_scale, {
                    "weight_loader": self.weight_loader,
                })
                set_weight_attrs(self.a2_scale, {
                    "weight_loader": self.weight_loader,
                })
Philipp Moritz's avatar
Philipp Moritz committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor,
                      weight_name: str, expert_id: int):
        tp_rank = get_tensor_model_parallel_rank()
        param_data = param.data
        shard_size = self.intermediate_size
        shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
        if weight_name.endswith("w1.weight"):
            param_data[expert_id, 0:shard_size, :] = loaded_weight[shard, :]
        if weight_name.endswith("w3.weight"):
            param_data[expert_id,
                       shard_size:2 * shard_size, :] = loaded_weight[shard, :]
        if weight_name.endswith("w2.weight"):
            param_data[expert_id, :, :] = loaded_weight[:, shard]
178
179
        if "act_scale" in weight_name or "weight_scale" in weight_name:
            param_data[expert_id] = loaded_weight
Philipp Moritz's avatar
Philipp Moritz committed
180

181
    def process_weights_after_loading(self):
182
183
184
185
186
187
188
189
190
191
        # Fp8 is the only case where we need to process after loading.
        if not self.use_fp8:
            return

        # If checkpoint is fp16, quantize here.
        if not self.quant_config.is_checkpoint_fp8_serialized:
            w13_weight = torch.empty_like(self.w13_weight.data,
                                          dtype=torch.float8_e4m3fn)
            w2_weight = torch.empty_like(self.w2_weight.data,
                                         dtype=torch.float8_e4m3fn)
192
            for expert in range(self.num_total_experts):
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
                w13_weight[expert, :, :], self.w13_scale[
                    expert] = ops.scaled_fp8_quant(
                        self.w13_weight.data[expert, :, :])
                w2_weight[expert, :, :], self.w2_scale[
                    expert] = ops.scaled_fp8_quant(
                        self.w2_weight.data[expert, :, :])
            self.w13_weight = nn.Parameter(w13_weight, requires_grad=False)
            self.w2_weight = nn.Parameter(w2_weight, requires_grad=False)

        # If checkpoint is fp8 + static, cleanup act_scales.
        #   Since state_dict has an act_scale per expert but our kernels
        #   are passed one act_scale shared across all experts.
        elif self.quant_config.activation_scheme == "static":
            if self.a13_scale is None or self.a2_scale is None:
                raise ValueError(
                    "QuantConfig has static quantization, but found "
                    "activation scales are None.")

            if (not all_close_1d(self.a13_scale)
                    or not all_close_1d(self.a2_scale)):
                print_warning_once(
                    "Found act_scales that are not equal for fp8 MoE layer. "
                    "Using the maximum across experts for each layer. ")

            self.a13_scale = nn.Parameter(self.a13_scale.max(),
                                          requires_grad=False)
            self.a2_scale = nn.Parameter(self.a2_scale.max(),
                                         requires_grad=False)
221

222
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
223
        num_tokens, hidden_size = hidden_states.shape
Philipp Moritz's avatar
Philipp Moritz committed
224
        hidden_states = hidden_states.view(-1, self.hidden_size)
225
        # router_logits: (num_tokens, n_experts)
226
        router_logits, _ = self.gate(hidden_states)
Philipp Moritz's avatar
Philipp Moritz committed
227
        final_hidden_states = fused_moe(hidden_states,
228
229
                                        self.w13_weight,
                                        self.w2_weight,
230
231
232
                                        router_logits,
                                        self.top_k,
                                        renormalize=True,
233
234
                                        inplace=True,
                                        use_fp8=self.use_fp8,
235
236
237
238
                                        w1_scale=self.w13_scale,
                                        w2_scale=self.w2_scale,
                                        a1_scale=self.a13_scale,
                                        a2_scale=self.a2_scale)
Philipp Moritz's avatar
Philipp Moritz committed
239

240
241
242
        if self.tp_size > 1:
            final_hidden_states = tensor_model_parallel_all_reduce(
                final_hidden_states)
243

244
        return final_hidden_states.view(num_tokens, hidden_size)
Pierre Stock's avatar
Pierre Stock committed
245
246
247
248


class MixtralAttention(nn.Module):

249
250
251
252
253
254
255
256
257
258
    def __init__(
        self,
        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,
    ) -> None:
Pierre Stock's avatar
Pierre Stock committed
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
        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 = hidden_size // self.total_num_heads
        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

281
282
283
        if isinstance(
                quant_config,
                Fp8Config) and not quant_config.is_checkpoint_fp8_serialized:
284
285
286
287
            print_warning_once(
                "For Mixtral FP8 quantization, we currently do not quantize "
                "the attention layers until their FP8 performance is improved."
            )
288
            quant_config = None
289

290
        self.qkv_proj = QKVParallelLinear(
Pierre Stock's avatar
Pierre Stock committed
291
292
293
294
295
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
296
            quant_config=quant_config,
Pierre Stock's avatar
Pierre Stock committed
297
        )
298
        self.o_proj = RowParallelLinear(
Pierre Stock's avatar
Pierre Stock committed
299
300
301
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
302
            quant_config=quant_config,
Pierre Stock's avatar
Pierre Stock committed
303
304
305
306
307
308
        )
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=int(self.rope_theta),
309
            is_neox_style=True,
Pierre Stock's avatar
Pierre Stock committed
310
        )
311
312
313
314
315
316
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
                              cache_config=cache_config,
                              quant_config=quant_config)
Pierre Stock's avatar
Pierre Stock committed
317
318
319
320
321

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
322
323
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Pierre Stock's avatar
Pierre Stock committed
324
    ) -> torch.Tensor:
325
        qkv, _ = self.qkv_proj(hidden_states)
Pierre Stock's avatar
Pierre Stock committed
326
327
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
328
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
329
        output, _ = self.o_proj(attn_output)
Pierre Stock's avatar
Pierre Stock committed
330
331
332
333
334
335
336
        return output


class MixtralDecoderLayer(nn.Module):

    def __init__(
        self,
337
        config: MixtralConfig,
338
        cache_config: Optional[CacheConfig] = None,
339
        quant_config: Optional[QuantizationConfig] = None,
Pierre Stock's avatar
Pierre Stock committed
340
341
342
343
344
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        # Requires transformers > 4.32.0
        rope_theta = getattr(config, "rope_theta", 10000)
345
        self.self_attn = MixtralAttention(
Pierre Stock's avatar
Pierre Stock committed
346
347
348
349
350
            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,
351
            cache_config=cache_config,
352
            quant_config=quant_config)
Philipp Moritz's avatar
Philipp Moritz committed
353
354
355
356
        self.block_sparse_moe = MixtralMoE(
            num_experts=config.num_local_experts,
            top_k=config.num_experts_per_tok,
            hidden_size=config.hidden_size,
357
            intermediate_size=config.intermediate_size,
358
            quant_config=quant_config)
359
360
361
362
        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
363
364
365
366

    def forward(
        self,
        positions: torch.Tensor,
367
        hidden_states: torch.Tensor,
368
369
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
370
        residual: Optional[torch.Tensor],
Pierre Stock's avatar
Pierre Stock committed
371
    ) -> torch.Tensor:
372
373
374
375
376
377
378
379
        # 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
380
            positions=positions,
381
            hidden_states=hidden_states,
Pierre Stock's avatar
Pierre Stock committed
382
            kv_cache=kv_cache,
383
            attn_metadata=attn_metadata,
Pierre Stock's avatar
Pierre Stock committed
384
385
        )

386
387
388
389
390
        # 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
391

392
393

class MixtralModel(nn.Module):
Pierre Stock's avatar
Pierre Stock committed
394
395
396

    def __init__(
        self,
397
        config: MixtralConfig,
398
        cache_config: Optional[CacheConfig] = None,
399
        quant_config: Optional[QuantizationConfig] = None,
400
        lora_config: Optional[LoRAConfig] = None,
Pierre Stock's avatar
Pierre Stock committed
401
402
403
    ) -> None:
        super().__init__()
        self.padding_idx = config.pad_token_id
404
405
406
407
        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
408
409

        self.embed_tokens = VocabParallelEmbedding(
410
            self.vocab_size,
Pierre Stock's avatar
Pierre Stock committed
411
            config.hidden_size,
412
            org_num_embeddings=config.vocab_size,
Pierre Stock's avatar
Pierre Stock committed
413
414
        )
        self.layers = nn.ModuleList([
415
416
417
            MixtralDecoderLayer(config,
                                cache_config,
                                quant_config=quant_config)
Pierre Stock's avatar
Pierre Stock committed
418
419
            for _ in range(config.num_hidden_layers)
        ])
420
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Pierre Stock's avatar
Pierre Stock committed
421
422
423
424
425

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
426
427
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
428
    ) -> torch.Tensor:
429
430
        hidden_states = self.embed_tokens(input_ids)
        residual = None
Pierre Stock's avatar
Pierre Stock committed
431
432
        for i in range(len(self.layers)):
            layer = self.layers[i]
433
            hidden_states, residual = layer(positions, hidden_states,
434
                                            kv_caches[i], attn_metadata,
435
                                            residual)
436
437
438
439
440
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


class MixtralForCausalLM(nn.Module):
441
442
    fall_back_to_pt_during_load = False

Terry's avatar
Terry committed
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
        "qkv_proj",
        "o_proj",
        "embed_tokens",
        "lm_head",
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }
    embedding_padding_modules = ["lm_head"]
463
464
465
466

    def __init__(
        self,
        config: MixtralConfig,
467
        cache_config: Optional[CacheConfig] = None,
468
        quant_config: Optional[QuantizationConfig] = None,
Terry's avatar
Terry committed
469
        lora_config: Optional[LoRAConfig] = None,
470
471
472
    ) -> None:
        super().__init__()
        self.config = config
473
        self.model = MixtralModel(config,
474
                                  cache_config,
475
                                  quant_config,
476
                                  lora_config=lora_config)
Terry's avatar
Terry committed
477
478
479
480
481
482
483
484
485
486
487
488
        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,
        )
489
490
491
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size)
        self.sampler = Sampler()
492
493
494
495
496

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
497
498
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
499
500
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
501
                                   attn_metadata)
Pierre Stock's avatar
Pierre Stock committed
502
503
        return hidden_states

504
505
506
507
508
509
    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

Pierre Stock's avatar
Pierre Stock committed
510
511
    def sample(
        self,
512
        logits: Optional[torch.Tensor],
Pierre Stock's avatar
Pierre Stock committed
513
        sampling_metadata: SamplingMetadata,
514
    ) -> Optional[SamplerOutput]:
515
        next_tokens = self.sampler(logits, sampling_metadata)
Pierre Stock's avatar
Pierre Stock committed
516
517
        return next_tokens

518
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
Pierre Stock's avatar
Pierre Stock committed
519
520
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
521
522
523
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
Pierre Stock's avatar
Pierre Stock committed
524
        ]
525

Philipp Moritz's avatar
Philipp Moritz committed
526
        expert_params_mapping = [
527
528
529
530
531
532
533
            # These are the weight scales for the experts
            # (param_name, weight_name, expert_id)
            ("w13_scale" if weight_name in ["w1", "w3"] else "w2_scale",
             f"experts.{expert_id}.{weight_name}.weight_scale", expert_id)
            for expert_id in range(self.config.num_local_experts)
            for weight_name in ["w1", "w2", "w3"]
        ] + [
534
            # These are the weights for the experts
Philipp Moritz's avatar
Philipp Moritz committed
535
            # (param_name, weight_name, expert_id)
536
            ("w13_weight" if weight_name in ["w1", "w3"] else "w2_weight",
Philipp Moritz's avatar
Philipp Moritz committed
537
538
539
             f"experts.{expert_id}.{weight_name}.weight", expert_id)
            for expert_id in range(self.config.num_local_experts)
            for weight_name in ["w1", "w2", "w3"]
540
541
542
        ] + [
            # These are the activation scales for the experts
            # (param_name, weight_name, expert_id)
543
            ("a13_scale" if weight_name in ["w1", "w3"] else "a2_scale",
544
545
546
             f"experts.{expert_id}.{weight_name}.act_scale", expert_id)
            for expert_id in range(self.config.num_local_experts)
            for weight_name in ["w1", "w2", "w3"]
Philipp Moritz's avatar
Philipp Moritz committed
547
548
        ]

Pierre Stock's avatar
Pierre Stock committed
549
        params_dict = dict(self.named_parameters())
550
        for name, loaded_weight in weights:
Pierre Stock's avatar
Pierre Stock committed
551
552
            if "rotary_emb.inv_freq" in name:
                continue
Philipp Moritz's avatar
Philipp Moritz committed
553

Pierre Stock's avatar
Pierre Stock committed
554
555
556
            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
557
558
559
560
561
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                param = params_dict[name]
Pierre Stock's avatar
Pierre Stock committed
562
563
564
565
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
Philipp Moritz's avatar
Philipp Moritz committed
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
                for param_name, weight_name, expert_id in expert_params_mapping:
                    if weight_name not in name:
                        continue
                    name = name.replace(weight_name, param_name)
                    param = params_dict[name]
                    weight_loader = param.weight_loader
                    weight_loader(param,
                                  loaded_weight,
                                  weight_name,
                                  expert_id=expert_id)
                    break
                else:
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue
581
582
583
584
585
586
587
588
589
590
591
592
593
594
                    # Remapping the name of FP8 kv-scale.
                    if name.endswith("kv_scale"):
                        remapped_kv_scale_name = name.replace(
                            ".kv_scale", ".attn.kv_scale")
                        if remapped_kv_scale_name not in params_dict:
                            print_warning_once(
                                "Found kv scale in the checkpoint "
                                f"(e.g. {name}), but not found the expected "
                                f"name in the model "
                                f"(e.g. {remapped_kv_scale_name}). "
                                "kv-scale is not loaded.")
                            continue
                        else:
                            name = remapped_kv_scale_name
Philipp Moritz's avatar
Philipp Moritz committed
595
596
597
598
                    param = params_dict[name]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
599
600
601
602
603


def all_close_1d(x: torch.Tensor) -> bool:
    assert len(x.shape) == 1
    return all(torch.allclose(x[0], x[i]) for i in range(x.shape[0]))