Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / moe.py: 68%
221 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""PyTorch MoE building blocks: router, experts and orchestrator.
17Provides :class:`FeedForward`, :class:`GroupedExperts`,
18:class:`TokenChoiceTopKRouter`, and the top-level :class:`MoE` orchestrator.
19Load-balancing utilities (expert bias update and auxiliary loss) are also
20included.
22All modules compose naturally with EP / TP parallel strategies via DTensor.
23Distributed collectives are handled by
24:mod:`hyper_parallel.core.expert_parallel.expert_parallel`; this module
25contains only single-device computation.
26"""
27from __future__ import annotations
29__all__ = [
30 "FeedForward",
31 "GroupedExperts",
32 "MoE",
33 "MoEAuxLossAutoScaler",
34 "TokenChoiceTopKRouter",
35 "update_expert_bias",
36]
38import math
39from typing import Any
41import torch
42import torch.nn.functional as F
43from torch import nn
45from hyper_parallel.components.functional.npu_grouped_swiglu import npu_grouped_swiglu
46from hyper_parallel.core.dtensor.dtensor import DTensor
49def _is_npu_tensor(tensor: torch.Tensor) -> bool:
50 """Return whether a tensor is stored on an Ascend NPU."""
51 return tensor.device.type == "npu"
54# ---------------------------------------------------------------------------
55# Grouped expert computation kernels
56# ---------------------------------------------------------------------------
58def _run_experts_for_loop(
59 w1: torch.Tensor,
60 w2: torch.Tensor,
61 w3: torch.Tensor,
62 x: torch.Tensor,
63 num_tokens_per_expert: torch.Tensor,
64 scores: torch.Tensor | None = None,
65) -> torch.Tensor:
66 """Run per-expert SwiGLU via a sequential loop (reference path).
68 Args:
69 w1: Shape ``[num_experts, hidden_dim, dim]``.
70 w2: Shape ``[num_experts, dim, hidden_dim]``.
71 w3: Shape ``[num_experts, hidden_dim, dim]``.
72 x: Routed tokens in expert-major order, shape
73 ``[total_routed_tokens, dim]``.
74 num_tokens_per_expert: 1-D integer tensor of length ``num_experts``.
75 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
76 Routing weights applied to intermediate activations (``silu(w1(x)) * w3(x)``)
77 before the w2 projection. When ``None``, no weighting is applied.
78 Defaults to ``None``.
80 Returns:
81 Expert output of shape ``[total_routed_tokens, dim]``.
82 """
83 # Use ``torch.cat`` instead of ``zeros_like + in-place slice assign``.
84 # On some backends (notably torch_npu) in-place slice assignment onto a
85 # ``requires_grad=False`` leaf tensor does not reliably upgrade it to a
86 # non-leaf with a ``grad_fn`` — the forward result may end up with
87 # ``grad_fn=None`` and downstream ``backward()`` fails with
88 # "element 0 of tensors does not require grad and does not have a grad_fn".
89 #
90 # Drain ``num_tokens_per_expert`` to host **once** via ``.tolist()``
91 # rather than calling ``int(n)`` per loop iteration — a single D2H
92 # copy instead of ``num_local_experts`` separate ones. Per-iter
93 # ``.item()`` would stall the host between expert kernels and shrink
94 # the dual-pipe overlap window.
95 counts_list = num_tokens_per_expert.tolist()
96 parts = []
97 offset = 0
98 for e, n in enumerate(counts_list):
99 if n == 0:
100 continue
101 x_e = x[offset:offset + n]
102 h = F.silu(x_e @ w1[e].T) * (x_e @ w3[e].T)
103 if scores is not None:
104 h = h * scores[offset:offset + n].unsqueeze(-1)
105 parts.append(h @ w2[e].T)
106 offset += n
107 if not parts:
108 # No routed tokens: return a grad-connected zero (not ``zeros_like``).
109 return x * 0.0
110 return torch.cat(parts, dim=0)
113def _run_experts_grouped_mm_gpu(
114 w1: torch.Tensor,
115 w2: torch.Tensor,
116 w3: torch.Tensor,
117 x: torch.Tensor,
118 num_tokens_per_expert: torch.Tensor,
119 scores: torch.Tensor | None = None,
120) -> torch.Tensor:
121 """Fused grouped matmul path for NVIDIA GPU using ``torch._grouped_mm``.
123 Args:
124 w1: Shape ``[num_experts, hidden_dim, dim]``.
125 w2: Shape ``[num_experts, dim, hidden_dim]``.
126 w3: Shape ``[num_experts, hidden_dim, dim]``.
127 x: Shape ``[total_routed_tokens, dim]``.
128 num_tokens_per_expert: 1-D integer tensor of length ``num_experts``.
129 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
130 Routing weights applied to intermediate activations (``silu(w1(x)) * w3(x)``)
131 before the w2 projection. When ``None``, no weighting is applied.
132 Defaults to ``None``.
134 Returns:
135 Expert output of shape ``[total_routed_tokens, dim]``.
136 """
137 # grouped_mm needs an end offset for every expert, including the last one.
138 offs = torch.cumsum(num_tokens_per_expert, dim=0).to(torch.int32)
139 # Column-major views let grouped_mm consume the existing parameter storage
140 # without retaining another full set of expert weights for backward.
141 w1_t = w1.transpose(1, 2)
142 w3_t = w3.transpose(1, 2)
143 w2_t = w2.transpose(1, 2)
144 h1 = torch._grouped_mm(x, w1_t, offs=offs) # pylint: disable=protected-access
145 h3 = torch._grouped_mm(x, w3_t, offs=offs) # pylint: disable=protected-access
146 h = F.silu(h1) * h3
147 if scores is not None:
148 h = h * scores.unsqueeze(-1)
149 return torch._grouped_mm(h, w2_t, offs=offs) # pylint: disable=protected-access
152def _run_experts_grouped_mm_npu(
153 w1: torch.Tensor,
154 w2: torch.Tensor,
155 w3: torch.Tensor,
156 x: torch.Tensor,
157 num_tokens_per_expert: torch.Tensor,
158 scores: torch.Tensor | None = None,
159) -> torch.Tensor:
160 """Fused grouped matmul path for Ascend NPU using ``torch_npu.npu_grouped_matmul``.
162 Requires ``torch_npu`` to be installed.
164 Args:
165 w1: Shape ``[num_experts, hidden_dim, dim]``.
166 w2: Shape ``[num_experts, dim, hidden_dim]``.
167 w3: Shape ``[num_experts, hidden_dim, dim]``.
168 x: Shape ``[total_routed_tokens, dim]``.
169 num_tokens_per_expert: 1-D integer tensor of length ``num_experts``.
170 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
171 Routing weights applied after the bias-free w2 projection.
172 When ``None``, no weighting is applied. Defaults to ``None``.
174 Returns:
175 Expert output of shape ``[total_routed_tokens, dim]``.
176 """
177 if x.shape[0] == 0:
178 parameter_zero = (w1.sum() + w2.sum() + w3.sum()).to(x.dtype) * 0.0
179 return x + parameter_zero
181 gate_up_weight = torch.cat((w1, w3), dim=1)
182 output = npu_grouped_swiglu(
183 x,
184 gate_up_weight,
185 w2,
186 num_tokens_per_expert,
187 )
188 if scores is not None:
189 output = output * scores.to(output.dtype).unsqueeze(-1)
190 return output
193# ---------------------------------------------------------------------------
194# FeedForward — shared expert / standard SwiGLU FFN
195# ---------------------------------------------------------------------------
197class FeedForward(nn.Module):
198 """SwiGLU feed-forward network, used as a shared (always-active) expert.
200 Implements: ``output = w2(silu(w1(x)) * w3(x))``
202 Args:
203 dim: Input embedding dimension.
204 hidden_dim: Intermediate hidden dimension.
205 bias: Whether to add a learnable bias. Defaults to ``False``.
207 Example::
208 >>> ff = FeedForward(dim=256, hidden_dim=512)
209 >>> out = ff(torch.randn(4, 16, 256))
210 >>> out.shape
211 torch.Size([4, 16, 256])
212 """
214 def __init__(self, dim: int, hidden_dim: int, bias: bool = False) -> None:
215 """Initialize FeedForward with three linear layers.
217 Args:
218 dim: Input embedding dimension.
219 hidden_dim: Intermediate hidden dimension.
220 bias: Whether to add a learnable bias. Defaults to ``False``.
221 """
222 super().__init__()
223 self.w1 = nn.Linear(dim, hidden_dim, bias=bias)
224 self.w2 = nn.Linear(hidden_dim, dim, bias=bias)
225 self.w3 = nn.Linear(dim, hidden_dim, bias=bias)
227 def forward(self, x: torch.Tensor) -> torch.Tensor:
228 """Compute SwiGLU feed-forward output.
230 Args:
231 x: Input tensor of shape ``(..., dim)``.
233 Returns:
234 Output tensor with the same leading shape and last dimension ``dim``.
235 """
236 return self.w2(F.silu(self.w1(x)) * self.w3(x))
239# ---------------------------------------------------------------------------
240# GroupedExperts
241# ---------------------------------------------------------------------------
243class GroupedExperts(nn.Module):
244 """Batch expert computation with optional grouped matrix-multiply.
246 All expert weights are stored in a single 3-D parameter so that EP / TP
247 sharding strategies can distribute the expert dimension via DTensor.
249 Args:
250 dim: Token embedding dimension.
251 hidden_dim: Expert hidden dimension (SwiGLU intermediate size).
252 num_experts: Total number of experts.
253 use_grouped_mm: If ``True``, uses a hardware-accelerated grouped
254 matmul kernel (``torch._grouped_mm`` on GPU,
255 ``torch_npu.npu_grouped_matmul`` on NPU). Falls back to the
256 for-loop path when neither is available. Defaults to ``False``.
258 Note:
259 When weights are DTensors (e.g. after TP sharding via
260 :class:`ExpertTensorParallel`), ``forward`` calls ``.to_local()``
261 before computation.
263 Example::
264 >>> experts = GroupedExperts(dim=8, hidden_dim=16, num_experts=4)
265 >>> x = torch.randn(10, 8)
266 >>> counts = torch.tensor([3, 2, 4, 1])
267 >>> out = experts(x, counts)
268 >>> out.shape
269 torch.Size([10, 8])
270 """
272 def __init__(
273 self,
274 dim: int,
275 hidden_dim: int,
276 num_experts: int,
277 use_grouped_mm: bool = False,
278 ) -> None:
279 """Initialize GroupedExperts with stacked expert weight matrices.
281 Args:
282 dim: Input embedding dimension.
283 hidden_dim: Intermediate hidden dimension.
284 num_experts: Number of experts.
285 use_grouped_mm: Whether to use grouped matrix multiplication.
286 """
287 super().__init__()
288 # Weight layout: [num_experts, out_dim, in_dim] so that the standard
289 # linear operation is x @ w[e].T.
290 self.w1 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
291 self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim))
292 self.w3 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
293 self.num_experts = num_experts
294 self.use_grouped_mm = use_grouped_mm
295 self._reset_parameters()
297 def _reset_parameters(self) -> None:
298 """Kaiming-uniform initialisation for all expert weight tensors."""
299 for weight in (self.w1, self.w2, self.w3):
300 nn.init.kaiming_uniform_(weight.view(weight.shape[0], -1), a=math.sqrt(5))
302 def forward(
303 self,
304 x: torch.Tensor,
305 num_tokens_per_expert: torch.Tensor,
306 scores: torch.Tensor | None = None,
307 ) -> torch.Tensor:
308 """Run all experts on their assigned tokens.
310 Args:
311 x: Routed tokens in expert-major order,
312 shape ``[total_routed_tokens, dim]``.
313 num_tokens_per_expert: 1-D integer tensor of length
314 ``num_local_experts`` with the token count per expert.
315 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
316 Routing weights for each token's expert output.
317 When ``None``, no weighting is applied. Defaults to ``None``.
319 Returns:
320 Expert output of shape ``[total_routed_tokens, dim]``.
321 """
322 # Extract local shard when parameters are DTensors (TP path).
323 w1 = self.w1.to_local() if isinstance(self.w1, DTensor) else self.w1
324 w2 = self.w2.to_local() if isinstance(self.w2, DTensor) else self.w2
325 w3 = self.w3.to_local() if isinstance(self.w3, DTensor) else self.w3
327 if not self.use_grouped_mm:
328 return _run_experts_for_loop(w1, w2, w3, x, num_tokens_per_expert, scores)
329 if _is_npu_tensor(x):
330 return _run_experts_grouped_mm_npu(w1, w2, w3, x, num_tokens_per_expert, scores)
331 if x.device.type == "cuda":
332 return _run_experts_grouped_mm_gpu(w1, w2, w3, x, num_tokens_per_expert, scores)
333 return _run_experts_for_loop(w1, w2, w3, x, num_tokens_per_expert, scores)
336# ---------------------------------------------------------------------------
337# TokenChoiceTopKRouter
338# ---------------------------------------------------------------------------
340class TokenChoiceTopKRouter(nn.Module):
341 """Top-K router: each token independently selects its top-K experts.
343 Args:
344 dim: Token embedding dimension (input to the gate).
345 num_experts: Total number of experts.
346 top_k: Experts selected per token. Defaults to ``1``.
347 score_func: Activation on gate logits before topk.
348 One of ``"sigmoid"`` (default) or ``"softmax"``.
349 num_expert_groups: For node-limited routing, number of expert
350 groups. ``None`` disables node-limited routing.
351 num_limited_groups: Groups to keep in node-limited routing. Required
352 when ``num_expert_groups`` is not ``None``.
353 route_scale: Scalar multiplier applied to routing scores.
354 Defaults to ``1.0``.
356 Example::
357 >>> router = TokenChoiceTopKRouter(dim=64, num_experts=8, top_k=2)
358 >>> scores, indices, counts = router(torch.randn(32, 64))
359 >>> scores.shape, indices.shape, counts.shape
360 (torch.Size([32, 2]), torch.Size([32, 2]), torch.Size([8]))
361 """
363 def __init__(
364 self,
365 dim: int,
366 num_experts: int,
367 top_k: int = 1,
368 score_func: str = "sigmoid",
369 num_expert_groups: int | None = None,
370 num_limited_groups: int | None = None,
371 route_scale: float = 1.0,
372 ) -> None:
373 """Initialize the top-K router.
375 Args:
376 dim: Input dimension of token representations.
377 num_experts: Total number of experts.
378 top_k: Number of experts to select per token.
379 score_func: Scoring function, ``"sigmoid"`` or ``"softmax"``.
380 num_expert_groups: Number of expert groups for node-limited routing.
381 num_limited_groups: Number of groups each token can route to.
382 route_scale: Scalar multiplier applied to routing scores.
383 """
384 super().__init__()
385 if score_func not in ("sigmoid", "softmax"):
386 raise ValueError(
387 f"score_func must be 'sigmoid' or 'softmax', got '{score_func}'."
388 )
389 if num_expert_groups is not None and num_limited_groups is None:
390 raise ValueError(
391 "num_limited_groups must be set when num_expert_groups is not None."
392 )
393 self.gate = nn.Linear(dim, num_experts, bias=False)
394 self.num_experts = num_experts
395 self.top_k = top_k
396 self.score_func = score_func
397 self.num_expert_groups = num_expert_groups
398 self.num_limited_groups = num_limited_groups
399 self.route_scale = route_scale
401 def forward(
402 self,
403 x: torch.Tensor,
404 expert_bias: torch.Tensor | None = None,
405 ) -> tuple:
406 """Compute routing scores and top-K expert assignments.
408 Args:
409 x: Token tensor of shape ``[num_tokens, dim]``.
410 expert_bias: Optional 1-D tensor of shape ``[num_experts]`` added
411 to gate logits for topk selection only (auxiliary-loss-free
412 load balancing). Does not affect the returned ``top_scores``.
414 Returns:
415 Tuple of:
417 - ``top_scores``: shape ``[num_tokens, top_k]`` — routing weights.
418 - ``selected_experts``: shape ``[num_tokens, top_k]`` — expert IDs.
419 - ``num_tokens_per_expert``: shape ``[num_experts]`` — load counts.
420 """
421 # Gate in float32 for numerical stability.
422 scores = self.gate(x).float()
424 if self.score_func == "sigmoid":
425 scores = torch.sigmoid(scores)
426 else:
427 scores = F.softmax(scores, dim=-1)
429 if self.route_scale != 1.0:
430 scores = scores * self.route_scale
432 # Node-limited routing — mask out low-scoring expert groups.
433 scores_for_topk = scores
434 if self.num_expert_groups is not None:
435 scores_for_topk = self._get_node_limited_routing_scores(scores)
437 # Add expert bias only for selection; returned scores remain unbiased.
438 scores_with_bias = scores_for_topk
439 if expert_bias is not None:
440 scores_with_bias = scores_for_topk + expert_bias.float()
442 top_scores, selected_experts = scores_with_bias.topk(self.top_k, dim=-1)
443 # Gather unbiased scores as the actual routing weights.
444 top_scores = scores.gather(1, selected_experts)
446 num_tokens_per_expert = torch.bincount(
447 selected_experts.flatten(), minlength=self.num_experts
448 )
449 return top_scores, selected_experts, num_tokens_per_expert
451 def _get_node_limited_routing_scores(self, scores: torch.Tensor) -> torch.Tensor:
452 """Mask out low-scoring expert groups (node-limited routing).
454 Args:
455 scores: Routing scores of shape ``[num_tokens, num_experts]``.
457 Returns:
458 Scores with non-selected groups masked to ``-inf``.
459 """
460 num_tokens, num_experts = scores.shape
461 experts_per_group = num_experts // self.num_expert_groups
462 group_scores = scores.view(
463 num_tokens, self.num_expert_groups, experts_per_group
464 ).max(dim=-1).values # [num_tokens, num_groups]
466 _, selected_groups = group_scores.topk(self.num_limited_groups, dim=-1)
468 mask = scores.new_zeros(num_tokens, self.num_expert_groups)
469 mask.scatter_(1, selected_groups, 1.0)
470 mask = (
471 mask.unsqueeze(-1)
472 .expand(num_tokens, self.num_expert_groups, experts_per_group)
473 .reshape(num_tokens, num_experts)
474 )
475 return scores.masked_fill(mask == 0, float("-inf"))
478# ---------------------------------------------------------------------------
479# Load-balance auxiliary loss
480# ---------------------------------------------------------------------------
482def _compute_load_balance_loss(
483 top_scores: torch.Tensor,
484 selected_experts: torch.Tensor,
485 num_experts: int,
486 sequence_partition_group: Any | None = None,
487) -> torch.Tensor:
488 """Compute load-balance auxiliary loss.
490 Standard formulation: ``loss = num_experts * Σ (expert_fraction_i * expert_prob_i)``
492 Where:
493 - expert_fraction_i = tokens routed to expert i / total routed tokens
494 - expert_prob_i = sum of routing scores for expert i / total tokens
496 When routing is perfectly balanced, loss ≈ 1.0.
497 When routing is imbalanced, loss > 1.0.
499 When ``sequence_partition_group`` is provided (e.g. the TP×CP or CP
500 process group), ``expert_fraction`` (token counts) is all-reduced across
501 the group so that the loss reflects the *global* token distribution over
502 the full sequence. ``expert_prob`` is kept local to preserve gradient
503 flow back to the router weights, and ``num_tokens`` is scaled by
504 ``num_sub_sequence`` (the group world size) to account for the sharded
505 sequence dimension. This follows the Megatron-LM
506 ``full_sequence_aux_loss`` pattern.
508 Args:
509 top_scores: Routing weights, shape ``[num_tokens, top_k]``.
510 selected_experts: Expert IDs, shape ``[num_tokens, top_k]``.
511 num_experts: Total number of experts.
512 sequence_partition_group: Optional process group spanning the
513 sequence-partition dimension (TP+SP or CP). When not ``None``,
514 ``expert_fraction`` is all-reduced across this group.
516 Returns:
517 Scalar loss tensor. Returns 0.0 for empty input.
518 """
519 import torch.distributed as dist # pylint: disable=C0415
521 num_tokens, top_k = top_scores.shape
523 if num_tokens == 0:
524 return torch.tensor(0.0, dtype=top_scores.dtype, device=top_scores.device)
526 flat_experts = selected_experts.flatten()
527 flat_scores = top_scores.flatten()
529 expert_fraction = torch.zeros(
530 num_experts, dtype=top_scores.dtype, device=top_scores.device
531 )
532 expert_fraction.scatter_add_(0, flat_experts, torch.ones_like(flat_scores))
534 num_sub_sequence = 1
535 if sequence_partition_group is not None:
536 num_sub_sequence = dist.get_world_size(sequence_partition_group)
537 dist.all_reduce(expert_fraction, group=sequence_partition_group)
539 expert_fraction = expert_fraction / (num_tokens * num_sub_sequence * top_k)
541 expert_prob = torch.zeros(
542 num_experts, dtype=top_scores.dtype, device=top_scores.device
543 )
544 expert_prob.scatter_add_(0, flat_experts, flat_scores)
545 expert_prob = expert_prob / (num_tokens * num_sub_sequence)
547 loss = num_experts * (expert_fraction * expert_prob).sum()
549 return loss
552# ---------------------------------------------------------------------------
553# MoEAuxLossAutoScaler — gradient injection for auxiliary loss
554# ---------------------------------------------------------------------------
557class MoEAuxLossAutoScaler(torch.autograd.Function):
558 """Autograd function that injects auxiliary-loss gradient into the backward chain.
560 Forward transparently passes through ``top_scores`` (values unchanged).
561 Backward injects a scaled gradient for ``aux_loss``, causing its gradients
562 to flow through the router weights without adding aux_loss to the main
563 loss explicitly.
565 The loss scale is managed via the class-level :meth:`set_loss_scale` method,
566 which should be called before ``loss.backward()`` in the training loop
567 (typically by the pipeline schedule or training framework).
569 Reference: Megatron-LM ``megatron/core/transformer/moe/moe_utils.py``
570 """
572 main_loss_backward_scale: torch.Tensor | None = None
574 @staticmethod
575 def forward(ctx: Any, output: torch.Tensor, aux_loss: torch.Tensor) -> torch.Tensor:
576 """Pass through ``output`` unchanged; save ``aux_loss`` for backward.
578 Args:
579 output: Router output (e.g. ``top_scores``), passed through unchanged.
580 aux_loss: Scalar auxiliary loss tensor to inject gradient for.
582 Returns:
583 The ``output`` tensor, identical in value but with an added
584 autograd edge to ``aux_loss``.
585 """
586 ctx.save_for_backward(aux_loss)
587 return output
589 @staticmethod
590 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
591 """Inject scaled aux_loss gradient into the backward chain.
593 Args:
594 grad_output: Gradient flowing back through ``output``.
596 Returns:
597 Tuple of (grad_output unchanged, scaled aux_loss gradient).
598 """
599 (aux_loss,) = ctx.saved_tensors
600 if MoEAuxLossAutoScaler.main_loss_backward_scale is None:
601 MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(
602 1.0, device=aux_loss.device,
603 )
604 aux_loss_backward_scale = MoEAuxLossAutoScaler.main_loss_backward_scale
605 scaled_aux_loss_grad = torch.ones_like(aux_loss) * aux_loss_backward_scale
606 return grad_output, scaled_aux_loss_grad
608 @staticmethod
609 def set_loss_scale(scale: torch.Tensor) -> None:
610 """Set the gradient scale for auxiliary loss.
612 Should be called before ``loss.backward()`` in the training loop.
613 Uses in-place copy on subsequent calls to avoid tensor reallocation.
615 Args:
616 scale: Tensor containing the loss scale value, typically
617 ``1 / (num_microbatches * dp_size)`` or similar.
618 """
619 if MoEAuxLossAutoScaler.main_loss_backward_scale is None:
620 MoEAuxLossAutoScaler.main_loss_backward_scale = scale
621 else:
622 MoEAuxLossAutoScaler.main_loss_backward_scale.copy_(scale)
625# ---------------------------------------------------------------------------
626# MoE orchestrator
627# ---------------------------------------------------------------------------
629class MoE(nn.Module):
630 """Mixture-of-Experts layer.
632 Orchestrates routing, token permutation, expert computation, and output
633 scatter-add. Supports shared experts, auxiliary-loss-free load balancing
634 via expert bias, node-limited routing, and auxiliary load-balance loss.
636 Args:
637 dim: Token embedding dimension.
638 hidden_dim: Expert hidden dimension.
639 num_experts: Total number of experts.
640 top_k: Experts selected per token. Defaults to ``1``.
641 score_before_experts: If ``True``, multiply routed tokens by routing
642 weights *before* expert computation; otherwise multiply expert
643 outputs *after*. Defaults to ``True``.
644 load_balance_coeff: When not ``None``, enables auxiliary load-balance
645 loss. The loss gradient is automatically injected into router
646 weights via :class:`MoEAuxLossAutoScaler` (no manual loss
647 addition needed). The scalar loss value is attached as
648 ``output._load_balance_loss`` for logging purposes.
649 sequence_partition_group: Optional process group spanning the
650 sequence-partition dimension (TP+SP or CP). When provided,
651 ``expert_fraction`` is all-reduced across this group so the
652 load-balance loss reflects global token distribution. Defaults
653 to ``None`` (no cross-rank synchronization).
654 shared_expert: Optional :class:`FeedForward` running on every token
655 in parallel; output added to routed-expert output.
656 router_kwargs: Extra keyword arguments forwarded to
657 :class:`TokenChoiceTopKRouter`.
658 use_grouped_mm: If ``True``, uses a hardware-accelerated grouped
659 matmul kernel (e.g. ``npu_grouped_matmul``) inside
660 :class:`GroupedExperts`. Defaults to ``False``.
662 Note:
663 *Auxiliary-loss-free load balancing*: call :func:`update_expert_bias`
664 once per optimiser step to adjust ``expert_bias`` from the accumulated
665 ``tokens_per_expert`` histogram.
667 Example::
668 >>> moe = MoE(dim=64, hidden_dim=128, num_experts=8, top_k=2)
669 >>> out = moe(torch.randn(2, 16, 64))
670 >>> out.shape
671 torch.Size([2, 16, 64])
672 """
674 def __init__(
675 self,
676 dim: int,
677 hidden_dim: int,
678 num_experts: int,
679 top_k: int = 1,
680 score_before_experts: bool = True,
681 load_balance_coeff: float | None = None,
682 sequence_partition_group: Any | None = None,
683 shared_expert: FeedForward | None = None,
684 router_kwargs: dict | None = None,
685 use_grouped_mm: bool = False,
686 ) -> None:
687 """Initialize MoE block with experts, router and optional shared expert.
689 Args:
690 dim: Input embedding dimension.
691 hidden_dim: Intermediate hidden dimension of each expert.
692 num_experts: Number of experts.
693 top_k: Number of experts to select per token.
694 score_before_experts: If ``True``, multiply routed hidden states by
695 the routing score.
696 load_balance_coeff: Load-balance loss coefficient.
697 shared_expert: Optional shared expert applied to all tokens.
698 router_kwargs: Additional keyword arguments for the router.
699 use_grouped_mm: Whether to use grouped matrix multiplication.
700 """
701 super().__init__()
702 router_kw = router_kwargs or {}
703 self.experts = GroupedExperts(
704 dim=dim, hidden_dim=hidden_dim, num_experts=num_experts,
705 use_grouped_mm=use_grouped_mm,
706 )
707 self.router = TokenChoiceTopKRouter(
708 dim=dim, num_experts=num_experts, top_k=top_k, **router_kw,
709 )
710 self.shared_expert = shared_expert
711 self.num_experts = num_experts
712 self.top_k = top_k
713 self.score_before_experts = score_before_experts
714 self.load_balance_coeff = load_balance_coeff
715 self.sequence_partition_group = sequence_partition_group
716 self.last_aux_loss: torch.Tensor | None = None
717 self.enable_expert_bias = True
719 # Auxiliary-loss-free load-balance buffers (no gradient).
720 self.register_buffer("expert_bias", torch.zeros(num_experts))
721 self.register_buffer("tokens_per_expert", torch.zeros(num_experts))
723 def permutation(
724 self,
725 selected_experts: torch.Tensor,
726 top_scores: torch.Tensor,
727 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
728 """Compute token-major → expert-major permutation indices.
730 Args:
731 selected_experts: Expert IDs of shape ``[num_tokens, top_k]``.
732 top_scores: Routing weights of shape ``[num_tokens, top_k]``.
734 Returns:
735 Tuple of ``(token_indices, top_scores_sorted,
736 num_tokens_per_expert)`` where *token_indices* maps each
737 expert-major slot back to its source token row,
738 *top_scores_sorted* are the scores in expert-major order, and
739 *num_tokens_per_expert* has shape ``[num_experts]``.
740 """
741 # flat_experts[i] is the expert ID for the i-th (token, top_k) slot.
742 flat_experts = selected_experts.flatten() # [num_tokens * top_k]
743 flat_indices = flat_experts.argsort(stable=True) # expert-major permutation
744 top_scores_sorted = top_scores.flatten()[flat_indices] # [num_tokens * top_k]
745 # Each entry in flat_indices maps to a position in [0, num_tokens * top_k);
746 # divide by top_k to recover the original token row index.
747 token_indices = flat_indices // self.top_k # [num_tokens * top_k]
748 num_tokens_per_expert = torch.bincount(
749 flat_experts, minlength=self.num_experts
750 )
751 return token_indices, top_scores_sorted, num_tokens_per_expert
753 def unpermutation(
754 self,
755 expert_out: torch.Tensor,
756 token_indices: torch.Tensor,
757 num_tokens: int,
758 dim: int,
759 ) -> torch.Tensor:
760 """Scatter expert outputs back to token-major order.
762 Args:
763 expert_out: Expert output tensor in expert-major order,
764 shape ``[total_routed_tokens, dim]``.
765 token_indices: Permutation indices from :meth:`permutation`,
766 shape ``[num_tokens * top_k]``.
767 num_tokens: Total number of tokens (unflattened).
768 dim: Feature dimension.
770 Returns:
771 Token-major output tensor of shape ``[num_tokens, dim]``.
772 """
773 # Use out-of-place ``scatter_add`` so autograd correctly records
774 # ``ScatterAddBackward``; ``new_zeros + scatter_add_`` on some
775 # backends (torch_npu) leaves the leaf un-upgraded and the result
776 # without a ``grad_fn``.
777 return torch.zeros(
778 num_tokens, dim, dtype=expert_out.dtype, device=expert_out.device,
779 ).scatter_add(
780 0,
781 token_indices.unsqueeze(1).expand(-1, dim),
782 expert_out,
783 )
785 def forward(self, x: torch.Tensor) -> torch.Tensor:
786 """Run the MoE layer.
788 Args:
789 x: Input tensor of shape ``[batch, seq_len, dim]``.
791 Returns:
792 Output tensor of shape ``[batch, seq_len, dim]``. When
793 ``load_balance_coeff`` is set, carries a ``_load_balance_loss``
794 attribute with the auxiliary loss scalar (for logging only;
795 gradient injection is handled automatically by
796 :class:`MoEAuxLossAutoScaler`).
797 """
798 bs, seq_len, dim = x.shape
799 num_tokens = bs * seq_len
800 x_flat = x.view(num_tokens, dim) # [num_tokens, dim]
802 # --- Routing ---
803 top_scores, selected_experts, token_counts = self.router(
804 x_flat, self.expert_bias
805 )
807 # Accumulate token histogram without creating gradient nodes.
808 with torch.no_grad():
809 self.tokens_per_expert.add_(token_counts.float())
811 # --- Auxiliary load-balance loss ---
812 # Compute aux_loss early and attach it to top_scores via
813 # MoEAuxLossAutoScaler so that backward through top_scores
814 # also triggers aux_loss gradient (injected into router weights).
815 if self.load_balance_coeff is not None:
816 lb_loss = self.load_balance_coeff * _compute_load_balance_loss(
817 top_scores, selected_experts, self.num_experts,
818 sequence_partition_group=self.sequence_partition_group,
819 )
820 # Apply AutoScaler *before* top_scores is used in the forward path
821 # so that the main backward through top_scores triggers aux_loss
822 # gradient injection. Forward values are unchanged.
823 top_scores = MoEAuxLossAutoScaler.apply(top_scores, lb_loss)
824 self.last_aux_loss = lb_loss.detach()
825 else:
826 lb_loss = None
827 self.last_aux_loss = None
829 # --- Token permutation: token-major → expert-major ---
830 token_indices, top_scores_sorted, num_tokens_per_expert = self.permutation(
831 selected_experts, top_scores
832 )
834 # Gather routed tokens in expert-major order.
835 routed_x = x_flat[token_indices] # [num_tokens * top_k, dim]
837 # --- Expert computation ---
838 if self.score_before_experts:
839 routed_x = routed_x * top_scores_sorted.unsqueeze(1)
840 expert_out = self.experts(routed_x, num_tokens_per_expert, scores=None)
841 else:
842 expert_out = self.experts(routed_x, num_tokens_per_expert, scores=top_scores_sorted)
844 # --- Shared expert (parallel with routed experts) ---
845 shared_out = None
846 if self.shared_expert is not None:
847 shared_out = self.shared_expert(x_flat)
850 # --- Scatter expert outputs back to token order ---
851 out = self.unpermutation(expert_out, token_indices, num_tokens, dim)
853 if shared_out is not None:
854 out = out + shared_out
856 result = out.view(bs, seq_len, dim)
858 # Attach auxiliary loss to the returned tensor for logging.
859 if lb_loss is not None:
860 result._load_balance_loss = lb_loss # pylint: disable=protected-access
862 return result
864 def update_expert_bias(
865 self,
866 lr: float = 1e-3,
867 num_recomputations: int = 1,
868 ) -> None:
869 """Update expert bias for auxiliary-loss-free load balancing.
871 Should be called once per training step after the optimizer step.
872 Adjusts ``expert_bias`` to push token load towards the mean, then
873 resets the ``tokens_per_expert`` accumulator.
875 The update delta is centered to have zero mean, preventing systematic
876 drift of all bias values over time.
878 When activation checkpoint is enabled, forward is re-executed during
879 backward, causing ``tokens_per_expert`` to accumulate twice (or more).
880 Use ``num_recomputations`` to correct for this double-counting.
882 Args:
883 lr: Step size for the bias update. Defaults to ``1e-3``.
884 num_recomputations: Number of times forward is executed per optimizer
885 step. Default ``1`` (normal training). Set to ``2`` when activation
886 checkpoint is enabled (forward + recompute during backward). For
887 nested checkpoint strategies, set to the actual execution count.
889 Example::
890 >>> # Single-card scenario without activation checkpoint:
891 >>> moe_layer.update_expert_bias(lr=1e-3)
892 >>>
893 >>> # With activation checkpoint (forward executed twice):
894 >>> moe_layer.update_expert_bias(lr=1e-3, num_recomputations=2)
895 """
896 with torch.no_grad():
897 if num_recomputations > 1:
898 self.tokens_per_expert.div_(num_recomputations)
899 avg = self.tokens_per_expert.float().mean()
900 delta = lr * (avg - self.tokens_per_expert.float()).sign()
901 delta = delta - delta.mean()
902 self.expert_bias.data += delta
903 self.tokens_per_expert.zero_()
906# ---------------------------------------------------------------------------
907# Expert bias update for auxiliary-loss-free load balancing
908# ---------------------------------------------------------------------------
910def update_expert_bias(
911 moe: MoE,
912 lr: float = 1e-3,
913 num_recomputations: int = 1,
914) -> None:
915 """Update expert bias for auxiliary-loss-free load balancing.
917 This is a module-level wrapper that calls ``moe.update_expert_bias()``.
918 Prefer calling the instance method directly.
920 Should be called once per training step after the optimizer step.
921 Adjusts ``moe.expert_bias`` to push token load towards the mean, then
922 resets the ``tokens_per_expert`` accumulator.
924 The update delta is centered to have zero mean, preventing systematic
925 drift of all bias values over time.
927 When activation checkpoint is enabled, forward is re-executed during
928 backward, causing ``tokens_per_expert`` to accumulate twice (or more).
929 Use ``num_recomputations`` to correct for this double-counting.
931 Args:
932 moe: The :class:`MoE` module whose bias should be updated.
933 lr: Step size for the bias update. Defaults to ``1e-3``.
934 num_recomputations: Number of times forward is executed per optimizer
935 step. Default ``1`` (normal training). Set to ``2`` when activation
936 checkpoint is enabled (forward + recompute during backward). For
937 nested checkpoint strategies, set to the actual execution count.
939 Example::
940 >>> # Single-card scenario without activation checkpoint:
941 >>> update_expert_bias(moe_layer, lr=1e-3)
942 >>>
943 >>> # With activation checkpoint (forward executed twice):
944 >>> update_expert_bias(moe_layer, lr=1e-3, num_recomputations=2)
945 """
946 moe.update_expert_bias(lr=lr, num_recomputations=num_recomputations)