glm4.py 11.1 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."""
zhuwenwen's avatar
zhuwenwen committed
25

zhuwenwen's avatar
zhuwenwen committed
26
import os
27
from collections.abc import Iterable
Yuxuan Zhang's avatar
Yuxuan Zhang committed
28
29
30
31
32

import torch
from torch import nn
from transformers import Glm4Config

33
import vllm.envs as envs
34
from vllm.attention.layer import Attention
Yuxuan Zhang's avatar
Yuxuan Zhang committed
35
36
37
38
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
39
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
Yuxuan Zhang's avatar
Yuxuan Zhang committed
40
41
42
43
44
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
45
from vllm.v1.attention.backend import AttentionType
Yuxuan Zhang's avatar
Yuxuan Zhang committed
46
47
48
49

from .interfaces import SupportsLoRA, SupportsPP
from .llama import LlamaMLP as Glm4MLP
from .llama import LlamaModel
zhuwenwen's avatar
zhuwenwen committed
50
from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
Yuxuan Zhang's avatar
Yuxuan Zhang committed
51

zhuwenwen's avatar
zhuwenwen committed
52
53
54
from vllm.utils import W8a8GetCacheJSON
from vllm import _custom_ops as ops
from vllm.model_executor.utils import pad_weight, gemm_bank_conf
Yuxuan Zhang's avatar
Yuxuan Zhang committed
55
56

class Glm4Attention(nn.Module):
57
58
59
60
61
62
63
    def __init__(
        self,
        config: Glm4Config,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
64
        head_dim: int | None = None,
65
        qkv_bias: bool = False,
66
67
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
68
69
70
        prefix: str = "",
        attn_type: str = AttentionType.DECODER,
    ) -> None:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
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
86
87

        rope_params = getattr(config, "rope_parameters", None)
zhuwenwen's avatar
zhuwenwen committed
88
        config.rope_parameters.setdefault("partial_rotary_factor", 0.5)
89

Yuxuan Zhang's avatar
Yuxuan Zhang committed
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.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.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,
            max_position=max_position,
114
            rope_parameters=config.rope_parameters,
zhuwenwen's avatar
zhuwenwen committed
115
            is_neox_style=False,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
116
        )
117
118
119
120
121
122
123
124
125
126
        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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

    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):
142
143
144
145
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
146
        config: Glm4Config | None = None,
147
    ) -> None:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
148
        super().__init__()
149
150
151
152
153

        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
154
155
156
157
158
159
160
161
        self.hidden_size = config.hidden_size

        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,
162
163
            qkv_bias=getattr(config, "attention_bias", False),
            head_dim=getattr(config, "head_dim", None),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
164
165
166
167
168
169
170
171
172
173
174
175
            cache_config=cache_config,
            quant_config=quant_config,
            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",
        )
176
177
178
179
180
181
182
183
        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
184
185
186
187
188

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

        hidden_states = self.post_self_attn_layernorm(hidden_states)

        # Fully Connected
205
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
        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,
223
224
    }
)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
225
226
class Glm4Model(LlamaModel):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
227
228
229
        super().__init__(
            vllm_config=vllm_config, prefix=prefix, layer_type=Glm4DecoderLayer
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252


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

        self.config = config

        self.quant_config = quant_config
253
254
255
        self.model = Glm4Model(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
256
257
258
259
260

        if get_pp_group().is_last_rank:
            if config.tie_word_embeddings:
                self.lm_head = self.model.embed_tokens
            else:
261
262
263
264
265
266
                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
267
268
269
270
271
272
        else:
            self.lm_head = PPMissingLayer()

        self.logits_processor = LogitsProcessor(config.vocab_size)

        self.make_empty_intermediate_tensors = (
273
274
            self.model.make_empty_intermediate_tensors
        )
zhuwenwen's avatar
zhuwenwen committed
275
276
277
278
279
280
281
282
283
284
285
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
            
        self.tritonsingleton= W8a8GetCacheJSON()      
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
        # self.use_lm_nn = os.environ.get('LM_NN') == '1'
        self.use_gemm_pad = os.environ.get('GEMM_PAD') == '1'
        self.use_fa_pad = os.environ.get('FA_PAD') == '1'
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
286
        self.w8a8_strategy = envs.VLLM_W8A8_BACKEND
Yuxuan Zhang's avatar
Yuxuan Zhang committed
287

288
289
    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
290
291
292

    def forward(
        self,
zhuwenwen's avatar
zhuwenwen committed
293
        input_ids: torch.Tensor,
Yuxuan Zhang's avatar
Yuxuan Zhang committed
294
        positions: torch.Tensor,
295
296
297
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
298
299
300
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
Yuxuan Zhang's avatar
Yuxuan Zhang committed
301
302
303
304
305
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
306
    ) -> torch.Tensor | None:
307
        logits = self.logits_processor(self.lm_head, hidden_states)
Yuxuan Zhang's avatar
Yuxuan Zhang committed
308
309
        return logits

310
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Yuxuan Zhang's avatar
Yuxuan Zhang committed
311
312
        loader = AutoWeightsLoader(
            self,
313
            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
Yuxuan Zhang's avatar
Yuxuan Zhang committed
314
315
        )
        return loader.load_weights(weights)