Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / common / moe.py: 84%
228 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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"""
27__all__ = [
28 "FeedForward",
29 "GroupedExperts",
30 "TokenChoiceTopKRouter",
31 "MoE",
32 "MoEAuxLossAutoScaler",
33 "update_expert_bias",
34]
36import math
37from typing import Any, Optional, Tuple
39import torch
40from torch import nn
41import torch.nn.functional as F
43from hyper_parallel.core.dtensor.dtensor import DTensor
46# ---------------------------------------------------------------------------
47# Grouped expert computation kernels
48# ---------------------------------------------------------------------------
50def _run_experts_for_loop(
51 w1: torch.Tensor,
52 w2: torch.Tensor,
53 w3: torch.Tensor,
54 x: torch.Tensor,
55 num_tokens_per_expert: torch.Tensor,
56 scores: Optional[torch.Tensor] = None,
57) -> torch.Tensor:
58 """Run per-expert SwiGLU via a sequential loop (reference path).
60 Args:
61 w1: Shape ``[num_experts, hidden_dim, dim]``.
62 w2: Shape ``[num_experts, dim, hidden_dim]``.
63 w3: Shape ``[num_experts, hidden_dim, dim]``.
64 x: Routed tokens in expert-major order, shape
65 ``[total_routed_tokens, dim]``.
66 num_tokens_per_expert: 1-D integer tensor of length ``num_experts``.
67 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
68 Routing weights applied to intermediate activations (``silu(w1(x)) * w3(x)``)
69 before the w2 projection. When ``None``, no weighting is applied.
70 Defaults to ``None``.
72 Returns:
73 Expert output of shape ``[total_routed_tokens, dim]``.
74 """
75 # Use ``torch.cat`` instead of ``zeros_like + in-place slice assign``.
76 # On some backends (notably torch_npu) in-place slice assignment onto a
77 # ``requires_grad=False`` leaf tensor does not reliably upgrade it to a
78 # non-leaf with a ``grad_fn`` — the forward result may end up with
79 # ``grad_fn=None`` and downstream ``backward()`` fails with
80 # "element 0 of tensors does not require grad and does not have a grad_fn".
81 #
82 # Drain ``num_tokens_per_expert`` to host **once** via ``.tolist()``
83 # rather than calling ``int(n)`` per loop iteration — a single D2H
84 # copy instead of ``num_local_experts`` separate ones. Per-iter
85 # ``.item()`` would stall the host between expert kernels and shrink
86 # the dual-pipe overlap window.
87 counts_list = num_tokens_per_expert.tolist()
88 parts = []
89 offset = 0
90 for e, n in enumerate(counts_list):
91 if n == 0:
92 continue
93 x_e = x[offset:offset + n]
94 h = F.silu(x_e @ w1[e].T) * (x_e @ w3[e].T)
95 if scores is not None:
96 h = h * scores[offset:offset + n].unsqueeze(-1)
97 parts.append(h @ w2[e].T)
98 offset += n
99 if not parts:
100 # No routed tokens: return a grad-connected zero (not ``zeros_like``).
101 return x * 0.0
102 return torch.cat(parts, dim=0)
105def _run_experts_grouped_mm_gpu(
106 w1: torch.Tensor,
107 w2: torch.Tensor,
108 w3: torch.Tensor,
109 x: torch.Tensor,
110 num_tokens_per_expert: torch.Tensor,
111 scores: Optional[torch.Tensor] = None,
112) -> torch.Tensor:
113 """Fused grouped matmul path for NVIDIA GPU using ``torch._grouped_mm``.
115 Args:
116 w1: Shape ``[num_experts, hidden_dim, dim]``.
117 w2: Shape ``[num_experts, dim, hidden_dim]``.
118 w3: Shape ``[num_experts, hidden_dim, dim]``.
119 x: Shape ``[total_routed_tokens, dim]``.
120 num_tokens_per_expert: 1-D integer tensor of length ``num_experts``.
121 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
122 Routing weights applied to intermediate activations (``silu(w1(x)) * w3(x)``)
123 before the w2 projection. When ``None``, no weighting is applied.
124 Defaults to ``None``.
126 Returns:
127 Expert output of shape ``[total_routed_tokens, dim]``.
128 """
129 # offs: cumulative split offsets (int32) for torch._grouped_mm.
130 offs = torch.cumsum(num_tokens_per_expert[:-1], dim=0).to(torch.int32)
131 # w1/w3 stored as [num_experts, hidden_dim, dim]; grouped_mm expects
132 # [num_experts, dim, hidden_dim], so transpose the inner two dims.
133 w1_t = w1.transpose(1, 2).contiguous() # [num_experts, dim, hidden_dim]
134 w3_t = w3.transpose(1, 2).contiguous()
135 w2_t = w2.transpose(1, 2).contiguous() # [num_experts, hidden_dim, dim]
136 h1 = torch._grouped_mm(x, w1_t, offs=offs) # pylint: disable=protected-access
137 h3 = torch._grouped_mm(x, w3_t, offs=offs) # pylint: disable=protected-access
138 h = F.silu(h1) * h3
139 if scores is not None:
140 h = h * scores.unsqueeze(-1)
141 return torch._grouped_mm(h, w2_t, offs=offs) # pylint: disable=protected-access
144def _run_experts_grouped_mm_npu(
145 w1: torch.Tensor,
146 w2: torch.Tensor,
147 w3: torch.Tensor,
148 x: torch.Tensor,
149 num_tokens_per_expert: torch.Tensor,
150 scores: Optional[torch.Tensor] = None,
151) -> torch.Tensor:
152 """Fused grouped matmul path for Ascend NPU using ``torch_npu.npu_grouped_matmul``.
154 Requires ``torch_npu`` to be installed.
156 Args:
157 w1: Shape ``[num_experts, hidden_dim, dim]``.
158 w2: Shape ``[num_experts, dim, hidden_dim]``.
159 w3: Shape ``[num_experts, hidden_dim, dim]``.
160 x: Shape ``[total_routed_tokens, dim]``.
161 num_tokens_per_expert: 1-D integer tensor of length ``num_experts``.
162 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
163 Routing weights applied to intermediate activations (``silu(w1(x)) * w3(x)``)
164 before the w2 projection. When ``None``, no weighting is applied.
165 Defaults to ``None``.
167 Returns:
168 Expert output of shape ``[total_routed_tokens, dim]``.
169 """
170 import torch_npu # pylint: disable=C0415
172 # npu_grouped_matmul computes y = x @ weight (no implicit transpose).
173 # Our weight storage is [num_experts, out_dim, in_dim], matching F.linear's
174 # convention (weight.T for y = x @ weight.T). Transpose each expert shard
175 # so the shapes satisfy: [tokens, in_dim] @ [in_dim, out_dim] = [tokens, out_dim].
176 num_experts = w1.shape[0]
177 counts = num_tokens_per_expert.tolist()
178 x_list = list(torch.split(x, counts, dim=0))
179 # w1, w3: [E, hidden_dim, dim] → transposed per-expert: [dim, hidden_dim]
180 # w2: [E, dim, hidden_dim] → transposed per-expert: [hidden_dim, dim]
181 w1_list = [w1[e].T.contiguous() for e in range(num_experts)]
182 w2_list = [w2[e].T.contiguous() for e in range(num_experts)]
183 w3_list = [w3[e].T.contiguous() for e in range(num_experts)]
185 # npu_grouped_matmul: multi-multi-multi mode (x[i] @ weight[i]).
186 # group_type=-1 selects independent per-expert matmul (no shared axis).
187 h1_list = torch_npu.npu_grouped_matmul(x_list, w1_list, group_type=-1)
188 h3_list = torch_npu.npu_grouped_matmul(x_list, w3_list, group_type=-1)
189 h_list = [F.silu(h1) * h3 for h1, h3 in zip(h1_list, h3_list)]
190 if scores is not None:
191 offset = 0
192 for i, h in enumerate(h_list):
193 n = counts[i]
194 if n > 0:
195 h_list[i] = h * scores[offset:offset + n].unsqueeze(-1)
196 offset += n
197 out_list = torch_npu.npu_grouped_matmul(h_list, w2_list, group_type=-1)
198 return torch.cat(out_list, dim=0)
201# ---------------------------------------------------------------------------
202# FeedForward — shared expert / standard SwiGLU FFN
203# ---------------------------------------------------------------------------
205class FeedForward(nn.Module):
206 """SwiGLU feed-forward network, used as a shared (always-active) expert.
208 Implements: ``output = w2(silu(w1(x)) * w3(x))``
210 Args:
211 dim: Input embedding dimension.
212 hidden_dim: Intermediate hidden dimension.
213 bias: Whether to add a learnable bias. Defaults to ``False``.
215 Example::
216 >>> ff = FeedForward(dim=256, hidden_dim=512)
217 >>> out = ff(torch.randn(4, 16, 256))
218 >>> out.shape
219 torch.Size([4, 16, 256])
220 """
222 def __init__(self, dim: int, hidden_dim: int, bias: bool = False) -> None:
223 """Initialize FeedForward with three linear layers.
225 Args:
226 dim: Input embedding dimension.
227 hidden_dim: Intermediate hidden dimension.
228 bias: Whether to add a learnable bias. Defaults to ``False``.
229 """
230 super().__init__()
231 self.w1 = nn.Linear(dim, hidden_dim, bias=bias)
232 self.w2 = nn.Linear(hidden_dim, dim, bias=bias)
233 self.w3 = nn.Linear(dim, hidden_dim, bias=bias)
235 def forward(self, x: torch.Tensor) -> torch.Tensor:
236 """Compute SwiGLU feed-forward output.
238 Args:
239 x: Input tensor of shape ``(..., dim)``.
241 Returns:
242 Output tensor with the same leading shape and last dimension ``dim``.
243 """
244 return self.w2(F.silu(self.w1(x)) * self.w3(x))
247# ---------------------------------------------------------------------------
248# GroupedExperts
249# ---------------------------------------------------------------------------
251class GroupedExperts(nn.Module):
252 """Batch expert computation with optional grouped matrix-multiply.
254 All expert weights are stored in a single 3-D parameter so that EP / TP
255 sharding strategies can distribute the expert dimension via DTensor.
257 Args:
258 dim: Token embedding dimension.
259 hidden_dim: Expert hidden dimension (SwiGLU intermediate size).
260 num_experts: Total number of experts.
261 use_grouped_mm: If ``True``, uses a hardware-accelerated grouped
262 matmul kernel (``torch._grouped_mm`` on GPU,
263 ``torch_npu.npu_grouped_matmul`` on NPU). Falls back to the
264 for-loop path when neither is available. Defaults to ``False``.
266 Note:
267 When weights are DTensors (e.g. after TP sharding via
268 :class:`ExpertTensorParallel`), ``forward`` calls ``.to_local()``
269 before computation.
271 Example::
272 >>> experts = GroupedExperts(dim=8, hidden_dim=16, num_experts=4)
273 >>> x = torch.randn(10, 8)
274 >>> counts = torch.tensor([3, 2, 4, 1])
275 >>> out = experts(x, counts)
276 >>> out.shape
277 torch.Size([10, 8])
278 """
280 def __init__(
281 self,
282 dim: int,
283 hidden_dim: int,
284 num_experts: int,
285 use_grouped_mm: bool = False,
286 ) -> None:
287 """Initialize GroupedExperts with stacked expert weight matrices.
289 Args:
290 dim: Input embedding dimension.
291 hidden_dim: Intermediate hidden dimension.
292 num_experts: Number of experts.
293 use_grouped_mm: Whether to use grouped matrix multiplication.
294 """
295 super().__init__()
296 # Weight layout: [num_experts, out_dim, in_dim] so that the standard
297 # linear operation is x @ w[e].T.
298 self.w1 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
299 self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim))
300 self.w3 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
301 self.num_experts = num_experts
302 self.use_grouped_mm = use_grouped_mm
303 self._reset_parameters()
305 def _reset_parameters(self) -> None:
306 """Kaiming-uniform initialisation for all expert weight tensors."""
307 for weight in (self.w1, self.w2, self.w3):
308 nn.init.kaiming_uniform_(weight.view(weight.shape[0], -1), a=math.sqrt(5))
310 def forward(
311 self,
312 x: torch.Tensor,
313 num_tokens_per_expert: torch.Tensor,
314 scores: Optional[torch.Tensor] = None,
315 ) -> torch.Tensor:
316 """Run all experts on their assigned tokens.
318 Args:
319 x: Routed tokens in expert-major order,
320 shape ``[total_routed_tokens, dim]``.
321 num_tokens_per_expert: 1-D integer tensor of length
322 ``num_local_experts`` with the token count per expert.
323 scores: Optional 1-D tensor of shape ``[total_routed_tokens]``.
324 Routing weights applied to intermediate activations before w2.
325 When ``None``, no weighting is applied. Defaults to ``None``.
327 Returns:
328 Expert output of shape ``[total_routed_tokens, dim]``.
329 """
330 # Extract local shard when parameters are DTensors (TP path).
331 w1 = self.w1.to_local() if isinstance(self.w1, DTensor) else self.w1
332 w2 = self.w2.to_local() if isinstance(self.w2, DTensor) else self.w2
333 w3 = self.w3.to_local() if isinstance(self.w3, DTensor) else self.w3
335 if not self.use_grouped_mm:
336 return _run_experts_for_loop(w1, w2, w3, x, num_tokens_per_expert, scores)
338 if hasattr(torch, 'npu') and torch.npu.is_available():
339 return _run_experts_grouped_mm_npu(w1, w2, w3, x, num_tokens_per_expert, scores)
340 if torch.cuda.is_available():
341 return _run_experts_grouped_mm_gpu(w1, w2, w3, x, num_tokens_per_expert, scores)
343 return _run_experts_for_loop(w1, w2, w3, x, num_tokens_per_expert, scores)
346# ---------------------------------------------------------------------------
347# TokenChoiceTopKRouter
348# ---------------------------------------------------------------------------
350class TokenChoiceTopKRouter(nn.Module):
351 """Top-K router: each token independently selects its top-K experts.
353 Args:
354 dim: Token embedding dimension (input to the gate).
355 num_experts: Total number of experts.
356 top_k: Experts selected per token. Defaults to ``1``.
357 score_func: Activation on gate logits before topk.
358 One of ``"sigmoid"`` (default) or ``"softmax"``.
359 num_expert_groups: For node-limited routing, number of expert
360 groups. ``None`` disables node-limited routing.
361 num_limited_groups: Groups to keep in node-limited routing. Required
362 when ``num_expert_groups`` is not ``None``.
363 route_scale: Scalar multiplier applied to routing scores.
364 Defaults to ``1.0``.
366 Example::
367 >>> router = TokenChoiceTopKRouter(dim=64, num_experts=8, top_k=2)
368 >>> scores, indices, counts = router(torch.randn(32, 64))
369 >>> scores.shape, indices.shape, counts.shape
370 (torch.Size([32, 2]), torch.Size([32, 2]), torch.Size([8]))
371 """
373 def __init__(
374 self,
375 dim: int,
376 num_experts: int,
377 top_k: int = 1,
378 score_func: str = "sigmoid",
379 num_expert_groups: Optional[int] = None,
380 num_limited_groups: Optional[int] = None,
381 route_scale: float = 1.0,
382 ) -> None:
383 """Initialize the top-K router.
385 Args:
386 dim: Input dimension of token representations.
387 num_experts: Total number of experts.
388 top_k: Number of experts to select per token.
389 score_func: Scoring function, ``"sigmoid"`` or ``"softmax"``.
390 num_expert_groups: Number of expert groups for node-limited routing.
391 num_limited_groups: Number of groups each token can route to.
392 route_scale: Scalar multiplier applied to routing scores.
393 """
394 super().__init__()
395 if score_func not in ("sigmoid", "softmax"):
396 raise ValueError(
397 f"score_func must be 'sigmoid' or 'softmax', got '{score_func}'."
398 )
399 if num_expert_groups is not None and num_limited_groups is None:
400 raise ValueError(
401 "num_limited_groups must be set when num_expert_groups is not None."
402 )
403 self.gate = nn.Linear(dim, num_experts, bias=False)
404 self.num_experts = num_experts
405 self.top_k = top_k
406 self.score_func = score_func
407 self.num_expert_groups = num_expert_groups
408 self.num_limited_groups = num_limited_groups
409 self.route_scale = route_scale
411 def forward(
412 self,
413 x: torch.Tensor,
414 expert_bias: Optional[torch.Tensor] = None,
415 ) -> tuple:
416 """Compute routing scores and top-K expert assignments.
418 Args:
419 x: Token tensor of shape ``[num_tokens, dim]``.
420 expert_bias: Optional 1-D tensor of shape ``[num_experts]`` added
421 to gate logits for topk selection only (auxiliary-loss-free
422 load balancing). Does not affect the returned ``top_scores``.
424 Returns:
425 Tuple of:
427 - ``top_scores``: shape ``[num_tokens, top_k]`` — routing weights.
428 - ``selected_experts``: shape ``[num_tokens, top_k]`` — expert IDs.
429 - ``num_tokens_per_expert``: shape ``[num_experts]`` — load counts.
430 """
431 # Gate in float32 for numerical stability.
432 scores = self.gate(x).float()
434 if self.score_func == "sigmoid":
435 scores = torch.sigmoid(scores)
436 else:
437 scores = F.softmax(scores, dim=-1)
439 if self.route_scale != 1.0:
440 scores = scores * self.route_scale
442 # Node-limited routing — mask out low-scoring expert groups.
443 scores_for_topk = scores
444 if self.num_expert_groups is not None:
445 scores_for_topk = self._get_node_limited_routing_scores(scores)
447 # Add expert bias only for selection; returned scores remain unbiased.
448 scores_with_bias = scores_for_topk
449 if expert_bias is not None:
450 scores_with_bias = scores_for_topk + expert_bias.float()
452 top_scores, selected_experts = scores_with_bias.topk(self.top_k, dim=-1)
453 # Gather unbiased scores as the actual routing weights.
454 top_scores = scores.gather(1, selected_experts)
456 num_tokens_per_expert = torch.bincount(
457 selected_experts.flatten(), minlength=self.num_experts
458 )
459 return top_scores, selected_experts, num_tokens_per_expert
461 def _get_node_limited_routing_scores(self, scores: torch.Tensor) -> torch.Tensor:
462 """Mask out low-scoring expert groups (node-limited routing).
464 Args:
465 scores: Routing scores of shape ``[num_tokens, num_experts]``.
467 Returns:
468 Scores with non-selected groups masked to ``-inf``.
469 """
470 num_tokens, num_experts = scores.shape
471 experts_per_group = num_experts // self.num_expert_groups
472 group_scores = scores.view(
473 num_tokens, self.num_expert_groups, experts_per_group
474 ).max(dim=-1).values # [num_tokens, num_groups]
476 _, selected_groups = group_scores.topk(self.num_limited_groups, dim=-1)
478 mask = scores.new_zeros(num_tokens, self.num_expert_groups)
479 mask.scatter_(1, selected_groups, 1.0)
480 mask = (
481 mask.unsqueeze(-1)
482 .expand(num_tokens, self.num_expert_groups, experts_per_group)
483 .reshape(num_tokens, num_experts)
484 )
485 return scores.masked_fill(mask == 0, float("-inf"))
488# ---------------------------------------------------------------------------
489# Load-balance auxiliary loss
490# ---------------------------------------------------------------------------
492def _compute_load_balance_loss(
493 top_scores: torch.Tensor,
494 selected_experts: torch.Tensor,
495 num_experts: int,
496 sequence_partition_group: Optional[Any] = None,
497) -> torch.Tensor:
498 """Compute load-balance auxiliary loss.
500 Standard formulation: ``loss = num_experts * Σ (expert_fraction_i * expert_prob_i)``
502 Where:
503 - expert_fraction_i = tokens routed to expert i / total routed tokens
504 - expert_prob_i = sum of routing scores for expert i / total tokens
506 When routing is perfectly balanced, loss ≈ 1.0.
507 When routing is imbalanced, loss > 1.0.
509 When ``sequence_partition_group`` is provided (e.g. the TP×CP or CP
510 process group), ``expert_fraction`` (token counts) is all-reduced across
511 the group so that the loss reflects the *global* token distribution over
512 the full sequence. ``expert_prob`` is kept local to preserve gradient
513 flow back to the router weights, and ``num_tokens`` is scaled by
514 ``num_sub_sequence`` (the group world size) to account for the sharded
515 sequence dimension. This follows the Megatron-LM
516 ``full_sequence_aux_loss`` pattern.
518 Args:
519 top_scores: Routing weights, shape ``[num_tokens, top_k]``.
520 selected_experts: Expert IDs, shape ``[num_tokens, top_k]``.
521 num_experts: Total number of experts.
522 sequence_partition_group: Optional process group spanning the
523 sequence-partition dimension (TP+SP or CP). When not ``None``,
524 ``expert_fraction`` is all-reduced across this group.
526 Returns:
527 Scalar loss tensor. Returns 0.0 for empty input.
528 """
529 import torch.distributed as dist # pylint: disable=C0415
531 num_tokens, top_k = top_scores.shape
533 if num_tokens == 0:
534 return torch.tensor(0.0, dtype=top_scores.dtype, device=top_scores.device)
536 flat_experts = selected_experts.flatten()
537 flat_scores = top_scores.flatten()
539 expert_fraction = torch.zeros(
540 num_experts, dtype=top_scores.dtype, device=top_scores.device
541 )
542 expert_fraction.scatter_add_(0, flat_experts, torch.ones_like(flat_scores))
544 num_sub_sequence = 1
545 if sequence_partition_group is not None:
546 num_sub_sequence = dist.get_world_size(sequence_partition_group)
547 dist.all_reduce(expert_fraction, group=sequence_partition_group)
549 expert_fraction = expert_fraction / (num_tokens * num_sub_sequence * top_k)
551 expert_prob = torch.zeros(
552 num_experts, dtype=top_scores.dtype, device=top_scores.device
553 )
554 expert_prob.scatter_add_(0, flat_experts, flat_scores)
555 expert_prob = expert_prob / (num_tokens * num_sub_sequence)
557 loss = num_experts * (expert_fraction * expert_prob).sum()
559 return loss
562# ---------------------------------------------------------------------------
563# MoEAuxLossAutoScaler — gradient injection for auxiliary loss
564# ---------------------------------------------------------------------------
567class MoEAuxLossAutoScaler(torch.autograd.Function):
568 """Autograd function that injects auxiliary-loss gradient into the backward chain.
570 Forward transparently passes through ``top_scores`` (values unchanged).
571 Backward injects a scaled gradient for ``aux_loss``, causing its gradients
572 to flow through the router weights without adding aux_loss to the main
573 loss explicitly.
575 The loss scale is managed via the class-level :meth:`set_loss_scale` method,
576 which should be called before ``loss.backward()`` in the training loop
577 (typically by the pipeline schedule or training framework).
579 Reference: Megatron-LM ``megatron/core/transformer/moe/moe_utils.py``
580 """
582 main_loss_backward_scale: Optional[torch.Tensor] = None
584 @staticmethod
585 def forward(ctx: Any, output: torch.Tensor, aux_loss: torch.Tensor) -> torch.Tensor:
586 """Pass through ``output`` unchanged; save ``aux_loss`` for backward.
588 Args:
589 output: Router output (e.g. ``top_scores``), passed through unchanged.
590 aux_loss: Scalar auxiliary loss tensor to inject gradient for.
592 Returns:
593 The ``output`` tensor, identical in value but with an added
594 autograd edge to ``aux_loss``.
595 """
596 ctx.save_for_backward(aux_loss)
597 return output
599 @staticmethod
600 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
601 """Inject scaled aux_loss gradient into the backward chain.
603 Args:
604 grad_output: Gradient flowing back through ``output``.
606 Returns:
607 Tuple of (grad_output unchanged, scaled aux_loss gradient).
608 """
609 (aux_loss,) = ctx.saved_tensors
610 if MoEAuxLossAutoScaler.main_loss_backward_scale is None:
611 MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(
612 1.0, device=aux_loss.device,
613 )
614 aux_loss_backward_scale = MoEAuxLossAutoScaler.main_loss_backward_scale
615 scaled_aux_loss_grad = torch.ones_like(aux_loss) * aux_loss_backward_scale
616 return grad_output, scaled_aux_loss_grad
618 @staticmethod
619 def set_loss_scale(scale: torch.Tensor) -> None:
620 """Set the gradient scale for auxiliary loss.
622 Should be called before ``loss.backward()`` in the training loop.
623 Uses in-place copy on subsequent calls to avoid tensor reallocation.
625 Args:
626 scale: Tensor containing the loss scale value, typically
627 ``1 / (num_microbatches * dp_size)`` or similar.
628 """
629 if MoEAuxLossAutoScaler.main_loss_backward_scale is None:
630 MoEAuxLossAutoScaler.main_loss_backward_scale = scale
631 else:
632 MoEAuxLossAutoScaler.main_loss_backward_scale.copy_(scale)
635# ---------------------------------------------------------------------------
636# MoE orchestrator
637# ---------------------------------------------------------------------------
639class MoE(nn.Module):
640 """Mixture-of-Experts layer.
642 Orchestrates routing, token permutation, expert computation, and output
643 scatter-add. Supports shared experts, auxiliary-loss-free load balancing
644 via expert bias, node-limited routing, and auxiliary load-balance loss.
646 Args:
647 dim: Token embedding dimension.
648 hidden_dim: Expert hidden dimension.
649 num_experts: Total number of experts.
650 top_k: Experts selected per token. Defaults to ``1``.
651 score_before_experts: If ``True``, multiply routed tokens by routing
652 weights *before* expert computation; otherwise multiply expert
653 outputs *after*. Defaults to ``True``.
654 load_balance_coeff: When not ``None``, enables auxiliary load-balance
655 loss. The loss gradient is automatically injected into router
656 weights via :class:`MoEAuxLossAutoScaler` (no manual loss
657 addition needed). The scalar loss value is attached as
658 ``output._load_balance_loss`` for logging purposes.
659 sequence_partition_group: Optional process group spanning the
660 sequence-partition dimension (TP+SP or CP). When provided,
661 ``expert_fraction`` is all-reduced across this group so the
662 load-balance loss reflects global token distribution. Defaults
663 to ``None`` (no cross-rank synchronization).
664 shared_expert: Optional :class:`FeedForward` running on every token
665 in parallel; output added to routed-expert output.
666 router_kwargs: Extra keyword arguments forwarded to
667 :class:`TokenChoiceTopKRouter`.
668 use_grouped_mm: If ``True``, uses a hardware-accelerated grouped
669 matmul kernel (e.g. ``npu_grouped_matmul``) inside
670 :class:`GroupedExperts`. Defaults to ``False``.
672 Note:
673 *Auxiliary-loss-free load balancing*: call :func:`update_expert_bias`
674 once per optimiser step to adjust ``expert_bias`` from the accumulated
675 ``tokens_per_expert`` histogram.
677 Example::
678 >>> moe = MoE(dim=64, hidden_dim=128, num_experts=8, top_k=2)
679 >>> out = moe(torch.randn(2, 16, 64))
680 >>> out.shape
681 torch.Size([2, 16, 64])
682 """
684 def __init__(
685 self,
686 dim: int,
687 hidden_dim: int,
688 num_experts: int,
689 top_k: int = 1,
690 score_before_experts: bool = True,
691 load_balance_coeff: Optional[float] = None,
692 sequence_partition_group: Optional[Any] = None,
693 shared_expert: Optional[FeedForward] = None,
694 router_kwargs: Optional[dict] = None,
695 use_grouped_mm: bool = False,
696 ) -> None:
697 """Initialize MoE block with experts, router and optional shared expert.
699 Args:
700 dim: Input embedding dimension.
701 hidden_dim: Intermediate hidden dimension of each expert.
702 num_experts: Number of experts.
703 top_k: Number of experts to select per token.
704 score_before_experts: If ``True``, multiply routed hidden states by
705 the routing score.
706 load_balance_coeff: Load-balance loss coefficient.
707 shared_expert: Optional shared expert applied to all tokens.
708 router_kwargs: Additional keyword arguments for the router.
709 use_grouped_mm: Whether to use grouped matrix multiplication.
710 """
711 super().__init__()
712 router_kw = router_kwargs or {}
713 self.experts = GroupedExperts(
714 dim=dim, hidden_dim=hidden_dim, num_experts=num_experts,
715 use_grouped_mm=use_grouped_mm,
716 )
717 self.router = TokenChoiceTopKRouter(
718 dim=dim, num_experts=num_experts, top_k=top_k, **router_kw,
719 )
720 self.shared_expert = shared_expert
721 self.num_experts = num_experts
722 self.top_k = top_k
723 self.score_before_experts = score_before_experts
724 self.load_balance_coeff = load_balance_coeff
725 self.sequence_partition_group = sequence_partition_group
726 self.last_aux_loss: Optional[torch.Tensor] = None
727 self.enable_expert_bias = True
729 # Auxiliary-loss-free load-balance buffers (no gradient).
730 self.register_buffer("expert_bias", torch.zeros(num_experts))
731 self.register_buffer("tokens_per_expert", torch.zeros(num_experts))
733 def permutation(
734 self,
735 selected_experts: torch.Tensor,
736 top_scores: torch.Tensor,
737 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
738 """Compute token-major → expert-major permutation indices.
740 Args:
741 selected_experts: Expert IDs of shape ``[num_tokens, top_k]``.
742 top_scores: Routing weights of shape ``[num_tokens, top_k]``.
744 Returns:
745 Tuple of ``(token_indices, top_scores_sorted,
746 num_tokens_per_expert)`` where *token_indices* maps each
747 expert-major slot back to its source token row,
748 *top_scores_sorted* are the scores in expert-major order, and
749 *num_tokens_per_expert* has shape ``[num_experts]``.
750 """
751 # flat_experts[i] is the expert ID for the i-th (token, top_k) slot.
752 flat_experts = selected_experts.flatten() # [num_tokens * top_k]
753 flat_indices = flat_experts.argsort(stable=True) # expert-major permutation
754 top_scores_sorted = top_scores.flatten()[flat_indices] # [num_tokens * top_k]
755 # Each entry in flat_indices maps to a position in [0, num_tokens * top_k);
756 # divide by top_k to recover the original token row index.
757 token_indices = flat_indices // self.top_k # [num_tokens * top_k]
758 num_tokens_per_expert = torch.bincount(
759 flat_experts, minlength=self.num_experts
760 )
761 return token_indices, top_scores_sorted, num_tokens_per_expert
763 def unpermutation(
764 self,
765 expert_out: torch.Tensor,
766 token_indices: torch.Tensor,
767 num_tokens: int,
768 dim: int,
769 ) -> torch.Tensor:
770 """Scatter expert outputs back to token-major order.
772 Args:
773 expert_out: Expert output tensor in expert-major order,
774 shape ``[total_routed_tokens, dim]``.
775 token_indices: Permutation indices from :meth:`permutation`,
776 shape ``[num_tokens * top_k]``.
777 num_tokens: Total number of tokens (unflattened).
778 dim: Feature dimension.
780 Returns:
781 Token-major output tensor of shape ``[num_tokens, dim]``.
782 """
783 # Use out-of-place ``scatter_add`` so autograd correctly records
784 # ``ScatterAddBackward``; ``new_zeros + scatter_add_`` on some
785 # backends (torch_npu) leaves the leaf un-upgraded and the result
786 # without a ``grad_fn``.
787 return torch.zeros(
788 num_tokens, dim, dtype=expert_out.dtype, device=expert_out.device,
789 ).scatter_add(
790 0,
791 token_indices.unsqueeze(1).expand(-1, dim),
792 expert_out,
793 )
795 def forward(self, x: torch.Tensor) -> torch.Tensor:
796 """Run the MoE layer.
798 Args:
799 x: Input tensor of shape ``[batch, seq_len, dim]``.
801 Returns:
802 Output tensor of shape ``[batch, seq_len, dim]``. When
803 ``load_balance_coeff`` is set, carries a ``_load_balance_loss``
804 attribute with the auxiliary loss scalar (for logging only;
805 gradient injection is handled automatically by
806 :class:`MoEAuxLossAutoScaler`).
807 """
808 bs, seq_len, dim = x.shape
809 num_tokens = bs * seq_len
810 x_flat = x.view(num_tokens, dim) # [num_tokens, dim]
812 # --- Routing ---
813 top_scores, selected_experts, token_counts = self.router(
814 x_flat, self.expert_bias
815 )
817 # Accumulate token histogram without creating gradient nodes.
818 with torch.no_grad():
819 self.tokens_per_expert.add_(token_counts.float())
821 # --- Auxiliary load-balance loss ---
822 # Compute aux_loss early and attach it to top_scores via
823 # MoEAuxLossAutoScaler so that backward through top_scores
824 # also triggers aux_loss gradient (injected into router weights).
825 if self.load_balance_coeff is not None:
826 lb_loss = self.load_balance_coeff * _compute_load_balance_loss(
827 top_scores, selected_experts, self.num_experts,
828 sequence_partition_group=self.sequence_partition_group,
829 )
830 # Apply AutoScaler *before* top_scores is used in the forward path
831 # so that the main backward through top_scores triggers aux_loss
832 # gradient injection. Forward values are unchanged.
833 top_scores = MoEAuxLossAutoScaler.apply(top_scores, lb_loss)
834 self.last_aux_loss = lb_loss.detach()
835 else:
836 lb_loss = None
837 self.last_aux_loss = None
839 # --- Token permutation: token-major → expert-major ---
840 token_indices, top_scores_sorted, num_tokens_per_expert = self.permutation(
841 selected_experts, top_scores
842 )
844 # Gather routed tokens in expert-major order.
845 routed_x = x_flat[token_indices] # [num_tokens * top_k, dim]
847 # --- Expert computation ---
848 if self.score_before_experts:
849 routed_x = routed_x * top_scores_sorted.unsqueeze(1)
850 expert_out = self.experts(routed_x, num_tokens_per_expert, scores=None)
851 else:
852 expert_out = self.experts(routed_x, num_tokens_per_expert, scores=top_scores_sorted)
854 # --- Shared expert (parallel with routed experts) ---
855 shared_out = None
856 if self.shared_expert is not None:
857 shared_out = self.shared_expert(x_flat)
860 # --- Scatter expert outputs back to token order ---
861 out = self.unpermutation(expert_out, token_indices, num_tokens, dim)
863 if shared_out is not None:
864 out = out + shared_out
866 result = out.view(bs, seq_len, dim)
868 # Attach auxiliary loss to the returned tensor for logging.
869 if lb_loss is not None:
870 result._load_balance_loss = lb_loss # pylint: disable=protected-access
872 return result
874 def update_expert_bias(
875 self,
876 lr: float = 1e-3,
877 num_recomputations: int = 1,
878 ) -> None:
879 """Update expert bias for auxiliary-loss-free load balancing.
881 Should be called once per training step after the optimizer step.
882 Adjusts ``expert_bias`` to push token load towards the mean, then
883 resets the ``tokens_per_expert`` accumulator.
885 The update delta is centered to have zero mean, preventing systematic
886 drift of all bias values over time.
888 When activation checkpoint is enabled, forward is re-executed during
889 backward, causing ``tokens_per_expert`` to accumulate twice (or more).
890 Use ``num_recomputations`` to correct for this double-counting.
892 Args:
893 lr: Step size for the bias update. Defaults to ``1e-3``.
894 num_recomputations: Number of times forward is executed per optimizer
895 step. Default ``1`` (normal training). Set to ``2`` when activation
896 checkpoint is enabled (forward + recompute during backward). For
897 nested checkpoint strategies, set to the actual execution count.
899 Example::
900 >>> # Single-card scenario without activation checkpoint:
901 >>> moe_layer.update_expert_bias(lr=1e-3)
902 >>>
903 >>> # With activation checkpoint (forward executed twice):
904 >>> moe_layer.update_expert_bias(lr=1e-3, num_recomputations=2)
905 """
906 with torch.no_grad():
907 if num_recomputations > 1:
908 self.tokens_per_expert.div_(num_recomputations)
909 avg = self.tokens_per_expert.float().mean()
910 delta = lr * (avg - self.tokens_per_expert.float()).sign()
911 delta = delta - delta.mean()
912 self.expert_bias.data += delta
913 self.tokens_per_expert.zero_()
916# ---------------------------------------------------------------------------
917# Expert bias update for auxiliary-loss-free load balancing
918# ---------------------------------------------------------------------------
920def update_expert_bias(
921 moe: "MoE",
922 lr: float = 1e-3,
923 num_recomputations: int = 1,
924) -> None:
925 """Update expert bias for auxiliary-loss-free load balancing.
927 This is a module-level wrapper that calls ``moe.update_expert_bias()``.
928 Prefer calling the instance method directly.
930 Should be called once per training step after the optimizer step.
931 Adjusts ``moe.expert_bias`` to push token load towards the mean, then
932 resets the ``tokens_per_expert`` accumulator.
934 The update delta is centered to have zero mean, preventing systematic
935 drift of all bias values over time.
937 When activation checkpoint is enabled, forward is re-executed during
938 backward, causing ``tokens_per_expert`` to accumulate twice (or more).
939 Use ``num_recomputations`` to correct for this double-counting.
941 Args:
942 moe: The :class:`MoE` module whose bias should be updated.
943 lr: Step size for the bias update. Defaults to ``1e-3``.
944 num_recomputations: Number of times forward is executed per optimizer
945 step. Default ``1`` (normal training). Set to ``2`` when activation
946 checkpoint is enabled (forward + recompute during backward). For
947 nested checkpoint strategies, set to the actual execution count.
949 Example::
950 >>> # Single-card scenario without activation checkpoint:
951 >>> update_expert_bias(moe_layer, lr=1e-3)
952 >>>
953 >>> # With activation checkpoint (forward executed twice):
954 >>> update_expert_bias(moe_layer, lr=1e-3, num_recomputations=2)
955 """
956 moe.update_expert_bias(lr=lr, num_recomputations=num_recomputations)