mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +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>
68 lines
3.0 KiB
Python
68 lines
3.0 KiB
Python
"""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()
|