vllm.model_executor.models.transformers.moe ¶

class MoEMixin(MixtureOfExperts):
    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
        self.check_version("5.0.0", "MoE models support")
        # Skip MixtureOfExperts.__init__ and call the next class in MRO
        super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix)

    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ):
        assert self.num_local_physical_experts == num_local_physical_experts
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
        for moe_block in self.mlp_layers:
            moe_block.n_local_physical_experts = num_local_physical_experts
            moe_block.n_physical_experts = num_physical_experts
            moe_block.n_redundant_experts = self.num_redundant_experts
            moe_block.experts.update_expert_map()

    def recursive_replace(self):
        """Initialize the MoE layers."""
        experts_name = "experts"
        text_config = self.text_config

        # Positional arguments
        num_experts = self.model_config.get_num_experts()
        top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"], None)
        assert top_k is not None
        hidden_size = text_config.hidden_size
        intermediate_size = getattr_iter(
            text_config, ["moe_intermediate_size", "intermediate_size"], None
        )
        assert intermediate_size is not None

        num_shared_experts = getattr_iter(
            text_config,
            [
                "n_shared_experts",  # DeepSeek, Docs, GLM
                "moe_num_shared_experts",  # Aria, Ernie
            ],
            0,
        )

        # Common kwargs
        renormalize = getattr(text_config, "norm_topk_prob", top_k > 1)

        # Routed scaling factor kwargs
        routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0)
        # aiter applies routed_scaling_factor internally
        apply_routed_scale_to_output = not rocm_aiter_ops.is_fused_moe_enabled()
        routed_scaling_factor_kwargs = dict(
            routed_scaling_factor=routed_scaling_factor,
            apply_routed_scale_to_output=apply_routed_scale_to_output,
        )

        # Dtype the router computes in, if it is not the activation dtype.
        config_router_dtype = getattr(text_config, "moe_router_dtype", None)
        config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype)

        # Grouped topk routing kwargs
        num_expert_group = getattr(text_config, "n_group", None)
        topk_group = getattr(text_config, "topk_group", None)
        use_grouped_topk = num_expert_group is not None and topk_group is not None
        grouped_topk_routing_kwargs = dict(
            use_grouped_topk=use_grouped_topk,
            num_expert_group=num_expert_group,
            topk_group=topk_group,
        )

        # MoE activation function
        activation = "silu"
        wrapped_arch = self.config.architectures[0].lower()
        if "gptoss" in wrapped_arch:
            activation = "swigluoai"

        # Expert parallel load balancing kwargs
        enable_eplb = self.parallel_config.enable_eplb
        num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts

        # MixtureOfExperts mixin settings
        ep_size = get_ep_group().world_size

        self.mlp_layers = []  # Used for MixtureOfExperts methods
        self.moe_layers = []
        self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
        self.num_logical_experts = num_experts
        self.num_physical_experts = num_experts + num_redundant_experts
        self.num_local_physical_experts = self.num_physical_experts // ep_size
        self.num_routed_experts = num_experts
        self.num_shared_experts = num_shared_experts
        self.num_redundant_experts = num_redundant_experts

        # Down projections of shared experts consumed by FusedMoE
        shared_down_projs: list[tuple[nn.Module, str]] = []

        # Recursively fuse MoE layers
        def _recursive_replace(module: nn.Module, prefix: str):
            for child_name, child_module in module.named_children():
                qual_name = maybe_prefix(prefix, child_name)
                # Naive implementations will have experts as ModuleList
                is_modulelist = isinstance(child_module, nn.ModuleList)
                # Packed implementations will have experts as 3D tensors of shapes like:
                # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size)
                # down_proj = (num_experts, intermediate_size, hidden_size)
                params = list(child_module.parameters())
                is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
                if child_name == experts_name and (is_modulelist or is_3d):
                    # Alias for readability
                    moe_block = module
                    experts = child_module
                    # Class of the fused block (parent of gate/experts/shared)
                    moe_block_cls = type(moe_block).__name__
                    experts_cls = type(experts).__name__
                    # Do the experts have biases
                    has_bias = False
                    for experts_param_name, _ in experts.named_parameters():
                        if "bias" in experts_param_name:
                            has_bias = True
                            break
                    # If the config does not specify num_shared_experts, but
                    # the model has shared experts, we assume there is one.
                    if self.num_shared_experts == 0:
                        for moe_block_param_name, _ in moe_block.named_parameters():
                            if "shared_expert" in moe_block_param_name:
                                self.num_shared_experts = 1
                                break

                    kwargs: dict[str, Any] = dict(
                        num_experts=num_experts,
                        top_k=top_k,
                        hidden_size=hidden_size,
                        intermediate_size=intermediate_size,
                        renormalize=renormalize,
                        quant_config=self.quant_config,
                        prefix=qual_name,
                        activation=activation,
                        enable_eplb=enable_eplb,
                        num_redundant_experts=num_redundant_experts,
                        has_bias=has_bias,
                        routed_experts_cls=TransformersRoutedExperts,
                    )
                    fuser = MoEBlockFuser.match(moe_block, experts_name)
                    # _maybe_apply_routed_scale_to_output edge case. Transformers
                    # decoder layers do not compensate for dividing by scaling factor.
                    reaches_fp16_trick = (
                        routed_scaling_factor != 1.0
                        and apply_routed_scale_to_output
                        and self.model_config.dtype == torch.float16
                        and fuser is not None
                        and fuser.shared_name is not None
                    )
                    if reaches_fp16_trick:
                        logger.warning_once(
                            "%s could be fused but routing it in vLLM would apply "
                            "`routed_scaling_factor` by dividing the shared expert "
                            "output, which only fp16 overflow protection expects the "
                            "decoder layer to compensate for. Falling back to routing "
                            "in Transformers; run in bfloat16 to fuse it.",
                            moe_block_cls,
                        )
                    if fuser is not None and not reaches_fp16_trick:
                        # MoE block forward is fully replaced.
                        # gate/router and shared expert (if any) runs in MoERunner.
                        shared_experts = fuser.shared_experts(moe_block, prefix)
                        # Store shared experts for later down projection adjustment
                        if shared_experts is not None:
                            hf_shared = shared_experts.shared_experts
                            glu_fuser = get_fuser(hf_shared)
                            down_name = getattr(glu_fuser, "down_name", None)
                            if down_name is not None:
                                shared_down_projs.append((hf_shared, down_name))
                        # Prefer config, otherwise read it from fuser.
                        router_dtype = config_router_dtype or fuser.router_dtype
                        gate = fuser.gate(moe_block, prefix, router_dtype)
                        kwargs |= dict(
                            scoring_func=fuser.scoring_func,
                            is_sequence_parallel=(
                                self.parallel_config.use_sequence_parallel_moe
                            ),
                            gate=gate,
                            shared_experts=shared_experts,
                        )
                        if router_dtype is not None:
                            kwargs["router_logits_dtype"] = router_dtype
                        if use_grouped_topk:
                            kwargs |= grouped_topk_routing_kwargs
                        if routed_scaling_factor != 1.0:
                            kwargs |= routed_scaling_factor_kwargs
                        bias = getattr(gate, "e_score_correction_bias", None)
                        if bias is not None:
                            kwargs["e_score_correction_bias"] = bias
                        fuser.rewrite_forward(moe_block)
                        routed = "gate + experts"
                        if fuser.shared_name:
                            routed += " + shared experts"
                        logger.info_once(
                            "Fused: %s (%s) -> MoERunner (internal routing)",
                            routed,
                            moe_block_cls,
                        )
                    else:
                        # MoE block forward is unmodified.
                        # gate/router and shared expert (if any) runs in Transformers.
                        # We then smuggle the topk_ids in using a custom op.
                        moe_state = TransformersMoEState()

                        def custom_routing_function(
                            hidden_states: torch.Tensor,
                            gating_output: torch.Tensor,
                            topk: int,
                            renormalize: bool,
                            moe_state: TransformersMoEState,
                        ):
                            """Return `topk_weights` from `gating_output` and the
                            `topk_ids` we stored in the layer earlier."""
                            topk_weights = gating_output
                            topk_ids = moe_state.topk_ids
                            assert topk_ids is not None
                            # Handle all gather in expert parallel
                            if topk_ids.size(0) != hidden_states.size(0):
                                dp_metadata = get_forward_context().dp_metadata
                                sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
                                is_sp = moe_state.is_sequence_parallel
                                group = get_ep_group() if is_sp else get_dp_group()
                                assert sizes[group.rank_in_group] == topk_ids.shape[0]
                                (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes)
                            return topk_weights, topk_ids

                        kwargs |= dict(
                            num_expert_group=num_expert_group,
                            topk_group=topk_group,
                            custom_routing_function=partial(
                                custom_routing_function, moe_state=moe_state
                            ),
                            runner_cls=TransformersMoERunner,
                            runner_args={"moe_state": moe_state},
                        )
                        logger.info_once(
                            "Fused: experts (%s) -> MoERunner (external routing)",
                            experts_cls,
                        )
                    fused_experts = FusedMoEFactory(**kwargs)
                    moe_block.experts = fused_experts
                    log_replacement(qual_name, experts, fused_experts)
                    # Update MixtureOfExperts mixin state
                    self.mlp_layers.append(moe_block)
                    self.moe_layers.append(fused_experts)
                else:
                    _recursive_replace(child_module, prefix=qual_name)

        _recursive_replace(self.model, prefix="model")
        self.num_moe_layers = len(self.moe_layers)
        # Continue with the replacement of layers in Base
        super().recursive_replace()
        # GLUFuser likely fused shared_experts. The down projection after the GLU
        # normally immediately reduces but we want FusedMoE to handle the reduction.
        for hf_shared, down_name in shared_down_projs:
            hf_shared.get_submodule(down_name).reduce_results = False

@PluggableLayer.register("transformers_fused_moe")
class TransformersMoERunner(MoERunner):
    """Custom MoERunner for the Transformers modeling backend."""

    # --8<-- [end:transformers_fused_moe]
    def __init__(self, *args, moe_state: TransformersMoEState, **kwargs):
        super().__init__(*args, **kwargs)
        self._moe_state = moe_state
        self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel

    def forward(
        self,
        hidden_states: torch.Tensor,
        topk_ids: torch.Tensor,
        topk_weights: torch.Tensor,
        **kwargs: Any,
    ) -> torch.Tensor:
        """In Transformers `experts.forward` will have this signature.

        We discard any extra kwargs because we cannot use them here."""
        # Note: we need to forward through a custom op so the topk_ids
        # can be transferred without interfering with cudagraphs.
        return torch.ops.vllm.transformers_moe_forward(
            hidden_states,
            topk_ids.to(torch.int32),
            topk_weights.to(torch.float32),
            self.layer_name,
        )

    def _forward_super(
        self,
        hidden_states: torch.Tensor,
        topk_weights: torch.Tensor,
    ) -> torch.Tensor:
        return super().forward(hidden_states, topk_weights)