mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
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 <noreply@anthropic.com>
This commit is contained in:
whn09
parent
e10220bed0
commit
bf447622c2
100
deepinfer/neuron_ops/vendor/deep_ep/comm.py
vendored
100
deepinfer/neuron_ops/vendor/deep_ep/comm.py
vendored
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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 <out>/<prompt>.pt, rank r>0 "
|
||||
"<out>/rank<r>/<prompt>.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)
|
||||
|
||||
|
||||
67
tests/neuron/deep_ep/dev.py
Normal file
67
tests/neuron/deep_ep/dev.py
Normal file
@ -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()
|
||||
@ -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:
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -49,10 +49,19 @@ run() { # run <stage> <command...>
|
||||
[ "$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 \
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user