wehub-resource-sync
190 lines
7.7 KiB
Python
190 lines
7.7 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
|
|
#
|
|
# Portions of this file are derived from TorchTitan.
|
|
# See THIRD_PARTY_NOTICES.md for the BSD-3-Clause notice.
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
Token-choice top-K router for expert parallelism.
|
|
|
|
Ported from TorchTitan's TokenChoiceTopKRouter with adaptations for DeepSpeed.
|
|
This module is self-contained: no imports from deepspeed.module_inject
|
|
or deepspeed.runtime.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from deepspeed.moe.ep_count import count_tokens_per_expert
|
|
|
|
|
|
class TokenChoiceTopKRouter(nn.Module):
|
|
"""Token-choice top-K routing for Mixture of Experts.
|
|
|
|
Each token is routed to top-K experts based on router scores.
|
|
Optionally supports node-limited (group-limited) routing where experts
|
|
are divided into groups (e.g., by node), and only ``num_limited_groups``
|
|
groups are considered before selecting top_k experts. This reduces
|
|
cross-node communication in distributed settings.
|
|
|
|
Args:
|
|
dim (int): Dimension of input tokens.
|
|
num_experts (int): Number of experts in each MoE layer.
|
|
num_expert_groups (int | None): Number of expert groups for
|
|
node-limited routing. If None, standard top-k routing is used.
|
|
Must be a divisor of num_experts.
|
|
num_limited_groups (int | None): Number of groups to select in
|
|
node-limited routing. Required when num_expert_groups is set.
|
|
top_k (int): Number of experts each token will be routed to.
|
|
score_func (str): ``"softmax"`` or ``"sigmoid"`` scoring function.
|
|
route_norm (bool): Whether to normalize routing scores.
|
|
route_scale (float): Scaling factor applied to routing scores.
|
|
gate_bias (bool): Whether to include a bias term in the gate linear.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_experts: int,
|
|
num_expert_groups: int | None,
|
|
num_limited_groups: int | None,
|
|
top_k: int,
|
|
score_func: str,
|
|
route_norm: bool,
|
|
route_scale: float,
|
|
gate_bias: bool,
|
|
group_score_func: str = "top2_sum",
|
|
):
|
|
super().__init__()
|
|
self.gate = nn.Linear(dim, num_experts, bias=gate_bias)
|
|
self.num_experts = num_experts
|
|
self.num_expert_groups = num_expert_groups
|
|
self.num_limited_groups = num_limited_groups
|
|
self.top_k = top_k
|
|
self.score_func = score_func
|
|
self.route_norm = route_norm
|
|
self.route_scale = route_scale
|
|
self.group_score_func = group_score_func
|
|
# Trainable expert score correction bias (e.g. DeepSeek-V3/Moonlight noaux_tc).
|
|
# Separate from the dynamic load-balancing expert_bias passed in forward().
|
|
self.e_score_correction_bias = None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Node-limited (group-limited) routing
|
|
# ------------------------------------------------------------------
|
|
|
|
def _get_node_limited_routing_scores(
|
|
self,
|
|
scores_for_choice: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Select ``num_limited_groups`` groups based on group scores and
|
|
mask out experts in non-selected groups.
|
|
|
|
Args:
|
|
scores_for_choice: Router scores with optional expert_bias,
|
|
shape ``(T, num_experts)``.
|
|
|
|
Returns:
|
|
Masked scores of the same shape, with non-selected group
|
|
entries set to ``-inf``.
|
|
"""
|
|
if self.num_limited_groups is None:
|
|
raise ValueError("num_limited_groups must be set when num_expert_groups is set")
|
|
assert self.num_expert_groups is not None
|
|
if self.num_limited_groups < 1:
|
|
raise ValueError(f"num_limited_groups must be >= 1, got {self.num_limited_groups}")
|
|
if self.num_experts % self.num_expert_groups != 0:
|
|
raise ValueError(f"num_experts ({self.num_experts}) must be divisible by "
|
|
f"num_expert_groups ({self.num_expert_groups})")
|
|
|
|
experts_per_group = self.num_experts // self.num_expert_groups
|
|
|
|
scores_grouped = scores_for_choice.view(-1, self.num_expert_groups, experts_per_group)
|
|
if self.group_score_func == "max":
|
|
group_scores = scores_grouped.max(dim=-1).values
|
|
elif self.group_score_func == "top2_sum":
|
|
if experts_per_group < 2:
|
|
raise ValueError(f"experts_per_group ({experts_per_group}) must be >= 2")
|
|
# DeepSeek-V3 scores each group by the sum of its top-2 experts.
|
|
top2_scores_in_group, _ = scores_grouped.topk(2, dim=-1)
|
|
group_scores = top2_scores_in_group.sum(dim=-1)
|
|
else:
|
|
raise NotImplementedError(f"Unknown group score function: {self.group_score_func}")
|
|
|
|
# Select top groups
|
|
_, group_idx = torch.topk(group_scores, k=self.num_limited_groups, dim=-1, sorted=False)
|
|
|
|
# Build mask: True = masked out (non-selected groups)
|
|
group_mask = torch.ones_like(group_scores, dtype=torch.bool)
|
|
group_mask.scatter_(1, group_idx, False)
|
|
|
|
scores_for_choice = scores_grouped.masked_fill(group_mask.unsqueeze(-1),
|
|
float("-inf")).view(-1, self.num_experts)
|
|
|
|
return scores_for_choice
|
|
|
|
# ------------------------------------------------------------------
|
|
# Forward
|
|
# ------------------------------------------------------------------
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
expert_bias: torch.Tensor | None = None,
|
|
) -> tuple:
|
|
"""
|
|
Args:
|
|
x: Input tensor of shape ``(T, dim)``.
|
|
expert_bias: Optional bias tensor of shape ``(num_experts,)``
|
|
used for load balancing.
|
|
|
|
Returns:
|
|
Tuple of:
|
|
- top_scores ``(T, top_k)``: routing weights for selected experts.
|
|
- selected_experts ``(T, top_k)``: expert indices per token.
|
|
- num_tokens_per_expert ``(num_experts,)``: histogram of token counts.
|
|
"""
|
|
# Gate projection -> (T, num_experts)
|
|
scores = self.gate(x)
|
|
|
|
# Scoring in float32 to avoid loss explosion
|
|
if self.score_func == "sigmoid":
|
|
scores = torch.sigmoid(scores.to(torch.float32))
|
|
elif self.score_func == "softmax":
|
|
scores = F.softmax(scores.to(torch.float32), dim=1)
|
|
else:
|
|
raise NotImplementedError(f"Unknown score function: {self.score_func}")
|
|
|
|
scores_for_choice = (scores if expert_bias is None else scores + expert_bias)
|
|
|
|
# Apply pre-trained score correction bias (e.g. DeepSeek-V3 noaux_tc routing)
|
|
if self.e_score_correction_bias is not None:
|
|
scores_for_choice = scores_for_choice + self.e_score_correction_bias.unsqueeze(0)
|
|
|
|
# Apply node-limited routing if configured
|
|
if self.num_expert_groups is not None:
|
|
scores_for_choice = self._get_node_limited_routing_scores(scores_for_choice)
|
|
|
|
# Select top-k experts per token
|
|
_, selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)
|
|
|
|
# Gather original (unbiased) scores for selected experts
|
|
top_scores = scores.gather(dim=1, index=selected_experts_indices)
|
|
|
|
# Optional normalization
|
|
if self.route_norm:
|
|
denominator = top_scores.sum(dim=-1, keepdim=True) + 1e-20
|
|
top_scores = top_scores / denominator
|
|
|
|
top_scores = top_scores * self.route_scale
|
|
|
|
num_tokens_per_expert = count_tokens_per_expert(selected_experts_indices, self.num_experts)
|
|
|
|
return top_scores, selected_experts_indices, num_tokens_per_expert
|