mixtral.py 24.3 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
        self.qkv_proj = QKVParallelLinear(
Pierre Stock's avatar
Pierre Stock committed
282
283
284
285
286
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
287
            quant_config=quant_config,
Pierre Stock's avatar
Pierre Stock committed
288
        )
289
        self.o_proj = RowParallelLinear(
Pierre Stock's avatar
Pierre Stock committed
290
291
292
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
293
            quant_config=quant_config,
Pierre Stock's avatar
Pierre Stock committed
294
295
296
297
298
299
        )
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=int(self.rope_theta),
300
            is_neox_style=True,
Pierre Stock's avatar
Pierre Stock committed
301
        )
302
303
304
305
306
307
        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
308
309
310
311
312

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
313
314
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Pierre Stock's avatar
Pierre Stock committed
315
    ) -> torch.Tensor:
316
        qkv, _ = self.qkv_proj(hidden_states)
Pierre Stock's avatar
Pierre Stock committed
317
318
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
319
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
320
        output, _ = self.o_proj(attn_output)
Pierre Stock's avatar
Pierre Stock committed
321
322
323
324
325
326
327
        return output


class MixtralDecoderLayer(nn.Module):

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

    def forward(
        self,
        positions: torch.Tensor,
358
        hidden_states: torch.Tensor,
359
360
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
361
        residual: Optional[torch.Tensor],
Pierre Stock's avatar
Pierre Stock committed
362
    ) -> torch.Tensor:
363
364
365
366
367
368
369
370
        # 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
371
            positions=positions,
372
            hidden_states=hidden_states,
Pierre Stock's avatar
Pierre Stock committed
373
            kv_cache=kv_cache,
374
            attn_metadata=attn_metadata,
Pierre Stock's avatar
Pierre Stock committed
375
376
        )

377
378
379
380
381
        # 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
382

383
384

class MixtralModel(nn.Module):
Pierre Stock's avatar
Pierre Stock committed
385
386
387

    def __init__(
        self,
388
        config: MixtralConfig,
389
        cache_config: Optional[CacheConfig] = None,
390
        quant_config: Optional[QuantizationConfig] = None,
391
        lora_config: Optional[LoRAConfig] = None,
Pierre Stock's avatar
Pierre Stock committed
392
393
394
    ) -> None:
        super().__init__()
        self.padding_idx = config.pad_token_id
395
396
397
398
        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
399
400

        self.embed_tokens = VocabParallelEmbedding(
401
            self.vocab_size,
Pierre Stock's avatar
Pierre Stock committed
402
            config.hidden_size,
403
            org_num_embeddings=config.vocab_size,
Pierre Stock's avatar
Pierre Stock committed
404
405
        )
        self.layers = nn.ModuleList([
406
407
408
            MixtralDecoderLayer(config,
                                cache_config,
                                quant_config=quant_config)
Pierre Stock's avatar
Pierre Stock committed
409
410
            for _ in range(config.num_hidden_layers)
        ])
411
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Pierre Stock's avatar
Pierre Stock committed
412
413
414
415
416

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
417
418
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
419
    ) -> torch.Tensor:
420
421
        hidden_states = self.embed_tokens(input_ids)
        residual = None
Pierre Stock's avatar
Pierre Stock committed
422
423
        for i in range(len(self.layers)):
            layer = self.layers[i]
424
            hidden_states, residual = layer(positions, hidden_states,
425
                                            kv_caches[i], attn_metadata,
426
                                            residual)
427
428
429
430
431
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


class MixtralForCausalLM(nn.Module):
432
433
    fall_back_to_pt_during_load = False

Terry's avatar
Terry committed
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
    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"]
454
455
456
457

    def __init__(
        self,
        config: MixtralConfig,
458
        cache_config: Optional[CacheConfig] = None,
459
        quant_config: Optional[QuantizationConfig] = None,
Terry's avatar
Terry committed
460
        lora_config: Optional[LoRAConfig] = None,
461
462
463
    ) -> None:
        super().__init__()
        self.config = config
464
        self.model = MixtralModel(config,
465
                                  cache_config,
466
                                  quant_config,
467
                                  lora_config=lora_config)
Terry's avatar
Terry committed
468
469
470
471
472
473
474
475
476
477
478
479
        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,
        )
480
481
482
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size)
        self.sampler = Sampler()
483
484
485
486
487

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
488
489
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
490
491
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
492
                                   attn_metadata)
Pierre Stock's avatar
Pierre Stock committed
493
494
        return hidden_states

495
496
497
498
499
500
    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
501
502
    def sample(
        self,
503
        logits: Optional[torch.Tensor],
Pierre Stock's avatar
Pierre Stock committed
504
        sampling_metadata: SamplingMetadata,
505
    ) -> Optional[SamplerOutput]:
506
        next_tokens = self.sampler(logits, sampling_metadata)
Pierre Stock's avatar
Pierre Stock committed
507
508
        return next_tokens

509
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
Pierre Stock's avatar
Pierre Stock committed
510
511
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
512
513
514
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
Pierre Stock's avatar
Pierre Stock committed
515
        ]
516

Philipp Moritz's avatar
Philipp Moritz committed
517
        expert_params_mapping = [
518
519
520
521
522
523
524
            # 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"]
        ] + [
525
            # These are the weights for the experts
Philipp Moritz's avatar
Philipp Moritz committed
526
            # (param_name, weight_name, expert_id)
527
            ("w13_weight" if weight_name in ["w1", "w3"] else "w2_weight",
Philipp Moritz's avatar
Philipp Moritz committed
528
529
530
             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"]
531
532
533
        ] + [
            # These are the activation scales for the experts
            # (param_name, weight_name, expert_id)
534
            ("a13_scale" if weight_name in ["w1", "w3"] else "a2_scale",
535
536
537
             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
538
539
        ]

Pierre Stock's avatar
Pierre Stock committed
540
        params_dict = dict(self.named_parameters())
541
        for name, loaded_weight in weights:
Pierre Stock's avatar
Pierre Stock committed
542
543
            if "rotary_emb.inv_freq" in name:
                continue
Philipp Moritz's avatar
Philipp Moritz committed
544

Pierre Stock's avatar
Pierre Stock committed
545
546
547
            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
548
549
550
551
552
                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
553
554
555
556
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
Philipp Moritz's avatar
Philipp Moritz committed
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
                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
572
573
574
575
576
577
578
579
580
581
582
583
584
585
                    # 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
586
587
588
589
                    param = params_dict[name]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
590
591
592
593
594


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