mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 15:27:26 +08:00
760e3cadc0
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
tilelang_ops interface tests (basic-mode stubs reported as skips via
nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
process group, device primitive checks, end-to-end vs golden on CPU and
on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
the port depends on or works around (fp8 strided-view indexing, uint64,
1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
104 lines
4.3 KiB
Python
104 lines
4.3 KiB
Python
"""BucketBuffer / BucketSession checks (subset of DeepEP-Ascend tests/bucket/*, CPU + gloo).
|
|
|
|
Covers what `model_executor/parallel_utils/cp.py` uses: a session, `session.allocate`,
|
|
in-place and out-of-place `all_gather`, and the NotImplementedError stubs.
|
|
"""
|
|
|
|
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
|
|
|
|
# DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank)
|
|
# with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo.
|
|
DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
|
|
|
|
|
|
def worker(local_rank: int, num_ranks: int, args):
|
|
os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
|
|
os.environ.setdefault('MASTER_PORT', str(args.port))
|
|
if DEVICE == 'neuron':
|
|
# Must be set before the Neuron runtime initializes (first device op) in this process.
|
|
os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
|
|
dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
|
|
group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD
|
|
rank = group.rank()
|
|
|
|
alignment = deep_ep.get_num_allocation_alignment()
|
|
assert alignment == 2 << 20 and deep_ep.get_num_rdma_alignment() == 512
|
|
|
|
# Plan-based construction (upstream BufferAllocator path).
|
|
plan = deep_ep.BufferAllocator()
|
|
gathered = plan.allocate((num_ranks, 2, 8), torch.float32)
|
|
buffer = deep_ep.BucketBuffer(group, plan, explicitly_destroy=True, device=torch.device(DEVICE))
|
|
assert buffer.num_bytes == plan.num_bytes
|
|
|
|
src = torch.arange(16, dtype=torch.float32, device=DEVICE).view(2, 8) + rank * 100
|
|
ref = torch.stack([src - rank * 100 + r * 100 for r in range(num_ranks)])
|
|
|
|
# Out-of-place: registered destination.
|
|
gathered.fill_(float('nan'))
|
|
result = buffer.all_gather(src, dsts=gathered, group=group).wait()
|
|
assert result.ndim == 1 and torch.equal(result, ref.reshape(-1))
|
|
|
|
# In-place: source is this rank's shard of a registered region.
|
|
gathered.fill_(float('nan'))
|
|
gathered[rank].copy_(src)
|
|
result = buffer.all_gather(gathered[rank], group=group).wait()
|
|
assert torch.equal(result, ref.reshape(-1))
|
|
buffer.barrier(group=group)
|
|
buffer.destroy()
|
|
|
|
# Session path (what CPManager uses).
|
|
buffer = deep_ep.BucketBuffer(group, 4 * alignment, explicitly_destroy=True,
|
|
device=torch.device(DEVICE))
|
|
with buffer.session() as session:
|
|
region = session.allocate((num_ranks, 2, 8), torch.float32, group=group)
|
|
shard = region[rank]
|
|
shard.copy_(src)
|
|
handle = buffer.all_gather(shard)
|
|
out = handle.wait()
|
|
assert torch.equal(out.view(num_ranks, 2, 8).flatten(0, 1), ref.reshape(-1, 8))
|
|
|
|
# Unregistered source is staged into the session storage.
|
|
unregistered = (src + 1000).contiguous()
|
|
out = buffer.all_gather(unregistered).wait()
|
|
ref2 = torch.stack([src - rank * 100 + r * 100 + 1000 for r in range(num_ranks)])
|
|
assert torch.equal(out, ref2.reshape(-1))
|
|
|
|
# Batched
|
|
outs = buffer.all_gather([shard, shard]).wait()
|
|
assert len(outs) == 2 and all(torch.equal(o, ref.reshape(-1)) for o in outs)
|
|
|
|
for name in ('reduce_scatter', 'all_reduce'):
|
|
try:
|
|
getattr(buffer, name)(shard)
|
|
except NotImplementedError:
|
|
pass
|
|
else:
|
|
raise AssertionError(f'{name} should raise NotImplementedError')
|
|
assert 'all_gather' not in buffer.__dict__
|
|
buffer.destroy()
|
|
|
|
if rank == 0:
|
|
print(f'Bucket 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('--port', type=int, default=29551)
|
|
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)
|