glm4_moe.py 26.4 KB
Newer Older
Yuxuan Zhang's avatar
Yuxuan Zhang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# Copyright 2025 The ZhipuAI Team.
# 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.
24
25
"""Inference-only GLM-4.5, GLM-4.6, GLM-4.7 model
compatible with HuggingFace weights."""
26

Yuxuan Zhang's avatar
Yuxuan Zhang committed
27
28
import typing
from collections.abc import Callable, Iterable
29
from itertools import islice
Yuxuan Zhang's avatar
Yuxuan Zhang committed
30
31
32

import torch
from torch import nn
33
from transformers.models.glm4_moe import Glm4MoeConfig
Yuxuan Zhang's avatar
Yuxuan Zhang committed
34
35
36

from vllm.compilation.decorators import support_torch_compile
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
37
38
39
40
41
from vllm.distributed import (
    get_ep_group,
    get_pp_group,
    get_tensor_model_parallel_world_size,
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
42
43
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import SiluAndMul
44
from vllm.model_executor.layers.attention import Attention
45
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
Yuxuan Zhang's avatar
Yuxuan Zhang committed
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
48
49
50
51
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
52
53
54
55
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
56
57
58
    ParallelLMHead,
    VocabParallelEmbedding,
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
59
from vllm.model_executor.model_loader.weight_utils import (
60
61
62
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
63
64
from vllm.sequence import IntermediateTensors

65
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
66
67
68
69
70
71
72
73
from .utils import (
    AutoWeightsLoader,
    PPMissingLayer,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
74
75
76
77
78
79
80
81
82
83

logger = init_logger(__name__)


class Glm4MoeMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
84
        quant_config: QuantizationConfig | None = None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
85
86
87
88
89
        reduce_results: bool = True,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
90
91
            hidden_size,
            [intermediate_size] * 2,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
92
93
            bias=False,
            quant_config=quant_config,
94
95
96
97
98
99
100
101
102
103
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            reduce_results=reduce_results,
            prefix=f"{prefix}.down_proj",
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
104
        if hidden_act != "silu":
105
106
107
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
108
109
110
111
112
113
114
115
116
117
118
119
        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class Glm4MoE(nn.Module):
    def __init__(
        self,
120
        config: Glm4MoeConfig,
121
        quant_config: QuantizationConfig | None = None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
122
123
124
125
126
127
128
129
        prefix: str = "",
        enable_eplb: bool = False,
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.routed_scaling_factor = config.routed_scaling_factor

        self.ep_group = get_ep_group().device_group
130
        self.ep_rank = get_ep_group().rank_in_group
Yuxuan Zhang's avatar
Yuxuan Zhang committed
131
132
133
134
135
        self.ep_size = self.ep_group.size()
        self.n_routed_experts: int = config.n_routed_experts
        self.n_shared_experts: int = config.n_shared_experts

        if config.hidden_act != "silu":
136
137
138
139
            raise ValueError(
                f"Unsupported activation: {config.hidden_act}. "
                "Only silu is supported for now."
            )
140
141
142
143
144
145
146
147
148
        # NOTE In the transformers implementation, the gate isn't an nn.Linear,
        # so we cannot use ReplicatedLinear here.
        # See: https://github.com/huggingface/transformers/blob/v4.55.1/src/transformers/models/glm4_moe/modeling_glm4_moe.py#L260
        self.gate = nn.Linear(
            config.hidden_size,
            config.n_routed_experts,
            bias=False,
            dtype=torch.float32,
        )
149
        self.gate.e_score_correction_bias = nn.Parameter(
150
151
            torch.empty(config.n_routed_experts, dtype=torch.float32)
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
152
153
154

        # Load balancing settings.
        vllm_config = get_current_vllm_config()
155
        eplb_config = vllm_config.parallel_config.eplb_config
Yuxuan Zhang's avatar
Yuxuan Zhang committed
156
157
        self.enable_eplb = enable_eplb

158
        self.n_redundant_experts = eplb_config.num_redundant_experts
Yuxuan Zhang's avatar
Yuxuan Zhang committed
159
        self.n_logical_experts = self.n_routed_experts
160
        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
Yuxuan Zhang's avatar
Yuxuan Zhang committed
161
162
        self.n_local_physical_experts = self.n_physical_experts // self.ep_size

163
164
165
166
        self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
        self.physical_expert_end = (
            self.physical_expert_start + self.n_local_physical_experts
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
167
168

        if config.n_shared_experts is not None:
169
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
Yuxuan Zhang's avatar
Yuxuan Zhang committed
170
171
172
173
174
            self.shared_experts = Glm4MoeMLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
175
                reduce_results=False,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
176
177
                prefix=f"{prefix}.shared_experts",
            )
178
        else:
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
            self.shared_experts = None

        self.experts = SharedFusedMoE(
            shared_experts=self.shared_experts,
            num_experts=config.n_routed_experts,
            top_k=config.num_experts_per_tok,
            hidden_size=config.hidden_size,
            intermediate_size=config.moe_intermediate_size,
            renormalize=config.norm_topk_prob,
            quant_config=quant_config,
            use_grouped_topk=True,
            num_expert_group=config.n_group,
            topk_group=config.topk_group,
            prefix=f"{prefix}.experts",
            scoring_func="sigmoid",
194
195
            routed_scaling_factor=self.routed_scaling_factor,
            apply_routed_scale_to_output=True,
196
197
198
            e_score_correction_bias=self.gate.e_score_correction_bias,
            enable_eplb=self.enable_eplb,
            num_redundant_experts=self.n_redundant_experts,
199
            router_logits_dtype=torch.float32,
200
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
201
202
203
204
205

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)

206
        # router_logits: (num_tokens, n_experts)
207
        router_logits = self.gate(hidden_states.to(dtype=torch.float32))
208

209
        final_hidden_states = self.experts(
210
211
            hidden_states=hidden_states, router_logits=router_logits
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
212
213
214
215
216
217
        return final_hidden_states.view(num_tokens, hidden_dim)


class Glm4MoeAttention(nn.Module):
    def __init__(
        self,
218
        config: Glm4MoeConfig,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
219
220
221
222
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position_embeddings: int = 131072,
223
        head_dim: int | None = None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
224
225
226
        rms_norm_eps: float = 1e-05,
        qkv_bias: bool = False,
        use_qk_norm: bool = False,
227
228
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
        prefix: str = "",
    ) -> None:
        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 = head_dim or (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.max_position_embeddings = max_position_embeddings
        self.use_qk_norm = use_qk_norm

254
255
256
257
258
259
260
261
262
        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=qkv_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
263

264
265
266
267
268
269
270
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
271

272
        config.rope_parameters.setdefault("partial_rotary_factor", 0.5)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
273
274
275
        self.rotary_emb = get_rope(
            self.head_dim,
            max_position=max_position_embeddings,
276
            rope_parameters=config.rope_parameters,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
        )
        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,
            prefix=f"{prefix}.attn",
        )

        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        if self.use_qk_norm:
300
301
302
303
304
305
            q = self.q_norm(q.reshape(-1, self.num_heads, self.head_dim)).reshape(
                q.shape
            )
            k = self.k_norm(k.reshape(-1, self.num_kv_heads, self.head_dim)).reshape(
                k.shape
            )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
306
307
308
309
310
311
312
313
314
315

        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output


class Glm4MoeDecoderLayer(nn.Module):
    def __init__(
        self,
316
        config: Glm4MoeConfig,
317
318
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
319
320
321
322
323
        prefix: str = "",
        enable_eplb: bool = False,
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
324
        max_position_embeddings = getattr(config, "max_position_embeddings", 131072)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
325
326
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
327
        layer_idx = int(prefix.split(sep=".")[-1])
Yuxuan Zhang's avatar
Yuxuan Zhang committed
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
        self.layer_idx = layer_idx

        self.self_attn = Glm4MoeAttention(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            max_position_embeddings=max_position_embeddings,
            head_dim=config.head_dim,
            rms_norm_eps=config.rms_norm_eps,
            qkv_bias=config.attention_bias,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
            use_qk_norm=config.use_qk_norm,
        )

345
346
347
348
        if (
            config.n_routed_experts is not None
            and layer_idx >= config.first_k_dense_replace
        ):
Yuxuan Zhang's avatar
Yuxuan Zhang committed
349
350
351
352
353
354
355
            self.mlp = Glm4MoE(
                config=config,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
                enable_eplb=enable_eplb,
            )
        else:
356
357
358
359
360
361
362
363
364
365
366
367
            self.mlp = Glm4MoeMLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )

        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
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
368
369
370
371
372
373
        self.routed_scaling_factor = config.routed_scaling_factor

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
374
        residual: torch.Tensor | None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
375
376
377
378
379
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
380
381
382
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
383
384
385
386
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


387
388
389
390
391
392
@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        "positions": -1,
        "intermediate_tensors": 0,
        "inputs_embeds": 0,
393
394
    }
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
395
396
397
398
399
400
401
402
403
404
405
406
407
408
class Glm4MoeModel(nn.Module):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()

        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        enable_eplb = vllm_config.parallel_config.enable_eplb
        self.config = config

        self.vocab_size = config.vocab_size

        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
409
410
                config.vocab_size, config.hidden_size, prefix=f"{prefix}.embed_tokens"
            )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
411
412
413
414
415
416
417
418
419
420
421
422
        else:
            self.embed_tokens = PPMissingLayer()

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: Glm4MoeDecoderLayer(
                config=config,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=prefix,
                enable_eplb=enable_eplb,
            ),
423
424
            prefix=f"{prefix}.layers",
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
425
426
427
428
429

        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
430
431
432
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
433

434
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
435
436
437
438
        return self.embed_tokens(input_ids)

    def forward(
        self,
439
        input_ids: torch.Tensor | None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
440
        positions: torch.Tensor,
441
442
443
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
444
445
446
447
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
448
                hidden_states = self.embed_input_ids(input_ids)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
449
450
451
452
453
454
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

455
        for layer in islice(self.layers, self.start_layer, self.end_layer):
Yuxuan Zhang's avatar
Yuxuan Zhang committed
456
457
458
            hidden_states, residual = layer(positions, hidden_states, residual)

        if not get_pp_group().is_last_rank:
459
460
461
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
462
463
464
465

        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

466
467
468
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
469
        return SharedFusedMoE.make_expert_params_mapping(
470
            self,
471
472
473
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
474
475
            num_experts=self.config.n_routed_experts,
        )
476

477
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
478
479
480
481
482
483
484
485
486
487
488
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
489
        expert_params_mapping = self.get_expert_mapping()
Yuxuan Zhang's avatar
Yuxuan Zhang committed
490
491
492
493
        for name, loaded_weight in weights:
            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is not None:
                continue
494
            for param_name, weight_name, shard_id in stacked_params_mapping:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
495
496
497
498
499
500
501
502
503
                # Skip non-stacked layers and experts (experts handled below).
                if weight_name not in name:
                    continue
                # We have mlp.experts[0].gate_proj in the checkpoint.
                # Since we handle the experts below in expert_params_mapping,
                # we need to skip here BEFORE we update the name, otherwise
                # name will be updated to mlp.experts[0].gate_up_proj, which
                # will then be updated below in expert_params_mapping
                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
504
                if ("mlp.experts." in name) and name not in params_dict:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
                    continue
                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
                if is_pp_missing_parameter(name, self):
                    continue

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                is_expert_weight = False
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue

                    # Anyway, this is an expert weight and should not be
                    # attempted to load as other weights later
                    is_expert_weight = True

                    # Do not modify `name` since the loop may continue here
                    # Instead, create a new variable
                    name_mapped = name.replace(weight_name, param_name)

                    if is_pp_missing_parameter(name_mapped, self):
                        continue

                    param = params_dict[name_mapped]
                    # We should ask the weight loader to return success or not
                    # here since otherwise we may skip experts with other
                    # available replicas.
539
540
541
542
543
544
545
546
547
548
549
                    weight_loader = typing.cast(
                        Callable[..., bool], param.weight_loader
                    )
                    success = weight_loader(
                        param,
                        loaded_weight,
                        name_mapped,
                        shard_id=shard_id,
                        expert_id=expert_id,
                        return_success=True,
                    )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
                    if success:
                        name = name_mapped
                        break
                else:
                    if is_expert_weight:
                        # We've checked that this is an expert weight
                        # However it's not mapped locally to this rank
                        # So we simply skip it
                        continue

                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue

                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

                    if is_pp_missing_parameter(name, self):
                        continue

                    param = params_dict[name]
573
574
575
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
576
577
578
579
580
581
                    weight_loader(param, loaded_weight)
            loaded_params.add(name)

        return loaded_params


582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
class Glm4MixtureOfExperts(MixtureOfExperts):
    def extract_moe_parameters(self, example_moe: Glm4MoE | None) -> None:
        if example_moe is None:
            raise RuntimeError("No Glm4MoE layer found in model.layers.")
        else:
            self.num_logical_experts = example_moe.n_logical_experts
            self.num_physical_experts = example_moe.n_physical_experts
            self.num_local_physical_experts = example_moe.n_local_physical_experts
            self.num_routed_experts = example_moe.n_routed_experts
            self.num_shared_experts = example_moe.n_shared_experts
            self.num_redundant_experts = example_moe.n_redundant_experts

    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None:
        assert self.num_local_physical_experts == num_local_physical_experts
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
        for moe in self.moe_mlp_layers:
            moe.n_local_physical_experts = num_local_physical_experts
            moe.n_physical_experts = num_physical_experts
            moe.n_redundant_experts = self.num_redundant_experts
            moe.experts.update_expert_map()


class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, Glm4MixtureOfExperts):
Yuxuan Zhang's avatar
Yuxuan Zhang committed
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    fall_back_to_pt_during_load = False

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.config = config
        self.quant_config = quant_config
631
632
633
        self.model = Glm4MoeModel(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
634
        if get_pp_group().is_last_rank:
635
636
637
638
639
640
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
641
642
643
644
        else:
            self.lm_head = PPMissingLayer()
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.make_empty_intermediate_tensors = (
645
646
            self.model.make_empty_intermediate_tensors
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
647
648
649
        self.expert_weights = []

        # Set MoE hyperparameters
650
        self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
Yuxuan Zhang's avatar
Yuxuan Zhang committed
651
652
        self.num_expert_groups = config.n_group

653
654
655
        self.moe_layers = []
        self.moe_mlp_layers: list[Glm4MoE] = []

656
        example_moe = None
Yuxuan Zhang's avatar
Yuxuan Zhang committed
657
        for layer in self.model.layers:
658
659
660
            if isinstance(layer, PPMissingLayer):
                continue

Yuxuan Zhang's avatar
Yuxuan Zhang committed
661
662
            assert isinstance(layer, Glm4MoeDecoderLayer)
            if isinstance(layer.mlp, Glm4MoE):
663
664
                # Pick last one layer since the first ones may be dense layers.
                example_moe = layer.mlp
665
                self.moe_mlp_layers.append(layer.mlp)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
666
667
                self.moe_layers.append(layer.mlp.experts)

668
        self.extract_moe_parameters(example_moe)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
669

670
671
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
672
673
674

    def forward(
        self,
675
        input_ids: torch.Tensor | None,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
676
        positions: torch.Tensor,
677
678
679
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
680
681
682
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
683
684
685
686
687
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
688
    ) -> torch.Tensor | None:
689
        logits = self.logits_processor(self.lm_head, hidden_states)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
690
691
        return logits

692
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
693
694
695
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)

696
697
698
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        return self.model.get_expert_mapping()

Yuxuan Zhang's avatar
Yuxuan Zhang committed
699

700
701
def get_spec_layer_idx_from_weight_name(
    config: Glm4MoeConfig, weight_name: str
702
) -> int | None:
703
704
705
    if hasattr(config, "num_nextn_predict_layers") and (
        config.num_nextn_predict_layers > 0
    ):
Yuxuan Zhang's avatar
Yuxuan Zhang committed
706
707
        layer_idx = config.num_hidden_layers
        for i in range(config.num_nextn_predict_layers):
708
            if f"layers.{layer_idx + i}." in weight_name:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
709
710
                return layer_idx + i
    return None