diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py index 93ba4c3..55d0bf9 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py @@ -127,20 +127,26 @@ def _plain_logprobs_and_greedy(logits, temperatures): # noqa: ANN001, ANN202 _COMPILED: dict = {} _FLAGS: dict = {} +_FLAGS_EPOCH: list = [None] -def reset_host_flags() -> None: - _FLAGS.clear() +def reset_host_flags(requests: tuple | None = None) -> None: + """Start a step. The flags are per-request sampling parameters, so they hold while the step's requests are those + of the previous step (``requests``: their ids): only a changed batch reads them again (a device read drains the + queue, and the first read of a step waits for the whole forward).""" + if requests is None or requests != _FLAGS_EPOCH[0]: + _FLAGS.clear() + _FLAGS_EPOCH[0] = requests def _host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]: # noqa: ANN001 """(any row filters by top-p / top-k, every row is greedy) for the step's sampling parameters. Reading them is a device sync, and a DSpark step samples the same parameter tensors six times (verification, then each draft), - so the flags are kept per tensor until the step ends (reset_host_flags: a later step's tensors may reuse the - address).""" + so the flags are kept until the batch changes (reset_host_flags).""" from deepinfer.sampling_params import SAMPLING_EPS - key = tuple((t.data_ptr(), t.shape[0]) for t in (temperatures, top_p, top_k)) + # By shape, not address: a step may hand over the same parameters in a fresh buffer. + key = ("filter", tuple(t.shape[0] for t in (temperatures, top_p, top_k)), vocab) flags = _FLAGS.get(key) if flags is None: t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu() @@ -193,7 +199,7 @@ def _greedy_verify( Sampled (non-greedy) verification is not ported (its int64 RNG reject draw).""" from deepinfer.sampling_params import SAMPLING_EPS - key = ("greedy", temperatures.data_ptr(), temperatures.shape[0]) + key = ("greedy", temperatures.shape[0]) if key not in _FLAGS: _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all()) if not _FLAGS[key]: diff --git a/deepinfer/neuron_ops/op_census.py b/deepinfer/neuron_ops/op_census.py index 8852ffe..4c21fc0 100644 --- a/deepinfer/neuron_ops/op_census.py +++ b/deepinfer/neuron_ops/op_census.py @@ -4,6 +4,7 @@ separate NEFF launch, so the census shows what still runs eagerly between and ar from __future__ import annotations import contextlib +import os from collections.abc import Iterator import torch @@ -33,3 +34,68 @@ def census(label: str, top: int = 40) -> Iterator[None]: graphs = [e for e in prof.key_averages() if "Compiled" in e.key or "neuron" in e.key.lower()] lines += [f"{e.count:6d} {e.cpu_time_total / 1e3:9.2f} ms {e.key}" for e in graphs[:10]] logger.info(f"Neuron op census ({label}): {total} aten ops outside bookkeeping\n" + "\n".join(lines)) + + +@contextlib.contextmanager +def trace_syncs(label: str) -> Iterator[None]: + """Log every device -> host read in the block (each drains the Neuron execution queue) with its caller.""" + import collections + import traceback + + hits: collections.Counter = collections.Counter() + saved = {} + + def site() -> str: + for frame in reversed(traceback.extract_stack()[:-2]): + if "/deepinfer/" in frame.filename and "op_census.py" not in frame.filename: + return f"{frame.filename.split('/deepinfer/')[-1]}:{frame.lineno} {frame.name}" + return "?" + + def wrap(name: str, check=None): # noqa: ANN001, ANN202 + orig = getattr(torch.Tensor, name) + saved[name] = orig + + def fn(self, *a, **k): # noqa: ANN001, ANN002, ANN003, ANN202 + if self.device.type == "neuron" and (check is None or check(*a, **k)): + key = f"{name:8s} {site()}" + if key not in hits and os.environ.get("DEEPINFER_NEURON_TRACE_SYNCS") == "2": + logger.info(f"host sync {key}:\n" + "".join(traceback.format_stack()[-14:-1])) + hits[key] += 1 + return orig(self, *a, **k) + + setattr(torch.Tensor, name, fn) + + def to_cpu(*a, **k) -> bool: # noqa: ANN002, ANN003 + dev = k.get("device", a[0] if a else None) + return isinstance(dev, (str, torch.device)) and str(dev).startswith("cpu") + + for name in ("item", "cpu", "tolist", "__bool__", "numpy"): + wrap(name) + wrap("to", to_cpu) + try: + yield + finally: + for name, orig in saved.items(): + setattr(torch.Tensor, name, orig) + logger.info(f"Neuron host syncs ({label}): {sum(hits.values())}\n" + "\n".join(f"{n:4d} {k}" for k, n in hits.most_common())) + + +@contextlib.contextmanager +def census_sites(label: str, top: int = 45) -> Iterator[None]: + """Count the ops run in the block (outermost aten op of each dispatch) by their innermost DeepInfer call site: + each eager op on a NeuronCore is its own NEFF launch.""" + import collections + + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CPU], with_stack=True) as prof: + yield + hits: collections.Counter = collections.Counter() + events = [e for e in prof.events() if e.name.startswith("aten::") and e.name not in _FREE] + inner = {id(e) for e in events for c in e.cpu_children if c.name.startswith("aten::")} + for e in events: + if e.cpu_parent is not None and e.cpu_parent.name.startswith("aten::"): + continue # count each dispatch once, at its outermost aten op + where = next((f for f in (e.stack or []) if "/deepinfer/" in f and "op_census.py" not in f), "?") + hits[where.split("/deepinfer/")[-1]] += 1 + logger.info(f"Neuron op sites ({label}): {sum(hits.values())} ops\n" + "\n".join(f"{n:5d} {k}" for k, n in hits.most_common(top))) diff --git a/deepinfer/utils.py b/deepinfer/utils.py index 491b1bf..3db7e9e 100644 --- a/deepinfer/utils.py +++ b/deepinfer/utils.py @@ -140,16 +140,59 @@ def allow_cuda_sync() -> AbstractContextManager[None]: return _sync_debug_arbiter.scope("allow") +def _neuron_bind_cpus(local_rank: int) -> list[int] | None: + """CPUs for the rank on logical NeuronCore ``local_rank`` (of CUDA_VISIBLE_DEVICES), on its device's NUMA node. + + The generic split (physical cores // ranks) gave each of 64 ranks one logical CPU, numbered by rank: the step's + dispatch thread and the runtime's threads shared it, and devices 0-3 (NUMA 1 on trn2.48xlarge) got CPUs of NUMA 0. + DEEPINFER_NEURON_CPU_BIND: ``node`` (default) the device's whole NUMA node, ``core`` one physical core and its + hyperthread on that node, ``legacy`` the generic split. None falls back to the generic split.""" + mode = os.environ.get("DEEPINFER_NEURON_CPU_BIND", "node") + if mode == "legacy": + return None + import json + import subprocess + + try: + devices = json.loads(subprocess.run(["neuron-ls", "--json-output"], capture_output=True, check=True, timeout=30).stdout) + except (OSError, subprocess.SubprocessError, ValueError): + logger.warning("neuron-ls unavailable: generic CPU binding") + return None + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + core = int(visible.split(",")[local_rank]) if visible else local_rank + device = next((d for d in devices if core in d["neuroncore_ids"]), None) + if device is None: + return None + + def cpus(spec: str) -> list[int]: + out = [] + for part in spec.split(","): + lo, _, hi = part.partition("-") + out.extend(range(int(lo), int(hi or lo) + 1)) + return out + + node_cpus = cpus(device["cpu_affinity"]) + if mode != "core": + return node_cpus + # One physical core per NeuronCore of the node, in core order, plus its hyperthread sibling. + node_cores = sorted(c for d in devices if d["numa_node"] == device["numa_node"] for c in d["neuroncore_ids"]) + physical = sorted({min(cpus(open(f"/sys/devices/system/cpu/cpu{c}/topology/thread_siblings_list").read().strip())) for c in node_cpus}) + first = physical[node_cores.index(core) % len(physical)] + return cpus(open(f"/sys/devices/system/cpu/cpu{first}/topology/thread_siblings_list").read().strip()) + + def bind_cpu(local_rank: int) -> None: from deepinfer import PLATFORM if PLATFORM.is_neuron and envs.DEEPINFER_NEURON_DEVICE == "cpu": # Neuron CPU mode computes on the host: a 1-core share (96 cores / 64 NeuronCores) would serialize it. return - phy_cpus = psutil.cpu_count(logical=False) - ngpus = get_actual_num_gpus() - num_cpus_per_gpu = phy_cpus // ngpus - bind_cpus = list(range(local_rank * num_cpus_per_gpu, (local_rank + 1) * num_cpus_per_gpu)) + bind_cpus = _neuron_bind_cpus(local_rank) if PLATFORM.is_neuron else None + if bind_cpus is None: + phy_cpus = psutil.cpu_count(logical=False) + ngpus = get_actual_num_gpus() + num_cpus_per_gpu = phy_cpus // ngpus + bind_cpus = list(range(local_rank * num_cpus_per_gpu, (local_rank + 1) * num_cpus_per_gpu)) logger.info(f"bind_cpu gpu_id: {local_rank}, cpus: {bind_cpus}.") import deep_cluster.topology diff --git a/deepinfer/worker/base.py b/deepinfer/worker/base.py index d13e3eb..2cff6a7 100644 --- a/deepinfer/worker/base.py +++ b/deepinfer/worker/base.py @@ -10,6 +10,7 @@ import functools import gc import math import os +import sys import signal import time from typing import Any @@ -148,6 +149,10 @@ class WorkerBase: dist.destroy_process_group() if not self.is_driver_worker: + if os.environ.get("NEURON_RT_INSPECT_ENABLE") == "1": + # A profiled Neuron follower (DEEPINFER_NEURON_PROFILE_RANKS): the runtime writes its trace in the + # teardown os._exit() skips. + sys.exit(0) os._exit(0) def maybe_gc_collect(self) -> None: diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index 50a9fc3..e0ea990 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -4,6 +4,7 @@ from __future__ import annotations import contextlib +import os import copy import time from dataclasses import dataclass @@ -447,11 +448,20 @@ class ModelRunner: if PLATFORM.is_neuron: from deepinfer.model_executor.layers.tilelang_ops.neuron.fused_sample import reset_host_flags - reset_host_flags() # sampling-parameter flags are read once a step (see _host_flags) + # Sampling-parameter flags are read once per batch of requests (see _host_flags). + reset_host_flags(tuple(seq.request_id for seq in seq_metadata_list) if seq_metadata_list else None) profiler = _maybe_profile_sample(self) from deepinfer.neuron_ops.op_census import census + from deepinfer.neuron_ops.op_census import census_sites, trace_syncs + + self._trace_sample_calls = getattr(self, "_trace_sample_calls", 0) + 1 + tracing = PLATFORM.is_neuron and os.environ.get("DEEPINFER_NEURON_TRACE_SYNCS", "0") != "0" and self._trace_sample_calls in (6, 7) with ( + trace_syncs(f"sample call {self._trace_sample_calls}") if tracing else contextlib.nullcontext(), + census_sites(f"sample call {self._trace_sample_calls}") + if tracing and os.environ.get("DEEPINFER_NEURON_TRACE_SYNCS") == "3" + else contextlib.nullcontext(), census("sample") if profiler is not None else contextlib.nullcontext(), cpu_timer.record("launch_sample"), cuda_timer.record("device_sample", scale=1),