mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +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>
96 lines
3.8 KiB
Python
96 lines
3.8 KiB
Python
"""BucketBuffer / BucketSession checks (subset of DeepEP-Ascend tests/bucket/*; CPU + gloo, or neuron, see dev.py).
|
|
|
|
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
|
|
from dev import DEVICE, init_dist, to_dev, to_host # noqa: E402
|
|
|
|
|
|
def worker(local_rank: int, num_ranks: int, args):
|
|
group = init_dist(local_rank, num_ranks, args.port)
|
|
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_host = torch.arange(16, dtype=torch.float32).view(2, 8) + rank * 100
|
|
src = to_dev(src_host)
|
|
ref = torch.stack([src_host - rank * 100 + r * 100 for r in range(num_ranks)]) # host reference
|
|
|
|
# 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(to_host(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(to_host(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(to_host(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_host - rank * 100 + r * 100 + 1000 for r in range(num_ranks)])
|
|
assert torch.equal(to_host(out), ref2.reshape(-1))
|
|
|
|
# Batched
|
|
outs = buffer.all_gather([shard, shard]).wait()
|
|
assert len(outs) == 2 and all(torch.equal(to_host(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)
|