moe.py

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
import torch.distributed as dist
import torch.distributed.nn.functional as dist_nn
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class MoERoutingStats:
    aux_loss: torch.Tensor
    z_loss: torch.Tensor
    expert_load: torch.Tensor
    expert_importance: torch.Tensor
    top1_load: torch.Tensor
    capacity: int | None
    dropped_tokens: int
    tokens_dispatched: int
    balance_strategy: str


class TopKRouter(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        num_experts_per_token: int = 2,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
        *,
        capacity_factor: float | None = None,
        capacity_factor_eval: float | None = None,
        min_capacity: int = 4,
        jitter_eps: float = 0.0,
        aux_loss_coef: float = 0.0,
        z_loss_coef: float = 0.0,
        router_fp32: bool = True,
        balance_strategy: str = "aux_loss",
        bias_update_rate: float = 0.01,
    ):
        super().__init__()
        if num_experts < 2:
            raise ValueError(f"num_experts must be at least 2 for MoE, got {num_experts}.")
        if not 1 <= num_experts_per_token <= num_experts:
            raise ValueError(
                "num_experts_per_token must be in [1, num_experts], "
                f"got num_experts_per_token={num_experts_per_token}, num_experts={num_experts}."
            )
        if capacity_factor is not None and capacity_factor <= 0:
            raise ValueError(f"capacity_factor must be positive when provided, got {capacity_factor}.")
        if capacity_factor_eval is not None and capacity_factor_eval <= 0:
            raise ValueError(f"capacity_factor_eval must be positive when provided, got {capacity_factor_eval}.")
        if min_capacity <= 0:
            raise ValueError(f"min_capacity must be positive, got {min_capacity}.")
        if jitter_eps < 0:
            raise ValueError(f"jitter_eps must be non-negative, got {jitter_eps}.")
        if aux_loss_coef < 0:
            raise ValueError(f"aux_loss_coef must be non-negative, got {aux_loss_coef}.")
        if z_loss_coef < 0:
            raise ValueError(f"z_loss_coef must be non-negative, got {z_loss_coef}.")
        if balance_strategy not in {"aux_loss", "bias", "none"}:
            raise ValueError(f"balance_strategy must be one of 'aux_loss', 'bias', or 'none', got {balance_strategy!r}.")
        if bias_update_rate < 0:
            raise ValueError(f"bias_update_rate must be non-negative, got {bias_update_rate}.")

        self.d_model = d_model
        self.num_experts = num_experts
        self.num_experts_per_token = num_experts_per_token
        self.capacity_factor = capacity_factor
        self.capacity_factor_eval = capacity_factor_eval if capacity_factor_eval is not None else capacity_factor
        self.min_capacity = min_capacity
        self.jitter_eps = jitter_eps
        self.aux_loss_coef = aux_loss_coef
        self.z_loss_coef = z_loss_coef
        self.router_fp32 = router_fp32
        self.balance_strategy = balance_strategy
        self.bias_update_rate = bias_update_rate

        self.W = nn.Parameter(torch.empty((num_experts, d_model), device=device, dtype=dtype))
        std = math.sqrt(2 / (d_model + num_experts))
        nn.init.trunc_normal_(self.W, mean=0.0, std=std, a=-3 * std, b=3 * std)
        self.register_buffer("expert_bias", torch.zeros(num_experts, device=device, dtype=torch.float32), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, routing_logits = self._compute_router_logits(x)
        return routing_logits

    def route_tokens(
        self,
        x: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        original_shape = x.shape[:-1]
        x_flat = x.reshape(-1, self.d_model)
        _, router_logits = self._compute_router_logits(x_flat)
        topk_logits, topk_indices = torch.topk(router_logits, k=self.num_experts_per_token, dim=-1)
        topk_weights = torch.softmax(topk_logits.to(torch.float32), dim=-1).to(x.dtype)
        return (
            router_logits.reshape(*original_shape, self.num_experts),
            topk_weights.reshape(*original_shape, self.num_experts_per_token),
            topk_indices.reshape(*original_shape, self.num_experts_per_token),
        )

    def build_routing(
        self,
        x_flat: torch.Tensor,
        *,
        output_dtype: torch.dtype,
    ) -> dict[str, torch.Tensor | MoERoutingStats | int | None]:
        num_tokens = x_flat.size(0)
        base_router_logits, router_logits = self._compute_router_logits(x_flat)
        topk_logits, topk_indices = torch.topk(router_logits, k=self.num_experts_per_token, dim=-1)
        topk_weights = torch.softmax(topk_logits.to(torch.float32), dim=-1).to(output_dtype)
        router_probs = torch.softmax(router_logits.to(torch.float32), dim=-1)

        token_indices = torch.arange(num_tokens, device=x_flat.device).unsqueeze(1).expand(-1, self.num_experts_per_token)
        route_ranks = torch.arange(self.num_experts_per_token, device=x_flat.device).view(1, -1).expand(num_tokens, -1)
        flat_tokens = token_indices.reshape(-1)
        flat_route_ranks = route_ranks.reshape(-1)
        flat_experts = topk_indices.reshape(-1)
        flat_weights = topk_weights.reshape(-1)

        capacity = self._compute_capacity(num_tokens)
        order = self._sorted_dispatch_order(flat_experts, flat_tokens, flat_route_ranks)
        sorted_experts = flat_experts[order]
        sorted_tokens = flat_tokens[order]
        sorted_weights = flat_weights[order]

        if sorted_experts.numel() == 0:
            kept_experts = sorted_experts
            kept_tokens = sorted_tokens
            kept_weights = sorted_weights
        else:
            expert_counts = torch.bincount(sorted_experts, minlength=self.num_experts)
            expert_offsets = torch.cumsum(expert_counts, dim=0) - expert_counts
            expert_positions = torch.arange(sorted_experts.numel(), device=x_flat.device) - expert_offsets[sorted_experts]
            if capacity is None:
                keep_mask = torch.ones_like(sorted_experts, dtype=torch.bool)
            else:
                keep_mask = expert_positions < capacity
            kept_experts = sorted_experts[keep_mask]
            kept_tokens = sorted_tokens[keep_mask]
            kept_weights = sorted_weights[keep_mask]

        stats = self._compute_routing_stats(
            base_router_logits=base_router_logits,
            router_probs=router_probs,
            topk_indices=topk_indices,
            kept_experts=kept_experts,
            num_tokens=num_tokens,
            capacity=capacity,
            total_routes=flat_experts.numel(),
        )
        self._maybe_update_bias(stats.top1_load)

        return {
            "router_logits": router_logits,
            "topk_weights": topk_weights,
            "topk_indices": topk_indices,
            "kept_experts": kept_experts,
            "kept_tokens": kept_tokens,
            "kept_weights": kept_weights,
            "stats": stats,
            "capacity": capacity,
        }

    def _compute_router_logits(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        router_input = x
        router_weight = self.W
        if self.router_fp32:
            router_input = x.to(torch.float32)
            router_weight = self.W.to(torch.float32)
        router_input = self._apply_jitter(router_input)
        base_router_logits = F.linear(router_input, router_weight)
        if self.balance_strategy == "bias":
            routing_logits = base_router_logits + self.expert_bias.to(base_router_logits.dtype)
        else:
            routing_logits = base_router_logits
        return base_router_logits, routing_logits

    def _apply_jitter(self, x: torch.Tensor) -> torch.Tensor:
        if not self.training or self.jitter_eps == 0:
            return x
        scale = torch.empty_like(x).uniform_(1 - self.jitter_eps, 1 + self.jitter_eps)
        return x * scale

    def _compute_capacity(self, num_tokens: int) -> int | None:
        capacity_factor = self.capacity_factor if self.training else self.capacity_factor_eval
        if capacity_factor is None:
            return None
        capacity = math.ceil(capacity_factor * num_tokens * self.num_experts_per_token / self.num_experts)
        return max(self.min_capacity, capacity)

    def _sorted_dispatch_order(
        self,
        expert_indices: torch.Tensor,
        token_indices: torch.Tensor,
        route_ranks: torch.Tensor,
    ) -> torch.Tensor:
        total_routes = token_indices.numel()
        sort_key = expert_indices * total_routes + route_ranks * max(token_indices.numel() // self.num_experts_per_token, 1) + token_indices
        return torch.argsort(sort_key)

    def _compute_routing_stats(
        self,
        *,
        base_router_logits: torch.Tensor,
        router_probs: torch.Tensor,
        topk_indices: torch.Tensor,
        kept_experts: torch.Tensor,
        num_tokens: int,
        capacity: int | None,
        total_routes: int,
    ) -> MoERoutingStats:
        device = base_router_logits.device
        top1_experts = topk_indices[:, 0]
        top1_counts = torch.bincount(top1_experts, minlength=self.num_experts).to(torch.float32)
        expert_load = torch.bincount(kept_experts, minlength=self.num_experts).to(torch.float32)
        global_top1_load, global_expert_importance, global_expert_load = self._synchronize_routing_stats(
            top1_counts=top1_counts,
            router_probs=router_probs,
            expert_load=expert_load,
            num_tokens=num_tokens,
        )
        aux_loss = base_router_logits.new_zeros(())
        if self.balance_strategy == "aux_loss":
            aux_loss = self.aux_loss_coef * self.num_experts * torch.sum(
                global_top1_load.to(torch.float32) * global_expert_importance.to(torch.float32)
            )
        z_loss = self.z_loss_coef * torch.mean(torch.logsumexp(base_router_logits.to(torch.float32), dim=-1) ** 2)
        dropped_tokens = int(total_routes - kept_experts.numel()) if capacity is not None else 0
        return MoERoutingStats(
            aux_loss=aux_loss.to(device=device),
            z_loss=z_loss.to(device=device),
            expert_load=global_expert_load.to(device=device),
            expert_importance=global_expert_importance.to(device=device),
            top1_load=global_top1_load.to(device=device),
            capacity=capacity,
            dropped_tokens=dropped_tokens,
            tokens_dispatched=int(kept_experts.numel()),
            balance_strategy=self.balance_strategy,
        )

    def _synchronize_routing_stats(
        self,
        *,
        top1_counts: torch.Tensor,
        router_probs: torch.Tensor,
        expert_load: torch.Tensor,
        num_tokens: int,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if not dist.is_available() or not dist.is_initialized():
            normalizer = max(num_tokens, 1)
            return top1_counts / normalizer, router_probs.mean(dim=0), expert_load

        top1_counts = top1_counts.to(torch.float32)
        expert_load = expert_load.to(torch.float32)
        importance_sum = router_probs.to(torch.float32).sum(dim=0)
        token_count = torch.tensor(float(num_tokens), device=router_probs.device, dtype=torch.float32)
        dist.all_reduce(top1_counts, op=dist.ReduceOp.SUM)
        dist.all_reduce(expert_load, op=dist.ReduceOp.SUM)
        importance_sum = dist_nn.all_reduce(importance_sum, op=dist.ReduceOp.SUM)
        dist.all_reduce(token_count, op=dist.ReduceOp.SUM)
        normalizer = max(float(token_count.item()), 1.0)
        return top1_counts / normalizer, importance_sum / normalizer, expert_load

    def _maybe_update_bias(self, top1_load: torch.Tensor) -> None:
        if self.balance_strategy != "bias" or self.bias_update_rate == 0 or not self.training:
            return
        with torch.no_grad():
            target = 1.0 / self.num_experts
            adjustment = self.bias_update_rate * (target - top1_load.to(torch.float32))
            self.expert_bias.add_(adjustment.to(self.expert_bias.device))
            self.expert_bias.clamp_(-2.0, 2.0)


class MoESwiGLU(nn.Module):
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int,
        num_experts_per_token: int = 2,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
        *,
        num_shared_experts: int = 0,
        capacity_factor: float | None = None,
        capacity_factor_eval: float | None = None,
        min_capacity: int = 4,
        jitter_eps: float = 0.0,
        aux_loss_coef: float = 0.0,
        z_loss_coef: float = 0.0,
        router_fp32: bool = True,
        balance_strategy: str = "aux_loss",
        bias_update_rate: float = 0.01,
        zero_init_output: bool = False,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_ff = d_ff
        self.num_experts = num_experts
        self.num_experts_per_token = num_experts_per_token
        self.num_shared_experts = num_shared_experts
        self.router = TopKRouter(
            d_model=d_model,
            num_experts=num_experts,
            num_experts_per_token=num_experts_per_token,
            device=device,
            dtype=dtype,
            capacity_factor=capacity_factor,
            capacity_factor_eval=capacity_factor_eval,
            min_capacity=min_capacity,
            jitter_eps=jitter_eps,
            aux_loss_coef=aux_loss_coef,
            z_loss_coef=z_loss_coef,
            router_fp32=router_fp32,
            balance_strategy=balance_strategy,
            bias_update_rate=bias_update_rate,
        )
        self.W1 = nn.Parameter(torch.empty((num_experts, d_ff, d_model), device=device, dtype=dtype))
        self.W3 = nn.Parameter(torch.empty((num_experts, d_ff, d_model), device=device, dtype=dtype))
        self.W2 = nn.Parameter(torch.empty((num_experts, d_model, d_ff), device=device, dtype=dtype))
        std = math.sqrt(2 / (d_ff + d_model))
        nn.init.trunc_normal_(self.W1, mean=0.0, std=std, a=-3 * std, b=3 * std)
        nn.init.trunc_normal_(self.W3, mean=0.0, std=std, a=-3 * std, b=3 * std)
        if zero_init_output:
            nn.init.zeros_(self.W2)
        else:
            nn.init.trunc_normal_(self.W2, mean=0.0, std=std, a=-3 * std, b=3 * std)
        if num_shared_experts > 0:
            self.shared_W1 = nn.Parameter(torch.empty((num_shared_experts, d_ff, d_model), device=device, dtype=dtype))
            self.shared_W3 = nn.Parameter(torch.empty((num_shared_experts, d_ff, d_model), device=device, dtype=dtype))
            self.shared_W2 = nn.Parameter(torch.empty((num_shared_experts, d_model, d_ff), device=device, dtype=dtype))
            nn.init.trunc_normal_(self.shared_W1, mean=0.0, std=std, a=-3 * std, b=3 * std)
            nn.init.trunc_normal_(self.shared_W3, mean=0.0, std=std, a=-3 * std, b=3 * std)
            if zero_init_output:
                nn.init.zeros_(self.shared_W2)
            else:
                nn.init.trunc_normal_(self.shared_W2, mean=0.0, std=std, a=-3 * std, b=3 * std)
        else:
            self.shared_W1 = None
            self.shared_W3 = None
            self.shared_W2 = None

        self.last_routing_stats: MoERoutingStats | None = None

    def forward(
        self,
        x: torch.Tensor,
        return_routing: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        original_shape = x.shape
        x_flat = x.reshape(-1, self.d_model)
        routing = self.router.build_routing(x_flat, output_dtype=x.dtype)
        router_logits = routing["router_logits"]
        topk_weights = routing["topk_weights"]
        topk_indices = routing["topk_indices"]
        kept_experts = routing["kept_experts"]
        kept_tokens = routing["kept_tokens"]
        kept_weights = routing["kept_weights"]

        output = x_flat.new_zeros((x_flat.size(0), self.d_model))
        if kept_experts.numel() > 0:
            expert_counts = torch.bincount(kept_experts, minlength=self.num_experts)
            max_tokens_per_expert = int(expert_counts.max().item())
            expert_offsets = torch.cumsum(expert_counts, dim=0) - expert_counts
            expert_positions = torch.arange(kept_experts.numel(), device=x.device) - expert_offsets[kept_experts]

            expert_inputs = x_flat.new_zeros((self.num_experts, max_tokens_per_expert, self.d_model))
            expert_inputs[kept_experts, expert_positions] = x_flat[kept_tokens]

            hidden_w1 = torch.einsum("efd,etd->etf", self.W1, expert_inputs)
            hidden_w3 = torch.einsum("efd,etd->etf", self.W3, expert_inputs)
            expert_hidden = F.silu(hidden_w1) * hidden_w3
            expert_outputs = torch.einsum("edf,etf->etd", self.W2, expert_hidden)

            kept_outputs = expert_outputs[kept_experts, expert_positions]
            output.index_add_(0, kept_tokens, kept_outputs * kept_weights.unsqueeze(-1))

        if self.num_shared_experts > 0 and self.shared_W1 is not None:
            shared_hidden_w1 = torch.einsum("efd,td->tef", self.shared_W1, x_flat)
            shared_hidden_w3 = torch.einsum("efd,td->tef", self.shared_W3, x_flat)
            shared_hidden = F.silu(shared_hidden_w1) * shared_hidden_w3
            shared_output = torch.einsum("edf,tef->ted", self.shared_W2, shared_hidden).mean(dim=1)
            output = output + shared_output

        output = output.reshape(*original_shape)
        self.last_routing_stats = routing["stats"]

        if return_routing:
            routed_shape = (*original_shape[:-1], self.num_experts_per_token)
            return (
                output,
                router_logits.reshape(*original_shape[:-1], self.num_experts),
                topk_weights.reshape(*routed_shape),
                topk_indices.reshape(*routed_shape),
            )
        return output


__all__ = [
    "MoERoutingStats",
    "MoESwiGLU",
    "TopKRouter",
]