chatglm.py 18 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
3
# Adapted from
4
5
# https://github.com/THUDM/ChatGLM2-6B
"""Inference-only ChatGLM model compatible with THUDM weights."""
6
import json
7
8
from collections.abc import Iterable
from typing import Optional, Union
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
9
10
11
12
13

import torch
from torch import nn
from torch.nn import LayerNorm

14
from vllm.attention import Attention
15
from vllm.compilation.decorators import support_torch_compile
16
from vllm.config import CacheConfig, VllmConfig
17
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
18
from vllm.model_executor.layers.activation import SiluAndMul
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
21
22
                                               QKVParallelLinear,
                                               RowParallelLinear)
23
from vllm.model_executor.layers.logits_processor import LogitsProcessor
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
from vllm.model_executor.layers.rotary_embedding import get_rope
26
from vllm.model_executor.layers.vocab_parallel_embedding import (
27
    ParallelLMHead, VocabParallelEmbedding)
28
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
29
from vllm.model_executor.sampling_metadata import SamplingMetadata
30
from vllm.sequence import IntermediateTensors
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
31
32
from vllm.transformers_utils.configs import ChatGLMConfig

33
from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant
34
from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
35
                    make_empty_intermediate_tensors_factory, make_layers,
36
                    maybe_prefix)
37

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
38
39
40

class GLMAttention(nn.Module):

41
42
    def __init__(
        self,
43
        config: ChatGLMConfig,
44
        cache_config: Optional[CacheConfig] = None,
45
        quant_config: Optional[QuantizationConfig] = None,
46
        prefix: str = "",
47
    ):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
48
49
50
51
52
53
54
55
56
57
        super().__init__()
        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.multi_query_attention = config.multi_query_attention
        self.total_num_kv_heads = (config.multi_query_group_num
                                   if config.multi_query_attention else
                                   config.num_attention_heads)
58
59
60
61
62
63
64
65
66
        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)
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
67
68
69
70
71
        self.head_dim = config.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

72
73
        self.query_key_value = QKVParallelLinear(
            self.hidden_size,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
74
            self.head_dim,
75
76
77
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=config.add_bias_linear or config.add_qkv_bias,
78
            quant_config=quant_config,
79
            prefix=f"{prefix}.query_key_value",
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
80
81
82
83
84
        )
        self.dense = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            config.hidden_size,
            bias=config.add_bias_linear,
85
            quant_config=quant_config,
86
            prefix=f"{prefix}.dense",
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
87
88
        )

89
90
91
        # https://huggingface.co/THUDM/chatglm3-6b-32k/blob/e210410255278dd9d74463cf396ba559c0ef801c/modeling_chatglm.py#L141
        rope_ratio = getattr(config, "rope_ratio", 1.0)
        max_positions = getattr(config, "seq_length", 8192)
92
93
94
        # NOTE: THUDM/cogagent-9b-20241220 uses original_rope=False,
        # which is equivalent to is_neox_style=True
        is_neox_style = not config.original_rope
Woosuk Kwon's avatar
Woosuk Kwon committed
95
        self.rotary_emb = get_rope(
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
96
97
            self.head_dim,
            rotary_dim=self.head_dim // 2,
98
99
            max_position=max_positions,
            base=10000 * rope_ratio,
100
            is_neox_style=is_neox_style,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
101
        )
102
103
104
105
106
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
                              cache_config=cache_config,
107
108
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
109
110
111
112
113
114
115
116

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.query_key_value(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
Woosuk Kwon's avatar
Woosuk Kwon committed
117
        q, k = self.rotary_emb(position_ids, q, k)
118
        context_layer = self.attn(q, k, v)
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
119
120
121
122
123
124
125
126
127
128
129
130
        attn_output, _ = self.dense(context_layer)
        return attn_output


class GLMMLP(nn.Module):
    """MLP.

    MLP will take the input with h hidden state, project it to 4*h
    hidden dimension, perform nonlinear transformation, and project the
    state back into h hidden dimension.
    """

131
132
    def __init__(
        self,
133
        config: ChatGLMConfig,
134
        quant_config: Optional[QuantizationConfig] = None,
135
        prefix: str = "",
136
    ):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
137
138
139
140
141
        super().__init__()

        self.add_bias = config.add_bias_linear

        # Project to 4h.
142
        self.dense_h_to_4h = MergedColumnParallelLinear(
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
143
            config.hidden_size,
144
            [config.ffn_hidden_size] * 2,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
145
            bias=config.add_bias_linear,
146
            quant_config=quant_config,
147
            prefix=f"{prefix}.dense_h_to_4h",
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
148
149
150
151
152
153
154
155
156
        )

        self.activation_func = SiluAndMul()

        # Project back to h.
        self.dense_4h_to_h = RowParallelLinear(
            config.ffn_hidden_size,
            config.hidden_size,
            bias=config.add_bias_linear,
157
            quant_config=quant_config,
158
            prefix=f"{prefix}.dense_4h_to_h",
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
        )

    def forward(self, hidden_states):
        # [s, b, 4hp]
        intermediate_parallel, _ = self.dense_h_to_4h(hidden_states)
        intermediate_parallel = self.activation_func(intermediate_parallel)
        # [s, b, h]
        output, _ = self.dense_4h_to_h(intermediate_parallel)
        return output


class GLMBlock(nn.Module):
    """A single transformer layer.

    Transformer layer takes input with size [s, b, h] and returns an
    output of the same size.
    """

    def __init__(
        self,
179
        config: ChatGLMConfig,
180
        cache_config: Optional[CacheConfig] = None,
181
        quant_config: Optional[QuantizationConfig] = None,
182
        prefix: str = "",
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
183
184
185
186
187
188
189
190
191
192
193
194
195
    ):
        super().__init__()
        self.apply_residual_connection_post_layernorm = (
            config.apply_residual_connection_post_layernorm)

        self.fp32_residual_connection = config.fp32_residual_connection

        layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm
        # Layernorm on the input data.
        self.input_layernorm = layer_norm_func(config.hidden_size,
                                               eps=config.layernorm_epsilon)

        # Self attention.
196
197
198
199
        self.self_attention = GLMAttention(config,
                                           cache_config,
                                           quant_config,
                                           prefix=f"{prefix}.self_attention")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
200
201
202
203
204
205
206
        self.hidden_dropout = config.hidden_dropout

        # Layernorm on the attention output
        self.post_attention_layernorm = layer_norm_func(
            config.hidden_size, eps=config.layernorm_epsilon)

        # MLP
207
        self.mlp = GLMMLP(config, quant_config, prefix=f"{prefix}.mlp")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
    ) -> torch.Tensor:
        # hidden_states: [num_tokens, h]
        # Layer norm at the beginning of the transformer layer.
        layernorm_output = self.input_layernorm(hidden_states)
        # Self attention.
        attention_output = self.self_attention(
            hidden_states=layernorm_output,
            position_ids=position_ids,
        )

        # Residual connection.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = hidden_states

        layernorm_input = residual + attention_output

        # Layer norm post the self attention.
        layernorm_output = self.post_attention_layernorm(layernorm_input)

        # Second residual connection.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = layernorm_input

        output = self.mlp(layernorm_output) + residual

        return output


class GLMTransformer(nn.Module):
    """Transformer class."""

248
249
    def __init__(
        self,
250
        config: ChatGLMConfig,
251
        cache_config: Optional[CacheConfig] = None,
252
        quant_config: Optional[QuantizationConfig] = None,
253
        prefix: str = "",
254
    ):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
255
256
257
258
259
260
261
        super().__init__()
        self.post_layer_norm = config.post_layer_norm

        # Number of layers.
        self.num_layers = config.num_layers

        # Transformer layers.
262
263
        self.start_layer, self.end_layer, self.layers = make_layers(
            self.num_layers,
264
265
            lambda prefix: GLMBlock(
                config, cache_config, quant_config, prefix=prefix),
266
267
            prefix=f"{prefix}.layers",
        )
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
268
269
270
271
272
273
274

        if self.post_layer_norm:
            layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm
            # Final layer norm before output.
            self.final_layernorm = layer_norm_func(
                config.hidden_size, eps=config.layernorm_epsilon)

275
276
277
278
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
279
280
281
282
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
283
    ) -> Union[torch.Tensor, IntermediateTensors]:
284
285
286
        for layer in self.layers[self.start_layer:self.end_layer]:
            hidden_states = layer(hidden_states=hidden_states,
                                  position_ids=position_ids)
287
288
289
290

        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
291
        # Final layer norm.
292
        if self.post_layer_norm:
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
293
294
295
296
297
            hidden_states = self.final_layernorm(hidden_states)

        return hidden_states


298
@support_torch_compile
299
300
301
302
303
class ChatGLMModel(nn.Module, SupportsQuant):
    packed_modules_mapping = {
        "linear_proj.merged_proj":
        ["linear_proj.gate_proj", "linear_proj.dense_h_to_4h"]
    }
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
304

305
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
306
307
        super().__init__()

308
309
310
311
        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

312
313
        self.config = config

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
314
        self.embedding = VocabParallelEmbedding(config.padded_vocab_size,
315
                                                config.hidden_size,
316
317
                                                quant_config=quant_config,
                                                prefix=f"{prefix}.embedding")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
318
319
320
321

        self.num_layers = config.num_layers
        self.multi_query_group_num = config.multi_query_group_num
        self.kv_channels = config.kv_channels
322
323
324
325
        self.encoder = GLMTransformer(config,
                                      cache_config,
                                      quant_config,
                                      prefix=f"{prefix}.encoder")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
326

327
        self.output_layer = ParallelLMHead(config.padded_vocab_size,
328
                                           config.hidden_size,
329
330
                                           quant_config=quant_config,
                                           prefix=f"{prefix}.output_layer")
331

332
333
334
        self.make_empty_intermediate_tensors = (
            self.encoder.make_empty_intermediate_tensors)

335
336
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embedding(input_ids)
337

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
338
339
340
    def forward(
        self,
        input_ids: torch.Tensor,
341
342
        positions: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors] = None,
343
        inputs_embeds: Optional[torch.Tensor] = None,
344
        **kwargs: object,
345
346
347
348
349
350
351
352
353
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
354

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
355
356
        # Run encoder.
        hidden_states = self.encoder(
357
            hidden_states=hidden_states,
358
            position_ids=positions,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
359
        )
360

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
361
362
        return hidden_states

363
364
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
365
366
367
368
369
370
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("linear_proj.merged_proj", "linear_proj.gate_proj", 0),
            ("linear_proj.merged_proj", "linear_proj.dense_h_to_4h", 1),
        ]
        params_dict = dict(self.named_parameters())
371
        loaded_params: set[str] = set()
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400

        for name, loaded_weight in weights:
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    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:
                if "rotary_pos_emb.inv_freq" in name:
                    continue
                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 = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
401

402
class ChatGLMBaseModel(nn.Module):
403
404
405
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_substr={".word_embeddings": ""}, )

406
407
408
409
410
411
412
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        transformer_type: type[ChatGLMModel] = ChatGLMModel,
    ) -> None:
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
413
        super().__init__()
414
415
416
417
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
        multimodal_config = vllm_config.model_config.multimodal_config
418
419
        self.config = config
        self.lora_config = lora_config
420
        self.multimodal_config = multimodal_config
421

422
        self.quant_config = quant_config
423
424
        self.max_position_embeddings = getattr(config, "max_sequence_length",
                                               8192)
425
426
427
        self.transformer = transformer_type(vllm_config=vllm_config,
                                            prefix=maybe_prefix(
                                                prefix, "transformer"))
428
429
430
        if self.config.tie_word_embeddings:
            self.transformer.output_layer.weight = (
                self.transformer.embedding.weight)
431
        self.lm_head = self.transformer.output_layer
432
        self.logits_processor = LogitsProcessor(config.padded_vocab_size)
433
434
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
435

436
437
438
439
440
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
441
        logits = self.logits_processor(self.lm_head, hidden_states,
442
443
444
                                       sampling_metadata)
        return logits

445
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
446
447
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
448
449


450
451
class ChatGLMForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,
                         SupportsQuant):
452
453
454
455
456
    packed_modules_mapping = {
        "query_key_value": ["query_key_value"],
        "dense_h_to_4h": ["dense_h_to_4h"]
    }

457
458
459
460
461
462
463
464
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        config = vllm_config.model_config.hf_config
        if hasattr(config, "vision_config"):
            hf_overrides = {"architectures": ["GLM4VForCausalLM"]}
            raise RuntimeError(
                "The configuration of this model indicates that it supports "
                "vision inputs, but you instantiated the text-only version "
                "of this model. Please use the vision model by setting "
465
                f"`--hf-overrides '{json.dumps(hf_overrides)}'`")
466

467
        super().__init__(vllm_config=vllm_config, prefix=prefix)
468

469
    def forward(
470
471
        self,
        input_ids: torch.Tensor,
472
473
474
475
        positions: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
476
477
        hidden_states = self.transformer(input_ids, positions,
                                         intermediate_tensors, inputs_embeds)
478
        return hidden_states