wehub-resource-sync
244 lines
12 KiB
Python
244 lines
12 KiB
Python
import math
|
|
import torch
|
|
import torch.distributed as dist
|
|
from typing import Optional, Tuple, Union
|
|
|
|
from .envs import get_global_seed
|
|
from .math import ceil_div
|
|
|
|
|
|
def dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
|
|
topk_idx: torch.Tensor, topk_weights: Optional[torch.Tensor],
|
|
num_max_tokens_per_rank: int, num_experts: int):
|
|
"""
|
|
The reference implementation of dispatching tokens to experts across multiple ranks.
|
|
|
|
Not expanded. Sorted by rank and then by token within each rank (i.e. sorted by `src_token_global_idx`).
|
|
|
|
Arguments:
|
|
- `x`: Input tokens, `[num_tokens, hidden]` or (`[num_tokens, hidden], [num_tokens, hidden_sf]`)
|
|
- `topk_idx`: Top-k expert indices for each token, `[num_tokens, num_topk]`
|
|
- `topk_weights`: Top-k weights for each token, can be None, `[num_tokens, num_topk]`
|
|
- `num_max_tokens_per_rank`: Maximum number of tokens per rank, must >= actual number of tokens per rank and aligned with DeepEP's `num_max_tokens_per_rank` since we're going to calculate `src_token_global_idx = src_rank_idx*num_max_tokens_per_rank + src_token_local_idx`
|
|
- `num_experts`: Total number of experts across all ranks
|
|
|
|
Returns:
|
|
- `recv_x`, `recv_topk_idx`, and `recv_topk_weights`: Received tokens, top-k indices, and top-k weights after dispatching. Out of range `recv_topk_idx` (i.e. that expert is not on the current rank) are set to -1.
|
|
- `recv_src_token_idx`: Received `src_token_global_idx` for each received token
|
|
- `num_recv_tokens_per_rank`: Number of received tokens from each rank, `[num_ranks]`
|
|
"""
|
|
# TODO: support forwarding
|
|
# TODO: make top-k weight fully optional
|
|
rank_idx = dist.get_rank()
|
|
num_ranks = dist.get_world_size()
|
|
|
|
assert num_experts % num_ranks == 0
|
|
num_experts_per_rank = num_experts // num_ranks
|
|
|
|
# Unpack SF
|
|
use_fp8 = isinstance(x, tuple)
|
|
x, sf = x if use_fp8 else (x, None)
|
|
|
|
# TODO: use SF bytes instead of hardcoded recipe
|
|
num_tokens, hidden = x.size()
|
|
num_tokens_, num_topk = topk_idx.size()
|
|
assert num_tokens == num_tokens_
|
|
if sf is not None:
|
|
num_tokens__, hidden_sf = sf.size()
|
|
assert num_tokens == num_tokens__
|
|
assert hidden_sf == ceil_div(hidden, 128)
|
|
if topk_weights is not None:
|
|
num_tokens__, num_topk_ = topk_weights.size()
|
|
assert num_tokens == num_tokens__
|
|
assert num_topk == num_topk_
|
|
|
|
# Prepare per-peer send buffers
|
|
send_x_list = []
|
|
send_sf_list = []
|
|
send_topk_idx_list = []
|
|
send_topk_weights_list = []
|
|
send_src_token_idx_list = []
|
|
num_send_tokens_per_rank = torch.zeros((num_ranks, ), dtype=torch.int, device=x.device)
|
|
for dst_rank_idx in range(num_ranks):
|
|
expert_start_idx = dst_rank_idx * num_experts_per_rank
|
|
expert_end_idx = expert_start_idx + num_experts_per_rank
|
|
|
|
# Get the indices of tokens
|
|
mask_to_send = ((expert_start_idx <= topk_idx) & (topk_idx < expert_end_idx)).any(dim=1)
|
|
indices_to_send = mask_to_send.nonzero(as_tuple=True)[0]
|
|
num_send_tokens_per_rank[dst_rank_idx] = indices_to_send.numel()
|
|
|
|
# Select the data for tokens
|
|
x_to_send = x[indices_to_send]
|
|
sf_to_send = sf[indices_to_send] if use_fp8 else None
|
|
topk_idx_to_send = topk_idx[indices_to_send]
|
|
topk_weights_to_send = topk_weights[indices_to_send]
|
|
masked_topk_idx = torch.where((expert_start_idx <= topk_idx_to_send) & (topk_idx_to_send < expert_end_idx),
|
|
topk_idx_to_send, torch.full_like(topk_idx_to_send, -1))
|
|
|
|
send_x_list.append(x_to_send)
|
|
send_sf_list.append(sf_to_send)
|
|
send_topk_idx_list.append(masked_topk_idx)
|
|
send_topk_weights_list.append(topk_weights_to_send)
|
|
send_src_token_idx_list.append(indices_to_send)
|
|
|
|
send_x = torch.cat(send_x_list, dim=0)
|
|
send_sf = torch.cat(send_sf_list, dim=0) if use_fp8 else None
|
|
send_topk_idx = torch.cat(send_topk_idx_list, dim=0)
|
|
send_topk_weights = torch.cat(send_topk_weights_list, dim=0)
|
|
send_src_token_idx = torch.cat(send_src_token_idx_list, dim=0).to(torch.int)
|
|
send_src_token_idx += rank_idx * num_max_tokens_per_rank
|
|
|
|
# Exchange size
|
|
num_recv_tokens_per_rank = torch.empty((num_ranks, ), dtype=torch.int, device=x.device)
|
|
dist.all_to_all_single(num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
|
num_recv_tokens = int(num_recv_tokens_per_rank.sum().item())
|
|
|
|
# Exchange main data
|
|
num_send_tokens_per_rank = num_send_tokens_per_rank.tolist()
|
|
num_recv_tokens_per_rank = num_recv_tokens_per_rank.tolist()
|
|
recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device=x.device)
|
|
recv_sf = torch.empty((num_recv_tokens, hidden_sf), dtype=sf.dtype, device=x.device) if use_fp8 else None
|
|
recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=topk_idx.dtype, device=x.device)
|
|
recv_topk_weights = torch.empty((num_recv_tokens, num_topk), dtype=topk_weights.dtype, device=x.device)
|
|
recv_src_token_idx = torch.empty((num_recv_tokens, ), dtype=torch.int, device=x.device)
|
|
dist.all_to_all_single(recv_x, send_x, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
|
if use_fp8:
|
|
dist.all_to_all_single(recv_sf, send_sf, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
|
dist.all_to_all_single(recv_topk_idx, send_topk_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
|
dist.all_to_all_single(recv_topk_weights, send_topk_weights, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
|
dist.all_to_all_single(recv_src_token_idx, send_src_token_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
|
|
|
|
# Mask top-k indices
|
|
expert_start_idx = rank_idx * num_experts_per_rank
|
|
expert_end_idx = expert_start_idx + num_experts_per_rank
|
|
mask = (expert_start_idx <= recv_topk_idx) & (recv_topk_idx < expert_end_idx)
|
|
recv_topk_idx = recv_topk_idx - expert_start_idx
|
|
recv_topk_idx.masked_fill_(~mask, -1)
|
|
|
|
# Pack SF
|
|
recv_x = (recv_x, recv_sf) if use_fp8 else recv_x
|
|
|
|
return (recv_x, recv_topk_idx, recv_topk_weights,
|
|
recv_src_token_idx, torch.tensor(num_recv_tokens_per_rank, dtype=torch.int))
|
|
|
|
|
|
def generate_pre_combine_data(src_token_global_idx: torch.Tensor,
|
|
num_max_tokens_per_rank: int, num_topk: int, hidden: int) -> torch.Tensor:
|
|
"""
|
|
Generate data needed for combine from `src_token_global_idx`.
|
|
Recall that `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
|
|
The generated data (denoted as `y`) of the i-th token has a shape of [num_topk, hidden], with
|
|
|
|
`y[j, k] = sin((token_seeds * P % max_seed + 1) / max_seed * (k + 1) + sin(seed))`
|
|
|
|
where `P=131071` is a large prime, `token_seeds` is calculated via `token_seeds = src_token_global_idx[i] * num_topk + j`,
|
|
and `max_seed = num_ranks * num_max_tokens_per_rank * num_topk`.
|
|
|
|
Arguments:
|
|
- `src_token_global_idx`: Source token global indices, `[num_tokens]`
|
|
|
|
Returns:
|
|
- Generated data, `[num_tokens, num_topk, hidden]`
|
|
"""
|
|
num_ranks = dist.get_world_size()
|
|
token_seeds = (src_token_global_idx.unsqueeze(1) * num_topk +
|
|
torch.arange(num_topk, device=src_token_global_idx.device).unsqueeze(0)) # [num_tokens, num_topk]
|
|
max_seed = num_ranks * num_max_tokens_per_rank * num_topk
|
|
result = torch.sin(
|
|
(((token_seeds * 131071 % max_seed).float() + 1) / max_seed).unsqueeze(-1) *
|
|
torch.arange(1, hidden + 1, device=src_token_global_idx.device, dtype=torch.float32).broadcast_to(1, 1, hidden) +
|
|
math.sin(float(get_global_seed()))
|
|
)
|
|
return result.to(torch.bfloat16)
|
|
|
|
|
|
def ordered_accumulate(data: torch.Tensor, initial_value: Optional[torch.Tensor] = None) -> torch.Tensor:
|
|
"""
|
|
Accumulate `data` in order along the num_topk dimension.
|
|
|
|
Arguments:
|
|
- `data`: Data to be accumulated, `[num_tokens, num_topk, hidden]`
|
|
- `initial_value`: Initial value for accumulation, `[num_tokens, hidden]`
|
|
|
|
Returns:
|
|
- Result, `[num_tokens, hidden]`
|
|
"""
|
|
num_topk = data.shape[1]
|
|
if initial_value is None:
|
|
result = torch.zeros((data.shape[0], data.shape[2]), dtype=torch.float32, device=data.device)
|
|
else:
|
|
result = initial_value.clone()
|
|
for i in range(num_topk):
|
|
result += data[:, i, :].float()
|
|
return result.to(data.dtype)
|
|
|
|
|
|
def combine(y: torch.Tensor, topk_idx: torch.Tensor,
|
|
num_scaleout_ranks: int, num_scaleup_ranks: int, num_experts: int,
|
|
bias: Optional[Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]],
|
|
reduce_in_local: bool, reduce_in_scaleup: bool) -> torch.Tensor:
|
|
"""
|
|
The reference implementation of (possibly multi-level reduction) combining tokens.
|
|
|
|
Arguments:
|
|
- `y`: Input tokens to be combined, `[num_tokens, num_topk, hidden]`
|
|
- `topk_idx`: `[num_tokens, num_topk]`
|
|
- `reduce_in_local` and `reduce_in_scaleup`: Whether to do reduction within rank and within scale-up group.
|
|
- `(True, True)` -> Hybrid combine
|
|
- `(True, False)` -> Non-hybrid combine
|
|
- `(False, False)` -> Equivalent to `allow_multiple_reduction` is False
|
|
Pay attention that `reduce_in_scaleup` = `True` or `False` is NOT equivalent even if `num_scaleout_ranks == 1` due to `bias` handling.
|
|
|
|
Returns:
|
|
- Combined result, `[num_tokens, hidden]`
|
|
"""
|
|
num_ranks = num_scaleout_ranks * num_scaleup_ranks
|
|
num_tokens, hidden = y.shape[0], y.shape[2]
|
|
num_topk = y.shape[1]
|
|
assert not (not reduce_in_local and reduce_in_scaleup), 'Invalid reduction configuration'
|
|
|
|
def grouped_reduce(data_to_reduce: torch.Tensor, group_id: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Perform in-place grouped reduction on `data_to_reduce` according to `group_id`.
|
|
The summation within each group are performed in strict order along the `num_topk` dimension.
|
|
The result for each group is stored at the rightmost token of that group, and other tokens are set to zero.
|
|
|
|
Arguments:
|
|
- `data_to_reduce`: Data to be reduced, `[num_tokens, num_topk, hidden]`
|
|
- `group_id`: group IDs for each token, `[num_tokens, num_topk]`
|
|
"""
|
|
# Shuffle to make tokens with the same group_id contiguous
|
|
group_id, src_indices = torch.sort(group_id, dim=-1, stable=True)
|
|
# transformed_src_indices[i, j] = i * num_topk + src_indices[i, j]
|
|
transformed_src_indices = (
|
|
(src_indices + torch.arange(0, num_tokens, device=y.device).unsqueeze(-1) * num_topk).flatten())
|
|
data_to_reduce = data_to_reduce.view(-1, hidden)[transformed_src_indices].view(num_tokens, num_topk, hidden)
|
|
# Perform segmented reduce within each group
|
|
cur_accum_buf = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=y.device)
|
|
for i in range(num_topk):
|
|
is_segment_break = torch.full((num_tokens, ), True, dtype=torch.bool, device=y.device) \
|
|
if i == num_topk - 1 else group_id[:, i] != group_id[:, i + 1]
|
|
cur_accum_buf += data_to_reduce[:, i, :].float()
|
|
# For one token, if `is_segment_break` is True,
|
|
# save the accumulated value and clear the buffer, otherwise, clear `data_to_reduce[:, i, :]`
|
|
segment_break_token_indices = torch.where(is_segment_break)[0]
|
|
data_to_reduce[segment_break_token_indices, i] = cur_accum_buf[segment_break_token_indices].to(data_to_reduce.dtype)
|
|
cur_accum_buf[segment_break_token_indices] = 0.0
|
|
non_segment_break_token_indices = torch.where(~is_segment_break)[0]
|
|
data_to_reduce[non_segment_break_token_indices, i] = 0.0
|
|
# Unshuffle
|
|
# noinspection PyShadowingNames
|
|
result = torch.empty_like(data_to_reduce)
|
|
result.view(-1, hidden)[transformed_src_indices] = data_to_reduce.view(-1, hidden)
|
|
return result.view(num_tokens, num_topk, hidden)
|
|
|
|
num_experts_per_rank = num_experts // num_ranks
|
|
if reduce_in_local:
|
|
y = grouped_reduce(y, topk_idx // num_experts_per_rank)
|
|
if reduce_in_scaleup:
|
|
y = grouped_reduce(y, topk_idx // (num_experts_per_rank * num_scaleup_ranks))
|
|
bias_sum = bias[0].float() + bias[1].float() if isinstance(bias, tuple) else bias.float() if bias is not None else None
|
|
result = ordered_accumulate(y, bias_sum)
|
|
return result
|