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

3
from typing import Iterable, Optional, Set, Tuple, Union
4
5
6
7

import torch
import torch.nn as nn

8
from vllm.attention import Attention
9
from vllm.config import CacheConfig, VllmConfig
10
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
11
12
                              get_tensor_model_parallel_world_size)
from vllm.model_executor.layers.fused_moe import FusedMoE
13
from vllm.model_executor.layers.linear import (QKVParallelLinear,
14
15
16
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
17
from vllm.model_executor.layers.quantization import QuantizationConfig
18
from vllm.model_executor.layers.rotary_embedding import get_rope
Joe Runde's avatar
Joe Runde committed
19
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
20
21
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
29
from .interfaces import SupportsPP
from .utils import (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
68
        params_dtype: Optional[torch.dtype] = None,
    ):
69
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(),
        )
        self.config = config
81
82
        self.tp_size = get_tensor_model_parallel_world_size()
        self.d_model = config.d_model
83
        self.intermediate_size = (self.config.ffn_config.ffn_hidden_size //
84
85
                                  self.tp_size)

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

128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

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,
    ):
        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,
                                   params_dtype=self.params_dtype)

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


class DbrxAttention(nn.Module):

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

    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)
239
        attn_output = self.attn(q, k, v)
240
241
242
243
244
245
246
247
248
        hidden_states, _ = self.out_proj(attn_output)
        return hidden_states


class DbrxFusedNormAttention(nn.Module):

    def __init__(
        self,
        config: DbrxConfig,
249
        cache_config: Optional[CacheConfig] = None,
250
        quant_config: Optional[QuantizationConfig] = None,
251
        prefix: str = "",
252
253
254
    ):
        super().__init__()
        self.d_model = config.d_model
255
256
257
258
        self.attn = DbrxAttention(config,
                                  cache_config,
                                  quant_config,
                                  prefix=f"{prefix}.attn")
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
        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,
284
        cache_config: Optional[CacheConfig] = None,
285
        quant_config: Optional[QuantizationConfig] = None,
286
        prefix: str = "",
287
288
    ):
        super().__init__()
289
290
291
292
293
        self.norm_attn_norm = DbrxFusedNormAttention(
            config,
            cache_config,
            quant_config,
            prefix=f"{prefix}.norm_attn_norm")
294
        self.ffn = DbrxMoE(config, quant_config)
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311

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

312
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
313
        super().__init__()
314
315
316
317
318

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

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

339
340
341
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.wte(input_ids)

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


365
class DbrxForCausalLM(nn.Module, SupportsPP):
366

367
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
368
        super().__init__()
369
370
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
371
        self.config = config
372
373
374
        if config.tie_word_embeddings:
            raise ValueError(
                "tie_word_embeddings is not supported for Dbrx models.")
375
        self.quant_config = quant_config
376
        self.unpadded_vocab_size = config.vocab_size
377
378
379
        self.transformer = DbrxModel(vllm_config=vllm_config,
                                     prefix=maybe_prefix(
                                         prefix, "transformer"))
380
381
382
383
384
        self.lm_head = ParallelLMHead(
            config.vocab_size,
            config.d_model,
            org_num_embeddings=config.vocab_size,
            padding_size=DEFAULT_VOCAB_PADDING_SIZE,
385
            quant_config=quant_config,
386
387
388
        )
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size)
Joe Runde's avatar
Joe Runde committed
389
        self.sampler = get_sampler()
390
391
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
392

393
394
395
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.transformer.get_input_embeddings(input_ids)

396
397
398
399
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
400
        intermediate_tensors: Optional[IntermediateTensors] = None,
401
        inputs_embeds: Optional[torch.Tensor] = None,
402
    ) -> Union[torch.Tensor, IntermediateTensors]:
403
404
        hidden_states = self.transformer(input_ids, positions,
                                         intermediate_tensors, inputs_embeds)
405
406
        return hidden_states

407
408
409
410
411
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
412
        logits = self.logits_processor(self.lm_head, hidden_states,
413
414
415
416
417
418
419
420
421
422
423
                                       sampling_metadata)
        return logits

    def sample(
        self,
        logits: Optional[torch.Tensor],
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        next_tokens = self.sampler(logits, sampling_metadata)
        return next_tokens

424
425
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
426
        expert_params_mapping = [(
427
            "w13" if weight_name in ["w1", "v1"] else "w2",
428
            f"mlp.{weight_name}",
429
430
        ) for weight_name in ["w1", "v1", "w2"]]
        params_dict = dict(self.named_parameters(remove_duplicate=False))
431
        loaded_params: Set[str] = set()
432

433
        for name, loaded_weight in weights:
434
435
            if (self.quant_config is not None and
                (scale_name := self.quant_config.get_cache_scale(name))):
436
                # Loading kv cache quantization scales
437
438
439
440
441
442
443
444
445
446
447
                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"
448
449
450
451
            for param_name, weight_name in expert_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
452
453
                if is_pp_missing_parameter(name, self):
                    continue
454
455
                param = params_dict[name]
                weight_loader = param.weight_loader
456
                weight_loader(param, loaded_weight, weight_name, name)
457
                break
458

459
            else:
460
461
462
463
464
                # Remapping the name of FP8 kv-scale.
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue

465
466
                if is_pp_missing_parameter(name, self):
                    continue
467
468
469
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
470
471
472
473
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
474
475
            loaded_params.add(name)
        return loaded_params