dbrx.py 17.9 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

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

import torch
import torch.nn as nn

9
from vllm.attention import Attention
10
from vllm.config import CacheConfig, VllmConfig
11
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
12
13
                              get_tensor_model_parallel_world_size)
from vllm.model_executor.layers.fused_moe import FusedMoE
14
from vllm.model_executor.layers.linear import (QKVParallelLinear,
15
16
17
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
18
from vllm.model_executor.layers.quantization import QuantizationConfig
19
20
21
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)
22
23
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
24
from vllm.model_executor.sampling_metadata import SamplingMetadata
25
from vllm.sequence import IntermediateTensors
26
27
from vllm.transformers_utils.configs.dbrx import DbrxConfig

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

33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

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

    def __init__(
        self,
        config: DbrxConfig,
        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,
53
            quant_config=None,
54
55
56
57
58
59
60
        )

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


61
class DbrxExperts(FusedMoE):
62
63
64
65

    def __init__(
        self,
        config: DbrxConfig,
66
        quant_config: Optional[QuantizationConfig] = None,
67
        params_dtype: Optional[torch.dtype] = None,
68
        prefix: str = "",
69
    ):
70
71
72
73
74
75
76
77
78
79
        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(),
80
            prefix=prefix,
81
82
        )
        self.config = config
83
        self.d_model = config.d_model
84
        self.intermediate_size = (self.config.ffn_config.ffn_hidden_size //
85
86
                                  self.tp_size)

87
    # Define custom weight loader for dbrx model
88
    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor,
89
                      weight_name: str, param_name: str):
90
91
92
93
94
95
96
        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"):
97
98
99
100
101
102
103
104
105
106
            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
107
        if weight_name.endswith("v1"):
108
109
110
111
112
113
114
115
116
117
118
            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
119
        if weight_name.endswith("w2"):
120
121
122
123
124
125
126
127
            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
128

129
130
131
132
133
134
135
136
137
138
139
140
141
142

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,
        config: DbrxConfig,
        quant_config: Optional[QuantizationConfig] = None,
        params_dtype: Optional[torch.dtype] = None,
143
        prefix: str = "",
144
145
146
147
148
149
150
151
152
153
154
    ):
        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,
155
156
                                   params_dtype=self.params_dtype,
                                   prefix=f"{prefix}.experts")
157

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


class DbrxAttention(nn.Module):

    def __init__(
        self,
        config: DbrxConfig,
172
        cache_config: Optional[CacheConfig] = None,
173
        quant_config: Optional[QuantizationConfig] = None,
174
        prefix: str = "",
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
    ):
        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,
192
            quant_config=quant_config,
193
194
195
196
197
        )
        self.out_proj = RowParallelLinear(
            self.d_model,
            self.d_model,
            bias=False,
198
            quant_config=quant_config,
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
        )
        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
224
225
226
227
228
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
                              cache_config=cache_config,
229
230
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
231
232
233
234
235
236
237
238
239
240
241

    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)
242
        attn_output = self.attn(q, k, v)
243
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,
        config: DbrxConfig,
252
        cache_config: Optional[CacheConfig] = None,
253
        quant_config: Optional[QuantizationConfig] = None,
254
        prefix: str = "",
255
256
257
    ):
        super().__init__()
        self.d_model = config.d_model
258
259
260
261
        self.attn = DbrxAttention(config,
                                  cache_config,
                                  quant_config,
                                  prefix=f"{prefix}.attn")
262
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,
        config: DbrxConfig,
287
        cache_config: Optional[CacheConfig] = None,
288
        quant_config: Optional[QuantizationConfig] = None,
289
        prefix: str = "",
290
291
    ):
        super().__init__()
292
293
294
295
296
        self.norm_attn_norm = DbrxFusedNormAttention(
            config,
            cache_config,
            quant_config,
            prefix=f"{prefix}.norm_attn_norm")
297
        self.ffn = DbrxMoE(config, quant_config, prefix=f"{prefix}.ffn")
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314

    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):

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

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

322
        self.quant_config = quant_config
323
324
325
326
        self.wte = VocabParallelEmbedding(
            config.vocab_size,
            config.d_model,
        )
327
328
        self.start_layer, self.end_layer, self.blocks = make_layers(
            config.n_layers,
329
330
            lambda prefix: DbrxBlock(
                config, cache_config, quant_config, prefix=prefix),
331
332
            prefix=f"{prefix}.blocks",
        )
333
334
335
336
337
338
        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)
339
340
341
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.d_model))
342

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

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

368
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
    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

417

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

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

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

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

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

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