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

# Copyright 2025 The Zhipu AI 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.
"""Inference-only GLM-4-0414 model compatible with HuggingFace weights."""
25

26
27
from collections.abc import Iterable
from typing import Optional, Union
Yuxuan Zhang's avatar
Yuxuan Zhang committed
28
29
30
31
32
33
34
35
36
37

import torch
from torch import nn
from transformers import Glm4Config

from vllm.attention import Attention, AttentionType
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
from vllm.model_executor.layers.layernorm import RMSNorm
38
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
Yuxuan Zhang's avatar
Yuxuan Zhang committed
39
40
41
42
43
44
45
46
47
48
49
50
51
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 ParallelLMHead
from vllm.sequence import IntermediateTensors

from .interfaces import SupportsLoRA, SupportsPP
from .llama import LlamaMLP as Glm4MLP
from .llama import LlamaModel
from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix


class Glm4Attention(nn.Module):
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
    def __init__(
        self,
        config: Glm4Config,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
        head_dim: Optional[int] = None,
        qkv_bias: bool = False,
        rope_theta: float = 10000,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        rope_scaling: Optional[tuple] = None,
        prefix: str = "",
        attn_type: str = AttentionType.DECODER,
    ) -> None:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
        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
        partial_rotary_factor = getattr(config, "partial_rotary_factor", 0.5)
        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
intervitens's avatar
intervitens committed
86
        self.rotary_dim = self.head_dim
Yuxuan Zhang's avatar
Yuxuan Zhang committed
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
        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
        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",
        )
        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,
            rotary_dim=self.rotary_dim,
            max_position=max_position,
            base=self.rope_theta,
            rope_scaling=rope_scaling,
            partial_rotary_factor=partial_rotary_factor,
intervitens's avatar
intervitens committed
114
            is_neox_style=False,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
115
        )
116
117
118
119
120
121
122
123
124
125
        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",
            attn_type=attn_type,
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

    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)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output


class Glm4DecoderLayer(nn.Module):
141
142
143
144
145
146
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
        config: Optional[Glm4Config] = None,
    ) -> None:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
147
        super().__init__()
148
149
150
151
152

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

Yuxuan Zhang's avatar
Yuxuan Zhang committed
153
154
155
156
157
158
159
160
161
162
163
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 1000000)
        rope_scaling = getattr(config, "rope_scaling", None)

        self.self_attn = Glm4Attention(
            config=config,
            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,
164
165
            qkv_bias=getattr(config, "attention_bias", False),
            head_dim=getattr(config, "head_dim", None),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
166
167
168
169
170
171
172
173
174
175
176
177
178
            cache_config=cache_config,
            quant_config=quant_config,
            rope_scaling=rope_scaling,
            prefix=f"{prefix}.self_attn",
            attn_type=AttentionType.DECODER,
        )
        self.mlp = Glm4MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
        )
179
180
181
182
183
184
185
186
        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
        )
        self.post_self_attn_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_mlp_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
187
188
189
190
191
192

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
193
    ) -> tuple[torch.Tensor, torch.Tensor]:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
194
195
196
197
198
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
199
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
200
201
202
203
204
205
206
207
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        hidden_states = self.post_self_attn_layernorm(hidden_states)

        # Fully Connected
208
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.post_mlp_layernorm(hidden_states)

        return hidden_states, residual


ALL_DECODER_LAYER_TYPES = {
    "attention": Glm4DecoderLayer,
}


@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        "positions": -1,
        "intermediate_tensors": 0,
        "inputs_embeds": 0,
226
227
    }
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
228
229
class Glm4Model(LlamaModel):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
230
231
232
        super().__init__(
            vllm_config=vllm_config, prefix=prefix, layer_type=Glm4DecoderLayer
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257


class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

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

        self.config = config
        self.lora_config = lora_config

        self.quant_config = quant_config
258
259
260
        self.model = Glm4Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
261
262
263
264
265

        if get_pp_group().is_last_rank:
            if config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
266
267
268
269
270
271
                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
272
273
274
275
276
277
        else:
            self.lm_head = PPMissingLayer()

        self.logits_processor = LogitsProcessor(config.vocab_size)

        self.make_empty_intermediate_tensors = (
278
279
            self.model.make_empty_intermediate_tensors
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
280
281
282
283
284
285
286
287
288
289
290

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

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
291
292
293
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
294
295
296
297
298
299
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> Optional[torch.Tensor]:
300
        logits = self.logits_processor(self.lm_head, hidden_states)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
301
302
        return logits

303
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
304
305
        loader = AutoWeightsLoader(
            self,
306
            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
307
308
        )
        return loader.load_weights(weights)