wehub-resource-sync
181 lines
8.6 KiB
Python
181 lines
8.6 KiB
Python
import torch
|
|
|
|
|
|
def generate_topk_idx(rank_count: torch.Tensor, num_tokens: int, num_experts: int, num_ranks: int, num_topk: int) -> torch.Tensor:
|
|
"""
|
|
Map rank count to expert indices
|
|
"""
|
|
assert torch.equal(torch.sum(rank_count, dim=1), torch.ones(num_tokens, dtype=torch.int, device='cuda') * num_topk)
|
|
assert (num_tokens, num_ranks) == rank_count.shape
|
|
num_experts_per_rank = num_experts // num_ranks
|
|
|
|
# Generate base value
|
|
base_vals = torch.arange(num_experts, device='cuda').view(1, num_ranks, num_experts_per_rank).expand(num_tokens, num_ranks, num_experts_per_rank)
|
|
|
|
# Randomize the ordering within each row
|
|
rand_vals = torch.rand(num_tokens, num_ranks, num_experts_per_rank, device='cuda')
|
|
perm_indices = torch.argsort(rand_vals, dim=-1)
|
|
permuted = torch.gather(base_vals, 2, perm_indices)
|
|
|
|
# Create the mask
|
|
k_idx = torch.arange(num_experts_per_rank, device='cuda').view(1, 1, num_experts_per_rank).expand(num_tokens, num_ranks, num_experts_per_rank)
|
|
rank_count_expanded = rank_count.unsqueeze(2).expand(num_tokens, num_ranks, num_experts_per_rank)
|
|
mask = k_idx < rank_count_expanded
|
|
|
|
# Get the final indices by masking and reshaping
|
|
selected = permuted[mask] # (num_tokens * num_topk,)
|
|
topk_idx = selected.view(num_tokens, num_topk)
|
|
|
|
return topk_idx
|
|
|
|
|
|
def generate_rank_count(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float) -> torch.Tensor:
|
|
"""
|
|
Generate rank count tensor for a given number of tokens, experts, ranks, and top-k.
|
|
|
|
This function generates a tensor of shape `(num_tokens, num_ranks)` where each element `[i, j]` represents
|
|
the number of topk experts that token `i` have on rank `j`. The distribution is such that
|
|
one special rank gets `ratio` times more traffic than the others.
|
|
"""
|
|
num_experts_per_rank = num_experts // num_ranks
|
|
num_normal_ranks = num_ranks - 1
|
|
|
|
assert ratio >= 1.0, 'ratio must be no less than 1.0'
|
|
|
|
# Generate rank count of each token from random distribution
|
|
random_scores = torch.rand(num_tokens, num_experts, device='cuda')
|
|
topk_weights_, topk_indices = torch.topk(random_scores, num_topk, dim=1, largest=True, sorted=False)
|
|
topk_indices //= num_experts_per_rank
|
|
sorted_topk_indices = torch.sort(topk_indices, dim=1)[0]
|
|
topk_indices_diff_mask = sorted_topk_indices[:, 1:] != sorted_topk_indices[:, :-1]
|
|
a = topk_indices_diff_mask.sum(dim=1) + 1
|
|
|
|
# Upper bound for this generating algorithm
|
|
upper_bound_per_token = int(num_normal_ranks / ratio) + 1
|
|
|
|
# Clamp the value in range [1, upper_bound_per_token] for each token
|
|
a = torch.clamp(a, None, upper_bound_per_token)
|
|
|
|
# Consider the special rank
|
|
sum_a = torch.sum(a).item()
|
|
normal_token_count = int(sum_a / (num_normal_ranks + ratio))
|
|
special_token_count = sum_a - normal_token_count * num_normal_ranks
|
|
special_token_count = min(special_token_count, int(normal_token_count * ratio) + 1)
|
|
|
|
# Tokens that the special rank must be in topk
|
|
must_mask = (a == num_ranks)
|
|
must_count = int(must_mask.sum().item())
|
|
special_token_count = max(must_count, special_token_count)
|
|
assert must_count <= special_token_count, 'Too many tokens with full rank assignment'
|
|
|
|
# Tokens that the special rank can optionally be in topk
|
|
optional_token_indices = torch.where(must_mask == 0)[0]
|
|
optional_token_indices = optional_token_indices[torch.randperm(num_tokens - must_count, device='cuda')][:special_token_count - must_count]
|
|
must_token_indices = torch.where(must_mask != 0)[0]
|
|
special_token_row_index = torch.cat(([must_token_indices, optional_token_indices]))
|
|
|
|
# Generate permutations for normal ranks
|
|
rank_perm = (torch.randperm(num_normal_ranks, device='cuda') + 1).repeat(num_tokens * num_topk // num_normal_ranks + 1)
|
|
|
|
# Compute cumulative sum of a to get starting indices in b for each row
|
|
a_cumsum = torch.cumsum(torch.cat((torch.tensor([0], device='cuda'), a)), dim=0)
|
|
row_starts = a_cumsum[:-1] # Starting indices for each row in b, shape (n,)
|
|
|
|
# Insert special rank index into the permutation for special tokens
|
|
rank_perm_with_special_rank = torch.zeros(num_tokens * num_topk, dtype=torch.long, device='cuda') # (n * k,)
|
|
special_token_mask = torch.zeros(num_tokens * num_topk, dtype=torch.bool, device='cuda')
|
|
special_token_flattened_row_index = row_starts[special_token_row_index]
|
|
special_token_mask[special_token_flattened_row_index] = 1
|
|
all_indices = torch.arange(num_tokens * num_topk, device='cuda')
|
|
non_special_indices = all_indices[special_token_mask != True]
|
|
rank_perm_with_special_rank[non_special_indices] = rank_perm[:len(non_special_indices)]
|
|
|
|
# Create column index grids
|
|
col_idx = torch.arange(num_topk, device='cuda').view(1, num_topk) # (1, num_topk)
|
|
|
|
# Compute modulo indices: col_idx % a[i] for each row
|
|
# torch.max is used to avoid zeros in case a[i] = 0 (which happens when the only topk rank is the special rank)
|
|
mod_idx = col_idx % a.view(num_tokens, 1) # (n, num_topk)
|
|
|
|
# Compute indices in b: row_start + (col % a[i])
|
|
b_indices = row_starts.view(num_tokens, 1) + mod_idx # (n, k)
|
|
|
|
# Gather values from b using computed indices
|
|
result = rank_perm_with_special_rank[b_indices]
|
|
|
|
# Shuffle rows randomly to avoid any pattern
|
|
shuffle_indices = torch.randperm(num_tokens, device='cuda')
|
|
result = result[shuffle_indices] # Shuffle rows
|
|
|
|
# Create rank count tensor
|
|
rank_count = torch.zeros((num_tokens, num_ranks), dtype=torch.int32, device='cuda')
|
|
rank_count.scatter_add_(dim=1, index=result, src=torch.ones_like(result, dtype=torch.int32))
|
|
return rank_count
|
|
|
|
|
|
def get_precise_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float):
|
|
"""
|
|
Generate precise unbalanced scores for testing.
|
|
|
|
Note that this function generates a distribution with precise unbalanced distribution,
|
|
which **differs from real distribution**.
|
|
"""
|
|
# Generate num topk experts for each rank
|
|
rank_count = generate_rank_count(num_tokens, num_experts, num_ranks, num_topk, ratio)
|
|
|
|
# Generate scores in a low distribution
|
|
threshold = 0.9
|
|
scores = torch.empty((num_tokens, num_experts), dtype=torch.float32, device='cuda')
|
|
scores.uniform_(to=threshold)
|
|
|
|
# Generate topk indices and change their scores to a high distribution
|
|
topk_idx = generate_topk_idx(rank_count, num_tokens, num_experts, num_ranks, num_topk)
|
|
topk_scores = torch.empty((num_tokens, num_topk), dtype=torch.float32, device='cuda')
|
|
topk_scores.uniform_(threshold + 1e-6, 1.0)
|
|
row_idx = torch.arange(num_tokens).unsqueeze(1).expand(num_tokens, num_topk)
|
|
scores[row_idx, topk_idx] = topk_scores
|
|
return scores
|
|
|
|
|
|
def get_scores_by_factor(num_tokens: int, num_experts: int, num_ranks: int, factor: float) -> torch.Tensor:
|
|
num_experts_per_rank = num_experts // num_ranks
|
|
scores = torch.empty((num_tokens, num_experts), dtype=torch.float32, device='cuda')
|
|
scores[:, :num_experts_per_rank].uniform_(to=factor)
|
|
scores[:, num_experts_per_rank:].uniform_(to=1)
|
|
return scores
|
|
|
|
|
|
def map_unbalanced_ratio_to_factor(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float) -> float:
|
|
num_iterations = 20
|
|
factor_l, factor_r = 1.0, 100.0
|
|
|
|
num_experts_per_rank = num_experts // num_ranks
|
|
for _i in range(num_iterations):
|
|
factor_mid = (factor_l + factor_r) / 2
|
|
scores = get_scores_by_factor(num_tokens, num_experts, num_ranks, factor_mid)
|
|
_, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
|
|
rank_idx = topk_idx // num_experts_per_rank
|
|
one_hot = torch.nn.functional.one_hot(rank_idx, num_ranks)
|
|
counts = one_hot.any(dim=1).to(torch.float).sum(dim=0)
|
|
if counts[0].item() > counts[1:].mean().item() * ratio:
|
|
factor_r = factor_mid
|
|
else:
|
|
factor_l = factor_mid
|
|
return factor_l
|
|
|
|
|
|
def get_random_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float):
|
|
"""Generate unbalanced scores with a given ratio.
|
|
"""
|
|
factor = 1.0
|
|
if ratio != 1.0:
|
|
factor = map_unbalanced_ratio_to_factor(num_tokens, num_experts, num_ranks, num_topk, ratio)
|
|
return get_scores_by_factor(num_tokens, num_experts, num_ranks, factor)
|
|
|
|
|
|
def get_unbalanced_scores(num_tokens: int, num_experts: int, num_ranks: int, num_topk: int, ratio: float, precise: bool):
|
|
if precise:
|
|
return get_precise_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio)
|
|
else:
|
|
return get_random_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio)
|