"""PyTorch layers for sparse mixture-of-experts models.""" from __future__ import annotations import dataclasses from collections.abc import Callable from contextlib import contextmanager, nullcontext from contextvars import ContextVar from functools import partial import torch import torch.nn as nn import torch.nn.functional as F _checkpoint_recompute = ContextVar("aliceai_t5_checkpoint_recompute", default=False) @contextmanager def _checkpoint_recompute_context(): token = _checkpoint_recompute.set(True) try: yield finally: _checkpoint_recompute.reset(token) def _checkpoint_contexts(): """Suppress router statistics updates when checkpointing repeats a forward.""" return nullcontext(), _checkpoint_recompute_context() @dataclasses.dataclass class Arguments: """Expert and router settings.""" hidden_size: int ffn_hidden_size: int activation_fn: Callable[[torch.Tensor], torch.Tensor] moe_num_experts: int moe_top_k: int bias: bool = True moe_normalize_expert_weights: int | float | bool | None = None routed_scaling_factor: float = 1.0 init_method: Callable[[torch.Tensor], torch.Tensor] = partial(nn.init.normal_, mean=0.0, std=0.02) def __post_init__(self) -> None: if self.moe_num_experts < 1: raise ValueError("moe_num_experts must be positive") if not 1 <= self.moe_top_k <= self.moe_num_experts: raise ValueError("moe_top_k must be in [1, moe_num_experts]") @dataclasses.dataclass class AdditionalArgs: moe_n_group: int | None = None moe_top_k_group: int | None = None group_routing: bool = False def normalize_expert_weights( expert_weights: torch.Tensor, norm_order: int | float | bool | None, routed_scaling_factor: float, ) -> torch.Tensor: """Apply optional Lp normalization and a routing scale; ``True`` selects L1.""" if norm_order: denominator = torch.norm(expert_weights, p=norm_order, dim=-1, keepdim=True) expert_weights = expert_weights / denominator.clamp_min(1e-20) return expert_weights * routed_scaling_factor class LearnedRouter(nn.Module): """Sigmoid top-k router with optional expert group selection.""" def __init__(self, args: Arguments, additional_args: AdditionalArgs): super().__init__() self.args = args self.additional_args = additional_args self.layer = nn.Linear(args.hidden_size, args.moe_num_experts, bias=False) self.register_buffer("top_k_bias", torch.zeros(1, args.moe_num_experts)) self.register_buffer("top_k_bias_update", torch.zeros(1, args.moe_num_experts), persistent=False) args.init_method(self.layer.weight) def _grouped_top_k(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: biased_scores = scores + self.top_k_bias if self.additional_args.group_routing: num_groups = self.additional_args.moe_n_group top_groups = self.additional_args.moe_top_k_group if num_groups is None or top_groups is None: raise ValueError("Grouped routing requires moe_n_group and moe_top_k_group") if self.args.moe_num_experts % num_groups: raise ValueError("moe_num_experts must be divisible by moe_n_group") grouped_scores = biased_scores.view(scores.shape[0], num_groups, -1) group_scores = grouped_scores.topk(min(2, grouped_scores.shape[-1]), dim=-1).values.sum(dim=-1) group_indices = group_scores.topk(top_groups, dim=-1, sorted=False).indices group_mask = torch.zeros_like(group_scores, dtype=torch.bool) group_mask.scatter_(1, group_indices, True) score_mask = group_mask.unsqueeze(-1).expand_as(grouped_scores).reshape_as(biased_scores) scores_for_choice = biased_scores.masked_fill(~score_mask, float("-inf")) else: scores_for_choice = biased_scores expert_indices = scores_for_choice.topk(self.args.moe_top_k, dim=-1, sorted=False).indices expert_weights = scores.gather(dim=-1, index=expert_indices) return expert_weights, expert_indices def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: router_logits = self.layer(hidden_states.reshape(-1, hidden_states.shape[-1])) scores = router_logits.float().sigmoid() expert_weights, expert_indices = self._grouped_top_k(scores) expert_weights = normalize_expert_weights( expert_weights, self.args.moe_normalize_expert_weights, self.args.routed_scaling_factor, ) return scores, expert_weights, expert_indices def _grouped_mm( inputs: torch.Tensor, weights: torch.Tensor, offsets: torch.Tensor, ) -> torch.Tensor: """Use a grouped CUDA operator when available, or per-expert ``torch.mm``.""" can_use_cuda_kernel = ( inputs.device.type == "cuda" and inputs.dtype in (torch.float16, torch.bfloat16) and weights.dtype == inputs.dtype and weights.device == inputs.device and torch.cuda.get_device_capability(inputs.device) >= (8, 0) ) if can_use_cuda_kernel: grouped_mm = getattr(F, "grouped_mm", None) if grouped_mm is None: grouped_mm = getattr(torch, "_grouped_mm", None) if grouped_mm is not None: return grouped_mm(inputs, weights, offs=offsets) outputs = [] start = 0 for expert_idx, end in enumerate(offsets.tolist()): if end > start: outputs.append(torch.mm(inputs[start:end], weights[expert_idx])) start = end if outputs: return torch.cat(outputs, dim=0) return torch.mm(inputs, weights[0]) def _unpermute_and_reduce( expert_output: torch.Tensor, inverse_permutation: torch.Tensor, expert_weights: torch.Tensor, *, num_tokens: int, top_k: int, ) -> torch.Tensor: """Combine routed expert outputs in token-major order. Routing weights are cast to the activation dtype before multiplication. Routing slots accumulate sequentially in that dtype; BF16 rounds after each multiplication and addition. """ output = expert_output.new_zeros((num_tokens, expert_output.shape[-1])) for top_k_position in range(top_k): grouped_rows = inverse_permutation[top_k_position::top_k] contribution = expert_output.index_select(0, grouped_rows) routing_weight = expert_weights[:, top_k_position, None].to(contribution.dtype) contribution.mul_(routing_weight) output.add_(contribution) return output class GroupedSwiGLU(nn.Module): """Expert weights of shape ``[num_experts * intermediate_size, hidden_size]``.""" def __init__(self, args: Arguments): super().__init__() self.args = args rows = args.moe_num_experts * args.ffn_hidden_size self.w1 = nn.Parameter(torch.empty(rows, args.hidden_size)) self.v1 = nn.Parameter(torch.empty(rows, args.hidden_size)) self.w2 = nn.Parameter(torch.empty(rows, args.hidden_size)) args.init_method(self.w1) args.init_method(self.v1) args.init_method(self.w2) def forward( self, hidden_states: torch.Tensor, expert_weights: torch.Tensor, expert_indices: torch.Tensor, ) -> torch.Tensor: num_tokens, hidden_size = hidden_states.shape num_experts = self.args.moe_num_experts intermediate_size = self.args.ffn_hidden_size if num_tokens == 0: # Keep every trainable branch in the autograd graph for empty # microbatches without reducing over the very large expert tensors. zero = expert_weights.sum() zero = zero + self.w1[0, 0] + self.v1[0, 0] + self.w2[0, 0] return hidden_states.reshape(0, hidden_size) + (zero * 0).to(hidden_states.dtype) w1 = self.w1.view(num_experts, intermediate_size, hidden_size) v1 = self.v1.view(num_experts, intermediate_size, hidden_size) w2 = self.w2.view(num_experts, intermediate_size, hidden_size) top_k = expert_indices.shape[-1] expert_ids = expert_indices.reshape(-1) # Group routed token slots by expert, then restore token-major order. expert_ids_grouped, permutation = torch.sort(expert_ids, stable=True) selected_states = hidden_states[permutation // top_k] histogram_input = ( expert_ids_grouped.float() if hidden_states.device.type in ("cpu", "mps") else expert_ids_grouped.int() ) tokens_per_expert = torch.histc( histogram_input, bins=num_experts, min=0, max=num_experts - 1, ) offsets = tokens_per_expert.cumsum(dim=0, dtype=torch.int32) gate = _grouped_mm(selected_states, w1.transpose(-2, -1), offsets) up = _grouped_mm(selected_states, v1.transpose(-2, -1), offsets) del selected_states gate = self.args.activation_fn(gate) gate.mul_(up) del up expert_output = _grouped_mm(gate, w2, offsets) del gate inverse_permutation = torch.empty_like(permutation) inverse_permutation[permutation] = torch.arange(permutation.numel(), device=permutation.device) return _unpermute_and_reduce( expert_output, inverse_permutation, expert_weights, num_tokens=num_tokens, top_k=top_k, ) class ExpertMLP(nn.Module): """Apply routed experts and an optional output bias.""" def __init__(self, experts_mlp: GroupedSwiGLU, args: Arguments): super().__init__() self.mlp = experts_mlp if args.bias: self.bias = nn.Parameter(torch.zeros(args.hidden_size)) else: self.register_parameter("bias", None) def forward( self, hidden_states: torch.Tensor, expert_weights: torch.Tensor, expert_indices: torch.Tensor, ) -> torch.Tensor: input_shape = hidden_states.shape flat_states = hidden_states.reshape(-1, input_shape[-1]) output = self.mlp(flat_states, expert_weights, expert_indices).view(input_shape) return output if self.bias is None else output + self.bias class dMoE(nn.Module): """Sparse MoE block.""" def __init__(self, args: Arguments, additional_args: AdditionalArgs): super().__init__() self.router = LearnedRouter(args, additional_args) self.experts = ExpertMLP(GroupedSwiGLU(args), args) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: _, expert_weights, expert_indices = self.router(hidden_states) if self.training and hidden_states.shape[0] > 0 and not _checkpoint_recompute.get(): with torch.no_grad(): tokens_per_expert = torch.bincount( expert_indices.flatten(), minlength=self.router.top_k_bias.shape[-1] ).to(torch.int32) self.router.top_k_bias_update.add_(tokens_per_expert / hidden_states.shape[0]) return self.experts(hidden_states, expert_weights, expert_indices) __all__ = [ "AdditionalArgs", "Arguments", "GroupedSwiGLU", "LearnedRouter", "ExpertMLP", "dMoE", ]