mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
b18af0027f
- fused_sample: _host_flags/_greedy_verify flags persist while the batch's requests are unchanged (each device read drained the queue and waited for the whole forward). Tiny golden 9/9; full model step time unchanged (660 ms), driver dflash launch 71 -> 10.5 ms host. - bind_cpu on Neuron: the device's NUMA node (DEEPINFER_NEURON_CPU_BIND=node|core|legacy); the generic split gave each rank one CPU and ranks 0-15 the wrong node. Neutral on TPOT (228.2 ms). - profiled followers exit through sys.exit so the runtime writes their trace. - op_census: trace_syncs / census_sites (DEEPINFER_NEURON_TRACE_SYNCS=1|2|3). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
102 lines
4.6 KiB
Python
102 lines
4.6 KiB
Python
"""Count the eager ops of one step on Neuron (issue #27): every aten op outside a compiled region graph is a
|
|
separate NEFF launch, so the census shows what still runs eagerly between and around the region graphs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import os
|
|
from collections.abc import Iterator
|
|
|
|
import torch
|
|
|
|
from deepinfer.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
#: Bookkeeping ops that launch nothing on the device.
|
|
_FREE = ("aten::view", "aten::_reshape_alias", "aten::as_strided", "aten::alias", "aten::detach", "aten::lift_fresh",
|
|
"aten::t", "aten::unsqueeze", "aten::squeeze", "aten::expand", "aten::permute", "aten::transpose",
|
|
"aten::slice", "aten::select", "aten::reshape", "aten::empty", "aten::empty_strided", "aten::resolve_conj",
|
|
"aten::resolve_neg", "aten::_unsafe_view", "aten::split", "aten::chunk", "aten::unbind", "aten::narrow")
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def census(label: str, top: int = 40) -> Iterator[None]:
|
|
"""Log the ops run inside the block, by count (the outermost aten op of each dispatch)."""
|
|
from torch.profiler import ProfilerActivity, profile
|
|
|
|
with profile(activities=[ProfilerActivity.CPU]) as prof:
|
|
yield
|
|
rows = [e for e in prof.key_averages() if e.key.startswith("aten::") and e.key not in _FREE]
|
|
rows.sort(key=lambda e: -e.count)
|
|
total = sum(e.count for e in rows)
|
|
lines = [f"{e.count:6d} {e.cpu_time_total / 1e3:9.2f} ms {e.key}" for e in rows[:top]]
|
|
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)))
|