mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
bf447622c2
- deep_ep comm: torch-neuronx picks the all_to_all lowering per rank from that rank's own splits (even / NKI a2a-v / 1-D-only padding fallback), so ranks could diverge. All-gather the split matrix first and choose one path for all ranks: even a2a, NKI a2a-v for groups of 4/8, otherwise a padded even a2a (2 ranks, 16 ranks where a2a-v fails on runtime 2.35). - PlatformNeuron: LOCAL_RANK=0 / LOCAL_WORLD_SIZE=1 so every one-core process is neuron:0 of a 1-device local world (otherwise rank > 0 died in _lazy_init at EP/DP > 1). - tests/neuron/deep_ep: references on host copies with the gloo group, buffers on the device; ep/engram/bucket pass at world size 2/4/8/16 on the neuron process group, and 1/2/4 on gloo. - run_deepinfer_tiny.py --ep N: one process per rank / NeuronCore; tiny model EP=2 runs end to end (CPU: within the noise floor). - run_all.sh: per-core busy check, bucket + world size 16 device stages. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
322 lines
16 KiB
Python
322 lines
16 KiB
Python
"""EPBuffer dispatch/combine correctness, ported from DeepEP-Ascend tests/ep/test_ep.py.
|
|
|
|
Default CPU + gloo. DEEP_EP_TEST_DEVICE=neuron: buffers on NeuronCores with a neuron process group; inputs and
|
|
references stay on the host (see dev.py). Run: python test_ep.py --num-processes {1,2,4,8}
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
|
|
import deep_ep # noqa: E402
|
|
from deep_ep.utils.math import align, per_token_cast_to_fp8 # noqa: E402
|
|
from refs import combine as ref_combine # noqa: E402
|
|
from refs import dispatch as ref_dispatch # noqa: E402
|
|
from refs import generate_pre_combine_data # noqa: E402
|
|
from dev import DEVICE, init_dist, to_dev, to_host # noqa: E402
|
|
|
|
# Inputs, references and checks live on the host; only the buffer calls run on DEVICE.
|
|
HOST = 'cpu'
|
|
|
|
|
|
def enumerate_modes():
|
|
for expert_alignment in (128, 1):
|
|
for use_fp8 in (True, False):
|
|
for num_bias in (0, 1, 2):
|
|
for col_major in ((False, True) if use_fp8 else (False,)):
|
|
yield expert_alignment, num_bias, use_fp8, col_major
|
|
|
|
|
|
def launch(buffer, name, defer_epilogue, **params):
|
|
result = getattr(buffer, name)(defer_epilogue=defer_epilogue, **params)
|
|
if defer_epilogue:
|
|
event = result
|
|
return *event.current_stream_wait(release_handle=True), event
|
|
result[-1].current_stream_wait(release_handle=True)
|
|
return result
|
|
|
|
|
|
class HostHandle:
|
|
"""Host copies of the tensors of an `EPHandle` (the device handle itself is what the buffer gets back)."""
|
|
|
|
def __init__(self, handle):
|
|
for name in ('psum_num_recv_tokens_per_rank', 'psum_num_recv_tokens_per_expert',
|
|
'num_unaligned_recv_tokens_per_expert', 'recv_src_metadata'):
|
|
setattr(self, name, to_host(getattr(handle, name)))
|
|
self.num_recv_tokens = handle.num_recv_tokens
|
|
self.num_recv_tokens_per_expert_list = list(handle.num_recv_tokens_per_expert_list)
|
|
self.num_sms = handle.num_sms
|
|
|
|
|
|
def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync):
|
|
expert_alignment, num_bias, use_fp8, col_major = mode
|
|
tag = f'mode={mode} defer={defer_epilogue} cpu_sync={do_cpu_sync}'
|
|
cap = args.num_tokens
|
|
num_tokens = max(1, args.num_tokens - rank)
|
|
hidden, num_topk, num_experts = args.hidden, args.num_topk, args.num_experts
|
|
num_local_experts = num_experts // num_ranks
|
|
|
|
torch.manual_seed(1234 + rank)
|
|
scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device=HOST)
|
|
topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
|
|
topk_idx = topk_idx.to(torch.int64)
|
|
if args.masked_ratio > 0:
|
|
rand_mask = torch.rand_like(topk_idx, dtype=torch.float)
|
|
topk_idx = topk_idx.masked_fill(rand_mask < args.masked_ratio, -1)
|
|
topk_weights = topk_weights.masked_fill(topk_idx < 0, 0)
|
|
topk_idx, topk_weights = topk_idx.contiguous(), topk_weights.contiguous()
|
|
|
|
x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST)
|
|
if use_fp8:
|
|
x_data, sf = per_token_cast_to_fp8(x)
|
|
sf = sf.view(torch.int16)
|
|
num_sf_packs = sf.size(1)
|
|
if col_major:
|
|
sf_storage = torch.empty((num_tokens * num_sf_packs,), dtype=torch.int16, device=HOST)
|
|
input_sf = torch.as_strided(sf_storage, (num_tokens, num_sf_packs), (1, num_tokens))
|
|
input_sf.copy_(sf)
|
|
else:
|
|
input_sf = sf
|
|
x_in = (x_data, input_sf)
|
|
else:
|
|
x_in = x
|
|
bias = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) if num_bias == 1 else None
|
|
if num_bias == 2:
|
|
bias = tuple(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) for _ in range(2))
|
|
|
|
d_stats = to_dev(torch.zeros((num_local_experts,), dtype=torch.int32, device=HOST))
|
|
d_topk_idx = to_dev(topk_idx)
|
|
dispatch_args = dict(
|
|
x=to_dev(x_in), topk_idx=d_topk_idx, topk_weights=to_dev(topk_weights),
|
|
num_sms=7, num_max_tokens_per_rank=cap, num_experts=num_experts,
|
|
expert_alignment=expert_alignment, async_with_compute_stream=True,
|
|
do_handle_copy=False, do_cpu_sync=do_cpu_sync, do_expand=True,
|
|
do_zero_padding=True, use_tma_aligned_col_major_sf=col_major)
|
|
|
|
ref_x = (x_data.view(torch.uint8), input_sf) if use_fp8 else x
|
|
(ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights,
|
|
ref_recv_src_token_idx, ref_num_recv_tokens_per_rank) = ref_dispatch(
|
|
ref_x, topk_idx, topk_weights, cap, num_experts)
|
|
|
|
d_recv_x, recv_topk_idx, d_recv_weights, d_handle, _ = launch(
|
|
buffer, 'dispatch', defer_epilogue, cumulative_local_expert_recv_stats=d_stats, **dispatch_args)
|
|
assert recv_topk_idx is None
|
|
assert d_recv_weights is not None
|
|
assert d_handle.num_sms == 7
|
|
if use_fp8: # strides are checked on the device tensor itself
|
|
sf_stride_on_device = d_recv_x[1].stride()
|
|
recv_x, recv_weights, handle, stats = to_host(d_recv_x), to_host(d_recv_weights), HostHandle(d_handle), to_host(d_stats)
|
|
|
|
recv_x_data, recv_sf = recv_x if use_fp8 else (recv_x, None)
|
|
check_x = recv_x_data.view(torch.uint8) if use_fp8 else recv_x_data
|
|
|
|
num_recv_tokens = handle.psum_num_recv_tokens_per_rank[-1].item()
|
|
num_expanded_tokens = align(handle.psum_num_recv_tokens_per_expert[-1].item(), expert_alignment)
|
|
expected_num_recv = num_recv_tokens if do_cpu_sync else num_ranks * cap
|
|
assert handle.num_recv_tokens == expected_num_recv, (handle.num_recv_tokens, expected_num_recv)
|
|
assert handle.recv_src_metadata.size(0) == expected_num_recv
|
|
if use_fp8:
|
|
expected_stride = (1, recv_x_data.size(0)) if col_major else (num_sf_packs, 1)
|
|
assert recv_sf.stride() == expected_stride, (recv_sf.stride(), expected_stride)
|
|
assert sf_stride_on_device == expected_stride, (sf_stride_on_device, expected_stride)
|
|
|
|
ref_recv_x_data, ref_recv_sf = ref_recv_x if use_fp8 else (ref_recv_x, None)
|
|
assert num_recv_tokens == ref_recv_x_data.size(0), (num_recv_tokens, ref_recv_x_data.size(0))
|
|
|
|
metadata = handle.recv_src_metadata[:num_recv_tokens]
|
|
recv_src_token_idx = metadata[:, -2]
|
|
recv_expanded_slots = metadata[:, :-2]
|
|
ref_order = torch.argsort(ref_recv_src_token_idx)
|
|
metadata_order = torch.argsort(recv_src_token_idx)
|
|
assert torch.equal(recv_src_token_idx[metadata_order], ref_recv_src_token_idx[ref_order])
|
|
|
|
def reorder(src):
|
|
out = torch.empty_like(src)
|
|
out[metadata_order] = src[ref_order]
|
|
return out
|
|
|
|
ref_x_by_meta = reorder(ref_recv_x_data)
|
|
ref_sf_by_meta = reorder(ref_recv_sf) if ref_recv_sf is not None else None
|
|
ref_idx_by_meta = reorder(ref_recv_topk_idx)
|
|
ref_w_by_meta = reorder(ref_recv_topk_weights)
|
|
|
|
recv_valid_mask = recv_expanded_slots >= 0
|
|
assert torch.equal(recv_valid_mask, ref_idx_by_meta >= 0)
|
|
|
|
valid_slots = recv_expanded_slots[recv_valid_mask].long()
|
|
assert valid_slots.numel() == handle.num_unaligned_recv_tokens_per_expert.sum().item()
|
|
if valid_slots.numel() > 0:
|
|
assert valid_slots.min().item() >= 0 and valid_slots.max().item() < num_expanded_tokens
|
|
assert torch.unique(valid_slots).numel() == valid_slots.numel()
|
|
expected_x = ref_x_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]
|
|
assert torch.equal(check_x[valid_slots], expected_x)
|
|
if ref_sf_by_meta is not None:
|
|
expected_sf = ref_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]
|
|
assert torch.equal(recv_sf[valid_slots], expected_sf)
|
|
assert torch.equal(recv_weights[valid_slots], ref_w_by_meta[recv_valid_mask])
|
|
|
|
# Master top-k slot column
|
|
topk_range = torch.arange(num_topk, device=HOST, dtype=metadata.dtype)
|
|
master = torch.where(recv_valid_mask, topk_range.view(1, num_topk),
|
|
torch.full_like(recv_expanded_slots, -1)).max(dim=1).values
|
|
assert torch.equal(metadata[:, -1], master)
|
|
|
|
# Per-rank and per-expert counts
|
|
psum_rank = [0] + handle.psum_num_recv_tokens_per_rank.tolist()
|
|
for r in range(num_ranks):
|
|
assert psum_rank[r + 1] - psum_rank[r] == ref_num_recv_tokens_per_rank[r].item()
|
|
psum_expert = [0] + handle.psum_num_recv_tokens_per_expert.tolist()
|
|
for e in range(num_local_experts):
|
|
count = psum_expert[e + 1] - align(psum_expert[e], expert_alignment)
|
|
ref_count = (ref_recv_topk_idx == e).sum().item()
|
|
assert count == ref_count, (e, count, ref_count)
|
|
assert stats[e].item() == ref_count
|
|
assert handle.num_unaligned_recv_tokens_per_expert[e].item() == ref_count
|
|
if do_cpu_sync:
|
|
assert handle.num_recv_tokens_per_expert_list[e] == align(ref_count, expert_alignment)
|
|
expert_start = align(psum_expert[e], expert_alignment)
|
|
expert_end = psum_expert[e + 1]
|
|
expert_slots = recv_expanded_slots[ref_idx_by_meta == e]
|
|
assert expert_slots.numel() == ref_count
|
|
if ref_count > 0:
|
|
assert torch.equal(torch.sort(expert_slots).values,
|
|
torch.arange(expert_start, expert_end, device=HOST,
|
|
dtype=expert_slots.dtype))
|
|
padding_end = align(expert_end, expert_alignment)
|
|
assert torch.count_nonzero(recv_x_data[expert_end:padding_end].view(torch.uint8)).item() == 0
|
|
if recv_sf is not None:
|
|
assert torch.count_nonzero(recv_sf[expert_end:padding_end]).item() == 0
|
|
|
|
# --- Combine ---
|
|
local_y = generate_pre_combine_data(metadata[:, -2], cap, num_topk, hidden)
|
|
input_for_combine = torch.empty((recv_x_data.size(0) + 1, hidden), dtype=torch.bfloat16, device=HOST)
|
|
input_for_combine[metadata[:, :-2].long().flatten()] = local_y.view(-1, hidden)
|
|
input_for_combine = input_for_combine[:-1].contiguous()
|
|
|
|
d_input_for_combine = to_dev(input_for_combine)
|
|
combined_x, combined_weights, _ = launch(
|
|
buffer, 'combine', defer_epilogue, x=d_input_for_combine, handle=d_handle,
|
|
topk_weights=d_recv_weights, bias=to_dev(bias), num_sms=7, async_with_compute_stream=True)
|
|
combined_x, combined_weights = to_host(combined_x), to_host(combined_weights)
|
|
|
|
ref_y = generate_pre_combine_data(rank * cap + torch.arange(num_tokens, device=HOST), cap, num_topk, hidden)
|
|
ref_y[topk_idx == -1] = 0
|
|
ref_combined_x = ref_combine(ref_y, topk_idx, num_experts, bias, reduce_in_local=True)
|
|
max_diff = (combined_x.float() - ref_combined_x.float()).abs().max().item()
|
|
assert torch.equal(combined_x, ref_combined_x), f'{tag} {max_diff=}'
|
|
assert torch.equal(combined_weights, topk_weights)
|
|
|
|
if num_bias == 0 and not col_major:
|
|
combined_x2, combined_weights2, _ = launch(
|
|
buffer, 'combine', defer_epilogue, x=d_input_for_combine, handle=d_handle,
|
|
bias=to_dev(bias), num_sms=7, async_with_compute_stream=True)
|
|
combined_x2 = to_host(combined_x2)
|
|
assert torch.equal(combined_x2, ref_combined_x)
|
|
assert combined_weights2 is None
|
|
|
|
# --- Cached dispatch replay ---
|
|
cached_args = dict(dispatch_args, num_sms=0, do_cpu_sync=False, handle=d_handle)
|
|
del cached_args['topk_idx']
|
|
cached_names = ('psum_num_recv_tokens_per_rank', 'psum_num_recv_tokens_per_expert',
|
|
'num_unaligned_recv_tokens_per_expert', 'recv_src_metadata')
|
|
cached_copies = [getattr(handle, name).clone() for name in cached_names]
|
|
padding_mask = torch.ones(num_expanded_tokens, dtype=torch.bool, device=HOST)
|
|
padding_mask[valid_slots] = False
|
|
for send_weights in (True, False):
|
|
updated_x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) * (8 if send_weights else 0.125)
|
|
if use_fp8:
|
|
updated_data, updated_sf = per_token_cast_to_fp8(updated_x)
|
|
updated_sf = updated_sf.view(torch.int16)
|
|
if col_major:
|
|
updated_sf = updated_sf.t().contiguous().t()
|
|
updated_in = (updated_data, updated_sf)
|
|
else:
|
|
updated_in = updated_x
|
|
updated_weights = torch.randn_like(topk_weights).masked_fill(topk_idx < 0, 0)
|
|
ref_updated = (updated_data.view(torch.uint8), updated_sf) if use_fp8 else updated_x
|
|
ref_cached_x, _, _, _, _ = ref_dispatch(ref_updated, topk_idx, updated_weights, cap, num_experts)
|
|
d_cached_x, cached_topk_idx, cached_w, cached_handle, _ = launch(
|
|
buffer, 'dispatch', defer_epilogue,
|
|
**dict(cached_args, x=to_dev(updated_in),
|
|
topk_weights=to_dev(updated_weights) if send_weights else None))
|
|
assert cached_handle is d_handle and cached_topk_idx is None
|
|
for name, saved in zip(cached_names, cached_copies):
|
|
assert torch.equal(to_host(getattr(d_handle, name)), saved), name
|
|
if use_fp8:
|
|
cached_sf_stride_on_device = d_cached_x[1].stride()
|
|
cached_x, cached_w = to_host(d_cached_x), to_host(cached_w)
|
|
cached_data, cached_sf = cached_x if use_fp8 else (cached_x, None)
|
|
ref_cached_data, ref_cached_sf = ref_cached_x if use_fp8 else (ref_cached_x, None)
|
|
cached_check = cached_data.view(torch.uint8) if use_fp8 else cached_data
|
|
ref_cached_by_meta = reorder(ref_cached_data)
|
|
assert torch.equal(cached_check[valid_slots],
|
|
ref_cached_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask])
|
|
assert torch.count_nonzero(cached_check[:num_expanded_tokens][padding_mask]).item() == 0
|
|
if cached_sf is not None:
|
|
assert cached_sf.stride() == expected_stride
|
|
assert cached_sf_stride_on_device == expected_stride
|
|
ref_cached_sf_by_meta = reorder(ref_cached_sf)
|
|
assert torch.equal(cached_sf[valid_slots],
|
|
ref_cached_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask])
|
|
if send_weights:
|
|
assert torch.equal(cached_w[valid_slots], reorder(updated_weights_ref(
|
|
updated_weights, topk_idx, cap, num_experts))[recv_valid_mask])
|
|
else:
|
|
assert cached_w is None
|
|
|
|
|
|
def updated_weights_ref(updated_weights, topk_idx, cap, num_experts):
|
|
_, _, ref_w, _, _ = ref_dispatch(
|
|
torch.zeros((topk_idx.size(0), 8), dtype=torch.bfloat16, device=HOST),
|
|
topk_idx, updated_weights, cap, num_experts)
|
|
return ref_w
|
|
|
|
|
|
def worker(local_rank: int, num_ranks: int, args: argparse.Namespace):
|
|
group = init_dist(local_rank, num_ranks, args.port)
|
|
rank = group.rank()
|
|
assert rank == dist.get_rank() # host references use the default (gloo) group
|
|
|
|
buffer = deep_ep.EPBuffer(
|
|
group, num_max_tokens_per_rank=args.num_tokens, hidden=args.hidden,
|
|
num_topk=args.num_topk, use_fp8_dispatch=True, allow_hybrid_mode=True,
|
|
allow_multiple_reduction=True, explicitly_destroy=True)
|
|
assert buffer.num_bytes > 0
|
|
assert deep_ep.get_physical_domain_size(group) == (1, num_ranks)
|
|
|
|
for mode in enumerate_modes():
|
|
for defer_epilogue in (True, False):
|
|
for do_cpu_sync in (True, False):
|
|
check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync)
|
|
if rank == 0:
|
|
print(f' ok: expert_alignment={mode[0]} num_bias={mode[1]} fp8={mode[2]} col_major={mode[3]}',
|
|
flush=True)
|
|
|
|
buffer.destroy()
|
|
if rank == 0:
|
|
print(f'EP tests passed (world_size={num_ranks})', flush=True)
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--num-processes', type=int, default=2)
|
|
parser.add_argument('--num-tokens', type=int, default=17)
|
|
parser.add_argument('--hidden', type=int, default=256)
|
|
parser.add_argument('--num-topk', type=int, default=4)
|
|
parser.add_argument('--num-experts', type=int, default=16)
|
|
parser.add_argument('--masked-ratio', type=float, default=0.2)
|
|
parser.add_argument('--port', type=int, default=29531)
|
|
args = parser.parse_args()
|
|
assert args.num_experts % args.num_processes == 0
|
|
if args.num_processes == 1:
|
|
worker(0, 1, args)
|
|
else:
|
|
mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
|