| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Iterable |
| 5 | from itertools import islice |
| 6 | from typing import Any |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | from transformers import DbrxConfig |
| 11 | |
| 12 | from vllm.config import CacheConfig, VllmConfig |
| 13 | from vllm.distributed import ( |
| 14 | get_pp_group, |
| 15 | get_tensor_model_parallel_rank, |
| 16 | get_tensor_model_parallel_world_size, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.attention import Attention |
| 19 | from vllm.model_executor.layers.fused_moe import ( |
| 20 | FusedMoEFactory, |
| 21 | GateLinear, |
| 22 | RoutedExperts, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.linear import ( |
| 25 | QKVParallelLinear, |
| 26 | RowParallelLinear, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 29 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 30 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 31 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 32 | ParallelLMHead, |
| 33 | VocabParallelEmbedding, |
| 34 | ) |
| 35 | from vllm.model_executor.model_loader.weight_utils import ( |
| 36 | default_weight_loader, |
| 37 | maybe_remap_kv_scale_name, |
| 38 | ) |
| 39 | from vllm.sequence import IntermediateTensors |
| 40 | |
| 41 | from .interfaces import SupportsPP |
| 42 | from .utils import ( |
| 43 | AutoWeightsLoader, |
| 44 | is_pp_missing_parameter, |
| 45 | make_empty_intermediate_tensors_factory, |
| 46 | make_layers, |
| 47 | maybe_prefix, |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | class DbrxRouter(nn.Module): |
| 52 | """A Router implementation for DBRX that returns logits for each expert |
| 53 | per token. |
| 54 | """ |
| 55 | |
| 56 | def __init__( |
| 57 | self, |
| 58 | config: DbrxConfig, |
| 59 | params_dtype: torch.dtype | None = None, |
| 60 | ): |
| 61 | super().__init__() |
| 62 | self.tp_size = get_tensor_model_parallel_world_size() |
| 63 | self.num_total_experts = config.ffn_config.moe_num_experts |
| 64 | self.d_model = config.d_model |
| 65 | self.layer = GateLinear( |
| 66 | self.d_model, |
| 67 | self.num_total_experts, |
| 68 | params_dtype=params_dtype, |
| 69 | ) |
| 70 | |
| 71 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 72 | router_logits, _ = self.layer(hidden_states) |
| 73 | return router_logits |
| 74 | |
| 75 | |
| 76 | class DbrxExperts(RoutedExperts): |
| 77 | def __init__( |
| 78 | self, |
| 79 | *args, |
| 80 | config: DbrxConfig, |
| 81 | **kwargs, |
| 82 | ): |
| 83 | super().__init__(*args, **kwargs) |
| 84 | self.config = config |
| 85 | self.d_model = config.d_model |
| 86 | self.tp_size = self.moe_config.tp_size |
| 87 | self.intermediate_size = self.config.ffn_config.ffn_hidden_size // self.tp_size |
| 88 | |
| 89 | # Define custom weight loader for dbrx model |
| 90 | def weight_loader( |
| 91 | self, |
| 92 | param: nn.Parameter, |
| 93 | loaded_weight: torch.Tensor, |
| 94 | weight_name: str, |
| 95 | shard_id: str, |
| 96 | expert_id: int | None = None, |
| 97 | return_success: Any = False, |
| 98 | ) -> Any: |
| 99 | param_name = shard_id |
| 100 | tp_rank = get_tensor_model_parallel_rank() |
| 101 | param_data = param.data |
| 102 | shard_size = self.intermediate_size |
| 103 | shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) |
| 104 | # DBRX uses GLU for each experts. |
| 105 | # GLU has 3 linear layers: w1, v1 and w2. |
| 106 | if weight_name.endswith("w1"): |
| 107 | if param_name.endswith("weight"): |
| 108 | loaded_weight = torch.reshape( |
| 109 | loaded_weight, |
| 110 | [-1, self.intermediate_size * self.tp_size, self.d_model], |
| 111 | ) |
| 112 | param_data[:, 0:shard_size, :] = loaded_weight[:, shard, :] |
| 113 | elif param_name.endswith("weight_scale"): |
| 114 | param_data[:, 0] = loaded_weight |
| 115 | else: |
| 116 | param_data = loaded_weight |
| 117 | if weight_name.endswith("v1"): |
| 118 | if param_name.endswith("weight"): |
| 119 | loaded_weight = torch.reshape( |
| 120 | loaded_weight, |
| 121 | [-1, self.intermediate_size * self.tp_size, self.d_model], |
| 122 | ) |
| 123 | param_data[:, shard_size : 2 * shard_size, :] = loaded_weight[ |
| 124 | :, shard, : |
| 125 | ] |
| 126 | elif param_name.endswith("weight_scale"): |
| 127 | param_data[:, 1] = loaded_weight |
| 128 | else: |
| 129 | param_data[:] = loaded_weight |
| 130 | if weight_name.endswith("w2"): |
| 131 | if param_name.endswith("weight"): |
| 132 | loaded_weight = torch.reshape( |
| 133 | loaded_weight, |
| 134 | [-1, self.intermediate_size * self.tp_size, self.d_model], |
| 135 | ).transpose(1, 2) |
| 136 | param_data[:] = loaded_weight[:, :, shard] |
| 137 | else: |
| 138 | param_data[:] = loaded_weight |
| 139 | return True if return_success else None |
| 140 | |
| 141 | |
| 142 | class DbrxMoE(nn.Module): |
| 143 | """A tensor-parallel MoE implementation for DBRX. |
| 144 | |
| 145 | Each expert's weights are sharded across all ranks and a fused MoE |
| 146 | kernel is used for the forward pass, and finally we reduce the outputs |
| 147 | across ranks. |
| 148 | """ |
| 149 | |
| 150 | def __init__( |
| 151 | self, |
| 152 | config: DbrxConfig, |
| 153 | quant_config: QuantizationConfig | None = None, |
| 154 | params_dtype: torch.dtype | None = None, |
| 155 | prefix: str = "", |
| 156 | ): |
| 157 | super().__init__() |
| 158 | self.d_model = config.d_model |
| 159 | if params_dtype is None: |
| 160 | params_dtype = torch.get_default_dtype() |
| 161 | self.params_dtype = params_dtype |
| 162 | |
| 163 | self.router = DbrxRouter(config, self.params_dtype) |
| 164 | |
| 165 | self.experts = FusedMoEFactory( |
| 166 | num_experts=config.ffn_config.moe_num_experts, |
| 167 | top_k=config.ffn_config.moe_top_k, |
| 168 | hidden_size=config.d_model, |
| 169 | intermediate_size=config.ffn_config.ffn_hidden_size, |
| 170 | params_dtype=self.params_dtype, |
| 171 | renormalize=True, |
| 172 | quant_config=quant_config, |
| 173 | tp_size=get_tensor_model_parallel_world_size(), |
| 174 | prefix=prefix, |
| 175 | routed_experts_cls=DbrxExperts, |
| 176 | routed_experts_args={"config": config}, |
| 177 | ) |
| 178 | |
| 179 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 180 | orig_shape = hidden_states.shape |
| 181 | hidden_states = hidden_states.view(-1, self.d_model) |
| 182 | # router_logits: (num_tokens, n_experts) |
| 183 | router_logits = self.router(hidden_states) |
| 184 | final_hidden_states = self.experts(hidden_states, router_logits) |
| 185 | return final_hidden_states.view(orig_shape) |
| 186 | |
| 187 | |
| 188 | class DbrxAttention(nn.Module): |
| 189 | def __init__( |
| 190 | self, |
| 191 | config: DbrxConfig, |
| 192 | cache_config: CacheConfig | None = None, |
| 193 | quant_config: QuantizationConfig | None = None, |
| 194 | prefix: str = "", |
| 195 | ): |
| 196 | super().__init__() |
| 197 | self.d_model = config.d_model |
| 198 | self.total_num_heads = config.n_heads |
| 199 | self.head_dim = self.d_model // self.total_num_heads |
| 200 | self.total_num_kv_heads = config.attn_config.kv_n_heads |
| 201 | self.clip_qkv = config.attn_config.clip_qkv |
| 202 | rope_parameters = { |
| 203 | "rope_type": "default", |
| 204 | "rope_theta": int(config.attn_config.rope_theta), |
| 205 | } |
| 206 | self.max_position = config.max_seq_len |
| 207 | |
| 208 | # pylint: disable=invalid-name |
| 209 | self.Wqkv = QKVParallelLinear( |
| 210 | self.d_model, |
| 211 | self.head_dim, |
| 212 | self.total_num_heads, |
| 213 | self.total_num_kv_heads, |
| 214 | bias=False, |
| 215 | quant_config=quant_config, |
| 216 | prefix=f"{prefix}.Wqkv", |
| 217 | ) |
| 218 | self.out_proj = RowParallelLinear( |
| 219 | self.d_model, |
| 220 | self.d_model, |
| 221 | bias=False, |
| 222 | quant_config=quant_config, |
| 223 | prefix=f"{prefix}.out_proj", |
| 224 | ) |
| 225 | self.rotary_emb = get_rope( |
| 226 | self.head_dim, |
| 227 | max_position=self.max_position, |
| 228 | rope_parameters=rope_parameters, |
| 229 | is_neox_style=True, |
| 230 | ) |
| 231 | |
| 232 | tp_world_size = get_tensor_model_parallel_world_size() |
| 233 | self.tp_size = tp_world_size |
| 234 | assert self.total_num_heads % tp_world_size == 0 |
| 235 | self.num_heads = self.total_num_heads // tp_world_size |
| 236 | if self.total_num_kv_heads >= tp_world_size: |
| 237 | # Number of KV heads is greater than TP size, so we partition |
| 238 | # the KV heads across multiple tensor parallel GPUs. |
| 239 | assert self.total_num_kv_heads % tp_world_size == 0 |
| 240 | else: |
| 241 | # Number of KV heads is less than TP size, so we replicate |
| 242 | # the KV heads across multiple tensor parallel GPUs. |
| 243 | assert tp_world_size % self.total_num_kv_heads == 0 |
| 244 | self.num_kv_heads = max(1, self.total_num_kv_heads // tp_world_size) |
| 245 | self.q_size = self.num_heads * self.head_dim |
| 246 | self.kv_size = self.num_kv_heads * self.head_dim |
| 247 | self.scaling = self.head_dim**-0.5 |
| 248 | self.attn = Attention( |
| 249 | self.num_heads, |
| 250 | self.head_dim, |
| 251 | self.scaling, |
| 252 | num_kv_heads=self.num_kv_heads, |
| 253 | cache_config=cache_config, |
| 254 | quant_config=quant_config, |
| 255 | prefix=f"{prefix}.attn", |
| 256 | ) |
| 257 | |
| 258 | def forward( |
| 259 | self, |
| 260 | position_ids: torch.Tensor, |
| 261 | hidden_states: torch.Tensor, |
| 262 | ) -> torch.Tensor: |
| 263 | qkv, _ = self.Wqkv(hidden_states) |
| 264 | if self.clip_qkv is not None: |
| 265 | qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv) |
| 266 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 267 | q, k = self.rotary_emb(position_ids, q, k) |
| 268 | attn_output = self.attn(q, k, v) |
| 269 | hidden_states, _ = self.out_proj(attn_output) |
| 270 | return hidden_states |
| 271 | |
| 272 | |
| 273 | class DbrxFusedNormAttention(nn.Module): |
| 274 | def __init__( |
| 275 | self, |
| 276 | config: DbrxConfig, |
| 277 | cache_config: CacheConfig | None = None, |
| 278 | quant_config: QuantizationConfig | None = None, |
| 279 | prefix: str = "", |
| 280 | ): |
| 281 | super().__init__() |
| 282 | self.d_model = config.d_model |
| 283 | self.attn = DbrxAttention( |
| 284 | config, cache_config, quant_config, prefix=f"{prefix}.attn" |
| 285 | ) |
| 286 | self.norm_1 = nn.LayerNorm(self.d_model) |
| 287 | self.norm_2 = nn.LayerNorm(self.d_model) |
| 288 | |
| 289 | def forward( |
| 290 | self, |
| 291 | position_ids: torch.Tensor, |
| 292 | hidden_states: torch.Tensor, |
| 293 | ) -> torch.Tensor: |
| 294 | residual = hidden_states |
| 295 | hidden_states = self.norm_1(hidden_states) |
| 296 | x = self.attn( |
| 297 | position_ids=position_ids, |
| 298 | hidden_states=hidden_states, |
| 299 | ) |
| 300 | hidden_states = residual + x |
| 301 | residual = hidden_states |
| 302 | hidden_states = self.norm_2(hidden_states) |
| 303 | return hidden_states, residual |
| 304 | |
| 305 | |
| 306 | class DbrxBlock(nn.Module): |
| 307 | def __init__( |
| 308 | self, |
| 309 | config: DbrxConfig, |
| 310 | cache_config: CacheConfig | None = None, |
| 311 | quant_config: QuantizationConfig | None = None, |
| 312 | prefix: str = "", |
| 313 | ): |
| 314 | super().__init__() |
| 315 | self.norm_attn_norm = DbrxFusedNormAttention( |
| 316 | config, cache_config, quant_config, prefix=f"{prefix}.norm_attn_norm" |
| 317 | ) |
| 318 | self.ffn = DbrxMoE(config, quant_config, prefix=f"{prefix}.ffn") |
| 319 | |
| 320 | def forward( |
| 321 | self, |
| 322 | position_ids: torch.Tensor, |
| 323 | hidden_states: torch.Tensor, |
| 324 | ) -> torch.Tensor: |
| 325 | hidden_states, residual = self.norm_attn_norm( |
| 326 | position_ids=position_ids, |
| 327 | hidden_states=hidden_states, |
| 328 | ) |
| 329 | hidden_states = self.ffn(hidden_states) |
| 330 | hidden_states = hidden_states + residual |
| 331 | return hidden_states |
| 332 | |
| 333 | |
| 334 | class DbrxModel(nn.Module): |
| 335 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 336 | super().__init__() |
| 337 | |
| 338 | config = vllm_config.model_config.hf_config |
| 339 | cache_config = vllm_config.cache_config |
| 340 | quant_config = vllm_config.quant_config |
| 341 | |
| 342 | self.quant_config = quant_config |
| 343 | self.wte = VocabParallelEmbedding( |
| 344 | config.vocab_size, |
| 345 | config.d_model, |
| 346 | ) |
| 347 | self.start_layer, self.end_layer, self.blocks = make_layers( |
| 348 | config.n_layers, |
| 349 | lambda prefix: DbrxBlock(config, cache_config, quant_config, prefix=prefix), |
| 350 | prefix=f"{prefix}.blocks", |
| 351 | ) |
| 352 | self.norm_f = nn.LayerNorm(config.d_model, eps=1e-5) |
| 353 | for module in self.modules(): |
| 354 | if hasattr(module, "bias") and isinstance(module.bias, nn.Parameter): |
| 355 | # Remove the bias term in Linear and LayerNorm. |
| 356 | module.register_parameter("bias", None) |
| 357 | self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( |
| 358 | ["hidden_states"], config.d_model |
| 359 | ) |
| 360 | |
| 361 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 362 | return self.wte(input_ids) |
| 363 | |
| 364 | def forward( |
| 365 | self, |
| 366 | input_ids: torch.Tensor | None, |
| 367 | position_ids: torch.Tensor, |
| 368 | intermediate_tensors: IntermediateTensors | None, |
| 369 | inputs_embeds: torch.Tensor | None = None, |
| 370 | ) -> torch.Tensor | IntermediateTensors: |
| 371 | if get_pp_group().is_first_rank: |
| 372 | if inputs_embeds is not None: |
| 373 | hidden_states = inputs_embeds |
| 374 | else: |
| 375 | hidden_states = self.embed_input_ids(input_ids) |
| 376 | else: |
| 377 | assert intermediate_tensors |
| 378 | hidden_states = intermediate_tensors["hidden_states"] |
| 379 | for block in islice(self.blocks, self.start_layer, self.end_layer): |
| 380 | hidden_states = block(position_ids, hidden_states) |
| 381 | if not get_pp_group().is_last_rank: |
| 382 | return IntermediateTensors({"hidden_states": hidden_states}) |
| 383 | hidden_states = self.norm_f(hidden_states) |
| 384 | return hidden_states |
| 385 | |
| 386 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 387 | expert_params_mapping = [ |
| 388 | ( |
| 389 | "w13" if weight_name in ["w1", "v1"] else "w2", |
| 390 | f"mlp.{weight_name}", |
| 391 | ) |
| 392 | for weight_name in ["w1", "v1", "w2"] |
| 393 | ] |
| 394 | params_dict = dict(self.named_parameters(remove_duplicate=False)) |
| 395 | loaded_params: set[str] = set() |
| 396 | |
| 397 | for name, loaded_weight in weights: |
| 398 | if name.endswith(("w1", "w2", "v1")): |
| 399 | name = name + "_weight" |
| 400 | for param_name, weight_name in expert_params_mapping: |
| 401 | if weight_name not in name: |
| 402 | continue |
| 403 | name = name.replace(weight_name, param_name) |
| 404 | if is_pp_missing_parameter(name, self): |
| 405 | continue |
| 406 | param = params_dict[name] |
| 407 | weight_loader = param.weight_loader |
| 408 | weight_loader(param, loaded_weight, weight_name, name) |
| 409 | break |
| 410 | |
| 411 | else: |
| 412 | if is_pp_missing_parameter(name, self): |
| 413 | continue |
| 414 | # Remapping the name of FP8 kv-scale. |
| 415 | remapped_name = maybe_remap_kv_scale_name(name, params_dict) |
| 416 | if remapped_name is None: |
| 417 | continue |
| 418 | name = remapped_name |
| 419 | param = params_dict[name] |
| 420 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 421 | weight_loader(param, loaded_weight) |
| 422 | loaded_params.add(name) |
| 423 | return loaded_params |
| 424 | |
| 425 | |
| 426 | class DbrxForCausalLM(nn.Module, SupportsPP): |
| 427 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 428 | super().__init__() |
| 429 | config = vllm_config.model_config.hf_config |
| 430 | quant_config = vllm_config.quant_config |
| 431 | self.config = config |
| 432 | if config.tie_word_embeddings: |
| 433 | raise ValueError("tie_word_embeddings is not supported for Dbrx models.") |
| 434 | self.quant_config = quant_config |
| 435 | |
| 436 | self.transformer = DbrxModel( |
| 437 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "transformer") |
| 438 | ) |
| 439 | self.lm_head = ParallelLMHead( |
| 440 | config.vocab_size, |
| 441 | config.d_model, |
| 442 | quant_config=quant_config, |
| 443 | prefix=maybe_prefix(prefix, "lm_head"), |
| 444 | ) |
| 445 | self.logits_processor = LogitsProcessor(config.vocab_size) |
| 446 | self.make_empty_intermediate_tensors = ( |
| 447 | self.transformer.make_empty_intermediate_tensors |
| 448 | ) |
| 449 | |
| 450 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 451 | return self.transformer.embed_input_ids(input_ids) |
| 452 | |
| 453 | def forward( |
| 454 | self, |
| 455 | input_ids: torch.Tensor | None, |
| 456 | positions: torch.Tensor, |
| 457 | intermediate_tensors: IntermediateTensors | None = None, |
| 458 | inputs_embeds: torch.Tensor | None = None, |
| 459 | ) -> torch.Tensor | IntermediateTensors: |
| 460 | hidden_states = self.transformer( |
| 461 | input_ids, positions, intermediate_tensors, inputs_embeds |
| 462 | ) |
| 463 | return hidden_states |
| 464 | |
| 465 | def compute_logits( |
| 466 | self, |
| 467 | hidden_states: torch.Tensor, |
| 468 | ) -> torch.Tensor | None: |
| 469 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 470 | return logits |
| 471 | |
| 472 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 473 | loader = AutoWeightsLoader(self) |
| 474 | return loader.load_weights(weights) |