mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
c142435472
Uneven all_to_all on the neuron process group is unreliable at scale (EP=64: "Invalid NEFF" in prefill; test_ep at 64 ranks hung after 9/18 modes). EPBuffer now moves tokens without all_to_all: - dispatch: one all_gather of x / sf / weights / topk_idx packed as uint8 rows (row count = global max via a 0-dim MAX all_reduce, or the static num_max_tokens_per_rank with DEEPINFER_NEURON_EP_AG_ROWS=cap); each rank selects the rows routed to its local experts in the same (source rank, source token) order as before, so layout / metadata / cached dispatch are unchanged. - combine: per-token fp32 partials over local experts (bf16-rounded as before) + weights, one fp32 SUM reduce_scatter (all_reduce fallback with DEEPINFER_NEURON_EP_REDUCE=ar), chunked under DEEPINFER_NEURON_EP_CHUNK_BYTES (64 MiB). - The a2a path stays behind DEEPINFER_NEURON_EP_ALGO=a2a. test_ep passes on the neuron group at 2/4/8/16/32/64 ranks (combine bit-identical to the reference in practice); new test_collectives.py found that int32 all_gather with full-range values is wrong on a 16-rank neuron group for >= 1 MiB outputs (uint8 / fp32 fine) -- payloads move as bytes. Tiny model EP=8 on NeuronCores: 25/27 argmax vs the reference, same as EP=1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
87 lines
4.4 KiB
Python
87 lines
4.4 KiB
Python
"""The collectives the all-gather EP path relies on (`deep_ep.comm`), checked against host math.
|
|
|
|
all_gather_into_tensor (uint8 / int32 / float32, 1-D), reduce_scatter_tensor and all_reduce (float32, SUM) on the
|
|
group handed to the buffers. DEEP_EP_TEST_DEVICE=neuron runs them on a torch-neuronx ``neuron`` group (see dev.py).
|
|
Run: python test_collectives.py --num-processes {2,4,...,64}
|
|
"""
|
|
|
|
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"))
|
|
from deep_ep import comm # noqa: E402
|
|
from dev import DEVICE, init_dist, to_dev, to_host # noqa: E402
|
|
|
|
|
|
def rank_data(r: int, n: int, dtype: torch.dtype, step: int) -> torch.Tensor:
|
|
g = torch.Generator().manual_seed(1000 * step + r)
|
|
if dtype == torch.uint8:
|
|
return torch.randint(0, 256, (n,), generator=g, dtype=torch.int32).to(torch.uint8)
|
|
if dtype == torch.int32:
|
|
# |v| < 2**24 only: a full-range int32 all-gather is WRONG on a 16-rank neuron group once the output reaches
|
|
# 1 MiB (every chunk corrupted; consistent with an fp32 round trip), while uint8 / fp32 of the same size are
|
|
# exact. The EP path therefore moves everything as uint8 bytes.
|
|
return torch.randint(-2**24, 2**24, (n,), generator=g, dtype=torch.int64).to(torch.int32)
|
|
return torch.randn((n,), generator=g)
|
|
|
|
|
|
def partial_data(r: int, idx: slice, step: int) -> torch.Tensor:
|
|
"""bf16-representable fp32 values, ~70% exact zeros, like the combine partials (closed form in the index)."""
|
|
i = torch.arange(idx.start, idx.stop, dtype=torch.int64)
|
|
v = (((i * 2654435761 + r * 97 + step * 13) % 255) - 127).float() / 64
|
|
return v.masked_fill((i * 31 + r + step) % 10 < 7, 0.0)
|
|
|
|
|
|
def worker(local_rank: int, num_ranks: int, args):
|
|
group = init_dist(local_rank, num_ranks, args.port)
|
|
rank = group.rank()
|
|
sizes = (1, 7, 4096 + 3, 3 * 7168 + 4, 65536 * 4 + 3)
|
|
for step, n in enumerate([n for _ in range(args.repeat) for n in sizes]):
|
|
# all_gather
|
|
for dtype in (torch.uint8, torch.int32, torch.float32):
|
|
local = rank_data(rank, n, dtype, step)
|
|
out = to_dev(torch.empty((num_ranks * n,), dtype=dtype))
|
|
comm.all_gather_into_tensor(out, to_dev(local), group)
|
|
ref = torch.cat([rank_data(r, n, dtype, step) for r in range(num_ranks)])
|
|
got = to_host(out)
|
|
if not torch.equal(got, ref):
|
|
bad = (got != ref).view(num_ranks, n)
|
|
rows = bad.any(1).nonzero().flatten().tolist()
|
|
first = bad.flatten().nonzero().flatten()[:4].tolist()
|
|
raise AssertionError(f'rank {rank}: all_gather {dtype} n={n}: wrong chunks from ranks {rows}, '
|
|
f'{int(bad.sum())} wrong elements, first flat idx {first}')
|
|
# reduce_scatter / all_reduce + slice (fp32 SUM). Inputs are a closed-form function of (rank, index), so
|
|
# the reference only materializes this rank's slice of every rank's input.
|
|
my = slice(rank * n, (rank + 1) * n)
|
|
terms = torch.stack([partial_data(r, my, step) for r in range(num_ranks)]).double()
|
|
my_ref, tol = terms.sum(0), 1e-6 * (terms.abs().sum(0).max().item() + 1)
|
|
for mode in ('rs', 'ar'):
|
|
out = to_dev(torch.empty((n,), dtype=torch.float32))
|
|
comm.reduce_scatter_tensor(out, to_dev(partial_data(rank, slice(0, num_ranks * n), step)), group,
|
|
mode=mode)
|
|
err = (to_host(out).double() - my_ref).abs().max().item()
|
|
assert err <= tol, f'reduce_scatter[{mode}] n={n} err={err} tol={tol}'
|
|
if rank == 0:
|
|
print(f' ok: n={n}', flush=True)
|
|
if rank == 0:
|
|
print(f'collective tests passed (world_size={num_ranks}, device={DEVICE})', flush=True)
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--num-processes', type=int, default=2)
|
|
parser.add_argument('--port', type=int, default=29541)
|
|
parser.add_argument('--repeat', type=int, default=1)
|
|
args = parser.parse_args()
|
|
if args.num_processes == 1:
|
|
worker(0, 1, args)
|
|
else:
|
|
mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
|