"vscode:/vscode.git/clone" did not exist on "9ddd0f975cb1470a7cfed025b51ed60199420f5a"
grok1.py 28.5 KB
Newer Older
Michael Goin's avatar
Michael Goin committed
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
Michael Goin's avatar
Michael Goin committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# Adapted from
# https://github.com/ROCm/vllm/blob/cea7419f151cc50293a05b7fac8547f8f887c9f6/vllm/model_executor/models/grok1.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.
Bijaya Dangol's avatar
Bijaya Dangol committed
24
"""Inference-only Grok (Grok1/Grok2) model."""
25

Bijaya Dangol's avatar
Bijaya Dangol committed
26
import math
27
from collections.abc import Iterable
28
from itertools import islice
29
from typing import Any
Michael Goin's avatar
Michael Goin committed
30
31
32
33
34
35
36
37

import torch
import torch.nn.functional as F
from torch import nn

from vllm.compilation.decorators import support_torch_compile
from vllm.config import CacheConfig, VllmConfig
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
Bijaya Dangol's avatar
Bijaya Dangol committed
38
39
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import GeluAndMul
40
from vllm.model_executor.layers.attention import Attention
Michael Goin's avatar
Michael Goin committed
41
42
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.layernorm import RMSNorm
43
from vllm.model_executor.layers.linear import (
Bijaya Dangol's avatar
Bijaya Dangol committed
44
    MergedColumnParallelLinear,
45
46
47
48
    QKVParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
Michael Goin's avatar
Michael Goin committed
49
50
51
52
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 (
53
54
55
    ParallelLMHead,
    VocabParallelEmbedding,
)
Michael Goin's avatar
Michael Goin committed
56
from vllm.model_executor.model_loader.weight_utils import (
57
58
59
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
Michael Goin's avatar
Michael Goin committed
60
61
62
from vllm.sequence import IntermediateTensors

from .interfaces import SupportsLoRA, SupportsPP
63
64
65
66
67
68
69
from .utils import (
    AutoWeightsLoader,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
Michael Goin's avatar
Michael Goin committed
70
71
72
73
74

# Default Grok1-specific constants, overridden by config values if present
DEFAULT_ATTN_OUTPUT_MULTIPLIER = 0.08838834764831845
DEFAULT_OUTPUT_MULTIPLIER_SCALE = 0.5773502691896257
DEFAULT_EMBEDDING_MULTIPLIER_SCALE = 78.38367176906169
Bijaya Dangol's avatar
Bijaya Dangol committed
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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
164
165
166
167
168
DEFAULT_ROUTER_LOGIT_SOFTCAP = 30.0

logger = init_logger(__name__)


def _get_num_experts(config) -> int:
    return getattr(config, "num_experts", getattr(config, "num_local_experts", 8))


def _get_moe_intermediate_size(config) -> int:
    return getattr(config, "moe_intermediate_size", config.intermediate_size)


def _get_grok_version(config) -> str:
    """Detect Grok version from HF config using multiple heuristics."""
    # Check for Grok2-specific attributes (both for robust detection)
    has_residual_moe = getattr(config, "residual_moe", False)
    has_moe_intermediate_size = hasattr(config, "moe_intermediate_size")

    if has_residual_moe or has_moe_intermediate_size:
        return "grok2"

    return "grok1"  # Default to Grok1


def _get_rope_parameters(config) -> dict[str, Any] | None:
    rope_parameters = getattr(config, "rope_parameters", None)
    if rope_parameters is None:
        rope_type = getattr(config, "rope_type", None)
        if rope_type is None:
            return None
        rope_parameters = {"rope_type": rope_type}
        rope_theta = getattr(config, "rope_theta", None)
        if rope_theta is not None:
            rope_parameters["rope_theta"] = rope_theta
        scaling_factor = getattr(config, "scaling_factor", None)
        if scaling_factor is not None:
            rope_parameters["factor"] = scaling_factor
        for name in (
            "original_max_position_embeddings",
            "extrapolation_factor",
            "attn_factor",
            "beta_fast",
            "beta_slow",
        ):
            value = getattr(config, name, None)
            if value is not None:
                rope_parameters[name] = value

    if rope_parameters.get("rope_type") == "original":
        rope_parameters = dict(rope_parameters)
        rope_parameters["rope_type"] = "default"
    return rope_parameters


def _get_moe_renormalize(config) -> bool:
    explicit_value = getattr(
        config, "moe_router_renormalize", getattr(config, "moe_renormalize", None)
    )
    if explicit_value is not None:
        return bool(explicit_value)
    return not getattr(config, "residual_moe", False)


class Grok1MLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
        self.act_fn = GeluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
        x, _ = self.down_proj(x)
        return x
Michael Goin's avatar
Michael Goin committed
169
170
171
172
173
174
175
176
177
178
179


class Grok1MoE(nn.Module):
    """A tensor-parallel MoE implementation for Grok1 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.
    """

180
181
182
183
184
185
    def __init__(
        self,
        num_experts: int,
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
Bijaya Dangol's avatar
Bijaya Dangol committed
186
        router_logit_soft_cap: float = 0.0,
187
188
189
        params_dtype: torch.dtype | None = None,
        quant_config: QuantizationConfig | None = None,
        tp_size: int | None = None,
Bijaya Dangol's avatar
Bijaya Dangol committed
190
        renormalize: bool = False,
191
192
        prefix: str = "",
    ):
Michael Goin's avatar
Michael Goin committed
193
194
195
196
        super().__init__()
        self.hidden_size = hidden_size

        # Gate always runs at half / full precision for now.
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
        self.gate = ReplicatedLinear(
            hidden_size,
            num_experts,
            bias=False,
            params_dtype=params_dtype,
            quant_config=None,
            prefix=f"{prefix}.gate",
        )

        self.experts = FusedMoE(
            num_experts=num_experts,
            top_k=top_k,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            params_dtype=params_dtype,
Bijaya Dangol's avatar
Bijaya Dangol committed
212
            renormalize=renormalize,
213
214
215
216
217
            quant_config=quant_config,
            tp_size=tp_size,
            activation="gelu",
            prefix=f"{prefix}.experts",
        )
Bijaya Dangol's avatar
Bijaya Dangol committed
218
        self.router_logit_soft_cap = router_logit_soft_cap
Michael Goin's avatar
Michael Goin committed
219
220
221
222
223
224
225

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # NOTE: hidden_states can have either 1D or 2D shape.
        orig_shape = hidden_states.shape
        hidden_states = hidden_states.view(-1, self.hidden_size)
        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
Bijaya Dangol's avatar
Bijaya Dangol committed
226
227
228
229
        if self.router_logit_soft_cap > 0:
            router_logits = self.router_logit_soft_cap * F.tanh(
                router_logits / self.router_logit_soft_cap
            )
Michael Goin's avatar
Michael Goin committed
230
231
232
233
234
235
        final_hidden_states = self.experts(hidden_states, router_logits)
        return final_hidden_states.view(orig_shape)


class Grok1Attention(nn.Module):
    def __init__(
236
237
238
239
240
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
241
        rope_parameters: dict[str, Any] | None = None,
242
243
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
244
245
        prefix: str = "",
        config=None,  # Added config parameter
Michael Goin's avatar
Michael Goin committed
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.config = config  # Store config reference
        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.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )
        self.rotary_emb = get_rope(
            self.head_dim,
            max_position=max_position,
288
            rope_parameters=rope_parameters,
Michael Goin's avatar
Michael Goin committed
289
290
291
            is_neox_style=True,
        )

292
        attn_logits_soft_cap = max(getattr(config, "attn_logit_softcapping", 30.0), 0.0)
Bijaya Dangol's avatar
Bijaya Dangol committed
293
294
295
296
297
298
299
300
301
        attn_logit_softcapping_method = getattr(
            config, "attn_logit_softcapping_method", None
        )
        if attn_logit_softcapping_method not in (None, "tanh"):
            logger.warning_once(
                "Grok attention logit softcapping method '%s' is not "
                "supported; falling back to default behavior.",
                attn_logit_softcapping_method,
            )
Michael Goin's avatar
Michael Goin committed
302

303
304
305
306
307
308
309
310
311
312
313
314
315
        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,
            logits_soft_cap=attn_logits_soft_cap,
            prefix=f"{prefix}.attn",
        )
        self.attn_multiplier = (
            getattr(self.config, "attn_output_multiplier", 1.0) if self.config else 1.0
        )
Michael Goin's avatar
Michael Goin committed
316
317
318
319
320
321
322
323
324

    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)
        q, k = self.rotary_emb(positions, q, k)
325
        attn_output = self.attn(q, k, v)
Michael Goin's avatar
Michael Goin committed
326
        output, _ = self.o_proj(attn_output)
327
        output *= self.attn_multiplier
Michael Goin's avatar
Michael Goin committed
328
329
330
331
332
333
334
        return output


class Grok1DecoderLayer(nn.Module):
    def __init__(
        self,
        config,
335
336
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
Michael Goin's avatar
Michael Goin committed
337
338
339
340
341
342
343
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        # Check for fp8 quantization
        self.use_fp8 = False
        if quant_config is not None:
344
            self.use_fp8 = getattr(quant_config, "is_fp8_w8a8", lambda: False)()
Michael Goin's avatar
Michael Goin committed
345
346
347
348
349
350
351
352
            if not self.use_fp8 and hasattr(quant_config, "is_fp8"):
                self.use_fp8 = quant_config.is_fp8

        self.attn = Grok1Attention(
            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,
Bijaya Dangol's avatar
Bijaya Dangol committed
353
            rope_parameters=_get_rope_parameters(config),
Michael Goin's avatar
Michael Goin committed
354
355
356
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
357
358
            config=config,
        )  # Pass config to Grok1Attention
Michael Goin's avatar
Michael Goin committed
359

Bijaya Dangol's avatar
Bijaya Dangol committed
360
        num_experts = _get_num_experts(config)
Michael Goin's avatar
Michael Goin committed
361
        num_experts_per_tok = getattr(config, "num_experts_per_tok", 2)
Bijaya Dangol's avatar
Bijaya Dangol committed
362
363
        moe_intermediate_size = _get_moe_intermediate_size(config)
        moe_renormalize = _get_moe_renormalize(config)
Michael Goin's avatar
Michael Goin committed
364

365
366
367
368
        self.moe_block = Grok1MoE(
            num_experts=num_experts,
            top_k=num_experts_per_tok,
            hidden_size=config.hidden_size,
Bijaya Dangol's avatar
Bijaya Dangol committed
369
370
371
372
373
374
375
376
377
            intermediate_size=moe_intermediate_size,
            router_logit_soft_cap=max(
                getattr(
                    config,
                    "router_logit_softcapping",
                    DEFAULT_ROUTER_LOGIT_SOFTCAP,
                ),
                0.0,
            ),
378
            quant_config=quant_config,
Bijaya Dangol's avatar
Bijaya Dangol committed
379
            renormalize=moe_renormalize,
380
381
            prefix=f"{prefix}.moe_block",
        )
Bijaya Dangol's avatar
Bijaya Dangol committed
382
383
        self.residual_moe = getattr(config, "residual_moe", False)
        self.residual_moe_scale = 1.0 / math.sqrt(2.0)
384
385
386
387
388

        self.pre_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.pre_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Bijaya Dangol's avatar
Bijaya Dangol committed
389
390
391
392
393
394
395
396
        self.mlp = None
        if self.residual_moe:
            self.mlp = Grok1MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )
Michael Goin's avatar
Michael Goin committed
397
398
399
400
401

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
402
        residual: torch.Tensor | None,
403
    ) -> tuple[torch.Tensor, torch.Tensor]:
Michael Goin's avatar
Michael Goin committed
404
405
406
407
408
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.pre_attn_norm(hidden_states)
        else:
409
            hidden_states, residual = self.pre_attn_norm(hidden_states, residual)
Michael Goin's avatar
Michael Goin committed
410
411
412
413
414
415
416
417
418
419
420

        hidden_states = self.attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        # Post attention normalization
        hidden_states = self.post_attn_norm(hidden_states)

        # MoE block with normalization
        hidden_states, residual = self.pre_moe_norm(hidden_states, residual)
Bijaya Dangol's avatar
Bijaya Dangol committed
421
422
423
424
425
426
427
        if self.residual_moe:
            assert self.mlp is not None
            hidden_states = (
                self.moe_block(hidden_states) + self.mlp(hidden_states)
            ) * self.residual_moe_scale
        else:
            hidden_states = self.moe_block(hidden_states)
Michael Goin's avatar
Michael Goin committed
428
429
430
431
432
433
434
        hidden_states = self.post_moe_norm(hidden_states)

        return hidden_states, residual


@support_torch_compile
class Grok1Model(nn.Module):
Bijaya Dangol's avatar
Bijaya Dangol committed
435
436
437
438
439
440
441
442
443
444
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        ckpt_gate_proj_name: str = "linear",
        ckpt_down_proj_name: str = "linear_1",
        ckpt_up_proj_name: str = "linear_v",
        weight_name_remapping: dict[str, str] | None = None,
    ):
Michael Goin's avatar
Michael Goin committed
445
446
447
448
449
450
        super().__init__()

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

451
452
        self.config = config
        self.quant_config = quant_config
453

Bijaya Dangol's avatar
Bijaya Dangol committed
454
455
456
457
458
459
        # Store expert naming for weight loading
        self.ckpt_gate_proj_name = ckpt_gate_proj_name
        self.ckpt_down_proj_name = ckpt_down_proj_name
        self.ckpt_up_proj_name = ckpt_up_proj_name
        self.weight_name_remapping = weight_name_remapping or {}

460
461
        self.vocab_size = config.vocab_size

Michael Goin's avatar
Michael Goin committed
462
        self.embedding_multiplier_scale = getattr(
463
464
            config, "embedding_multiplier_scale", DEFAULT_EMBEDDING_MULTIPLIER_SCALE
        )
Michael Goin's avatar
Michael Goin committed
465
466
467
468
469
470
471
472
473
474
475
476

        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            quant_config=quant_config,
        )

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: Grok1DecoderLayer(
                config, cache_config, quant_config=quant_config, prefix=prefix
            ),
477
478
            prefix=f"{prefix}.layers",
        )
Michael Goin's avatar
Michael Goin committed
479
480

        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
481
482
483
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
Michael Goin's avatar
Michael Goin committed
484

485
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
Michael Goin's avatar
Michael Goin committed
486
487
488
489
490
491
        hidden_states = self.embed_tokens(input_ids)
        hidden_states = hidden_states * self.embedding_multiplier_scale
        return hidden_states

    def forward(
        self,
492
        input_ids: torch.Tensor | None,
Michael Goin's avatar
Michael Goin committed
493
        positions: torch.Tensor,
494
495
496
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
Michael Goin's avatar
Michael Goin committed
497
498
499
500
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
501
                hidden_states = self.embed_input_ids(input_ids)
Michael Goin's avatar
Michael Goin committed
502
503
504
505
506
507
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

508
        for layer in islice(self.layers, self.start_layer, self.end_layer):
509
            hidden_states, residual = layer(positions, hidden_states, residual)
Michael Goin's avatar
Michael Goin committed
510
511

        if not get_pp_group().is_last_rank:
512
513
514
            return IntermediateTensors(
                {"hidden_states": hidden_states, "residual": residual}
            )
Michael Goin's avatar
Michael Goin committed
515
516
517
518

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

519
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
Bijaya Dangol's avatar
Bijaya Dangol committed
520
521
        # Map expert parameter names to standard names
        num_experts = _get_num_experts(self.config)
522
        return FusedMoE.make_expert_params_mapping(
523
            self,
Bijaya Dangol's avatar
Bijaya Dangol committed
524
525
526
            ckpt_gate_proj_name=self.ckpt_gate_proj_name,
            ckpt_down_proj_name=self.ckpt_down_proj_name,
            ckpt_up_proj_name=self.ckpt_up_proj_name,
527
528
            num_experts=num_experts,
        )
529

530
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Michael Goin's avatar
Michael Goin committed
531
532
533
534
535
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
Bijaya Dangol's avatar
Bijaya Dangol committed
536
537
            ("mlp.gate_up_proj", "mlp.gate_proj", 0),
            ("mlp.gate_up_proj", "mlp.up_proj", 1),
Michael Goin's avatar
Michael Goin committed
538
539
540
        ]

        params_dict = dict(self.named_parameters())
541
        loaded_params: set[str] = set()
542
        expert_params_mapping = self.get_expert_mapping()
Michael Goin's avatar
Michael Goin committed
543
        for name, loaded_weight in weights:
Bijaya Dangol's avatar
Bijaya Dangol committed
544
545
546
547
            # Apply version-specific weight name remapping
            for old_pattern, new_pattern in self.weight_name_remapping.items():
                if old_pattern in name:
                    name = name.replace(old_pattern, new_pattern)
548
549
550
            if self.quant_config is not None and (
                scale_name := self.quant_config.get_cache_scale(name)
            ):
Michael Goin's avatar
Michael Goin committed
551
552
                # Loading kv cache quantization scales
                param = params_dict[scale_name]
553
554
555
556
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                loaded_weight = (
                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
                )
Michael Goin's avatar
Michael Goin committed
557
558
559
560
                weight_loader(param, loaded_weight)
                loaded_params.add(scale_name)
                continue

561
            for param_name, weight_name, shard_id in stacked_params_mapping:
Michael Goin's avatar
Michael Goin committed
562
563
564
565
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
566
567
568
                if (
                    name.endswith(".bias") or name.endswith("_bias")
                ) and name not in params_dict:
Michael Goin's avatar
Michael Goin committed
569
570
571
572
573
574
575
576
577
                    continue
                # Skip layers on other devices.
                if is_pp_missing_parameter(name, self):
                    continue
                if name.endswith("scale"):
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue
Bijaya Dangol's avatar
Bijaya Dangol committed
578
579
                if name not in params_dict:
                    continue
Michael Goin's avatar
Michael Goin committed
580
581
582
583
584
585
586
587
588
589
590
591
592
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue
                    name = name.replace(weight_name, param_name)
                    # Skip layers on other devices.
                    if is_pp_missing_parameter(name, self):
                        continue
593
594
595
                    if (
                        name.endswith(".bias") or name.endswith("_bias")
                    ) and name not in params_dict:
Michael Goin's avatar
Michael Goin committed
596
597
598
                        continue
                    param = params_dict[name]
                    weight_loader = param.weight_loader
599
600
601
602
603
604
605
                    weight_loader(
                        param,
                        loaded_weight,
                        name,
                        shard_id=shard_id,
                        expert_id=expert_id,
                    )
Michael Goin's avatar
Michael Goin committed
606
607
608
                    break
                else:
                    # Skip loading extra bias for GPTQ models.
609
610
611
                    if (
                        name.endswith(".bias") or name.endswith("_bias")
                    ) and name not in params_dict:
Michael Goin's avatar
Michael Goin committed
612
613
614
615
616
617
618
619
620
621
622
623
624
625
                        continue
                    # Skip layers on other devices.
                    if is_pp_missing_parameter(name, self):
                        continue

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

                    # Handle Grok1-specific norm.scale naming
                    if "norm.scale" in name:
                        name = name.replace("scale", "weight")

Bijaya Dangol's avatar
Bijaya Dangol committed
626
627
                    if name not in params_dict:
                        continue
Michael Goin's avatar
Michael Goin committed
628
                    param = params_dict[name]
629
630
631
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
Michael Goin's avatar
Michael Goin committed
632
633
634
                    weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params
635
636


Bijaya Dangol's avatar
Bijaya Dangol committed
637
638
639
class GrokBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
    """Base class for Grok models with shared logic."""

640
641
    fall_back_to_pt_during_load = False

Bijaya Dangol's avatar
Bijaya Dangol committed
642
    # Subclasses should override these
643
644
645
646
647
648
649
650
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
    }

Bijaya Dangol's avatar
Bijaya Dangol committed
651
652
653
654
655
656
657
658
659
    # Expert weight naming - subclasses override these
    ckpt_gate_proj_name: str = "linear"
    ckpt_down_proj_name: str = "linear_1"
    ckpt_up_proj_name: str = "linear_v"

    def get_weight_name_remapping(self) -> dict[str, str]:
        """Return weight name remapping for this version. Override in subclasses."""
        return {}

660
661
662
663
664
665
666
667
668
    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

669
        self.model = Grok1Model(
Bijaya Dangol's avatar
Bijaya Dangol committed
670
671
672
673
674
675
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "model"),
            ckpt_gate_proj_name=self.ckpt_gate_proj_name,
            ckpt_down_proj_name=self.ckpt_down_proj_name,
            ckpt_up_proj_name=self.ckpt_up_proj_name,
            weight_name_remapping=self.get_weight_name_remapping(),
676
        )
677
678

        self.lm_head = ParallelLMHead(
679
            config.vocab_size,
680
681
682
683
684
685
686
687
688
            config.hidden_size,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "lm_head"),
        )

        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

        self.output_multiplier_scale = getattr(
689
690
691
            config, "output_multiplier_scale", DEFAULT_OUTPUT_MULTIPLIER_SCALE
        )
        self.logits_processor = LogitsProcessor(
Bijaya Dangol's avatar
Bijaya Dangol committed
692
693
694
            config.vocab_size,
            scale=self.output_multiplier_scale,
            soft_cap=getattr(config, "final_logit_softcapping", None),
695
        )
696
697

        self.make_empty_intermediate_tensors = (
698
699
            self.model.make_empty_intermediate_tensors
        )
700

701
702
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
703
704
705

    def forward(
        self,
706
        input_ids: torch.Tensor | None,
707
        positions: torch.Tensor,
708
709
710
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
711
712
713
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
714
715
716
717
718
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
719
    ) -> torch.Tensor | None:
720
        logits = self.logits_processor(self.lm_head, hidden_states)
721
722
        return logits

723
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
724
        # Skip lm_head when tie_word_embeddings is True
725
        skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else None
726

727
728
729
730
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=skip_prefixes,
        )
731
        return loader.load_weights(weights)
732
733
734

    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        return self.model.get_expert_mapping()
Bijaya Dangol's avatar
Bijaya Dangol committed
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801


class Grok1ForCausalLM(GrokBaseForCausalLM):
    """Grok1-specific implementation."""

    # Grok1 expert weight naming
    ckpt_gate_proj_name = "linear"
    ckpt_down_proj_name = "linear_1"
    ckpt_up_proj_name = "linear_v"

    def get_weight_name_remapping(self) -> dict[str, str]:
        # Grok1 uses standard naming, no remapping needed
        return {}


class Grok2ForCausalLM(GrokBaseForCausalLM):
    """Grok2-specific implementation."""

    # Grok2 has additional packed modules for MLP
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # Grok2 expert weight naming
    ckpt_gate_proj_name = "w1"
    ckpt_down_proj_name = "w2"
    ckpt_up_proj_name = "w3"

    def get_weight_name_remapping(self) -> dict[str, str]:
        # Grok2 checkpoint uses different naming conventions
        return {
            ".self_attn.": ".attn.",
            ".block_sparse_moe.": ".moe_block.",
        }


# Version dispatch mapping
_GROK_VERSIONS: dict[str, type[GrokBaseForCausalLM]] = {
    "grok1": Grok1ForCausalLM,
    "grok2": Grok2ForCausalLM,
}


class GrokForCausalLM(GrokBaseForCausalLM):
    """Factory class that dispatches to version-specific implementation."""

    def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""):
        config = vllm_config.model_config.hf_config
        version = _get_grok_version(config)

        instance_cls = _GROK_VERSIONS.get(version)
        if instance_cls is None:
            raise ValueError(f"Unsupported Grok version: {version}")

        # Merge class attributes for LoRA/quantization compatibility
        cls.packed_modules_mapping = dict(cls.packed_modules_mapping)
        cls.packed_modules_mapping.update(instance_cls.packed_modules_mapping)

        return instance_cls(vllm_config=vllm_config, prefix=prefix)