diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index bbcb94f32c..1b3c0d39bf 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -486,6 +486,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | | `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | ✅︎ | ✅︎ | +| `Kolibri1ForCausalLM` | Kolibri 1 | `Aleph-Alpha/Kolibri-1`, `Aleph-Alpha/Kolibri-1-BF16`, etc. | ✅︎ | ✅︎ | | `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | ✅︎ | ✅︎ | | `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ | | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index b08efd9fb0..37605f984e 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -360,6 +360,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "KimiLinearForCausalLM": _HfExamplesInfo( "moonshotai/Kimi-Linear-48B-A3B-Instruct", trust_remote_code=True ), + "Kolibri1ForCausalLM": _HfExamplesInfo( + "Aleph-Alpha/Kolibri-1", min_transformers_version="5.20.0" + ), "LagunaForCausalLM": _HfExamplesInfo("poolside/Laguna-XS.2"), "Lfm2ForCausalLM": _HfExamplesInfo("LiquidAI/LFM2-1.2B"), "Lfm2MoeForCausalLM": _HfExamplesInfo( diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index f0e70c896b..284b1d6a09 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -144,6 +144,16 @@ class DisconnectedRouter(TopKRouter): return logits, value, index +class ExtraInputRouter(TopKRouter): + """Router taking a block-owned correction bias as a second input (HYV3).""" + + def forward(self, hidden_states, bias=0.0): + logits = F.linear(hidden_states, self.weight) + _, index = torch.topk(logits.sigmoid() + bias, self.top_k, dim=-1) + value = logits.sigmoid().gather(1, index) + return logits, value, index + + class MoEExperts(nn.Module): """Packed experts (3D weights); only its name (`experts`) matters here.""" @@ -166,17 +176,19 @@ class MoEBlock(nn.Module): def _shared(self, x, logits): """The term added to the experts' output (none for a plain block).""" - return 0 + return None def forward(self, hidden_states): x = hidden_states.reshape(-1, hidden_states.shape[-1]) logits, weights, index = self.gate(x) - out = self.experts(x, index, weights) + self._shared(x, logits) + out = self.experts(x, index, weights) + if (shared := self._shared(x, logits)) is not None: + out = out + shared return out.reshape(hidden_states.shape) class MoEBlockNoShared(MoEBlock): - """No shared-expert child but a gate-derived add -> trace skipped, still fuses.""" + """No shared-expert child but a gate-derived add -> declined, not misread.""" def _shared(self, x, logits): return logits.sum() @@ -293,6 +305,115 @@ class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted): self.extra = BufferScale() +class MoEBlockExpertsKeywords(MoEBlock): + """Passes the experts' routing by keyword -> still fuses.""" + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + return self.experts(x, weights=weights, index=index) + + +class ScaledExperts(MoEExperts): + def forward(self, hidden_states, index, weights, scale=1.0): + return hidden_states * scale + + +class MoEBlockExpertsExtraArg(MoEBlock): + """Passes the experts an extra argument the fused forward would drop.""" + + def __init__(self): + super().__init__() + self.experts = ScaledExperts() + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + return self.experts(x, index, weights, scale=2.0) + + +class MoEBlockSharedUnflattened(MoEBlockSharedNoGate): + """Gate and shared expert read the unflattened input (Kolibri-1).""" + + def forward(self, hidden_states): + _, weights, index = self.gate(hidden_states) + x = hidden_states.view(-1, hidden_states.shape[-1]) + out = self.experts(x, index, weights).view(*hidden_states.shape) + return out + self.shared_expert(hidden_states) + + +class MoEBlockSharedAveraged(MoEBlockSharedUnflattened): + """Averages the experts' and shared expert's outputs (Cohere2-MoE).""" + + def forward(self, hidden_states): + return super().forward(hidden_states) / 2 + + +class MoEBlockScaled(MoEBlockSharedNoGate): + """Scales the experts' output by `routed_scaling_factor` (Laguna).""" + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + return self.experts(x, index, weights) * 2.5 + self.shared_expert(x) + + +class MoEBlockScaledBiasedRouter(MoEBlockScaled): + """An unrecognized router with a scaled experts' output.""" + + def __init__(self): + super().__init__() + self.gate = BiasedRouter() + + +class MoEBlockSharedBiasedRouter(MoEBlockSharedUnflattened): + """An unrecognized router with a shared expert.""" + + def __init__(self): + super().__init__() + self.gate = BiasedRouter() + + +class MoEBlockGateExtraInput(MoEBlock): + """Passes the gate a block-owned bias the fused forward would drop (HYV3).""" + + def __init__(self): + super().__init__(ExtraInputRouter) + self.register_buffer("e_score_correction_bias", torch.zeros(8)) + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x, self.e_score_correction_bias) + return self.experts(x, index, weights) + + +class MoEBlockCastCombine(MoEBlockSharedNoGate): + """Combines the outputs in fp32 (HYV3).""" + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + out = self.experts(x, index, weights).float() + self.shared_expert(x).float() + return out.to(x.dtype) + + +class MoEBlockIdentityLatent(MoEBlockSharedNoGate): + """Wraps the experts in unused (identity) latent projections (NemotronH).""" + + def __init__(self): + super().__init__() + self.fc1_latent_proj = nn.Identity() + self.fc2_latent_proj = nn.Identity() + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + out = self.fc2_latent_proj( + self.experts(self.fc1_latent_proj(x), index, weights) + ) + return out + self.shared_expert(x) + + @pytest.mark.parametrize("sigmoid", [False, True]) def test_moe_fuser_detects_router(sigmoid): with torch.device("meta"): @@ -385,16 +506,6 @@ def test_moe_fuser_detects_shared_experts(): assert fuser.shared_gate_name == "shared_expert_gate" -def test_moe_fuser_skips_shared_detection_without_extra_children(): - """With only experts and gate, shared-expert detection (and its block trace) - is skipped, so a gate-derived add is not misread as a shared expert.""" - with torch.device("meta"): - block = MoEBlockNoShared() - fuser = MoEBlockFuser.match(block, "experts") - assert isinstance(fuser, MoEBlockFuser) - assert fuser.shared_name is None and fuser.shared_gate_name is None - - def test_moe_fuser_shared_without_gate(): with torch.device("meta"): block = MoEBlockSharedNoGate() @@ -422,6 +533,10 @@ def test_moe_fuser_detects_non_glu_shared_expert(): MoEBlockTupleVar, # tuple returned via a name binding, not a literal MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow MoEBlockUnaccountedBuffer, # buffer-only child outside the fused dataflow + MoEBlockExpertsExtraArg, # extra experts argument the fused forward drops + MoEBlockNoShared, # gate-derived term added to the experts' output + MoEBlockSharedAveraged, # experts' and shared expert's outputs averaged + MoEBlockGateExtraInput, # gate input the fused forward drops ], ) def test_moe_fuser_declines_unsupported(block_cls): @@ -444,3 +559,53 @@ def test_moe_fuser_router_requires_connected_dataflow(): with torch.device("meta"): block = MoEBlock(DisconnectedRouter) assert MoEBlockFuser.match(block, "experts") is None + + +def test_moe_fuser_unrecognized_router_moves_shared_expert(): + """Without vLLM routing, the shared expert still moves to vLLM.""" + with torch.device("meta"): + block = MoEBlockSharedBiasedRouter() + fuser = MoEBlockFuser.match(block, "experts") + assert fuser.scoring_func is None + assert fuser.shared_name == "shared_expert" + + +@pytest.mark.parametrize( + ("block_cls", "scoring_func"), + [ + (MoEBlockSharedUnflattened, None), # routing falls back to Transformers + (MoEBlock, "softmax"), # no shared output to divide + ], +) +def test_moe_fuser_divided_shared_output(block_cls, scoring_func): + with torch.device("meta"): + fuser = MoEBlockFuser.match(block_cls(), "experts", divides_shared_output=True) + assert fuser.scoring_func == scoring_func + + +def test_moe_fuser_scaled_output_needs_vllm_routing(): + """Only vLLM routing applies the routed scale, so nothing else moves without it.""" + with torch.device("meta"): + block = MoEBlockScaledBiasedRouter() + assert MoEBlockFuser.match(block, "experts") is None + + +def test_moe_fuser_matches_experts_called_by_keyword(): + with torch.device("meta"): + block = MoEBlockExpertsKeywords() + assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) + + +def test_moe_fuser_scaled_output_with_vllm_routing(): + """A scaled block fuses because vLLM routing applies `routed_scaling_factor`.""" + with torch.device("meta"): + fuser = MoEBlockFuser.match(MoEBlockScaled(), "experts") + assert fuser.scoring_func == "softmax" + assert fuser.shared_name == "shared_expert" + + +@pytest.mark.parametrize("block_cls", [MoEBlockCastCombine, MoEBlockIdentityLatent]) +def test_moe_fuser_looks_through_casts_and_identity(block_cls): + with torch.device("meta"): + fuser = MoEBlockFuser.match(block_cls(), "experts") + assert fuser.shared_name == "shared_expert" diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 09e5025c4e..afc9218c6b 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -717,6 +717,7 @@ _TRANSFORMERS_SUPPORTED_MODELS = { "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"), "Jais2ForCausalLM": ("transformers", "TransformersForCausalLM"), + "Kolibri1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"), "MellumForCausalLM": ("transformers", "TransformersMoEForCausalLM"), "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"), "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"), diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index 06fdac109c..46a1ce33bc 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing.""" +"""MoE fuser: run an HF MoE block's gate, experts and shared experts in `FusedMoE`.""" import ast +import functools import inspect import textwrap import types @@ -11,20 +12,25 @@ from dataclasses import dataclass import torch from torch import fx, nn +from typing_extensions import TypeIs from vllm.distributed import tensor_model_parallel_all_gather +from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import GateLinear from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.transformers.fx_utils import ( find_node, is_complete, is_op, + output_value, peel, trace, ) from vllm.model_executor.models.transformers.utils import named_state from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk +logger = init_logger(__name__) + def _own_returns(node: ast.AST) -> Iterator[ast.Return]: """`return` statements in `node`'s own scope, not in nested functions.""" @@ -126,6 +132,35 @@ def _renormalizes(topk: fx.Node) -> bool: return False +def _unwrap(node: object, moe_block: nn.Module) -> object: + """Strip reshapes, dtype casts and `nn.Identity` calls.""" + while True: + node = peel(node) + if ( + is_op(node, "reshape") + or is_op(node, "view") + or ( + isinstance(node, fx.Node) + and node.op == "call_module" + and isinstance(moe_block.get_submodule(node.target), nn.Identity) + ) + ): + node = node.args[0] + else: + return node + + +def _reads_input(node: object, x: fx.Node, moe_block: nn.Module) -> TypeIs[fx.Node]: + """Is `node` a child module of `moe_block` called on just `x`?""" + return ( + isinstance(node, fx.Node) + and node.op == "call_module" + and len(node.args) == 1 + and not node.kwargs + and _unwrap(node.args[0], moe_block) is x + ) + + class SharedExpertMLP(nn.Module): """Wraps an HF shared expert, applying the output gating it is paired with.""" @@ -162,7 +197,8 @@ class MoEBlockFuser: """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional).""" gate_name: str - scoring_func: str + scoring_func: str | None + """`None` if routing stays in the Transformers gate.""" shared_name: str | None shared_gate_name: str | None router_dtype: torch.dtype | None = None @@ -215,83 +251,142 @@ class MoEBlockFuser: scoring_func = "softmax" if is_softmax else "sigmoid" return scoring_func, _forced_dtype(logits_cone), renormalize + @staticmethod + def _match_experts( + graph: fx.Graph, experts_name: str, moe_block: nn.Module + ) -> tuple[fx.Node, fx.Node, fx.Node] | None: + """Matches `experts(x, ids, weights)` with `_, weights, ids = gate(x)`, `x` + being the block's input. + + Returns the experts call, `x` and the gate call.""" + experts = find_node( + graph, lambda n: n.op == "call_module" and n.target == experts_name + ) + if experts is None: + return None + # Bind by signature so keyword calls match too. Fail closed on any other + # argument: the rewritten forward would drop it. + forward = inspect.signature(moe_block.get_submodule(experts_name).forward) + try: + bound = forward.bind(*experts.args, **experts.kwargs) + except TypeError: + return None + if len(bound.arguments) != 3: + return None + x, ids, weights = (_unwrap(a, moe_block) for a in bound.arguments.values()) + if not (isinstance(x, fx.Node) and x.op == "placeholder"): + return None + if not (is_op(ids, "getitem") and is_op(weights, "getitem")): + return None + gate = ids.args[0] + if not ( + _reads_input(gate, x, moe_block) + and ids.args == (gate, 2) + and weights.args == (gate, 1) + ): + return None + return experts, x, gate + @staticmethod def _match_shared_experts( - graph: fx.Graph, experts: str - ) -> tuple[str | None, str | None]: - """Detects the shared expert and its optional gate by dataflow.""" - experts_predicate = lambda n: n.op == "call_module" and n.target == experts - if (experts_node := find_node(graph, experts_predicate)) is None: - return None, None - from_experts = _reaches(experts_node, "users") - for add in graph.nodes: - if not is_op(add, "add"): + node: object, x: fx.Node, moe_block: nn.Module + ) -> tuple[str, str | None] | None: + """Matches `shared(x)` or `sigmoid(shared_gate(x)) * shared(x)`. + + Returns the shared expert's name and its optional gate's name.""" + if _reads_input(node, x, moe_block): + return node.target, None + if not (is_op(node, "mul") and len(node.args) == 2): + return None + for gated, shared in (node.args, node.args[::-1]): + gated, shared = peel(gated), peel(shared) + if not (is_op(gated, "sigmoid") and _reads_input(shared, x, moe_block)): continue - operands = [a for a in add.args if isinstance(a, fx.Node)] - # Exactly one side is the experts' output; the other is the shared path. - sides = [a in from_experts for a in operands] - if len(operands) != 2 or sides.count(True) != 1: - continue - cone = _reaches(operands[sides.index(False)], "all_input_nodes") - modules = [n for n in cone if n.op == "call_module" and n.target != experts] - # A sigmoid wrapping one of those modules marks the shared-expert gate. - gate = next( - ( - src - for n in cone - if is_op(n, "sigmoid") - and isinstance(src := peel(n.args[0]), fx.Node) - and src in modules - ), - None, - ) - shared = [n for n in modules if n is not gate] - if len(shared) != 1: - return None, None - return shared[0].target, (gate.target if gate is not None else None) - return None, None + gate = peel(gated.args[0]) + if _reads_input(gate, x, moe_block) and _is_scalar_gate( + moe_block.get_submodule(gate.target) + ): + return shared.target, gate.target + return None @classmethod - def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None": + def match( + cls, + moe_block: nn.Module, + experts_name: str, + divides_shared_output: bool = False, + ) -> "MoEBlockFuser | None": + """Matches a block returning `experts(x, ids, weights) [* c] [+ shared(x)]` + with `_, weights, ids = gate(x)`, `x` being the block's input, up to + reshapes, casts and `nn.Identity`. + + Args: + moe_block: The MoE block. + experts_name: The name of the experts in `moe_block`. + divides_shared_output: Whether vLLM routing would divide the shared + expert's output by `routed_scaling_factor` (fp16 overflow + protection), which Transformers decoder layers do not undo. + + Returns: + A fuser if more than the experts can run in vLLM, else `None`. + + """ # Standard MoE block returns a single tensor. if _returns_tuple(type(moe_block)): return None - # Router: the child that scores + top-k selects. - gate_name = scoring_func = router_dtype = renormalize = None - for name, child in moe_block.named_children(): - if ( - name != experts_name - and (router := cls._match_router(child)) is not None - ): - gate_name = name - scoring_func, router_dtype, renormalize = router - break - if gate_name is None or scoring_func is None: + if (graph := trace(moe_block)) is None: return None - # Shared expert: a child the block adds to the experts' output. - shared_name = shared_gate_name = None - others = [ - n - for n, _ in moe_block.named_children() - if n not in {experts_name, gate_name} - ] - if others: - graph = trace(moe_block) - if graph is None: - return None - shared_name, shared_gate_name = cls._match_shared_experts( - graph, experts_name + if (matched := cls._match_experts(graph, experts_name, moe_block)) is None: + return None + experts, x, gate = matched + unwrap = functools.partial(_unwrap, moe_block=moe_block) + gate_name = gate.target + router = cls._match_router(getattr(moe_block, gate_name)) + scoring_func, router_dtype, renormalize = router or (None, None, None) + + def is_routed(n: object) -> bool: + return n is experts or ( + is_op(n, "mul") + and not isinstance(n.args[1], fx.Node) + and unwrap(n.args[0]) is experts ) - if shared_gate_name is not None and not _is_scalar_gate( - getattr(moe_block, shared_gate_name) - ): + + # Output: the experts' output, plus the shared expert's if any. + routed = out = unwrap(output_value(graph)) + shared = None + if not is_routed(out): + if not (is_op(out, "add") and len(out.args) == 2): return None + routed, shared = map(unwrap, out.args) + if is_routed(shared): + routed, shared = shared, routed + if not is_routed(routed): + return None + shared_name = shared_gate_name = None + if shared is not None: + if (match := cls._match_shared_experts(shared, x, moe_block)) is None: + return None + shared_name, shared_gate_name = match # Fail closed: `rewrite_forward` runs only the experts and the detected # shared expert, so any other stateful child would be dropped. accounted = {experts_name, gate_name, shared_name, shared_gate_name} for name, child in moe_block.named_children(): if name not in accounted and next(named_state(child), None) is not None: return None + if scoring_func is not None and shared_name and divides_shared_output: + 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.", + type(moe_block).__name__, + ) + scoring_func = router_dtype = renormalize = None + # Without vLLM routing, only a shared expert can move to vLLM, and only + # if the block does not scale the experts' output (vLLM routing does). + if scoring_func is None and (shared_name is None or routed is not experts): + return None return cls( gate_name, scoring_func, diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 953459172d..a426a2c9d9 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -127,6 +127,43 @@ direct_register_custom_op( ) +def _hf_topk_router_precomputed( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + moe_state: TransformersMoEState, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return `(topk_weights, topk_ids)` the HF block forward already computed: + weights arrive as `gating_output`, ids via `moe_state`.""" + 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 + assert dp_metadata is not None + sizes = dp_metadata.get_chunk_sizes_across_dp_rank() + assert sizes is not None + 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 + + +def _hf_topk_router( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + gate: nn.Module, +) -> tuple[torch.Tensor, torch.Tensor]: + """Route with a Transformers gate returning `(_, topk_weights, topk_ids)`.""" + _, topk_weights, topk_ids = gate(hidden_states) + return topk_weights, topk_ids + + class TransformersRoutedExperts(RoutedExperts): def get_expert_mapping( self, @@ -310,26 +347,17 @@ class MoEMixin(MixtureOfExperts, Base): is_fused_checkpoint_transposed=is_transposed, 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 = ( + divides_shared_output = ( 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: + fuser = MoEBlockFuser.match( + moe_block, experts_name, divides_shared_output + ) + if fuser is not None: # MoE block forward is fully replaced. # gate/router and shared expert (if any) runs in MoERunner. shared_experts = fuser.shared_experts(moe_block, prefix) @@ -341,39 +369,52 @@ class MoEMixin(MixtureOfExperts, Base): down_name = glu_fuser.down_name 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) + if scoring_func := fuser.scoring_func: + # Internal routing kwargs. + # 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=scoring_func, + renormalize=( + layer_renormalize + if fuser.renormalize is None + else fuser.renormalize + ), + gate=gate, + ) + 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 + else: + # External routing kwargs + kwargs |= dict( + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=partial( + _hf_topk_router, + gate=getattr(moe_block, fuser.gate_name), + ), + ) + kwargs |= dict( - scoring_func=fuser.scoring_func, - renormalize=( - layer_renormalize - if fuser.renormalize is None - else fuser.renormalize - ), 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, + "Fused: gate + experts%s (%s) -> MoERunner (%s routing)", + " + shared experts" if fuser.shared_name else "", moe_block_cls, + "external" if fuser.scoring_func is None else "internal", ) else: # MoE block forward is unmodified. @@ -381,35 +422,11 @@ class MoEMixin(MixtureOfExperts, Base): # 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 - assert dp_metadata is not None - sizes = dp_metadata.get_chunk_sizes_across_dp_rank() - assert sizes is not None - 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 + _hf_topk_router_precomputed, moe_state=moe_state ), runner_cls=TransformersMoERunner, runner_args={"moe_state": moe_state},