dbrx.py 18 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
from collections.abc import Iterable
from typing import Optional, Union
6
7
8

import torch
import torch.nn as nn
9
from transformers import PretrainedConfig
10

11
from vllm.attention import Attention
12
from vllm.config import CacheConfig, VllmConfig
13
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
14
15
                              get_tensor_model_parallel_world_size)
from vllm.model_executor.layers.fused_moe import FusedMoE
16
from vllm.model_executor.layers.linear import (QKVParallelLinear,
17
18
19
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
20
from vllm.model_executor.layers.quantization import QuantizationConfig
21
22
23
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
24
25
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
26
from vllm.model_executor.sampling_metadata import SamplingMetadata
27
from vllm.sequence import IntermediateTensors
28

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

34
35
36
37
38
39
40
41

class DbrxRouter(nn.Module):
    """A Router implementation for DBRX that returns logits for each expert
    per token.
    """

    def __init__(
        self,
42
        config: PretrainedConfig,
43
44
45
46
47
48
49
50
51
52
53
        params_dtype: Optional[torch.dtype] = None,
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.num_total_experts = config.ffn_config.moe_num_experts
        self.d_model = config.d_model
        self.layer = ReplicatedLinear(
            self.d_model,
            self.num_total_experts,
            bias=False,
            params_dtype=params_dtype,
54
            quant_config=None,
55
56
57
58
59
60
61
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        router_logits, _ = self.layer(hidden_states)
        return router_logits


62
class DbrxExperts(FusedMoE):
63
64
65

    def __init__(
        self,
66
        config: PretrainedConfig,
67
        quant_config: Optional[QuantizationConfig] = None,
68
        params_dtype: Optional[torch.dtype] = None,
69
        prefix: str = "",
70
    ):
71
72
73
74
75
76
77
78
79
80
        super().__init__(
            num_experts=config.ffn_config.moe_num_experts,
            top_k=config.ffn_config.moe_top_k,
            hidden_size=config.d_model,
            intermediate_size=config.ffn_config.ffn_hidden_size,
            params_dtype=params_dtype,
            reduce_results=True,
            renormalize=True,
            quant_config=quant_config,
            tp_size=get_tensor_model_parallel_world_size(),
81
            prefix=prefix,
82
83
        )
        self.config = config
84
        self.d_model = config.d_model
85
        self.intermediate_size = (self.config.ffn_config.ffn_hidden_size //
86
87
                                  self.tp_size)

88
    # Define custom weight loader for dbrx model
89
    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor,
90
                      weight_name: str, param_name: str):
91
92
93
94
95
96
97
        tp_rank = get_tensor_model_parallel_rank()
        param_data = param.data
        shard_size = self.intermediate_size
        shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
        # DBRX uses GLU for each experts.
        # GLU has 3 linear layers: w1, v1 and w2.
        if weight_name.endswith("w1"):
98
99
100
101
102
103
104
105
106
107
            if param_name.endswith("weight"):
                loaded_weight = torch.reshape(
                    loaded_weight,
                    [-1, self.intermediate_size * self.tp_size, self.d_model],
                )
                param_data[:, 0:shard_size, :] = loaded_weight[:, shard, :]
            elif param_name.endswith("weight_scale"):
                param_data[:, 0] = loaded_weight
            else:
                param_data = loaded_weight
108
        if weight_name.endswith("v1"):
109
110
111
112
113
114
115
116
117
118
119
            if param_name.endswith("weight"):
                loaded_weight = torch.reshape(
                    loaded_weight,
                    [-1, self.intermediate_size * self.tp_size, self.d_model],
                )
                param_data[:, shard_size:2 *
                           shard_size, :] = loaded_weight[:, shard, :]
            elif param_name.endswith("weight_scale"):
                param_data[:, 1] = loaded_weight
            else:
                param_data[:] = loaded_weight
120
        if weight_name.endswith("w2"):
121
122
123
124
125
126
127
128
            if param_name.endswith("weight"):
                loaded_weight = torch.reshape(
                    loaded_weight,
                    [-1, self.intermediate_size * self.tp_size, self.d_model],
                ).transpose(1, 2)
                param_data[:] = loaded_weight[:, :, shard]
            else:
                param_data[:] = loaded_weight
129

130
131
132
133
134
135
136
137
138
139
140

class DbrxMoE(nn.Module):
    """A tensor-parallel MoE implementation for DBRX.

    Each expert's weights are sharded across all ranks and a fused MoE
    kernel is used for the forward pass, and finally we reduce the outputs
    across ranks.
    """

    def __init__(
        self,
141
        config: PretrainedConfig,
142
143
        quant_config: Optional[QuantizationConfig] = None,
        params_dtype: Optional[torch.dtype] = None,
144
        prefix: str = "",
145
146
147
148
149
150
151
152
153
154
155
    ):
        super().__init__()
        self.d_model = config.d_model
        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        self.params_dtype = params_dtype

        self.router = DbrxRouter(config, self.params_dtype)

        self.experts = DbrxExperts(config=config,
                                   quant_config=quant_config,
156
157
                                   params_dtype=self.params_dtype,
                                   prefix=f"{prefix}.experts")
158

159
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
160
        orig_shape = hidden_states.shape
161
162
163
        hidden_states = hidden_states.view(-1, self.d_model)
        # router_logits: (num_tokens, n_experts)
        router_logits = self.router(hidden_states)
164
165
        final_hidden_states = self.experts(hidden_states, router_logits)
        return final_hidden_states.view(orig_shape)
166
167
168
169
170
171


class DbrxAttention(nn.Module):

    def __init__(
        self,
172
        config: PretrainedConfig,
173
        cache_config: Optional[CacheConfig] = None,
174
        quant_config: Optional[QuantizationConfig] = None,
175
        prefix: str = "",
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
    ):
        super().__init__()
        self.d_model = config.d_model
        self.total_num_heads = config.n_heads
        self.head_dim = self.d_model // self.total_num_heads
        self.total_num_kv_heads = config.attn_config.kv_n_heads
        self.clip_qkv = config.attn_config.clip_qkv
        self.rope_theta = config.attn_config.rope_theta
        self.max_position = config.max_seq_len

        # pylint: disable=invalid-name
        self.Wqkv = QKVParallelLinear(
            self.d_model,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
193
            quant_config=quant_config,
194
195
196
197
198
        )
        self.out_proj = RowParallelLinear(
            self.d_model,
            self.d_model,
            bias=False,
199
            quant_config=quant_config,
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
        )
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=self.max_position,
            base=int(self.rope_theta),
            is_neox_style=True,
        )

        tp_world_size = get_tensor_model_parallel_world_size()
        self.tp_size = tp_world_size
        assert self.total_num_heads % tp_world_size == 0
        self.num_heads = self.total_num_heads // tp_world_size
        if self.total_num_kv_heads >= tp_world_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_world_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_world_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_world_size)
        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
225
226
227
228
229
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
                              cache_config=cache_config,
230
231
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
232
233
234
235
236
237
238
239
240
241
242

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.Wqkv(hidden_states)
        if self.clip_qkv is not None:
            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(position_ids, q, k)
243
        attn_output = self.attn(q, k, v)
244
245
246
247
248
249
250
251
        hidden_states, _ = self.out_proj(attn_output)
        return hidden_states


class DbrxFusedNormAttention(nn.Module):

    def __init__(
        self,
252
        config: PretrainedConfig,
253
        cache_config: Optional[CacheConfig] = None,
254
        quant_config: Optional[QuantizationConfig] = None,
255
        prefix: str = "",
256
257
258
    ):
        super().__init__()
        self.d_model = config.d_model
259
260
261
262
        self.attn = DbrxAttention(config,
                                  cache_config,
                                  quant_config,
                                  prefix=f"{prefix}.attn")
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        self.norm_1 = nn.LayerNorm(self.d_model)
        self.norm_2 = nn.LayerNorm(self.d_model)

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.norm_1(hidden_states)
        x = self.attn(
            position_ids=position_ids,
            hidden_states=hidden_states,
        )
        hidden_states = residual + x
        residual = hidden_states
        hidden_states = self.norm_2(hidden_states)
        return hidden_states, residual


class DbrxBlock(nn.Module):

    def __init__(
        self,
287
        config: PretrainedConfig,
288
        cache_config: Optional[CacheConfig] = None,
289
        quant_config: Optional[QuantizationConfig] = None,
290
        prefix: str = "",
291
292
    ):
        super().__init__()
293
294
295
296
297
        self.norm_attn_norm = DbrxFusedNormAttention(
            config,
            cache_config,
            quant_config,
            prefix=f"{prefix}.norm_attn_norm")
298
        self.ffn = DbrxMoE(config, quant_config, prefix=f"{prefix}.ffn")
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states, residual = self.norm_attn_norm(
            position_ids=position_ids,
            hidden_states=hidden_states,
        )
        hidden_states = self.ffn(hidden_states)
        hidden_states = hidden_states + residual
        return hidden_states


class DbrxModel(nn.Module):

316
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
317
        super().__init__()
318
319
320
321
322

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

323
        self.quant_config = quant_config
324
325
326
327
        self.wte = VocabParallelEmbedding(
            config.vocab_size,
            config.d_model,
        )
328
329
        self.start_layer, self.end_layer, self.blocks = make_layers(
            config.n_layers,
330
331
            lambda prefix: DbrxBlock(
                config, cache_config, quant_config, prefix=prefix),
332
333
            prefix=f"{prefix}.blocks",
        )
334
335
336
337
338
339
        self.norm_f = nn.LayerNorm(config.d_model, eps=1e-5)
        for module in self.modules():
            if hasattr(module, "bias") and isinstance(module.bias,
                                                      nn.Parameter):
                # Remove the bias term in Linear and LayerNorm.
                module.register_parameter("bias", None)
340
341
342
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.d_model))
343

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

347
348
349
350
    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
351
        intermediate_tensors: Optional[IntermediateTensors],
352
        inputs_embeds: Optional[torch.Tensor] = None,
353
354
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
355
356
357
358
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
359
360
361
        else:
            assert intermediate_tensors
            hidden_states = intermediate_tensors["hidden_states"]
362
363
        for block in self.blocks[self.start_layer:self.end_layer]:
            hidden_states = block(position_ids, hidden_states)
364
365
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
366
367
368
        hidden_states = self.norm_f(hidden_states)
        return hidden_states

369
370
371
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
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
        expert_params_mapping = [(
            "w13" if weight_name in ["w1", "v1"] else "w2",
            f"mlp.{weight_name}",
        ) for weight_name in ["w1", "v1", "w2"]]
        params_dict = dict(self.named_parameters(remove_duplicate=False))
        loaded_params: set[str] = set()

        for name, loaded_weight in weights:
            if (self.quant_config is not None and
                (scale_name := self.quant_config.get_cache_scale(name))):
                # Loading kv cache quantization scales
                param = params_dict[scale_name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
                                 loaded_weight[0])
                weight_loader(param, loaded_weight)
                loaded_params.add(scale_name)
                continue

            if name.endswith(("w1", "w2", "v1")):
                name = name + "_weight"
            for param_name, weight_name in expert_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, weight_name, name)
                break

            else:
                if is_pp_missing_parameter(name, self):
                    continue
                # Remapping the name of FP8 kv-scale.
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    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

418

419
class DbrxForCausalLM(nn.Module, SupportsPP):
420

421
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
422
        super().__init__()
423
424
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
425
        self.config = config
426
427
428
        if config.tie_word_embeddings:
            raise ValueError(
                "tie_word_embeddings is not supported for Dbrx models.")
429
        self.quant_config = quant_config
430
        self.unpadded_vocab_size = config.vocab_size
431
432
433
        self.transformer = DbrxModel(vllm_config=vllm_config,
                                     prefix=maybe_prefix(
                                         prefix, "transformer"))
434
435
436
437
438
        self.lm_head = ParallelLMHead(
            config.vocab_size,
            config.d_model,
            org_num_embeddings=config.vocab_size,
            padding_size=DEFAULT_VOCAB_PADDING_SIZE,
439
            quant_config=quant_config,
440
441
442
        )
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size)
443
444
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
445

446
447
448
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.transformer.get_input_embeddings(input_ids)

449
450
451
452
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
453
        intermediate_tensors: Optional[IntermediateTensors] = None,
454
        inputs_embeds: Optional[torch.Tensor] = None,
455
    ) -> Union[torch.Tensor, IntermediateTensors]:
456
457
        hidden_states = self.transformer(input_ids, positions,
                                         intermediate_tensors, inputs_embeds)
458
459
        return hidden_states

460
461
462
463
464
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
465
        logits = self.logits_processor(self.lm_head, hidden_states,
466
467
468
                                       sampling_metadata)
        return logits

469
470
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
471
472
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)