| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2024 The vLLM team. |
| 5 | # |
| 6 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 7 | # you may not use this file except in compliance with the License. |
| 8 | # You may obtain a copy of the License at |
| 9 | # |
| 10 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | # |
| 12 | # Unless required by applicable law or agreed to in writing, software |
| 13 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | # See the License for the specific language governing permissions and |
| 16 | # limitations under the License. |
| 17 | """Transformers modeling backend mixin for Mixture of Experts (MoE) models.""" |
| 18 | |
| 19 | from dataclasses import dataclass |
| 20 | from functools import partial |
| 21 | from typing import TYPE_CHECKING, Any |
| 22 | |
| 23 | import torch |
| 24 | import torch.nn as nn |
| 25 | |
| 26 | from vllm._aiter_ops import rocm_aiter_ops |
| 27 | from vllm.config.utils import getattr_iter |
| 28 | from vllm.distributed import get_dp_group, get_ep_group |
| 29 | from vllm.forward_context import ForwardContext, get_forward_context |
| 30 | from vllm.logger import init_logger |
| 31 | from vllm.model_executor.custom_op import PluggableLayer |
| 32 | from vllm.model_executor.layers.fused_moe import ( |
| 33 | FusedMoEFactory, |
| 34 | MoERunner, |
| 35 | RoutedExperts, |
| 36 | ) |
| 37 | from vllm.model_executor.models.interfaces import MixtureOfExperts |
| 38 | from vllm.model_executor.models.transformers.fuser import get_fuser |
| 39 | from vllm.model_executor.models.transformers.fusers.glu import GLUFuser |
| 40 | from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser |
| 41 | from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix |
| 42 | from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op |
| 43 | |
| 44 | from .base import Base |
| 45 | from .utils import log_replacement, maybe_per_layer |
| 46 | |
| 47 | if TYPE_CHECKING: |
| 48 | from vllm.config import VllmConfig |
| 49 | |
| 50 | logger = init_logger(__name__) |
| 51 | |
| 52 | |
| 53 | @dataclass |
| 54 | class TransformersMoEState: |
| 55 | topk_ids: torch.Tensor | None = None |
| 56 | is_sequence_parallel: bool = False |
| 57 | |
| 58 | |
| 59 | # --8<-- [start:transformers_fused_moe] |
| 60 | @PluggableLayer.register("transformers_fused_moe") |
| 61 | class TransformersMoERunner(MoERunner): |
| 62 | """Custom MoERunner for the Transformers modeling backend.""" |
| 63 | |
| 64 | # --8<-- [end:transformers_fused_moe] |
| 65 | def __init__(self, *args, moe_state: TransformersMoEState, **kwargs): |
| 66 | super().__init__(*args, **kwargs) |
| 67 | self._moe_state = moe_state |
| 68 | self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel |
| 69 | |
| 70 | def forward( # type: ignore[override] |
| 71 | self, |
| 72 | hidden_states: torch.Tensor, |
| 73 | topk_ids: torch.Tensor, |
| 74 | topk_weights: torch.Tensor, |
| 75 | **kwargs: Any, |
| 76 | ) -> torch.Tensor: |
| 77 | """In Transformers `experts.forward` will have this signature, which is why |
| 78 | it deliberately does not match `MoERunner.forward`. |
| 79 | |
| 80 | We discard any extra kwargs because we cannot use them here.""" |
| 81 | # Note: we need to forward through a custom op so the topk_ids |
| 82 | # can be transferred without interfering with cudagraphs. |
| 83 | return torch.ops.vllm.transformers_moe_forward( |
| 84 | hidden_states, |
| 85 | topk_ids.to(torch.int32), |
| 86 | topk_weights.to(torch.float32), |
| 87 | self.layer_name, |
| 88 | ) |
| 89 | |
| 90 | def _forward_super( |
| 91 | self, |
| 92 | hidden_states: torch.Tensor, |
| 93 | topk_weights: torch.Tensor, |
| 94 | ) -> torch.Tensor: |
| 95 | return super().forward(hidden_states, topk_weights) |
| 96 | |
| 97 | |
| 98 | def _transformers_moe_forward( |
| 99 | hidden_states: torch.Tensor, |
| 100 | topk_ids: torch.Tensor, |
| 101 | topk_weights: torch.Tensor, |
| 102 | layer_name: str, |
| 103 | ) -> torch.Tensor: |
| 104 | """Store the `topk_ids` in the layer and call the actual forward.""" |
| 105 | forward_context: ForwardContext = get_forward_context() |
| 106 | self = forward_context.no_compile_layers[layer_name] |
| 107 | self._moe_state.topk_ids = topk_ids |
| 108 | return self._forward_super(hidden_states, topk_weights) |
| 109 | |
| 110 | |
| 111 | def _transformers_moe_forward_fake( |
| 112 | hidden_states: torch.Tensor, |
| 113 | topk_ids: torch.Tensor, |
| 114 | topk_weights: torch.Tensor, |
| 115 | layer_name: str, |
| 116 | ) -> torch.Tensor: |
| 117 | return torch.empty_like(hidden_states) |
| 118 | |
| 119 | |
| 120 | direct_register_custom_op( |
| 121 | op_name="transformers_moe_forward", |
| 122 | op_func=_transformers_moe_forward, |
| 123 | mutates_args=["hidden_states"], |
| 124 | fake_impl=_transformers_moe_forward_fake, |
| 125 | tags=(torch.Tag.needs_fixed_stride_order,), |
| 126 | ) |
| 127 | |
| 128 | |
| 129 | class TransformersRoutedExperts(RoutedExperts): |
| 130 | def get_expert_mapping( |
| 131 | self, |
| 132 | ckpt_gate_proj_name: str | None = None, |
| 133 | ckpt_down_proj_name: str | None = None, |
| 134 | ckpt_up_proj_name: str | None = None, |
| 135 | include_fused: bool = False, |
| 136 | ) -> list[tuple[str, str, int, str]]: |
| 137 | common_map = super().get_expert_mapping( |
| 138 | ckpt_gate_proj_name or "gate_proj", |
| 139 | ckpt_down_proj_name or "down_proj", |
| 140 | ckpt_up_proj_name or "up_proj", |
| 141 | include_fused, |
| 142 | ) |
| 143 | mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused) |
| 144 | if not include_fused: |
| 145 | return common_map + mixtral_map |
| 146 | common_fused, common_unfused = common_map[:3], common_map[3:] |
| 147 | mixtral_fused, mixtral_unfused = mixtral_map[:3], mixtral_map[3:] |
| 148 | return common_fused + mixtral_fused + common_unfused + mixtral_unfused |
| 149 | |
| 150 | |
| 151 | class MoEMixin(MixtureOfExperts, Base): |
| 152 | mlp_layers: list[nn.Module] |
| 153 | """MoE blocks whose experts were replaced, for the `MixtureOfExperts` methods.""" |
| 154 | |
| 155 | # Redeclared because `recursive_replace` assigns it below the first read |
| 156 | num_local_physical_experts: int |
| 157 | |
| 158 | def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): |
| 159 | # Skip MixtureOfExperts.__init__ and call the next class in MRO |
| 160 | super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix) |
| 161 | |
| 162 | def update_physical_experts_metadata( |
| 163 | self, |
| 164 | num_physical_experts: int, |
| 165 | num_local_physical_experts: int, |
| 166 | ): |
| 167 | assert self.num_local_physical_experts == num_local_physical_experts |
| 168 | self.num_physical_experts = num_physical_experts |
| 169 | self.num_local_physical_experts = num_local_physical_experts |
| 170 | self.num_redundant_experts = num_physical_experts - self.num_logical_experts |
| 171 | for moe_block in self.mlp_layers: |
| 172 | moe_block.n_local_physical_experts = num_local_physical_experts |
| 173 | moe_block.n_physical_experts = num_physical_experts |
| 174 | moe_block.n_redundant_experts = self.num_redundant_experts |
| 175 | moe_block.experts.update_expert_map() |
| 176 | |
| 177 | def recursive_replace(self): |
| 178 | """Initialize the MoE layers.""" |
| 179 | experts_name = "experts" |
| 180 | text_config = self.text_config |
| 181 | |
| 182 | # Positional arguments |
| 183 | num_experts = self.model_config.get_num_experts() |
| 184 | top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"], None) |
| 185 | assert top_k is not None |
| 186 | hidden_size = text_config.hidden_size |
| 187 | intermediate_size = getattr_iter( |
| 188 | text_config, ["moe_intermediate_size", "intermediate_size"], None |
| 189 | ) |
| 190 | assert intermediate_size is not None |
| 191 | |
| 192 | num_shared_experts = getattr_iter( |
| 193 | text_config, |
| 194 | [ |
| 195 | "n_shared_experts", # DeepSeek, Docs, GLM |
| 196 | "moe_num_shared_experts", # Aria, Ernie |
| 197 | ], |
| 198 | 0, |
| 199 | ) |
| 200 | |
| 201 | # Common kwargs |
| 202 | norm_topk_prob = getattr(text_config, "norm_topk_prob", None) |
| 203 | |
| 204 | # Routed scaling factor kwargs |
| 205 | routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0) |
| 206 | # aiter applies routed_scaling_factor internally |
| 207 | apply_routed_scale_to_output = not rocm_aiter_ops.is_fused_moe_enabled() |
| 208 | routed_scaling_factor_kwargs = dict( |
| 209 | routed_scaling_factor=routed_scaling_factor, |
| 210 | apply_routed_scale_to_output=apply_routed_scale_to_output, |
| 211 | ) |
| 212 | |
| 213 | # Dtype the router computes in, if it is not the activation dtype. |
| 214 | config_router_dtype = getattr(text_config, "moe_router_dtype", None) |
| 215 | if config_router_dtype is not None: |
| 216 | config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype) |
| 217 | |
| 218 | # Grouped topk routing kwargs |
| 219 | num_expert_group = getattr(text_config, "n_group", None) |
| 220 | topk_group = getattr(text_config, "topk_group", None) |
| 221 | use_grouped_topk = num_expert_group is not None and topk_group is not None |
| 222 | grouped_topk_routing_kwargs = dict( |
| 223 | use_grouped_topk=use_grouped_topk, |
| 224 | num_expert_group=num_expert_group, |
| 225 | topk_group=topk_group, |
| 226 | ) |
| 227 | |
| 228 | # MoE activation function |
| 229 | activation = "silu" |
| 230 | wrapped_arch = self.config.architectures[0].lower() |
| 231 | if "gptoss" in wrapped_arch: |
| 232 | activation = "swigluoai" |
| 233 | |
| 234 | # Expert parallel load balancing kwargs |
| 235 | enable_eplb = self.parallel_config.enable_eplb |
| 236 | num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts |
| 237 | |
| 238 | # MixtureOfExperts mixin settings |
| 239 | ep_size = get_ep_group().world_size |
| 240 | |
| 241 | self.mlp_layers = [] |
| 242 | self.moe_layers = [] |
| 243 | self.num_expert_groups = 1 if num_expert_group is None else num_expert_group |
| 244 | self.num_logical_experts = num_experts |
| 245 | self.num_physical_experts = num_experts + num_redundant_experts |
| 246 | self.num_local_physical_experts = self.num_physical_experts // ep_size |
| 247 | self.num_routed_experts = num_experts |
| 248 | self.num_shared_experts = num_shared_experts |
| 249 | self.num_redundant_experts = num_redundant_experts |
| 250 | |
| 251 | # Down projections of shared experts consumed by FusedMoE |
| 252 | shared_down_projs: list[tuple[nn.Module, str]] = [] |
| 253 | |
| 254 | # Recursively fuse MoE layers |
| 255 | def _recursive_replace(module: nn.Module, prefix: str): |
| 256 | for child_name, child_module in module.named_children(): |
| 257 | qual_name = maybe_prefix(prefix, child_name) |
| 258 | # Naive implementations will have experts as ModuleList |
| 259 | is_modulelist = isinstance(child_module, nn.ModuleList) |
| 260 | # Packed implementations will have experts as 3D tensors of shapes like: |
| 261 | # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size) |
| 262 | # down_proj = (num_experts, intermediate_size, hidden_size) |
| 263 | params = list(child_module.parameters()) |
| 264 | is_3d = len(params) > 0 and all(p.ndim == 3 for p in params) |
| 265 | if child_name == experts_name and (is_modulelist or is_3d): |
| 266 | # Alias for readability |
| 267 | moe_block = module |
| 268 | experts = child_module |
| 269 | # Class of the fused block (parent of gate/experts/shared) |
| 270 | moe_block_cls = type(moe_block).__name__ |
| 271 | experts_cls = type(experts).__name__ |
| 272 | # Do the experts have biases |
| 273 | has_bias = False |
| 274 | for experts_param_name, _ in experts.named_parameters(): |
| 275 | if "bias" in experts_param_name: |
| 276 | has_bias = True |
| 277 | break |
| 278 | # If the config does not specify num_shared_experts, but |
| 279 | # the model has shared experts, we assume there is one. |
| 280 | if self.num_shared_experts == 0: |
| 281 | for moe_block_param_name, _ in moe_block.named_parameters(): |
| 282 | if "shared_expert" in moe_block_param_name: |
| 283 | self.num_shared_experts = 1 |
| 284 | break |
| 285 | |
| 286 | # Only pay for the layer index if something is per-layer; |
| 287 | # `extract_layer_index` raises for unnumbered module paths. |
| 288 | per_layer = (top_k, intermediate_size, norm_topk_prob) |
| 289 | layer_idx = ( |
| 290 | extract_layer_index(qual_name) |
| 291 | if any(isinstance(v, list) for v in per_layer) |
| 292 | else 0 |
| 293 | ) |
| 294 | layer_top_k = maybe_per_layer(top_k, layer_idx) |
| 295 | layer_renormalize = maybe_per_layer(norm_topk_prob, layer_idx) |
| 296 | if layer_renormalize is None: |
| 297 | layer_renormalize = layer_top_k > 1 |
| 298 | |
| 299 | kwargs: dict[str, Any] = dict( |
| 300 | num_experts=num_experts, |
| 301 | top_k=layer_top_k, |
| 302 | hidden_size=hidden_size, |
| 303 | intermediate_size=maybe_per_layer(intermediate_size, layer_idx), |
| 304 | renormalize=layer_renormalize, |
| 305 | quant_config=self.quant_config, |
| 306 | prefix=qual_name, |
| 307 | activation=activation, |
| 308 | enable_eplb=enable_eplb, |
| 309 | num_redundant_experts=num_redundant_experts, |
| 310 | has_bias=has_bias, |
| 311 | routed_experts_cls=TransformersRoutedExperts, |
| 312 | ) |
| 313 | fuser = MoEBlockFuser.match(moe_block, experts_name) |
| 314 | # _maybe_apply_routed_scale_to_output edge case. Transformers |
| 315 | # decoder layers do not compensate for dividing by scaling factor. |
| 316 | reaches_fp16_trick = ( |
| 317 | routed_scaling_factor != 1.0 |
| 318 | and apply_routed_scale_to_output |
| 319 | and self.model_config.dtype == torch.float16 |
| 320 | and fuser is not None |
| 321 | and fuser.shared_name is not None |
| 322 | ) |
| 323 | if reaches_fp16_trick: |
| 324 | logger.warning_once( |
| 325 | "%s could be fused but routing it in vLLM would apply " |
| 326 | "`routed_scaling_factor` by dividing the shared expert " |
| 327 | "output, which only fp16 overflow protection expects the " |
| 328 | "decoder layer to compensate for. Falling back to routing " |
| 329 | "in Transformers; run in bfloat16 to fuse it.", |
| 330 | moe_block_cls, |
| 331 | ) |
| 332 | if fuser is not None and not reaches_fp16_trick: |
| 333 | # MoE block forward is fully replaced. |
| 334 | # gate/router and shared expert (if any) runs in MoERunner. |
| 335 | shared_experts = fuser.shared_experts(moe_block, prefix) |
| 336 | # Store shared experts for later down projection adjustment |
| 337 | if shared_experts is not None: |
| 338 | hf_shared = shared_experts.shared_experts |
| 339 | glu_fuser = get_fuser(hf_shared, GLUFuser) |
| 340 | if glu_fuser is not None: |
| 341 | down_name = glu_fuser.down_name |
| 342 | if down_name is not None: |
| 343 | shared_down_projs.append((hf_shared, down_name)) |
| 344 | # Prefer config, otherwise read it from fuser. |
| 345 | router_dtype = config_router_dtype or fuser.router_dtype |
| 346 | gate = fuser.gate(moe_block, prefix, router_dtype) |
| 347 | kwargs |= dict( |
| 348 | scoring_func=fuser.scoring_func, |
| 349 | is_sequence_parallel=( |
| 350 | self.parallel_config.use_sequence_parallel_moe |
| 351 | ), |
| 352 | gate=gate, |
| 353 | shared_experts=shared_experts, |
| 354 | ) |
| 355 | if router_dtype is not None: |
| 356 | kwargs["router_logits_dtype"] = router_dtype |
| 357 | if use_grouped_topk: |
| 358 | kwargs |= grouped_topk_routing_kwargs |
| 359 | if routed_scaling_factor != 1.0: |
| 360 | kwargs |= routed_scaling_factor_kwargs |
| 361 | bias = getattr(gate, "e_score_correction_bias", None) |
| 362 | if bias is not None: |
| 363 | kwargs["e_score_correction_bias"] = bias |
| 364 | fuser.rewrite_forward(moe_block) |
| 365 | routed = "gate + experts" |
| 366 | if fuser.shared_name: |
| 367 | routed += " + shared experts" |
| 368 | logger.info_once( |
| 369 | "Fused: %s (%s) -> MoERunner (internal routing)", |
| 370 | routed, |
| 371 | moe_block_cls, |
| 372 | ) |
| 373 | else: |
| 374 | # MoE block forward is unmodified. |
| 375 | # gate/router and shared expert (if any) runs in Transformers. |
| 376 | # We then smuggle the topk_ids in using a custom op. |
| 377 | moe_state = TransformersMoEState() |
| 378 | |
| 379 | def custom_routing_function( |
| 380 | hidden_states: torch.Tensor, |
| 381 | gating_output: torch.Tensor, |
| 382 | topk: int, |
| 383 | renormalize: bool, |
| 384 | moe_state: TransformersMoEState, |
| 385 | ): |
| 386 | """Return `topk_weights` from `gating_output` and the |
| 387 | `topk_ids` we stored in the layer earlier.""" |
| 388 | topk_weights = gating_output |
| 389 | topk_ids = moe_state.topk_ids |
| 390 | assert topk_ids is not None |
| 391 | # Handle all gather in expert parallel |
| 392 | if topk_ids.size(0) != hidden_states.size(0): |
| 393 | dp_metadata = get_forward_context().dp_metadata |
| 394 | assert dp_metadata is not None |
| 395 | sizes = dp_metadata.get_chunk_sizes_across_dp_rank() |
| 396 | assert sizes is not None |
| 397 | is_sp = moe_state.is_sequence_parallel |
| 398 | group = get_ep_group() if is_sp else get_dp_group() |
| 399 | assert sizes[group.rank_in_group] == topk_ids.shape[0] |
| 400 | (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes) |
| 401 | return topk_weights, topk_ids |
| 402 | |
| 403 | kwargs |= dict( |
| 404 | num_expert_group=num_expert_group, |
| 405 | topk_group=topk_group, |
| 406 | custom_routing_function=partial( |
| 407 | custom_routing_function, moe_state=moe_state |
| 408 | ), |
| 409 | runner_cls=TransformersMoERunner, |
| 410 | runner_args={"moe_state": moe_state}, |
| 411 | ) |
| 412 | logger.info_once( |
| 413 | "Fused: experts (%s) -> MoERunner (external routing)", |
| 414 | experts_cls, |
| 415 | ) |
| 416 | fused_experts = FusedMoEFactory(**kwargs) |
| 417 | moe_block.experts = fused_experts |
| 418 | log_replacement(qual_name, experts, fused_experts) |
| 419 | # Update MixtureOfExperts mixin state |
| 420 | self.mlp_layers.append(moe_block) |
| 421 | self.moe_layers.append(fused_experts) |
| 422 | else: |
| 423 | _recursive_replace(child_module, prefix=qual_name) |
| 424 | |
| 425 | _recursive_replace(self.model, prefix="model") |
| 426 | self.num_moe_layers = len(self.moe_layers) |
| 427 | # Continue with the replacement of layers in Base |
| 428 | super().recursive_replace() |
| 429 | # GLUFuser likely fused shared_experts. The down projection after the GLU |
| 430 | # normally immediately reduces but we want FusedMoE to handle the reduction. |
| 431 | for hf_shared, down_name in shared_down_projs: |
| 432 | hf_shared.get_submodule(down_name).reduce_results = False |