wehub-resource-sync
58 lines
3.5 KiB
Python
58 lines
3.5 KiB
Python
import torch
|
|
|
|
from deep_ep.utils.math import ceil_div
|
|
from deep_ep.utils.gate import get_unbalanced_scores
|
|
|
|
|
|
def test_unbalanced_scores():
|
|
print('Testing gate score generation (Output with num_tokens = 4096, num_experts = 512):')
|
|
for num_tokens in [1, 4096]:
|
|
for num_experts_per_rank in [1, 4, 8, 16]:
|
|
for num_ranks in [2, 4, 8, 16, 64, 72]:
|
|
num_experts = num_experts_per_rank * num_ranks
|
|
for num_topk in [1, 2, 4, 6, 8, 9]:
|
|
if num_topk > num_experts:
|
|
continue
|
|
for ratio in [1.0, 2.0, 4.0]:
|
|
for precise in [1, 0]:
|
|
total_rank_count = torch.zeros(num_ranks, device='cuda')
|
|
|
|
# This is the requirement from precise generation algorithm
|
|
lower_bound_per_token = max(1, ceil_div(num_topk, num_experts_per_rank))
|
|
upper_bound_per_token = min(min(num_topk, num_ranks), int((num_ranks - 1) / ratio) + 1)
|
|
if lower_bound_per_token > upper_bound_per_token:
|
|
continue
|
|
|
|
# Repeat for each rank
|
|
for rank_idx in range(num_ranks):
|
|
scores = get_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio, precise)
|
|
_topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
|
|
topk_idx = topk_idx // num_experts_per_rank
|
|
row_indices = torch.arange(num_tokens).unsqueeze(1).expand(num_tokens, num_topk).flatten()
|
|
topk_idx = topk_idx.flatten()
|
|
rank_count = torch.zeros((num_tokens, num_ranks), device='cuda')
|
|
rank_count[row_indices, topk_idx] = 1
|
|
rank_count = rank_count.sum(dim=0)
|
|
total_rank_count += rank_count
|
|
|
|
# Calculate the actual ratio and inequality
|
|
practical_ratio = total_rank_count[0].item() / max(total_rank_count[1:].min().item(), 1)
|
|
inequality = total_rank_count[1:].max().item() / max(total_rank_count[1:].min().item(), 1)
|
|
total_sent_tokens = int(total_rank_count.sum().item())
|
|
if num_tokens > 1000:
|
|
if num_ranks in [8, 64] and num_experts_per_rank == 8:
|
|
print(f' > {precise=}, {num_ranks=:2d}, {num_topk=}, expected_ratio={ratio} | '
|
|
f'ratio={practical_ratio:6.3f}, {inequality=:6.3f}, {total_sent_tokens=:7d}')
|
|
|
|
# Only check the ratio and inequality in precise mode
|
|
if precise:
|
|
assert abs(practical_ratio - ratio) / ratio < 0.1 and inequality < 1.02, \
|
|
f'Failed to generate unbalanced scores with following config: \n' \
|
|
f'{precise=}, {num_tokens=}, {num_experts=:3d}, {num_ranks=:2d}, {num_topk=}, expected_ratio={ratio} | ' \
|
|
f'ratio={practical_ratio:6.3f}, {inequality=:6.3f}, {total_sent_tokens=:7d}'
|
|
print()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_unbalanced_scores()
|