internlm2.py 17.2 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from functools import partial
4
from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union
Fengzhe Zhou's avatar
Fengzhe Zhou committed
5
6

import torch
7
from torch import nn
Fengzhe Zhou's avatar
Fengzhe Zhou committed
8
9
from transformers import PretrainedConfig

10
from vllm.attention import Attention
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, VllmConfig
13
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
14
15
16
                              get_tensor_model_parallel_world_size,
                              split_tensor_along_last_dim,
                              tensor_model_parallel_all_gather)
17
18
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
20
21
                                               QKVParallelLinear,
                                               RowParallelLinear)
22
from vllm.model_executor.layers.logits_processor import LogitsProcessor
23
from vllm.model_executor.layers.pooler import Pooler, PoolingType
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
from vllm.model_executor.layers.rotary_embedding import get_rope
Joe Runde's avatar
Joe Runde committed
26
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
27
from vllm.model_executor.layers.vocab_parallel_embedding import (
28
    ParallelLMHead, VocabParallelEmbedding)
29
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
30
from vllm.model_executor.pooling_metadata import PoolingMetadata
31
from vllm.model_executor.sampling_metadata import SamplingMetadata
32
from vllm.sequence import IntermediateTensors, PoolerOutput
Fengzhe Zhou's avatar
Fengzhe Zhou committed
33

34
from .interfaces import SupportsLoRA, SupportsPP
35
from .utils import (is_pp_missing_parameter,
36
37
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
38

39
40
41
42
43
44
45
46

class InternLM2MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
47
        quant_config: Optional[QuantizationConfig] = None,
48
        prefix: str = "",
49
50
51
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
52
53
54
55
56
57
58
59
60
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.w2 = RowParallelLinear(
            intermediate_size,
            hidden_size,
61
            bias=False,
62
63
64
            quant_config=quant_config,
            prefix=f"{prefix}.w2",
        )
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.w2(x)
        return x


class InternLM2Attention(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        rope_theta: float = 10000,
        rope_scaling: Optional[Dict[str, Any]] = None,
        max_position_embeddings: int = 8192,
87
        cache_config: Optional[CacheConfig] = None,
88
        quant_config: Optional[QuantizationConfig] = None,
89
        prefix: str = "",
90
91
92
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
93
94
        self.tp_size = get_tensor_model_parallel_world_size()
        self.tp_rank = get_tensor_model_parallel_rank()
95
        self.total_num_heads = num_heads
96
97
        assert self.total_num_heads % self.tp_size == 0
        self.num_heads = self.total_num_heads // self.tp_size
98
        self.total_num_kv_heads = num_kv_heads
99
        if self.total_num_kv_heads >= self.tp_size:
100
101
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
102
            assert self.total_num_kv_heads % self.tp_size == 0
103
104
105
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
106
107
            assert self.tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size)
108
109
110
        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
111
        self.key_value_groups = int(self.num_heads / self.num_kv_heads)
112
113
114
115
116
117
118
119
120
121
        self.scaling = self.head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

        self.wqkv = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
122
            quant_config=quant_config,
123
            prefix=f"{prefix}.wqkv",
124
125
126
127
128
        )
        self.wo = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
129
            quant_config=quant_config,
130
            prefix=f"{prefix}.wo",
131
132
133
134
135
136
137
138
139
        )

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position_embeddings,
            base=rope_theta,
            rope_scaling=rope_scaling,
        )
140
141
142
143
144
145
146
147
148
        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",
        )
149

150
    def split_qkv(self, qkv: torch.Tensor):
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
        seq_len = qkv.shape[0]
        if self.tp_size > 1:
            qkv_map = [self.q_size, self.kv_size, self.kv_size] * self.tp_size
            qkv = tensor_model_parallel_all_gather(qkv)
            qkv = torch.split(qkv, qkv_map, dim=-1)
            qkv = qkv[::3] + qkv[1::3] + qkv[2::3]
            qkv = torch.cat(qkv, dim=-1)

        qkv = qkv.view(seq_len, self.total_num_kv_heads,
                       self.key_value_groups + 2, self.head_dim)
        q, k, v = torch.split(qkv, [self.key_value_groups, 1, 1], dim=-2)
        q = q.reshape(seq_len, self.q_size * self.tp_size)
        k = k.reshape(seq_len, self.kv_size * self.tp_size)
        v = v.reshape(seq_len, self.kv_size * self.tp_size)

        if self.tp_size > 1:
            splitter = partial(split_tensor_along_last_dim,
                               num_partitions=self.tp_size)
            q = splitter(q)[self.tp_rank]
            k = splitter(k)[self.tp_rank]
            v = splitter(v)[self.tp_rank]
172
173
        return q, k, v

174
175
176
177
178
179
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.wqkv(hidden_states)
180
        q, k, v = self.split_qkv(qkv)
181
        q, k = self.rotary_emb(positions, q, k)
182
        attn_output = self.attn(q, k, v)
183
184
185
186
187
188
189
190
191
        output, _ = self.wo(attn_output)
        return output


class InternLMDecoderLayer(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
192
        cache_config: Optional[CacheConfig] = None,
193
        quant_config: Optional[QuantizationConfig] = None,
194
        prefix: str = "",
195
196
197
198
199
200
201
202
203
204
205
206
207
208
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
        self.attention = InternLM2Attention(
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
209
            cache_config=cache_config,
210
            quant_config=quant_config,
211
            prefix=f"{prefix}.attention",
212
213
214
215
216
        )
        self.feed_forward = InternLM2MLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
217
            quant_config=quant_config,
218
            prefix=f"{prefix}.feed_forward",
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
        )
        self.attention_norm = RMSNorm(config.hidden_size,
                                      eps=config.rms_norm_eps)
        self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.attention_norm(hidden_states)
        else:
            hidden_states, residual = self.attention_norm(
                hidden_states, residual)
        hidden_states = self.attention(
            positions=positions,
            hidden_states=hidden_states,
        )

        # Fully Connected
        hidden_states, residual = self.ffn_norm(hidden_states, residual)
        hidden_states = self.feed_forward(hidden_states)
        return hidden_states, residual


248
@support_torch_compile
249
class InternLM2Model(nn.Module):
Fengzhe Zhou's avatar
Fengzhe Zhou committed
250

251
252
253
254
255
256
    def __init__(
            self,
            *,
            vllm_config: VllmConfig,
            prefix: str = "",
            layer_type: Type[InternLMDecoderLayer] = InternLMDecoderLayer):
257
        super().__init__()
258
259
260
261
262

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

263
264
265
266
267
268
269
        self.config = config
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.tok_embeddings = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
270
271
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
272
            lambda prefix: layer_type(
273
                config, cache_config, quant_config, prefix=prefix),
274
            prefix=f"{prefix}.layers")
275
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
276
277
278
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
279

280
281
282
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.tok_embeddings(input_ids)

283
284
285
286
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
287
        intermediate_tensors: Optional[IntermediateTensors] = None,
288
        inputs_embeds: Optional[torch.Tensor] = None,
289
290
291
292
293
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
294
                hidden_states = self.get_input_embeddings(input_ids)
295
            residual = None
296
        else:
297
298
299
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
300
301
        for layer in self.layers[self.start_layer:self.end_layer]:
            hidden_states, residual = layer(positions, hidden_states, residual)
302
303
304
305
306
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })
307
308
309
310
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


311
312
313
314
315
316
class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
    packed_modules_mapping = {
        "wqkv": ["wqkv"],
        "gate_up_proj": ["w1", "w3"],
    }

317
318
319
320
321
    def __init__(self,
                 *,
                 vllm_config: VllmConfig,
                 prefix: str = "",
                 model_type: Type[InternLM2Model] = InternLM2Model):
322
        super().__init__()
323
324
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
325
326
        lora_config = vllm_config.lora_config

327
        self.config = config
328
        self.quant_config = quant_config
329
330
        self.lora_config = lora_config

331
332
        self.model = model_type(vllm_config=vllm_config,
                                prefix=maybe_prefix(prefix, "model"))
333
334
        self.output = ParallelLMHead(config.vocab_size,
                                     config.hidden_size,
335
336
                                     quant_config=quant_config,
                                     prefix=maybe_prefix(prefix, "output"))
337
338
        if self.config.tie_word_embeddings:
            self.output.weight = self.model.tok_embeddings.weight
339
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
340
        self.sampler = get_sampler()
341
342
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
343

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

347
348
349
350
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
351
        intermediate_tensors: Optional[IntermediateTensors],
352
        inputs_embeds: Optional[torch.Tensor] = None,
353
    ) -> torch.Tensor:
354
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
355
                                   inputs_embeds)
356
357
        return hidden_states

358
359
360
361
362
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
363
        logits = self.logits_processor(self.output, hidden_states,
364
365
366
                                       sampling_metadata)
        return logits

367
368
    def sample(
        self,
369
        logits: torch.Tensor,
370
371
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
372
        next_tokens = self.sampler(logits, sampling_metadata)
373
        return next_tokens
Fengzhe Zhou's avatar
Fengzhe Zhou committed
374

375
376
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
Fengzhe Zhou's avatar
Fengzhe Zhou committed
377
378
379
380
381
382
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "w1", 0),
            ("gate_up_proj", "w3", 1),
        ]
        params_dict = dict(self.named_parameters())
383
        loaded_params: Set[str] = set()
384
        for name, loaded_weight in weights:
Fengzhe Zhou's avatar
Fengzhe Zhou committed
385
386
387
388
389
390
391
392
393
            if "rotary_emb.inv_freq" in name:
                continue
            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
394
395
                if is_pp_missing_parameter(name, self):
                    continue
Fengzhe Zhou's avatar
Fengzhe Zhou committed
396
397
398
399
400
401
402
403
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
404
405
                if is_pp_missing_parameter(name, self):
                    continue
Fengzhe Zhou's avatar
Fengzhe Zhou committed
406
                param = params_dict[name]
407
408
409
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
410
411
            loaded_params.add(name)
        return loaded_params
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453


class InternLM2ForRewardModel(InternLM2ForCausalLM):

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        model_type: Type[InternLM2Model] = InternLM2Model,
    ):
        super().__init__(vllm_config=vllm_config,
                         prefix=prefix,
                         model_type=model_type)

        for attr in ("output", "logits_processor", "sampler"):
            delattr(self, attr)

        config = vllm_config.model_config.hf_config
        self.v_head = RowParallelLinear(
            config.hidden_size,
            1,
            bias=False,
            input_is_parallel=False,
            prefix=maybe_prefix(prefix, "v_head"),
        )

        pooler_config = vllm_config.model_config.pooler_config
        self._pooler = Pooler.from_config_with_defaults(
            pooler_config,
            pooling_type=PoolingType.ALL,
            normalize=False,
            softmax=False,
        )

    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]:
454
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
455
456
457
458
459
460
461
462
463
464
                                   inputs_embeds)
        logits, _ = self.v_head(hidden_states)
        return logits

    def pooler(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Optional[PoolerOutput]:
        return self._pooler(hidden_states, pooling_metadata)