chatglm.py 19 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
2
# Adapted from
3
4
5
# https://github.com/THUDM/ChatGLM2-6B
"""Inference-only ChatGLM model compatible with THUDM weights."""
from typing import Iterable, List, Optional, Set, Tuple, Union
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
6
7
8
9
10

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

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

30
from .interfaces import SupportsLoRA, SupportsPP
31
from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
32
                    make_empty_intermediate_tensors_factory, make_layers,
33
                    maybe_prefix)
34

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
35
36
37

class GLMAttention(nn.Module):

38
39
    def __init__(
        self,
40
        config: ChatGLMConfig,
41
        cache_config: Optional[CacheConfig] = None,
42
        quant_config: Optional[QuantizationConfig] = None,
43
        prefix: str = "",
44
    ):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
45
46
47
48
49
50
51
52
53
54
        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)
55
56
57
58
59
60
61
62
63
        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
64
65
66
67
68
        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

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

86
87
88
        # 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)
89
90
91
        # 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
92
        self.rotary_emb = get_rope(
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
93
94
            self.head_dim,
            rotary_dim=self.head_dim // 2,
95
96
            max_position=max_positions,
            base=10000 * rope_ratio,
97
            is_neox_style=is_neox_style,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
98
        )
99
100
101
102
103
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
                              cache_config=cache_config,
104
105
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
106
107
108
109
110

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
111
112
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
113
114
115
    ) -> 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
116
        q, k = self.rotary_emb(position_ids, q, k)
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
117
118
119
120
        context_layer = self.attn(
            q,
            k,
            v,
121
122
            kv_cache,
            attn_metadata,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
123
124
125
126
127
128
129
130
131
132
133
134
135
        )
        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.
    """

136
137
    def __init__(
        self,
138
        config: ChatGLMConfig,
139
        quant_config: Optional[QuantizationConfig] = None,
140
        prefix: str = "",
141
    ):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
142
143
144
145
146
        super().__init__()

        self.add_bias = config.add_bias_linear

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

        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,
162
            quant_config=quant_config,
163
            prefix=f"{prefix}.dense_4h_to_h",
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
        )

    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,
184
        config: ChatGLMConfig,
185
        cache_config: Optional[CacheConfig] = None,
186
        quant_config: Optional[QuantizationConfig] = None,
187
        prefix: str = "",
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
188
189
190
191
192
193
194
195
196
197
198
199
200
    ):
        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.
201
202
203
204
        self.self_attention = GLMAttention(config,
                                           cache_config,
                                           quant_config,
                                           prefix=f"{prefix}.self_attention")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
205
206
207
208
209
210
211
        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
212
        self.mlp = GLMMLP(config, quant_config, prefix=f"{prefix}.mlp")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
213
214
215
216
217

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
218
219
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
220
221
222
223
224
225
226
227
228
    ) -> 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,
            kv_cache=kv_cache,
229
            attn_metadata=attn_metadata,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
        )

        # 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."""

257
258
    def __init__(
        self,
259
        config: ChatGLMConfig,
260
        cache_config: Optional[CacheConfig] = None,
261
        quant_config: Optional[QuantizationConfig] = None,
262
        prefix: str = "",
263
    ):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
264
265
266
267
268
269
270
        super().__init__()
        self.post_layer_norm = config.post_layer_norm

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

        # Transformer layers.
271
272
        self.start_layer, self.end_layer, self.layers = make_layers(
            self.num_layers,
273
274
            lambda prefix: GLMBlock(
                config, cache_config, quant_config, prefix=prefix),
275
276
            prefix=f"{prefix}.layers",
        )
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
277
278
279
280
281
282
283

        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)

284
285
286
287
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
288
289
290
291
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
292
293
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
294
    ) -> Union[torch.Tensor, IntermediateTensors]:
295
        for i in range(self.start_layer, self.end_layer):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
296
297
298
299
            layer = self.layers[i]
            hidden_states = layer(
                hidden_states=hidden_states,
                position_ids=position_ids,
300
                kv_cache=kv_caches[i - self.start_layer],
301
                attn_metadata=attn_metadata,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
302
            )
303
304
305
306

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

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
307
        # Final layer norm.
308
        if self.post_layer_norm:
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
309
310
311
312
313
314
315
            hidden_states = self.final_layernorm(hidden_states)

        return hidden_states


class ChatGLMModel(nn.Module):

316
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
317
318
        super().__init__()

319
320
321
322
        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

323
324
        self.config = config

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
325
        self.embedding = VocabParallelEmbedding(config.padded_vocab_size,
326
                                                config.hidden_size,
327
328
                                                quant_config=quant_config,
                                                prefix=f"{prefix}.embedding")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
329
330
331
332

        self.num_layers = config.num_layers
        self.multi_query_group_num = config.multi_query_group_num
        self.kv_channels = config.kv_channels
333
334
335
336
        self.encoder = GLMTransformer(config,
                                      cache_config,
                                      quant_config,
                                      prefix=f"{prefix}.encoder")
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
337

338
        self.output_layer = ParallelLMHead(config.padded_vocab_size,
339
                                           config.hidden_size,
340
341
                                           quant_config=quant_config,
                                           prefix=f"{prefix}.output_layer")
342

343
344
345
        self.make_empty_intermediate_tensors = (
            self.encoder.make_empty_intermediate_tensors)

346
347
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embedding(input_ids)
348

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
349
350
351
    def forward(
        self,
        input_ids: torch.Tensor,
352
        positions: torch.Tensor,
353
354
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
355
        intermediate_tensors: Optional[IntermediateTensors] = None,
356
        inputs_embeds: Optional[torch.Tensor] = None,
357
        **kwargs: object,
358
359
360
361
362
363
364
365
366
    ) -> 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"]
367

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
368
369
        # Run encoder.
        hidden_states = self.encoder(
370
            hidden_states=hidden_states,
371
            position_ids=positions,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
372
            kv_caches=kv_caches,
373
            attn_metadata=attn_metadata,
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
374
        )
375

GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
376
377
        return hidden_states

378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
        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())
        loaded_params: Set[str] = set()

        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
416

417
class ChatGLMBaseModel(nn.Module):
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
418

419
420
421
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_substr={".word_embeddings": ""}, )

422
423
424
425
426
427
428
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        transformer_type: type[ChatGLMModel] = ChatGLMModel,
    ) -> None:
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
429
        super().__init__()
430
431
432
433
        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
434
435
        self.config = config
        self.lora_config = lora_config
436
        self.multimodal_config = multimodal_config
437

438
        self.quant_config = quant_config
439
440
        self.max_position_embeddings = getattr(config, "max_sequence_length",
                                               8192)
441
442
443
        self.transformer = transformer_type(vllm_config=vllm_config,
                                            prefix=maybe_prefix(
                                                prefix, "transformer"))
444
445
446
        if self.config.tie_word_embeddings:
            self.transformer.output_layer.weight = (
                self.transformer.embedding.weight)
447
        self.lm_head = self.transformer.output_layer
448
        self.logits_processor = LogitsProcessor(config.padded_vocab_size)
Joe Runde's avatar
Joe Runde committed
449
        self.sampler = get_sampler()
450
451
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
452

453
454
455
456
457
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
458
        logits = self.logits_processor(self.lm_head, hidden_states,
459
460
461
                                       sampling_metadata)
        return logits

462
463
    def sample(
        self,
464
        logits: torch.Tensor,
465
        sampling_metadata: SamplingMetadata,
466
    ) -> Optional[SamplerOutput]:
467
        next_tokens = self.sampler(logits, sampling_metadata)
GoHomeToMacDonal's avatar
GoHomeToMacDonal committed
468
469
        return next_tokens

470
471
472
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
473
474


475
class ChatGLMForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP):
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
    packed_modules_mapping = {
        "query_key_value": ["query_key_value"],
        "dense_h_to_4h": ["dense_h_to_4h"]
    }
    # LoRA specific attributes
    supported_lora_modules = [
        "query_key_value",
        "dense",
        "dense_h_to_4h",
        "dense_4h_to_h",
    ]

    embedding_modules = {}
    embedding_padding_modules = []

491
492
493
494
495
496
497
498
499
    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 "
                f"`--hf-overrides {hf_overrides!r}`")
500

501
        super().__init__(vllm_config=vllm_config, prefix=prefix)
502

503
    def forward(
504
505
        self,
        input_ids: torch.Tensor,
506
507
508
509
510
511
512
513
514
515
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        hidden_states = self.transformer(input_ids, positions, kv_caches,
                                         attn_metadata, intermediate_tensors,
                                         inputs_embeds)
        return hidden_states