From bf447622c29c61fe24b230a0b53efec6ca9712dc Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 21 Oct 2026 11:39:58 +0800 Subject: [PATCH] Neuron: EP > 1 with real neuron collectives (issue #6) - 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 --- deepinfer/neuron_ops/vendor/deep_ep/comm.py | 100 +++++++++++++++++- deepinfer/platform.py | 7 ++ examples/neuron/run_deepinfer_tiny.py | 61 ++++++++++- tests/neuron/deep_ep/dev.py | 67 +++++++++++++ tests/neuron/deep_ep/test_bucket.py | 32 +++--- tests/neuron/deep_ep/test_engram.py | 41 ++++---- tests/neuron/deep_ep/test_ep.py | 106 ++++++++++++-------- tests/neuron/run_all.sh | 27 +++-- 8 files changed, 342 insertions(+), 99 deletions(-) create mode 100644 tests/neuron/deep_ep/dev.py diff --git a/deepinfer/neuron_ops/vendor/deep_ep/comm.py b/deepinfer/neuron_ops/vendor/deep_ep/comm.py index 1158275..824fff2 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/comm.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/comm.py @@ -12,6 +12,7 @@ All ranks are reported as one "direct" (NVLink-equivalent) domain, which is what from __future__ import annotations +import os from typing import Any, List, Optional, Sequence, Tuple import torch @@ -81,6 +82,57 @@ def barrier(buffer: Any, wait_comm_stream: bool = True, # --- Collective helpers (world-size-1 groups never touch `torch.distributed`) --- +def _nki_a2av_group_size(n: int) -> bool: + """Group sizes routed to the torch-neuronx NKI all_to_all_v kernel. + + torch-neuronx accepts 4/8/16/32k (whole trn2 devices), but on trn2.48xlarge with runtime 2.35 only 4 and 8 work: + a 16-rank group (4 devices) fails at load with ``ENC:__select_mesh_subtype not a supported operation for mesh + (op_type 12)`` / "NRT model scheduling failed". Everything else takes the padded even all-to-all. + """ + return n in (4, 8) + + +def _use_nki_a2av() -> bool: + # Same rule as torch-neuronx `UseNkiAllToAllV()`: enabled when unset or "1". + return os.environ.get('NEURON_USE_NKI_ALLTOALLV', '1') == '1' + + +def _split_matrix(input_split_sizes: Sequence[int], device: torch.device, + group: dist.ProcessGroup) -> List[List[int]]: + """`m[src][dst]` = rows rank `src` sends to rank `dst`, for the whole group. + + NEURON: host sync — one all-gather of the per-rank split sizes, so that every rank takes + the *same* collective path (see `all_to_all_single`). + """ + num_ranks = group.size() + local = torch.tensor(list(input_split_sizes), dtype=torch.int32).to(device) + gathered = torch.empty((num_ranks * num_ranks,), dtype=torch.int32, device=device) + dist.all_gather_into_tensor(gathered, local, group=group) + flat = gathered.cpu().tolist() + return [flat[r * num_ranks:(r + 1) * num_ranks] for r in range(num_ranks)] + + +def _padded_all_to_all(output: torch.Tensor, input_: torch.Tensor, + output_split_sizes: Sequence[int], input_split_sizes: Sequence[int], + max_rows: int, group: dist.ProcessGroup) -> None: + """Uneven all-to-all as one *even* all-to-all over chunks padded to the global max.""" + num_ranks = group.size() + trailing = tuple(input_.shape[1:]) + send = torch.zeros((num_ranks * max_rows,) + trailing, dtype=input_.dtype, device=input_.device) + offset = 0 + for dst, rows in enumerate(input_split_sizes): + if rows: + send.narrow(0, dst * max_rows, rows).copy_(input_.narrow(0, offset, rows)) + offset += rows + recv = torch.empty_like(send) + dist.all_to_all_single(recv, send, group=group) + offset = 0 + for src, rows in enumerate(output_split_sizes): + if rows: + output.narrow(0, offset, rows).copy_(recv.narrow(0, src * max_rows, rows)) + offset += rows + + def all_to_all_single(output: torch.Tensor, input_: torch.Tensor, output_split_sizes: Optional[Sequence[int]], input_split_sizes: Optional[Sequence[int]], @@ -89,14 +141,54 @@ def all_to_all_single(output: torch.Tensor, input_: torch.Tensor, Split sizes are host values; computing them from routing data forces a device->host sync in the caller (see `EPBuffer.dispatch`). + + On the torch-neuronx backend, uneven splits need care: `NeuronBackend::alltoall_base` + picks its path from the *local* split sizes (locally even -> even ALLTOALL, otherwise NKI + all_to_all_v for group sizes 4/8/16/32k, otherwise a 1-D-only padding fallback that is + wrong unless every rank has the same largest chunk). A rank whose splits happen to be + even while a peer's are not would enter a different collective. So the global split + matrix is gathered first and the path is chosen identically on every rank: + * globally uniform splits -> even all-to-all; + * 4/8-rank group and no rank locally even -> NKI all_to_all_v (via torch-neuronx); + * otherwise (2-rank groups, 16+, NEURON_USE_NKI_ALLTOALLV=0) -> even all-to-all over chunks padded + to the global max. """ if group.size() == 1: output.view(-1).copy_(input_.reshape(-1)[:output.numel()]) return output - dist.all_to_all_single(output, input_.contiguous(), - None if output_split_sizes is None else list(output_split_sizes), - None if input_split_sizes is None else list(input_split_sizes), - group=group) + input_ = input_.contiguous() + if output.device.type == 'cpu' or (output_split_sizes is None and input_split_sizes is None): + dist.all_to_all_single(output, input_, + None if output_split_sizes is None else list(output_split_sizes), + None if input_split_sizes is None else list(input_split_sizes), + group=group) + return output + + num_ranks, rank = group.size(), group.rank() + assert output_split_sizes is not None and input_split_sizes is not None + matrix = _split_matrix(input_split_sizes, output.device, group) + assert list(matrix[rank]) == list(input_split_sizes) + assert [matrix[src][rank] for src in range(num_ranks)] == list(output_split_sizes), \ + 'all_to_all_single: output split sizes disagree with the peers\' input split sizes' + values = {v for row in matrix for v in row} + max_rows = max(values) + if max_rows == 0: + return output # globally empty exchange (every rank sees the same matrix) + if len(values) == 1: + dist.all_to_all_single(output, input_, group=group) + return output + + def locally_even(r: int) -> bool: + sends = matrix[r] + recvs = [matrix[src][r] for src in range(num_ranks)] + return len(set(sends)) == 1 and len(set(recvs)) == 1 and (sends[0] + recvs[0]) > 0 + + if _use_nki_a2av() and _nki_a2av_group_size(num_ranks) and \ + not any(locally_even(r) for r in range(num_ranks)): + dist.all_to_all_single(output, input_, list(output_split_sizes), list(input_split_sizes), + group=group) + else: + _padded_all_to_all(output, input_, output_split_sizes, input_split_sizes, max_rows, group) return output diff --git a/deepinfer/platform.py b/deepinfer/platform.py index 42efbf3..15c923c 100644 --- a/deepinfer/platform.py +++ b/deepinfer/platform.py @@ -250,6 +250,13 @@ class PlatformNeuron(Platform): os.environ["CUDA_VISIBLE_DEVICES"] = cores local_rank = int(os.getenv("DEEPINFER_NEURON_LOCAL_RANK", "0")) os.environ["NEURON_RT_VISIBLE_CORES"] = cores.split(",")[local_rank] + # torch-neuronx numbers a process's torch device by LOCAL_RANK out of LOCAL_WORLD_SIZE when the first + # "neuron" process group is created (distributed/backend.py, else it derives them from the hosts of all + # ranks). With one core per process every rank must be device 0 of a 1-device world: `cuda:N` maps to + # neuron:0 here, and the per-device RNG table is sized before init_dist (seeding in init_model), so + # otherwise rank r > 0 dies in _lazy_init ("idx < num_devices") at EP/DP > 1. + os.environ["LOCAL_RANK"] = "0" + os.environ["LOCAL_WORLD_SIZE"] = "1" else: os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0") self.distributed_backend = "gloo" diff --git a/examples/neuron/run_deepinfer_tiny.py b/examples/neuron/run_deepinfer_tiny.py index 1e3ce86..5b1dad8 100644 --- a/examples/neuron/run_deepinfer_tiny.py +++ b/examples/neuron/run_deepinfer_tiny.py @@ -46,6 +46,11 @@ def parse_args(): ap.add_argument("--hidden", action="store_true", help="also record per-layer block outputs (hc residual stream)") ap.add_argument("--compare", action="store_true", help="run compare_golden on every prompt afterwards") ap.add_argument("--attn-p-bf16", action="store_true", help="experiment: patch the flash_mla shim with the proposed bf16-P fix") + ap.add_argument("--ep", type=int, default=1, + help="expert parallel size (= data parallel size): one engine process per rank, every rank runs the " + "same teacher-forced prompts. Rank r uses entry r of CUDA_VISIBLE_DEVICES (default 0..ep-1), " + "i.e. one NeuronCore per process on --device neuron. Rank 0 writes /.pt, rank r>0 " + "/rank/.pt") return ap.parse_args() @@ -194,6 +199,54 @@ def main(): args = parse_args() for k in ("hf_ckpt", "golden_dir", "out"): setattr(args, k, os.path.abspath(os.path.expanduser(getattr(args, k)))) + if args.ep > 1: + launch_ep(args) + return + run(args) + + +def launch_ep(args) -> None: + """EP=DP=args.ep: one spawned engine process per rank (what examples/_launch.py does for the V4 scripts).""" + import multiprocessing as mp + + cores = os.environ.get("CUDA_VISIBLE_DEVICES") or ",".join(str(i) for i in range(args.ep)) + cores = cores.split(",") + assert len(cores) >= args.ep, f"--ep {args.ep} needs {args.ep} entries in CUDA_VISIBLE_DEVICES, got {cores}" + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "23333") + ctx = mp.get_context("spawn") + procs = [ctx.Process(target=_ep_rank_main, args=(args, rank, cores[rank]), daemon=False) for rank in range(args.ep)] + for p in procs: + p.start() + # Fail fast: a dead rank leaves its peers blocked in a collective forever. + while any(p.is_alive() for p in procs): + if any(p.exitcode not in (None, 0) for p in procs): + for p in procs: + if p.is_alive(): + p.terminate() + break + time.sleep(1) + for p in procs: + p.join() + failed = [rank for rank, p in enumerate(procs) if p.exitcode != 0] + if failed: + raise RuntimeError(f"EP ranks failed: {failed}") + + +def _ep_rank_main(args, rank: int, core: str) -> None: + # Before deepinfer / the Neuron runtime initialize in this process: one NeuronCore per rank. + os.environ["CUDA_VISIBLE_DEVICES"] = core + os.environ.pop("NEURON_RT_VISIBLE_CORES", None) # PlatformNeuron derives it from CUDA_VISIBLE_DEVICES + if rank > 0: + args.out = os.path.join(args.out, f"rank{rank}") + try: + run(args, rank) + except Exception: + traceback.print_exc() + sys.exit(1) + + +def run(args, rank: int = 0): setup_env(args.device) os.makedirs(args.out, exist_ok=True) @@ -218,10 +271,10 @@ def main(): llm = LLM( model=args.hf_ckpt, trust_remote_code=True, - expert_parallel_size=1, + expert_parallel_size=args.ep, context_parallel_size=1, - data_parallel_size=1, - engine_rank=0, + data_parallel_size=args.ep, + engine_rank=rank, chunked_prefill=True, decode_sequence=False, max_model_len=args.max_model_len, @@ -229,7 +282,7 @@ def main(): max_num_seqs=1, gpu_memory_utilization=0.5, ) - print(f"[run] engine up in {time.time() - t0:.1f}s", flush=True) + print(f"[run] engine up in {time.time() - t0:.1f}s (ep={args.ep}, rank={rank})", flush=True) status = install_hidden_hooks(rec, llm) if args.hidden else "hidden not requested" print(f"[run] {status}", flush=True) diff --git a/tests/neuron/deep_ep/dev.py b/tests/neuron/deep_ep/dev.py new file mode 100644 index 0000000..eb0d69a --- /dev/null +++ b/tests/neuron/deep_ep/dev.py @@ -0,0 +1,67 @@ +"""Device plumbing shared by the deep_ep tests. + +`DEEP_EP_TEST_DEVICE=neuron` runs every rank on its own NeuronCore (core `DEEP_EP_TEST_CORE_BASE + rank`) and +hands the buffers a torch-neuronx ``neuron`` process group, exercising the real collectives. Inputs are generated +and references are computed on the **host** (gloo default group, CPU copies), then compared against the device +results copied back, so a device bug cannot cancel out between the code under test and its reference. +Default (`cpu`): everything on the host with gloo, as before. + +FP8 tensors cross the host/device boundary as `uint8` views (indexing / cat on float8 is wrong on the neuron +device, see tests/neuron/device/test_device_primitives.py); strides are preserved (column-major SF inputs). +""" + +import os + +import torch +import torch.distributed as dist + +DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu') +_FP8 = (torch.float8_e4m3fn, torch.float8_e5m2) + + +def init_dist(local_rank: int, num_ranks: int, port: int): + """Default gloo PG (host references) + the group handed to the buffers. + + Mirrors DeepInfer's own layout (`worker/base.py`: gloo default PG, `PLATFORM.distributed_backend` subgroups). + torch-neuronx sets up the runtime when the first ``neuron`` group is created and asserts that no device op ran + before, so NEURON_RT_VISIBLE_CORES is pinned here and nothing touches the device until the group exists. + """ + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ.setdefault('MASTER_PORT', str(port)) + if DEVICE == 'neuron': + os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank) + # Same process model as DeepInfer (deepinfer/platform.py PlatformNeuron): one core per process, which + # torch-neuronx then exposes as neuron:0 of a 1-device local world on every rank. + os.environ['LOCAL_RANK'] = '0' + os.environ['LOCAL_WORLD_SIZE'] = '1' + dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks) + if DEVICE == 'neuron': + import torch_neuronx # noqa: F401 (registers the "neuron" backend) + group = dist.new_group(list(range(num_ranks)), backend='neuron') + else: + group = dist.group.WORLD + return group + + +def to_dev(t): + """Host tensor -> DEVICE, same dtype/shape/strides (identity on cpu).""" + if t is None or DEVICE == 'cpu': + return t + if isinstance(t, tuple): + return tuple(to_dev(x) for x in t) + is_fp8 = t.dtype in _FP8 + src = t.view(torch.uint8) if is_fp8 else t + out = torch.empty_strided(src.size(), src.stride(), dtype=src.dtype, device=DEVICE) + out.copy_(src) + return out.view(t.dtype) if is_fp8 else out + + +def to_host(t): + """DEVICE tensor -> host (identity on cpu); FP8 through a uint8 view.""" + if t is None or DEVICE == 'cpu': + return t + if isinstance(t, tuple): + return tuple(to_host(x) for x in t) + if t.dtype in _FP8: + return t.view(torch.uint8).cpu().view(t.dtype) + return t.cpu() diff --git a/tests/neuron/deep_ep/test_bucket.py b/tests/neuron/deep_ep/test_bucket.py index 09d0c82..3a8845f 100644 --- a/tests/neuron/deep_ep/test_bucket.py +++ b/tests/neuron/deep_ep/test_bucket.py @@ -1,4 +1,4 @@ -"""BucketBuffer / BucketSession checks (subset of DeepEP-Ascend tests/bucket/*, CPU + gloo). +"""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. @@ -15,20 +15,11 @@ 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') +from dev import DEVICE, init_dist, to_dev, to_host # noqa: E402 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 + group = init_dist(local_rank, num_ranks, args.port) rank = group.rank() alignment = deep_ep.get_num_allocation_alignment() @@ -40,19 +31,20 @@ def worker(local_rank: int, num_ranks: int, args): 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)]) + 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(result, ref.reshape(-1)) + 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(result, ref.reshape(-1)) + assert torch.equal(to_host(result), ref.reshape(-1)) buffer.barrier(group=group) buffer.destroy() @@ -65,17 +57,17 @@ def worker(local_rank: int, num_ranks: int, args): 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)) + 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 - rank * 100 + r * 100 + 1000 for r in range(num_ranks)]) - assert torch.equal(out, ref2.reshape(-1)) + 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(o, ref.reshape(-1)) for o in outs) + 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: diff --git a/tests/neuron/deep_ep/test_engram.py b/tests/neuron/deep_ep/test_engram.py index 53574a8..779e3e0 100644 --- a/tests/neuron/deep_ep/test_engram.py +++ b/tests/neuron/deep_ep/test_engram.py @@ -1,4 +1,8 @@ -"""EngramBuffer correctness, ported from DeepEP-Ascend tests/engram/test_engram.py (CPU + gloo).""" +"""EngramBuffer correctness, ported from DeepEP-Ascend tests/engram/test_engram.py. + +Default CPU + gloo; DEEP_EP_TEST_DEVICE=neuron runs the buffer on NeuronCores with a neuron process group while the +tables, indices and reference lookups stay on the host (see dev.py). +""" import argparse import os @@ -12,10 +16,9 @@ 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 ceil_div, per_token_cast_to_fp8 # noqa: E402 +from dev import init_dist, to_dev, to_host # 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') +HOST = 'cpu' # inputs + references; the buffer runs on dev.DEVICE def run(rank: int, num_ranks: int, group, args, use_fp8: bool): @@ -41,39 +44,40 @@ def run(rank: int, num_ranks: int, group, args, use_fp8: bool): torch.manual_seed(7 + rank) local_storages, sfs = [], ([] if use_fp8 else None) for num_entries in num_entries_per_layer: - local_bf16 = torch.randn((num_entries, args.hidden), dtype=torch.bfloat16, device=DEVICE) + local_bf16 = torch.randn((num_entries, args.hidden), dtype=torch.bfloat16, device=HOST) if use_fp8: local_storage, local_sf = per_token_cast_to_fp8(local_bf16) local_sf = local_sf.view(torch.int16).contiguous() - sf = torch.empty((num_ranks * num_entries, num_sf_packs), dtype=local_sf.dtype, device=DEVICE) + sf = torch.empty((num_ranks * num_entries, num_sf_packs), dtype=local_sf.dtype, device=HOST) if num_ranks == 1: sf.copy_(local_sf) else: # gloo has no int16 collective; gather the bytes instead (DeepInfer's # NgramEmbedding.gather_sf gathers int16 directly on the accelerator). - dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.view(torch.uint8), group=group) + # Host reference on the default (gloo) group. + dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.view(torch.uint8)) sfs.append(sf) else: local_storage = local_bf16 local_storages.append(local_storage.contiguous()) - buffer.write(local_storages, sfs=sfs) + buffer.write([to_dev(t) for t in local_storages], sfs=None if sfs is None else [to_dev(t) for t in sfs]) # Reference global tables global_bits = [] for local_storage in local_storages: bits = torch.empty((num_ranks * local_storage.shape[0], args.hidden * dtype.itemsize), - dtype=torch.uint8, device=DEVICE) + dtype=torch.uint8, device=HOST) local_bits = local_storage.view(torch.uint8).view(local_storage.shape[0], -1).contiguous() if num_ranks == 1: bits.copy_(local_bits) else: - dist.all_gather_into_tensor(bits, local_bits, group=group) + dist.all_gather_into_tensor(bits, local_bits) # host, default (gloo) group global_bits.append(bits) def generate_indices(): return torch.stack([ torch.randint(0, num_ranks * num_entries, (args.num_tokens, args.num_entries_per_token), - dtype=torch.int, device=DEVICE) + dtype=torch.int, device=HOST) for num_entries in num_entries_per_layer]) for col_major in ((False, True) if use_fp8 else (False,)): @@ -84,29 +88,28 @@ def run(rank: int, num_ranks: int, group, args, use_fp8: bool): ref_sfs = [sfs[l][indices[l].reshape(-1).long()].reshape(args.num_tokens, -1) for l in range(num_layers)] if use_fp8 else None - hooks = buffer.fetch(indices, use_tma_aligned_col_major_sf=col_major) + hooks = buffer.fetch(to_dev(indices), use_tma_aligned_col_major_sf=col_major) assert len(hooks) == num_layers for l in layer_order: fetched = hooks[l]() data, fetched_sf = fetched if use_fp8 else (fetched, None) + if fetched_sf is not None: + sf_stride_on_device = fetched_sf.stride() + data, fetched_sf = to_host(data), to_host(fetched_sf) assert data.shape == (args.num_tokens, args.num_entries_per_token * args.hidden) assert torch.equal(ref_data[l], data.reshape(args.num_tokens, -1).contiguous().view(torch.uint8)), \ f'data mismatch (layer {l}, fp8={use_fp8}, col_major={col_major})' if use_fp8: expected_stride = (1, args.num_tokens) if col_major else (fetched_sf.size(1), 1) assert fetched_sf.stride() == expected_stride, (fetched_sf.stride(), expected_stride) + assert sf_stride_on_device == expected_stride, (sf_stride_on_device, expected_stride) assert torch.equal(ref_sfs[l], fetched_sf), f'sf mismatch (layer {l})' buffer.destroy() 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 + group = init_dist(local_rank, num_ranks, args.port) + assert group.rank() == dist.get_rank() for use_fp8 in (False, True): run(group.rank(), num_ranks, group, args, use_fp8) if group.rank() == 0: diff --git a/tests/neuron/deep_ep/test_ep.py b/tests/neuron/deep_ep/test_ep.py index b72bdc3..60aaa7a 100644 --- a/tests/neuron/deep_ep/test_ep.py +++ b/tests/neuron/deep_ep/test_ep.py @@ -1,6 +1,7 @@ """EPBuffer dispatch/combine correctness, ported from DeepEP-Ascend tests/ep/test_ep.py. -CPU + gloo, no Neuron device. Run: python test_ep.py --num-processes {1,2,4} +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 @@ -18,10 +19,10 @@ 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 -# 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') +# Inputs, references and checks live on the host; only the buffer calls run on DEVICE. +HOST = 'cpu' def enumerate_modes(): @@ -41,6 +42,18 @@ def launch(buffer, name, defer_epilogue, **params): 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}' @@ -50,7 +63,7 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): num_local_experts = num_experts // num_ranks torch.manual_seed(1234 + rank) - scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device=DEVICE) + 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: @@ -59,13 +72,13 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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=DEVICE) + 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=DEVICE) + 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: @@ -73,13 +86,14 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): x_in = (x_data, input_sf) else: x_in = x - bias = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) if num_bias == 1 else None + 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=DEVICE) for _ in range(2)) + bias = tuple(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) for _ in range(2)) - stats = torch.zeros((num_local_experts,), dtype=torch.int32, device=DEVICE) + 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=x_in, topk_idx=topk_idx, topk_weights=topk_weights, + 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, @@ -90,11 +104,14 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): ref_recv_src_token_idx, ref_num_recv_tokens_per_rank) = ref_dispatch( ref_x, topk_idx, topk_weights, cap, num_experts) - recv_x, recv_topk_idx, recv_weights, handle, _ = launch( - buffer, 'dispatch', defer_epilogue, cumulative_local_expert_recv_stats=stats, **dispatch_args) + 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 recv_weights is not None - assert handle.num_sms == 7 + 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 @@ -107,6 +124,7 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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)) @@ -144,7 +162,7 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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=DEVICE, dtype=metadata.dtype) + 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) @@ -168,7 +186,7 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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=DEVICE, + 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 @@ -177,15 +195,17 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): # --- 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=DEVICE) + 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=input_for_combine, handle=handle, - topk_weights=recv_weights, bias=bias, num_sms=7, async_with_compute_stream=True) + 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=DEVICE), cap, num_topk, hidden) + 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() @@ -194,21 +214,22 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): if num_bias == 0 and not col_major: combined_x2, combined_weights2, _ = launch( - buffer, 'combine', defer_epilogue, x=input_for_combine, handle=handle, - bias=bias, num_sms=7, async_with_compute_stream=True) + 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=handle) + cached_args = dict(dispatch_args, num_sms=0, do_cpu_sync=False, handle=d_handle) del cached_args['topk_idx'] - cached_tensors = (handle.psum_num_recv_tokens_per_rank, handle.psum_num_recv_tokens_per_expert, - handle.num_unaligned_recv_tokens_per_expert, handle.recv_src_metadata) - cached_copies = [t.clone() for t in cached_tensors] - padding_mask = torch.ones(num_expanded_tokens, dtype=torch.bool, device=DEVICE) + 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=DEVICE) * (8 if send_weights else 0.125) + 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) @@ -220,13 +241,16 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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) - cached_x, cached_topk_idx, cached_w, cached_handle, _ = launch( + d_cached_x, cached_topk_idx, cached_w, cached_handle, _ = launch( buffer, 'dispatch', defer_epilogue, - **dict(cached_args, x=updated_in, - topk_weights=updated_weights if send_weights else None)) - assert cached_handle is handle and cached_topk_idx is None - for tensor, saved in zip(cached_tensors, cached_copies): - assert torch.equal(tensor, saved) + **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 @@ -236,6 +260,7 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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]) @@ -248,20 +273,15 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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=DEVICE), + 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): - 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 + 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, diff --git a/tests/neuron/run_all.sh b/tests/neuron/run_all.sh index 7fc2a59..77b917e 100755 --- a/tests/neuron/run_all.sh +++ b/tests/neuron/run_all.sh @@ -49,10 +49,19 @@ run() { # run [ "$r" = FAIL ] && tail -5 "$LOGDIR/$stage.log" | sed 's/^/ | /' } -cores_free() { - local busy - busy=$(neuron-ls 2>/dev/null | grep -cE "python|vLLM|Worker|torchrun") - [ "$busy" = "0" ] || { echo "NeuronCores busy ($busy processes); not starting device stages"; return 1; } +cores_free() { # are the NeuronCores in $CORES free? (other cores may be busy) + python3 - "$CORES" <<'EOF' +import json, subprocess, sys +want = set() +for part in sys.argv[1].split(","): + a, _, b = part.partition("-") + want.update(range(int(a), int(b) + 1) if b else [int(a)]) +busy = [d["neuron_device"] for d in json.loads(subprocess.check_output(["neuron-ls", "-j"])) + if d.get("neuron_processes") and want & set(d["neuroncore_ids"])] +if busy: + print(f"NeuronCores in {sys.argv[1]} are in use (devices {busy}); not starting device stages") +sys.exit(1 if busy else 0) +EOF } cpu_stages() { @@ -87,12 +96,12 @@ device_stages() { c0="$(first_core)"; n="$(core_count)" export DEEPINFER_NEURON_DEVICE=neuron run device_primitives env NEURON_RT_VISIBLE_CORES="$c0" python "$T/device/test_device_primitives.py" - for ws in 2 4 8; do + for ws in 2 4 8 16; do [ "$ws" -le "$n" ] || continue - run "deep_ep_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \ - python "$T/deep_ep/test_ep.py" --num-processes $ws --port $((29631 + ws)) - run "engram_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \ - python "$T/deep_ep/test_engram.py" --num-processes $ws --port $((29641 + ws)) + for t in ep engram bucket; do + run "${t}_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \ + python "$T/deep_ep/test_$t.py" --num-processes $ws --port $((29600 + 10 * ws + ${#t})) + done done run e2e_neuron env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29733 \ python examples/neuron/run_deepinfer_tiny.py --device neuron --num-kv-blocks 32 \