wehub-resource-sync
754 lines
32 KiB
Python
754 lines
32 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
|
|
"""AutoEP MoE Layer: drop-in replacement for HF MoE blocks with EP support.
|
|
|
|
Contains AutoEPMoELayer, compute_split_plan, _AllToAllV, and helper functions.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Literal, NamedTuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import deepspeed.comm as dist
|
|
from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults
|
|
from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter
|
|
from deepspeed.utils import logger
|
|
from deepspeed.moe.ep_router import TokenChoiceTopKRouter
|
|
from deepspeed.moe.ep_count import count_tokens_per_expert
|
|
from deepspeed.moe.ep_experts import GroupedExperts
|
|
from deepspeed.moe.ep_kernels import TokenReorderer
|
|
from deepspeed.moe.ep_repack import _gather_source_zero_params, repack_expert_requires_grad_flags, repack_expert_weights
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Named tuples
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class RouterOutput(NamedTuple):
|
|
top_scores: torch.Tensor # [T, K]
|
|
selected_experts: torch.Tensor # [T, K]
|
|
num_tokens_per_expert: torch.Tensor # [E_global]
|
|
|
|
|
|
class SplitPlan(NamedTuple):
|
|
input_splits: list[int] # len=ep_size
|
|
output_splits: list[int] # len=ep_size
|
|
local_counts: torch.Tensor # [E_local]
|
|
local_counts_by_source: torch.Tensor # [ep_size, E_local]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper functions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def resolve_score_apply_mode(
|
|
spec: MoELayerSpec,
|
|
config_override: Literal["auto", "pre", "post"],
|
|
) -> Literal["pre", "post"]:
|
|
"""Resolve score-application mode from config override or preset default."""
|
|
if config_override != "auto":
|
|
return config_override
|
|
return spec.score_apply
|
|
|
|
|
|
def resolve_combine_impl(
|
|
config_override: Literal["auto", "weighted_sum", "legacy_bmm"], ) -> Literal["weighted_sum", "legacy_bmm"]:
|
|
"""Resolve combine implementation from config override or default."""
|
|
if config_override != "auto":
|
|
return config_override
|
|
return "weighted_sum"
|
|
|
|
|
|
def _copy_parameter_data(target: nn.Parameter, source: torch.Tensor) -> None:
|
|
full_shape = torch.Size(getattr(source, "ds_shape", source.shape))
|
|
with torch.no_grad():
|
|
source_data = source.data
|
|
if torch.Size(source_data.shape) != full_shape:
|
|
raise RuntimeError("AutoEP source parameter must be gathered before copying: "
|
|
f"expected full shape {tuple(full_shape)}, got {tuple(source_data.shape)}")
|
|
if (torch.Size(target.data.shape) != full_shape or target.data.dtype != source_data.dtype
|
|
or target.data.device != source_data.device):
|
|
target.data = torch.empty(full_shape, dtype=source_data.dtype, device=source_data.device)
|
|
target.data.copy_(source_data)
|
|
|
|
|
|
def apply_scores_before_experts_if_enabled(
|
|
routed_input: torch.Tensor,
|
|
top_scores: torch.Tensor,
|
|
score_apply: Literal["pre", "post"],
|
|
) -> torch.Tensor:
|
|
"""Pre-multiply token representations by router scores before expert compute."""
|
|
if score_apply == "pre":
|
|
return (routed_input.to(torch.float32) * top_scores.reshape(-1, 1)).to(routed_input.dtype)
|
|
return routed_input
|
|
|
|
|
|
def compute_split_plan(
|
|
selected_experts: torch.Tensor, # [T, K]
|
|
num_experts: int,
|
|
ep_size: int,
|
|
num_local_experts: int,
|
|
ep_group: dist.ProcessGroup | None,
|
|
) -> SplitPlan:
|
|
"""Compute AllToAllV split sizes for token dispatch/combine.
|
|
|
|
Returns SplitPlan with input_splits, output_splits, local_counts, and
|
|
local_counts_by_source.
|
|
"""
|
|
T_K = selected_experts.numel()
|
|
|
|
if ep_size == 1:
|
|
# No dispatch needed - all tokens stay local
|
|
num_tokens_per_expert = count_tokens_per_expert(
|
|
selected_experts,
|
|
num_experts,
|
|
out_dtype=torch.int32,
|
|
)
|
|
return SplitPlan(
|
|
input_splits=[T_K],
|
|
output_splits=[T_K],
|
|
local_counts=num_tokens_per_expert,
|
|
local_counts_by_source=num_tokens_per_expert.view(1, num_local_experts),
|
|
)
|
|
|
|
# Count tokens per expert globally
|
|
num_tokens_per_expert = count_tokens_per_expert(
|
|
selected_experts,
|
|
num_experts,
|
|
out_dtype=torch.int32,
|
|
)
|
|
|
|
# Reshape to [ep_size, num_local_experts] to get per-rank counts
|
|
count_matrix = num_tokens_per_expert.view(ep_size, num_local_experts)
|
|
|
|
# input_splits: how many tokens THIS rank sends to each destination rank
|
|
input_splits = count_matrix.sum(dim=1).cpu().tolist()
|
|
|
|
# Exchange counts with all ranks to get output_splits
|
|
# Each rank tells every other rank how many tokens it will send
|
|
local_counts_tensor = count_matrix.sum(dim=1).clone() # [ep_size]
|
|
remote_counts_tensor = torch.zeros_like(local_counts_tensor)
|
|
|
|
dist.all_to_all_single(
|
|
remote_counts_tensor,
|
|
local_counts_tensor,
|
|
group=ep_group,
|
|
)
|
|
output_splits = remote_counts_tensor.cpu().tolist()
|
|
|
|
# local_counts: how many tokens this rank will process for each local expert
|
|
# After receiving tokens, we need per-expert counts for this rank
|
|
local_expert_counts = count_matrix[:, :].clone() # [ep_size, E_local]
|
|
|
|
# Exchange the detailed per-expert counts
|
|
# Each rank needs to know, for its local experts, how many tokens come from each source
|
|
local_expert_counts_flat = local_expert_counts.view(-1).contiguous() # [ep_size * E_local]
|
|
received_counts_flat = torch.zeros_like(local_expert_counts_flat)
|
|
|
|
dist.all_to_all_single(
|
|
received_counts_flat,
|
|
local_expert_counts_flat,
|
|
group=ep_group,
|
|
)
|
|
|
|
# Sum over source ranks to get total per local expert
|
|
received_counts = received_counts_flat.view(ep_size, num_local_experts)
|
|
local_counts = received_counts.sum(dim=0) # [E_local]
|
|
|
|
return SplitPlan(
|
|
input_splits=input_splits,
|
|
output_splits=output_splits,
|
|
local_counts=local_counts,
|
|
local_counts_by_source=received_counts,
|
|
)
|
|
|
|
|
|
def compute_split_plan_from_expert_indices(
|
|
expert_indices: torch.Tensor,
|
|
num_experts: int,
|
|
ep_size: int,
|
|
num_local_experts: int,
|
|
ep_group: dist.ProcessGroup | None,
|
|
) -> SplitPlan:
|
|
"""Compute EP AllToAllV splits for an already partitioned assignment list."""
|
|
if ep_size == 1:
|
|
counts = count_tokens_per_expert(expert_indices, num_experts, out_dtype=torch.int32)
|
|
return SplitPlan([int(expert_indices.numel())], [int(expert_indices.numel())], counts,
|
|
counts.view(1, num_local_experts))
|
|
|
|
counts = count_tokens_per_expert(expert_indices, num_experts, out_dtype=torch.int32)
|
|
count_matrix = counts.view(ep_size, num_local_experts)
|
|
input_splits = count_matrix.sum(dim=1).cpu().tolist()
|
|
local_counts_tensor = count_matrix.sum(dim=1).clone()
|
|
remote_counts_tensor = torch.zeros_like(local_counts_tensor)
|
|
dist.all_to_all_single(remote_counts_tensor, local_counts_tensor, group=ep_group)
|
|
output_splits = remote_counts_tensor.cpu().tolist()
|
|
|
|
local_expert_counts_flat = count_matrix.reshape(-1).contiguous()
|
|
received_counts_flat = torch.zeros_like(local_expert_counts_flat)
|
|
dist.all_to_all_single(received_counts_flat, local_expert_counts_flat, group=ep_group)
|
|
received_counts = received_counts_flat.view(ep_size, num_local_experts)
|
|
local_counts = received_counts.sum(dim=0)
|
|
return SplitPlan(input_splits, output_splits, local_counts, received_counts)
|
|
|
|
|
|
class _AllToAllV(torch.autograd.Function):
|
|
"""Autograd-compatible all-to-all with variable split sizes."""
|
|
|
|
@staticmethod
|
|
def forward(ctx, group, x, input_splits, output_splits):
|
|
ctx.group = group
|
|
ctx.input_splits = input_splits
|
|
ctx.output_splits = output_splits
|
|
|
|
output_size = sum(output_splits)
|
|
output = torch.empty(
|
|
(output_size, x.shape[1]),
|
|
dtype=x.dtype,
|
|
device=x.device,
|
|
)
|
|
|
|
dist.all_to_all_single(
|
|
output,
|
|
x.contiguous(),
|
|
output_split_sizes=output_splits,
|
|
input_split_sizes=input_splits,
|
|
group=group,
|
|
)
|
|
return output
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_out):
|
|
# Reverse the splits for backward
|
|
grad_out = grad_out.contiguous()
|
|
input_size = sum(ctx.input_splits)
|
|
grad_input = torch.empty(
|
|
(input_size, grad_out.shape[1]),
|
|
dtype=grad_out.dtype,
|
|
device=grad_out.device,
|
|
)
|
|
|
|
dist.all_to_all_single(
|
|
grad_input,
|
|
grad_out,
|
|
output_split_sizes=ctx.input_splits,
|
|
input_split_sizes=ctx.output_splits,
|
|
group=ctx.group,
|
|
)
|
|
return None, grad_input, None, None
|
|
|
|
|
|
def permute_by_local_expert(
|
|
tokens: torch.Tensor,
|
|
local_counts: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
|
|
"""Reorder tokens so they are grouped contiguously by local expert ID.
|
|
|
|
Uses TorchTitan's Triton kernel for permutation index generation.
|
|
|
|
Returns:
|
|
tokens_permuted: [N_padded, H] (alignment-padded)
|
|
permuted_indices: [N_padded] (maps padded positions -> original positions)
|
|
aligned_counts: [E_local] aligned token counts per expert (for expert computation)
|
|
n_tokens: original token count before padding (for unpermute)
|
|
"""
|
|
from deepspeed.moe.ep_kernels import generate_permute_indices, TOKEN_GROUP_ALIGN_SIZE_M
|
|
|
|
if local_counts.ndim == 1:
|
|
# [E_local]: already aggregated over sources (ep_degree=1)
|
|
ep_degree = 1
|
|
num_local_experts = local_counts.shape[0]
|
|
local_counts_flat = local_counts
|
|
elif local_counts.ndim == 2:
|
|
# [ep_size, E_local]: preserve per-source layout for correct regrouping
|
|
ep_degree, num_local_experts = local_counts.shape
|
|
local_counts_flat = local_counts.reshape(-1)
|
|
else:
|
|
raise ValueError(
|
|
f"local_counts must have shape [E_local] or [ep_degree, E_local], got {tuple(local_counts.shape)}")
|
|
|
|
n_tokens = tokens.shape[0]
|
|
alignment = TOKEN_GROUP_ALIGN_SIZE_M
|
|
|
|
# Compute padded max length
|
|
x_padded_per_expert = n_tokens + num_local_experts * alignment
|
|
padded_max_len = ((x_padded_per_expert + alignment - 1) // alignment) * alignment
|
|
|
|
# Use the pure-PyTorch path for host tensors. The CPU accelerator reports
|
|
# CPU tensors as "on accelerator", but Triton still requires a GPU driver.
|
|
use_cpu = tokens.device.type == "cpu"
|
|
counts_for_permute = local_counts_flat.cpu() if use_cpu else local_counts_flat
|
|
with torch.no_grad():
|
|
permuted_indices, m_sizes, _offsets = generate_permute_indices(
|
|
counts_for_permute,
|
|
num_local_experts,
|
|
ep_degree,
|
|
padded_max_len,
|
|
alignment,
|
|
use_cpu=use_cpu,
|
|
)
|
|
if not use_cpu:
|
|
permuted_indices = permuted_indices.to(tokens.device)
|
|
m_sizes = m_sizes.to(tokens.device)
|
|
|
|
# Add padding row for out-of-bounds indices (index n_tokens -> zero row)
|
|
tokens_padded = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1], ))))
|
|
tokens_permuted = tokens_padded[permuted_indices, :]
|
|
|
|
return tokens_permuted, permuted_indices, m_sizes, n_tokens
|
|
|
|
|
|
def unpermute_by_local_expert(
|
|
expert_output: torch.Tensor,
|
|
permuted_indices: torch.Tensor,
|
|
n_tokens: int,
|
|
) -> torch.Tensor:
|
|
"""Reverse permute_by_local_expert: restore original token order and strip padding.
|
|
|
|
Args:
|
|
expert_output: [N_padded, H] from expert computation
|
|
permuted_indices: [N_padded] index mapping from permute_by_local_expert
|
|
n_tokens: original token count before alignment padding
|
|
"""
|
|
# Scatter expert outputs back to original positions.
|
|
# permuted_indices values range 0..n_tokens, where n_tokens is the zero-padding row.
|
|
out_unpermuted = expert_output.new_zeros((n_tokens + 1, expert_output.shape[-1]))
|
|
out_unpermuted[permuted_indices, :] = expert_output
|
|
# Strip the zero-padding row to get [n_tokens, H]
|
|
return out_unpermuted[:-1]
|
|
|
|
|
|
def combine_from_routed(
|
|
expert_output: torch.Tensor, # [N, H]
|
|
top_scores: torch.Tensor, # [T, K]
|
|
token_indices_sorted: torch.Tensor, # [N]
|
|
top_k: int,
|
|
score_apply: Literal["pre", "post"],
|
|
combine_impl: Literal["weighted_sum", "legacy_bmm"],
|
|
shape: tuple[int, int, int], # (B, S, H)
|
|
) -> torch.Tensor:
|
|
"""Scatter-add expert outputs back to original token positions."""
|
|
bsz, seqlen, hdim = shape
|
|
T = bsz * seqlen
|
|
|
|
# Create output tensor
|
|
output = torch.zeros(T * top_k, hdim, dtype=expert_output.dtype, device=expert_output.device)
|
|
|
|
# Place expert outputs back in unsorted order
|
|
output[token_indices_sorted] = expert_output
|
|
|
|
# Reshape to [T, K, H]
|
|
output = output.reshape(T, top_k, hdim)
|
|
|
|
if score_apply == "post":
|
|
if combine_impl == "legacy_bmm":
|
|
# Legacy reduction path retained as a debug option for model-family
|
|
# verification. The weighted-sum path is the default.
|
|
output = torch.bmm(
|
|
top_scores.reshape(-1, 1, top_k).float(),
|
|
output.float(),
|
|
).to(expert_output.dtype).squeeze(1)
|
|
else:
|
|
# Match the runtime HF grouped-mm path: apply routing weights per
|
|
# token-slot sample, then reduce over top-k.
|
|
output = (output.float() * top_scores.reshape(T, top_k, 1).float()).sum(dim=1).to(expert_output.dtype)
|
|
else:
|
|
# Scores already applied pre-experts, just sum over top_k
|
|
output = output.sum(dim=1)
|
|
|
|
return output.reshape(bsz, seqlen, hdim)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AutoEPMoELayer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class AutoEPMoELayer(nn.Module):
|
|
"""Drop-in replacement for HF MoE blocks with Expert Parallelism support."""
|
|
|
|
_is_autoep_layer = True # Marker for AutoTP skip handshake
|
|
|
|
def __init__(
|
|
self,
|
|
spec: MoELayerSpec,
|
|
source_module: nn.Module,
|
|
ep_size: int,
|
|
ep_rank: int,
|
|
config: AutoEPConfig,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.model_family = spec.model_family
|
|
self.return_router_logits = spec.return_router_logits
|
|
self.router_logits_capture_target = spec.router_logits_capture_target
|
|
self.router_logits_capture_index = spec.router_logits_capture_index
|
|
self.router_logits_capture_mode = spec.router_logits_capture_mode
|
|
self.moe_output_shape = spec.moe_output_shape
|
|
self.top_k = spec.top_k
|
|
self.score_apply = resolve_score_apply_mode(spec, config.score_apply)
|
|
self.combine_impl = resolve_combine_impl(config.combine_impl)
|
|
route_norm = spec.route_norm if config.route_norm is None else config.route_norm
|
|
self.ep_size = ep_size
|
|
self.ep_rank = ep_rank
|
|
self.num_experts = spec.num_experts
|
|
self.num_local_experts = spec.num_experts // ep_size
|
|
self.hidden_size = spec.hidden_size
|
|
self.ep_group_name = f"ep_size_{ep_size}"
|
|
self.ep_group = None # Set by set_deepspeed_parallelism()
|
|
self.folding_group_handles = None
|
|
self.tp_group = None
|
|
resolved_config = resolve_autoep_config_defaults(config, spec.model_family)
|
|
self.validate_folding_routing = bool(resolved_config.validate_folding_routing)
|
|
|
|
# Router: copy gate weights from source
|
|
source_gate = getattr(source_module, spec.router_name)
|
|
source_gate_bias = getattr(source_gate, 'bias', None)
|
|
source_ecb = getattr(source_gate, 'e_score_correction_bias', None)
|
|
unsupported_router_biases = [
|
|
getattr(source_gate, bias_name, None) for bias_name in spec.unsupported_router_bias_names
|
|
]
|
|
if not spec.supports_expert_bias and resolved_config.load_balance_coeff is not None:
|
|
raise ValueError(f"AutoEP preset '{spec.model_family}' does not support load_balance_coeff/expert_bias "
|
|
"yet. Set load_balance_coeff=None.")
|
|
with _gather_source_zero_params([source_gate.weight, source_gate_bias, source_ecb,
|
|
*unsupported_router_biases]):
|
|
for bias_name, router_bias in zip(spec.unsupported_router_bias_names, unsupported_router_biases):
|
|
if router_bias is None:
|
|
continue
|
|
if torch.is_tensor(router_bias) and torch.count_nonzero(router_bias.detach()).item() == 0:
|
|
continue
|
|
raise ValueError(f"AutoEP preset '{spec.model_family}' does not support nonzero router bias "
|
|
f"'{bias_name}' yet.")
|
|
self.router = TokenChoiceTopKRouter(
|
|
dim=spec.hidden_size,
|
|
num_experts=spec.num_experts,
|
|
num_expert_groups=spec.num_expert_groups,
|
|
num_limited_groups=spec.num_limited_groups,
|
|
top_k=spec.top_k,
|
|
score_func=spec.score_func,
|
|
route_norm=route_norm,
|
|
route_scale=spec.route_scale,
|
|
gate_bias=spec.gate_bias,
|
|
group_score_func=spec.group_score_func,
|
|
)
|
|
# Copy gate weights
|
|
_copy_parameter_data(self.router.gate.weight, source_gate.weight)
|
|
self.router.gate.weight.requires_grad_(source_gate.weight.requires_grad)
|
|
if spec.gate_bias and source_gate_bias is not None:
|
|
_copy_parameter_data(self.router.gate.bias, source_gate_bias)
|
|
self.router.gate.bias.requires_grad_(source_gate_bias.requires_grad)
|
|
|
|
# Copy pre-trained score correction bias (DeepSeek-V3/Moonlight noaux_tc routing)
|
|
if source_ecb is not None and isinstance(source_ecb, nn.Parameter):
|
|
self.router.e_score_correction_bias = nn.Parameter(source_ecb.data.clone(),
|
|
requires_grad=source_ecb.requires_grad)
|
|
logger.info('AutoEP: copied e_score_correction_bias from source gate '
|
|
'(shape=%s)', source_ecb.shape)
|
|
|
|
# Alias router under the name OutputRecorder expects (layer_name if provided),
|
|
# but only when OutputRecorder captures from the router child and the alias is safe.
|
|
alias_target = spec.router_logits_capture_layer_name or spec.router_name
|
|
if spec.router_logits_capture_target == "router" and alias_target != "router":
|
|
if "." in alias_target or alias_target in ("experts", "shared_experts") or hasattr(self, alias_target):
|
|
logger.warning(f"Skipping router alias '{alias_target}' to avoid name collision.")
|
|
else:
|
|
setattr(self, alias_target, self.router)
|
|
|
|
# Experts: extract local expert weights
|
|
w1, w2, w3 = repack_expert_weights(
|
|
experts_source=getattr(source_module, spec.experts_name),
|
|
spec=spec,
|
|
ep_rank=ep_rank,
|
|
ep_size=ep_size,
|
|
)
|
|
w1_requires_grad, w2_requires_grad, w3_requires_grad = repack_expert_requires_grad_flags(
|
|
experts_source=getattr(source_module, spec.experts_name),
|
|
spec=spec,
|
|
ep_rank=ep_rank,
|
|
ep_size=ep_size,
|
|
)
|
|
self.experts = GroupedExperts(
|
|
dim=spec.hidden_size,
|
|
hidden_dim=spec.ffn_hidden_size,
|
|
num_experts=self.num_local_experts,
|
|
use_grouped_mm=config.use_grouped_mm,
|
|
)
|
|
_copy_parameter_data(self.experts.w1, w1)
|
|
_copy_parameter_data(self.experts.w2, w2)
|
|
_copy_parameter_data(self.experts.w3, w3)
|
|
self.experts.w1.requires_grad_(w1_requires_grad)
|
|
self.experts.w2.requires_grad_(w2_requires_grad)
|
|
self.experts.w3.requires_grad_(w3_requires_grad)
|
|
|
|
self.reorderer = TokenReorderer(num_experts=self.num_experts, top_k=self.top_k)
|
|
self.shared_experts = getattr(source_module, spec.shared_experts_name,
|
|
None) if spec.has_shared_experts else None
|
|
self.shared_experts_gate = getattr(source_module, spec.shared_experts_gate_name,
|
|
None) if spec.shared_experts_gate_name else None
|
|
|
|
# Mark expert params for EDP gradient reduction
|
|
for param in self.experts.parameters():
|
|
param.allreduce = False
|
|
param.group_name = self.ep_group_name
|
|
param.ds_zero_placement_family = "autoep_expert"
|
|
param.ds_zero_partition_group_name = self.ep_group_name
|
|
|
|
# Mark shared expert and router params for global DP reduction.
|
|
# The router runs redundantly on every TP peer and its gradient is
|
|
# rebuilt into a replicated full view by the restore all-gather, so it
|
|
# is tagged as the replicated family (AVERAGE TP reduction); a SUM would
|
|
# double it under tp_size=2. See mark_autoep_folding_router_parameter.
|
|
for param in self.router.parameters():
|
|
param.allreduce = True
|
|
mark_autoep_folding_router_parameter(param)
|
|
param.ds_zero_placement_family = "replicated"
|
|
if self.shared_experts is not None:
|
|
for param in self.shared_experts.parameters():
|
|
param.allreduce = True
|
|
param.ds_zero_placement_family = "replicated"
|
|
if self.shared_experts_gate is not None:
|
|
for param in self.shared_experts_gate.parameters():
|
|
param.allreduce = True
|
|
param.ds_zero_placement_family = "replicated"
|
|
|
|
# Load balancing buffers
|
|
self.load_balance_coeff = resolved_config.load_balance_coeff
|
|
buf_device = source_gate.weight.device
|
|
if self.load_balance_coeff is not None:
|
|
self.register_buffer(
|
|
"expert_bias",
|
|
torch.zeros(spec.num_experts, dtype=torch.float32, device=buf_device),
|
|
persistent=True,
|
|
)
|
|
else:
|
|
self.expert_bias = None
|
|
self.register_buffer(
|
|
"tokens_per_expert",
|
|
torch.zeros(spec.num_experts, dtype=torch.float32, device=buf_device),
|
|
persistent=False,
|
|
)
|
|
|
|
# Router-logit cache
|
|
self._cached_router_logits = None
|
|
self._register_logit_hook()
|
|
|
|
def _register_logit_hook(self):
|
|
"""Register a forward hook that caches gate logits for OutputRecorder capture."""
|
|
if self.router_logits_capture_target != "router":
|
|
return
|
|
|
|
def hook_fn(module, input, output):
|
|
x = input[0] # [T, H]
|
|
logits = module.gate(x) # [T, E_global]
|
|
if self.router_logits_capture_mode == "post_score":
|
|
if self.router.score_func == "softmax":
|
|
logits = torch.softmax(logits.float(), dim=-1).to(logits.dtype)
|
|
elif self.router.score_func == "sigmoid":
|
|
logits = torch.sigmoid(logits.float()).to(logits.dtype)
|
|
self._cached_router_logits = logits
|
|
|
|
self.router.register_forward_hook(hook_fn)
|
|
|
|
def set_deepspeed_parallelism(
|
|
self,
|
|
use_data_before_expert_parallel_: bool = False,
|
|
folding_group_handles=None,
|
|
) -> None:
|
|
"""Bind EP group handle to this module."""
|
|
from deepspeed.utils import groups
|
|
from deepspeed.utils.bwc import bwc_pipeline_parallel_world_size
|
|
|
|
if folding_group_handles is not None:
|
|
self.folding_group_handles = folding_group_handles
|
|
self.ep_group_name = folding_group_handles.ep_group_name
|
|
self.ep_group = folding_group_handles.ep_group
|
|
self.tp_group = folding_group_handles.tp_group
|
|
self.ep_rank = dist.get_rank(group=self.ep_group)
|
|
return
|
|
|
|
if self.ep_group_name not in groups._get_expert_parallel_group_dict():
|
|
mp_size = max(
|
|
getattr(groups, '_get_model_parallel_world_size', lambda: 1)(),
|
|
getattr(groups, '_get_sequence_parallel_world_size', lambda: 1)(),
|
|
)
|
|
mp_mode = "tp" if getattr(groups, '_get_model_parallel_world_size', lambda: 1)() > 1 else "sp"
|
|
pp_size = 1 if groups.mpu is None else bwc_pipeline_parallel_world_size(groups.mpu)
|
|
groups._create_expert_and_data_parallel(
|
|
expert_parallel_size_=self.ep_size,
|
|
mp_size=mp_size,
|
|
pp_size=pp_size,
|
|
mp_mode=mp_mode,
|
|
use_data_before_expert_parallel_=use_data_before_expert_parallel_,
|
|
)
|
|
self.ep_group = groups._get_expert_parallel_group(self.ep_group_name)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
|
"""Forward pass.
|
|
|
|
Args:
|
|
hidden_states: [B, S, H]
|
|
|
|
Returns:
|
|
[B, S, H] or ([B, S, H], [T, E]) if return_router_logits.
|
|
Some HF MoE contracts return ([T, H], [T, E]) instead.
|
|
"""
|
|
bsz, seqlen, hdim = hidden_states.shape
|
|
x = hidden_states.reshape(-1, hdim) # [T, H]
|
|
|
|
# Router
|
|
ro: RouterOutput = RouterOutput(*self.router(x, self.expert_bias))
|
|
|
|
# Accumulate expert utilization
|
|
with torch.no_grad():
|
|
self.tokens_per_expert.add_(ro.num_tokens_per_expert)
|
|
|
|
# Reorder tokens by expert
|
|
top_scores_sorted, token_indices_sorted, _ = self.reorderer(ro.top_scores, ro.selected_experts)
|
|
expert_indices_sorted = ro.selected_experts.reshape(-1).index_select(0, token_indices_sorted)
|
|
|
|
folded_tp = self.folding_group_handles is not None and self.folding_group_handles.spec.tp_size > 1
|
|
restore_ctx = None
|
|
if folded_tp:
|
|
from deepspeed.moe.ep_tp_dispatch import (
|
|
RoutedAssignmentPayload,
|
|
assignment_ordinals_by_expert,
|
|
assert_tp_payload_consistent,
|
|
dispatch_counters,
|
|
partition_assignments,
|
|
restore_combined,
|
|
)
|
|
payload = RoutedAssignmentPayload(
|
|
token_indices=(token_indices_sorted // self.top_k).to(torch.long),
|
|
expert_indices=expert_indices_sorted.to(torch.long),
|
|
assignment_indices=assignment_ordinals_by_expert(expert_indices_sorted.to(torch.long)),
|
|
capacity_slots=(token_indices_sorted % self.top_k).to(torch.long),
|
|
combine_weights=top_scores_sorted
|
|
if self.score_apply == "post" else torch.ones_like(top_scores_sorted),
|
|
drop_mask=torch.zeros_like(top_scores_sorted, dtype=torch.bool),
|
|
pad_mask=torch.zeros_like(top_scores_sorted, dtype=torch.bool),
|
|
input_splits=[0 for _ in range(self.ep_size)],
|
|
output_splits=[0 for _ in range(self.ep_size)],
|
|
extra={
|
|
"destination_ranks": (expert_indices_sorted // self.num_local_experts).to(torch.long),
|
|
"top_scores": top_scores_sorted,
|
|
"num_tokens": torch.tensor(bsz * seqlen, device=hidden_states.device, dtype=torch.long),
|
|
},
|
|
)
|
|
if self.validate_folding_routing:
|
|
assert_tp_payload_consistent(payload,
|
|
tp_group=self.tp_group,
|
|
tp_size=self.folding_group_handles.spec.tp_size)
|
|
tp_rank = dist.get_rank(group=self.tp_group)
|
|
local_payload, restore_ctx = partition_assignments(payload,
|
|
tp_group=self.tp_group,
|
|
tp_rank=tp_rank,
|
|
tp_size=self.folding_group_handles.spec.tp_size)
|
|
token_indices_for_compute = token_indices_sorted.index_select(0, restore_ctx.local_indices)
|
|
top_scores_for_compute = top_scores_sorted.index_select(0, restore_ctx.local_indices)
|
|
expert_indices_for_plan = local_payload.expert_indices
|
|
else:
|
|
token_indices_for_compute = token_indices_sorted
|
|
top_scores_for_compute = top_scores_sorted
|
|
expert_indices_for_plan = expert_indices_sorted
|
|
|
|
routed_input = x[token_indices_for_compute // self.top_k] # [N, H]
|
|
routed_input = apply_scores_before_experts_if_enabled(routed_input,
|
|
top_scores_for_compute,
|
|
score_apply=self.score_apply)
|
|
|
|
if self.ep_size == 1:
|
|
# No AllToAll needed - local computation only
|
|
local_counts = count_tokens_per_expert(
|
|
ro.selected_experts,
|
|
self.num_local_experts,
|
|
out_dtype=torch.int32,
|
|
)
|
|
|
|
routed_input_permuted, perm_indices, aligned_counts, n_tokens = permute_by_local_expert(
|
|
routed_input, local_counts)
|
|
expert_output = self.experts(routed_input_permuted, aligned_counts)
|
|
expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens)
|
|
else:
|
|
# EP dispatch/compute/combine
|
|
if folded_tp:
|
|
plan = compute_split_plan_from_expert_indices(
|
|
expert_indices=expert_indices_for_plan,
|
|
num_experts=self.num_experts,
|
|
ep_size=self.ep_size,
|
|
num_local_experts=self.num_local_experts,
|
|
ep_group=self.ep_group,
|
|
)
|
|
else:
|
|
plan = compute_split_plan(
|
|
selected_experts=ro.selected_experts,
|
|
num_experts=self.num_experts,
|
|
ep_size=self.ep_size,
|
|
num_local_experts=self.num_local_experts,
|
|
ep_group=self.ep_group,
|
|
)
|
|
|
|
routed_input = _AllToAllV.apply(self.ep_group, routed_input, plan.input_splits, plan.output_splits)
|
|
|
|
routed_input, perm_indices, aligned_counts, n_tokens = permute_by_local_expert(
|
|
routed_input, plan.local_counts_by_source)
|
|
expert_output = self.experts(routed_input, aligned_counts)
|
|
expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens)
|
|
|
|
expert_output = _AllToAllV.apply(self.ep_group, expert_output, plan.output_splits, plan.input_splits)
|
|
|
|
if folded_tp:
|
|
output = restore_combined(expert_output,
|
|
restore_ctx,
|
|
tp_group=self.tp_group,
|
|
validate_coverage=self.validate_folding_routing).reshape(bsz, seqlen, hdim)
|
|
self._last_folding_dispatch_counters = dispatch_counters(restore_ctx)
|
|
else:
|
|
output = combine_from_routed(
|
|
expert_output,
|
|
top_scores=ro.top_scores,
|
|
token_indices_sorted=token_indices_sorted,
|
|
top_k=self.top_k,
|
|
score_apply=self.score_apply,
|
|
combine_impl=self.combine_impl,
|
|
shape=(bsz, seqlen, hdim),
|
|
)
|
|
|
|
if self.moe_output_shape == "flat":
|
|
output = output.reshape(-1, hdim)
|
|
shared_expert_input = x
|
|
elif self.shared_experts_gate is not None:
|
|
shared_expert_input = x
|
|
else:
|
|
shared_expert_input = hidden_states
|
|
|
|
if self.shared_experts is not None:
|
|
shared_expert_output = self.shared_experts(shared_expert_input)
|
|
if self.shared_experts_gate is not None:
|
|
shared_expert_gate = torch.sigmoid(self.shared_experts_gate(shared_expert_input))
|
|
shared_expert_output = shared_expert_gate * shared_expert_output
|
|
if shared_expert_output.shape != output.shape:
|
|
shared_expert_output = shared_expert_output.reshape_as(output)
|
|
output = output + shared_expert_output
|
|
|
|
if self.return_router_logits:
|
|
logits = self._cached_router_logits
|
|
self._cached_router_logits = None
|
|
return output, logits
|
|
|
|
self._cached_router_logits = None
|
|
return output
|