wehub-resource-sync
382 lines
12 KiB
Python
382 lines
12 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 reordering and permutation utilities for expert parallelism.
|
|
|
|
Ported from TorchTitan's TokenReorderer, Triton kernels, and alignment
|
|
utilities with adaptations for DeepSpeed:
|
|
- Triton import guarded with try/except; pure-PyTorch fallback provided
|
|
- Alignment config exposed as TOKEN_GROUP_ALIGN_SIZE_M
|
|
|
|
This module is self-contained: no imports from deepspeed.module_inject
|
|
or deepspeed.runtime.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Callable
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from deepspeed.moe.ep_count import count_tokens_per_expert
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Try to import Triton; fall back gracefully
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_TRITON_AVAILABLE = False
|
|
try:
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
_TRITON_AVAILABLE = True
|
|
except ImportError:
|
|
logger.info("Triton not available; using pure-PyTorch CPU fallback for "
|
|
"permutation index generation.")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Alignment constant
|
|
# ---------------------------------------------------------------------------
|
|
|
|
TOKEN_GROUP_ALIGN_SIZE_M = 8
|
|
"""Alignment granularity for token groups in grouped GEMM.
|
|
|
|
- bf16: 8 (16 bytes / 2 bytes per elem)
|
|
- fp8: 16 (16 bytes / 1 byte per elem)
|
|
- mxfp8: 32 (scaling block size)
|
|
"""
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Utility: round up
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _round_up(x: int, y: int) -> int:
|
|
"""Round *x* up to the nearest multiple of *y*."""
|
|
return ((x + y - 1) // y) * y
|
|
|
|
|
|
# ===================================================================
|
|
# Triton kernel for filling permutation indices
|
|
# ===================================================================
|
|
|
|
if _TRITON_AVAILABLE:
|
|
|
|
@triton.jit
|
|
def _fill_indices_kernel(
|
|
tokens_per_expert_group_ptr,
|
|
start_index_values_ptr,
|
|
write_offsets_ptr,
|
|
output_ptr,
|
|
experts_per_rank: tl.constexpr,
|
|
num_ranks: tl.constexpr,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(axis=0)
|
|
num_programs = tl.num_programs(axis=0)
|
|
|
|
for expert_id in range(pid, experts_per_rank, num_programs):
|
|
write_offset = tl.load(write_offsets_ptr + expert_id)
|
|
|
|
for r in range(num_ranks):
|
|
i = r * experts_per_rank + expert_id
|
|
start_index = tl.load(start_index_values_ptr + i)
|
|
length = tl.load(tokens_per_expert_group_ptr + i)
|
|
|
|
offsets = tl.arange(0, BLOCK_SIZE)
|
|
for chunk_start in range(0, length, BLOCK_SIZE):
|
|
chunk_offsets = chunk_start + offsets
|
|
mask = chunk_offsets < length
|
|
values = start_index + chunk_offsets
|
|
dest_indices = write_offset + chunk_offsets
|
|
tl.store(output_ptr + dest_indices, values, mask=mask)
|
|
|
|
write_offset += length
|
|
|
|
|
|
# ===================================================================
|
|
# Triton wrapper
|
|
# ===================================================================
|
|
|
|
|
|
def fill_indices_wrapper(
|
|
tokens_per_expert_group: torch.Tensor,
|
|
start_index_values: torch.Tensor,
|
|
write_offsets: torch.Tensor,
|
|
experts_per_rank: int,
|
|
num_ranks: int,
|
|
max_len: int,
|
|
block_size: int = 128,
|
|
max_blocks: int = 1024,
|
|
) -> torch.Tensor:
|
|
"""Launch the Triton kernel to fill permutation indices.
|
|
|
|
Falls back to :func:`fill_indices_cpu` when Triton is unavailable.
|
|
"""
|
|
if not _TRITON_AVAILABLE:
|
|
return fill_indices_cpu(
|
|
tokens_per_expert_group,
|
|
start_index_values,
|
|
write_offsets,
|
|
experts_per_rank,
|
|
num_ranks,
|
|
max_len,
|
|
)
|
|
|
|
permuted_indices = torch.full((max_len, ), -1, dtype=torch.int32, device=tokens_per_expert_group.device)
|
|
|
|
num_blocks = min(experts_per_rank, max_blocks)
|
|
grid = (num_blocks, )
|
|
|
|
_fill_indices_kernel[grid](
|
|
tokens_per_expert_group,
|
|
start_index_values,
|
|
write_offsets,
|
|
permuted_indices,
|
|
experts_per_rank,
|
|
num_ranks,
|
|
BLOCK_SIZE=block_size,
|
|
)
|
|
return permuted_indices
|
|
|
|
|
|
# ===================================================================
|
|
# CPU reference implementation (always available)
|
|
# ===================================================================
|
|
|
|
|
|
def fill_indices_cpu(
|
|
tokens_per_expert_group: torch.Tensor,
|
|
start_index_values: torch.Tensor,
|
|
write_offsets: torch.Tensor,
|
|
experts_per_rank: int,
|
|
num_ranks: int,
|
|
max_len: int,
|
|
) -> torch.Tensor:
|
|
"""Pure-PyTorch CPU reference for filling permutation indices."""
|
|
permuted_indices = torch.full(
|
|
(max_len, ),
|
|
-1,
|
|
dtype=torch.int32,
|
|
)
|
|
for e in range(experts_per_rank):
|
|
write_start = write_offsets[e].item()
|
|
for r in range(num_ranks):
|
|
i = r * experts_per_rank + e
|
|
start_index = start_index_values[i].item()
|
|
length = tokens_per_expert_group[i].item()
|
|
if length > 0:
|
|
end_idx = min(write_start + length, max_len)
|
|
permuted_indices[write_start:end_idx] = torch.arange(
|
|
start_index,
|
|
start_index + (end_idx - write_start),
|
|
dtype=torch.int32,
|
|
)
|
|
write_start += length
|
|
return permuted_indices
|
|
|
|
|
|
# ===================================================================
|
|
# generate_permute_indices
|
|
# ===================================================================
|
|
|
|
|
|
def generate_permute_indices(
|
|
tokens_per_expert_group: torch.Tensor,
|
|
experts_per_rank: int,
|
|
num_ranks: int,
|
|
max_len: int,
|
|
alignment: int,
|
|
use_cpu: bool = False,
|
|
) -> tuple:
|
|
"""Prepare permutation indices and aligned token counts per expert.
|
|
|
|
Args:
|
|
tokens_per_expert_group: Token counts for each expert from all ranks,
|
|
shape ``(num_ranks * experts_per_rank,)``.
|
|
experts_per_rank: Number of experts per rank.
|
|
num_ranks: Number of ranks.
|
|
max_len: Maximum length of the output index vector.
|
|
alignment: Alignment for ``m_sizes`` and padding minimum.
|
|
use_cpu: Whether to force the CPU implementation.
|
|
|
|
Returns:
|
|
Tuple of:
|
|
- permuted_indices: Index mapping from original to expert-grouped order.
|
|
- m_sizes: Aligned token counts per expert.
|
|
- m_offsets: Cumulative sum of m_sizes.
|
|
"""
|
|
# Prefix sum for start indices
|
|
start_index_values = (torch.cumsum(tokens_per_expert_group, 0) - tokens_per_expert_group)
|
|
|
|
# Total tokens per expert across all ranks
|
|
total_tokens_per_expert = tokens_per_expert_group.view(num_ranks, -1).sum(0)
|
|
|
|
# Pad empty experts to alignment minimum
|
|
total_tokens_per_expert = torch.clamp_min(total_tokens_per_expert, alignment)
|
|
|
|
# Align chunk sizes (ceiling division * alignment)
|
|
m_sizes = ((total_tokens_per_expert + alignment - 1) // alignment * alignment).to(torch.int32)
|
|
|
|
# Write offsets per local expert
|
|
m_offsets = torch.cumsum(m_sizes, 0)
|
|
write_offsets = m_offsets - m_sizes
|
|
|
|
if use_cpu:
|
|
permuted_indices = fill_indices_cpu(
|
|
tokens_per_expert_group,
|
|
start_index_values,
|
|
write_offsets,
|
|
experts_per_rank,
|
|
num_ranks,
|
|
max_len,
|
|
)
|
|
else:
|
|
permuted_indices = fill_indices_wrapper(
|
|
tokens_per_expert_group,
|
|
start_index_values,
|
|
write_offsets,
|
|
experts_per_rank,
|
|
num_ranks,
|
|
max_len,
|
|
)
|
|
|
|
return permuted_indices, m_sizes, m_offsets.to(torch.int32)
|
|
|
|
|
|
# ===================================================================
|
|
# _permute / _unpermute / indices_padding_wrapper
|
|
# ===================================================================
|
|
|
|
|
|
def _permute(
|
|
x: torch.Tensor,
|
|
num_tokens_per_expert: torch.Tensor,
|
|
ep_degree: int,
|
|
num_local_experts: int,
|
|
) -> tuple:
|
|
"""Permute tokens into expert-grouped order with alignment padding.
|
|
|
|
Returns:
|
|
Tuple of (input_shape, permuted_x, permuted_indices, aligned_counts).
|
|
"""
|
|
global TOKEN_GROUP_ALIGN_SIZE_M
|
|
x_padded_per_expert = x.shape[0] + num_local_experts * TOKEN_GROUP_ALIGN_SIZE_M
|
|
padded_max_len = _round_up(x_padded_per_expert, TOKEN_GROUP_ALIGN_SIZE_M)
|
|
|
|
with torch.no_grad():
|
|
permuted_indices, num_tokens_per_expert, _offsets = generate_permute_indices(
|
|
num_tokens_per_expert,
|
|
num_local_experts,
|
|
ep_degree,
|
|
padded_max_len,
|
|
TOKEN_GROUP_ALIGN_SIZE_M,
|
|
)
|
|
|
|
# Append a single zero-row for safe indexing of padding slots
|
|
x = torch.vstack((x, x.new_zeros((x.shape[-1]))))
|
|
input_shape = x.shape
|
|
x = x[permuted_indices, :]
|
|
|
|
return input_shape, x, permuted_indices, num_tokens_per_expert
|
|
|
|
|
|
def _unpermute(
|
|
out: torch.Tensor,
|
|
input_shape: torch.Size,
|
|
permuted_indices: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Reverse the permutation produced by :func:`_permute`."""
|
|
out_unpermuted = out.new_empty(input_shape)
|
|
out_unpermuted[permuted_indices, :] = out
|
|
# Strip the extra zero-row appended during _permute
|
|
out = out_unpermuted[:-1]
|
|
return out
|
|
|
|
|
|
def indices_padding_wrapper(func: Callable) -> Callable:
|
|
"""Decorator that pads / aligns token groups for ``torch._grouped_mm``.
|
|
|
|
Wraps an expert-computation function so that each expert's token
|
|
count is a multiple of ``TOKEN_GROUP_ALIGN_SIZE_M``.
|
|
"""
|
|
|
|
def wrapper(
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
w3: torch.Tensor,
|
|
x: torch.Tensor,
|
|
num_tokens_per_expert: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
num_local_experts = w1.shape[0]
|
|
ep_degree = num_tokens_per_expert.shape[0] // num_local_experts
|
|
|
|
input_shape, x, permuted_indices, num_tokens_per_expert = _permute(x, num_tokens_per_expert, ep_degree,
|
|
num_local_experts)
|
|
|
|
out = func(w1, w2, w3, x, num_tokens_per_expert)
|
|
|
|
out = _unpermute(out, input_shape, permuted_indices)
|
|
return out
|
|
|
|
return wrapper
|
|
|
|
|
|
# ===================================================================
|
|
# TokenReorderer
|
|
# ===================================================================
|
|
|
|
|
|
class TokenReorderer(nn.Module):
|
|
"""Reorder token indices to match expert order for efficient parallel
|
|
processing.
|
|
|
|
Args:
|
|
num_experts (int): Number of experts in the MoE layer.
|
|
top_k (int): Number of experts each token is routed to.
|
|
"""
|
|
|
|
def __init__(self, num_experts: int, top_k: int):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
self.top_k = top_k
|
|
|
|
def forward(
|
|
self,
|
|
top_scores: torch.Tensor,
|
|
selected_experts_indices: torch.Tensor,
|
|
) -> tuple:
|
|
"""
|
|
Args:
|
|
top_scores: Routing scores, shape ``(T, top_k)``.
|
|
selected_experts_indices: Expert indices, shape ``(T, top_k)``.
|
|
|
|
Returns:
|
|
Tuple of:
|
|
- top_scores_experts_sorted ``(T * top_k,)``: scores in
|
|
expert-sorted order.
|
|
- token_indices_experts_sorted ``(T * top_k,)``: flattened
|
|
token-slot indices sorted by expert.
|
|
- num_tokens_per_expert ``(num_experts,)``: histogram.
|
|
"""
|
|
num_tokens_per_expert = count_tokens_per_expert(selected_experts_indices, self.num_experts)
|
|
|
|
token_indices_experts_sorted = torch.argsort(selected_experts_indices.view(-1), stable=True)
|
|
|
|
top_scores_experts_sorted = top_scores.view(-1)[token_indices_experts_sorted]
|
|
|
|
return (
|
|
top_scores_experts_sorted,
|
|
token_indices_experts_sorted,
|
|
num_tokens_per_expert,
|
|
)
|