From 0f3911afb5231edb76aeae502fab5eb982a86054 Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Thu, 29 Oct 2026 08:39:28 +0100 Subject: [PATCH] prain: span profiler with flame graphs (PRAIN_PROFILE: engine step, every model module, MoE route/plan/copy/round, call counts; python -m prain.profile); prain web --host --- README.md | 12 +- prain/cli.py | 5 +- prain/experts.py | 12 +- prain/profile.py | 277 ++++++++++++++++++++++++++++++++++++++++++ prain/tiered.py | 52 ++++---- prain/web.py | 11 +- tests/test_cli.py | 7 ++ tests/test_profile.py | 100 +++++++++++++++ 8 files changed, 444 insertions(+), 32 deletions(-) create mode 100644 prain/profile.py create mode 100644 tests/test_profile.py diff --git a/README.md b/README.md index 7c284bf..f768d35 100644 --- a/README.md +++ b/README.md @@ -35,8 +35,8 @@ The first start writes the bf16 pack (42.5 GiB, ~3 min); with `quant = "int8"` o then quantizes that pack once on the GPU (~3 min). Once the int4 pack is written, the bf16 files (`layerNNN.w13`, `.w2`, `.done`) are not read and can be deleted; a later bf16 or int8 start writes them again. Chat with `prain tui`, or `prain web --config FILE` for several users with -drafts. Tool calls work through vLLM's `gemma4` parser, so coding agents and MCP clients can use -the OpenAI-compatible API. +drafts (`--host 0.0.0.0` serves it to other machines; it has no login). Tool calls work through +vLLM's `gemma4` parser, so coding agents and MCP clients can use the OpenAI-compatible API. ## Expert formats @@ -123,6 +123,14 @@ code at a negative log-likelihood of 5-9 per token on a correct engine. The simulator treats every disk-tier expert as a disk read, so it under-predicts bf16 (0.6 vs 1.04 tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s). +## Profile + +Serve with `PRAIN_PROFILE=/abs/dir`. Writing a number N into `dir/start` records the next N engine +steps: spans around the engine step, every module of the model and prain's MoE phases (route, +plan, copy: h2d/repack/d2d, round), each synced with the GPU. `python -m prain.profile dir` +writes `decode.svg`/`prefill.svg` flame graphs and prints self time and calls per step. The +syncs remove the CPU/GPU overlap: a decode step took 395 ms profiled, 307 ms without. + ## vLLM patches - `vllm-0.32-rocm-no-cuda-families`: gfx1032 reports capability 10.3, and vLLM's CUDA "family" diff --git a/prain/cli.py b/prain/cli.py index 7aed86c..e44b023 100644 --- a/prain/cli.py +++ b/prain/cli.py @@ -167,6 +167,8 @@ def main(argv=None): "--scheduling-policy priority") if name == "web": ui.add_argument("--port", type=int, default=8080) + ui.add_argument("--host", default="127.0.0.1", + help="address to bind; 0.0.0.0 for every interface (no login!)") ui.add_argument("--config", help="placement file: [serve] max_requests and the " "[draft] table (default: their defaults)") args = parser.parse_args(argv) @@ -199,7 +201,8 @@ def main(argv=None): from prain.web import serve as web config = load_config(args.config) if args.config else Config(pack_dir="") # defaults - web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft) + web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft, + host=args.host) if __name__ == "__main__": diff --git a/prain/experts.py b/prain/experts.py index 102c242..a91ff45 100644 --- a/prain/experts.py +++ b/prain/experts.py @@ -32,7 +32,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts -from prain import drafts +from prain import drafts, profile from prain.config import layer_index, load_config from prain.policy import SlotCache, write_stats from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4, @@ -127,6 +127,8 @@ class TieredRoutedExperts(RoutedExperts): raise NotImplementedError("prain: one GPU only (no TP, DP or EP)") if vllm_config.use_v2_model_runner: _install_draft_hook() + if profile.DIR: + profile.install() elif not LAYERS: logger.warning("prain: drafts need vLLM's V2 model runner; with " "VLLM_USE_V2_MODEL_RUNNER=0 every request runs with all its experts") @@ -308,8 +310,12 @@ class TieredRoutedExperts(RoutedExperts): for dsts, srcs in zip(self._slots, self._source(expert)): for i, (dst, src) in enumerate(zip(dsts, srcs)): if i == 0 and self.quant == "int4": # the pack keeps the CPU kernel's order - src = to_gpu_int4(src.to(dst.device, non_blocking=True)) - dst[slot].copy_(src, non_blocking=True) + with profile.span("h2d"): + src = src.to(dst.device, non_blocking=True) + with profile.span("repack"): + src = to_gpu_int4(src) + with profile.span("h2d" if src.device.type == "cpu" else "d2d"): + dst[slot].copy_(src, non_blocking=True) def _load(self, expert, slot): """Demand copy on the compute stream.""" diff --git a/prain/profile.py b/prain/profile.py new file mode 100644 index 0000000..22c7544 --- /dev/null +++ b/prain/profile.py @@ -0,0 +1,277 @@ +"""Span profiler for the vLLM engine process: flame-graph data per engine step. + +Active only with PRAIN_PROFILE=DIR. Writing a number N into DIR/start records the next N engine +steps: the step, the runner's execute_model and sample_tokens, every module forward of the model +and prain's MoE phases become nested spans. Each span syncs the GPU at both ends, so GPU work is +charged to the span that queued it (the overlap of CPU and GPU work is lost; a sync costs ~10 us). +Then it writes DIR/decode.folded and DIR/prefill.folded ("frame;frame;frame microseconds" lines: +time spent in the last frame itself, summed over the steps of that kind; a step with more than +DECODE_TOKENS tokens is a prefill step), DIR/steps.json, and removes DIR/start. +scripts/flamegraph.py draws them. +""" + +import contextlib +import json +import os +import threading +import time + +import torch + +DIR = os.environ.get("PRAIN_PROFILE") +SYNC = [torch.cuda.synchronize] # tests replace it +DECODE_TOKENS = 4 + + +class _Recorder: + def __init__(self): + self.active, self.left, self.thread = False, 0, None + self.stack = [] # [frame, start, time in children] + self.totals, self.steps, self.hooks = {}, {}, [] + self.kind = "decode" + + +_R = _Recorder() +_NULL = contextlib.nullcontext() + + +def _enter(frame): + SYNC[0]() + _R.stack.append([frame, time.perf_counter(), 0.0]) + + +def _exit(): + SYNC[0]() + frame, start, children = _R.stack.pop() + took = time.perf_counter() - start + stack = (*(f[0] for f in _R.stack), frame) + _R.step[stack] = _R.step.get(stack, 0.0) + took - children + _R.step_calls[stack] = _R.step_calls.get(stack, 0) + 1 + if _R.stack: + _R.stack[-1][2] += took + + +def _recording(): + return _R.active and threading.get_ident() == _R.thread + + +def _pre_hook(frame): + def hook(module, args): # must return None: anything else replaces the module's inputs + if _recording(): + _enter(frame) + return hook + + +def _post_hook(module, args, out): # must return None, or it replaces the output + if _recording(): + _exit() + + +@contextlib.contextmanager +def _span(name): + _enter(name) + try: + yield + finally: + _exit() + + +def span(name): + """with span("copy"): ... -- a no-op unless a recording runs.""" + return _span(name) if _recording() else _NULL + + +def step_begin(): + if not _R.active and DIR and os.path.exists(os.path.join(DIR, "start")): + with open(os.path.join(DIR, "start")) as f: + text = f.read().strip() + _R.left = int(text) if text.isdigit() else 1 + _R.active, _R.thread, _R.totals, _R.steps = True, threading.get_ident(), {}, {} + _R.calls = {} + if _recording(): + _R.step, _R.step_calls, _R.kind = {}, {}, "decode" + _enter("step") + + +def note_step(model, tokens): + """In execute_model: the step's token count, and the module hooks on first use.""" + if not _recording(): + return + if tokens > DECODE_TOKENS: + _R.kind = "prefill" + if not _R.hooks: + for name, module in model.named_modules(): + attr, cls = name.rsplit(".", 1)[-1], type(module).__name__ + frame = cls if not name or attr.isdigit() else f"{attr}:{cls}" + _R.hooks.append(module.register_forward_pre_hook(_pre_hook(frame))) + _R.hooks.append(module.register_forward_hook(_post_hook)) + + +def step_end(): + if not _recording(): + return + while _R.stack: # also closes spans an exception left open + _exit() + for stack, seconds in _R.step.items(): + key = (_R.kind, *stack) + _R.totals[key] = _R.totals.get(key, 0.0) + seconds + _R.calls[key] = _R.calls.get(key, 0) + _R.step_calls[stack] + _R.steps[_R.kind] = _R.steps.get(_R.kind, 0) + 1 + _R.left -= 1 + if _R.left <= 0: + _finish() + + +def _finish(): + _R.active = False + for hook in _R.hooks: + hook.remove() + _R.hooks = [] + for kind in _R.steps: + for suffix, values, fmt in (("folded", _R.totals, lambda s: round(s * 1e6)), + ("calls", _R.calls, int)): + with open(os.path.join(DIR, f"{kind}.{suffix}"), "w") as f: + for key, value in sorted(values.items()): + if key[0] == kind: + f.write(f"{';'.join(key[1:])} {fmt(value)}\n") + with open(os.path.join(DIR, "steps.json"), "w") as f: + json.dump({"steps": _R.steps, "decode_tokens": DECODE_TOKENS}, f) + with contextlib.suppress(OSError): + os.remove(os.path.join(DIR, "start")) + + +def install(): + """Spans around vLLM's engine step and the model runner (the vLLM process with the GPU).""" + from vllm.v1.engine.core import EngineCoreProc + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + def wrap(cls, name, around): + original = getattr(cls, name) + if not getattr(original, "prain_profile", False): + wrapped = around(original) + wrapped.prain_profile = True + setattr(cls, name, wrapped) + + def engine_step(original): + def step(self, *args, **kwargs): + step_begin() + try: + return original(self, *args, **kwargs) + finally: + step_end() + return step + + def execute_model(original): + def execute(self, scheduler_output, *args, **kwargs): + with span("execute_model"): + note_step(self.model, getattr(scheduler_output, "total_num_scheduled_tokens", 0)) + return original(self, scheduler_output, *args, **kwargs) + return execute + + def sample_tokens(original): + def sample(self, *args, **kwargs): + with span("sample_tokens"): + return original(self, *args, **kwargs) + return sample + + wrap(EngineCoreProc, "_process_engine_step", engine_step) + wrap(GPUModelRunner, "execute_model", execute_model) + wrap(GPUModelRunner, "sample_tokens", sample_tokens) + + +_COLORS = (("copy", "#e8603c"), ("round", "#f0a030"), ("expert", "#f0a030"), ("moe", "#f0a030"), + ("route", "#e8c440"), ("plan", "#e8c440"), ("degrade", "#e8c440"), ("cpu", "#e8c440"), + ("attn", "#5b9bd5"), ("rotary", "#5b9bd5"), ("mlp", "#6abf69"), ("gate_up", "#6abf69"), + ("down_proj", "#6abf69"), ("norm", "#a0a0a0"), ("logits", "#a070c0"), + ("lm_head", "#a070c0"), ("sample", "#a070c0"), ("embed", "#a070c0")) + + +def _tree(lines): + root = {"name": "all", "self": 0.0, "kids": {}} + for line in lines: + stack, _, us = line.strip().rpartition(" ") + node = root + for frame in stack.split(";"): + node = node["kids"].setdefault(frame, {"name": frame, "self": 0.0, "kids": {}}) + node["self"] += float(us) + + def total(node): + node["total"] = node["self"] + sum(total(k) for k in node["kids"].values()) + return node["total"] + total(root) + return root + + +def top(lines, steps=1, n=25, calls=None): + """[(frame, ms per step, % of all, calls per step or None)] by time in the frame itself; + calls: the .calls lines of the same recording.""" + own, count = {}, {} + for line in lines: + stack, _, us = line.strip().rpartition(" ") + frame = stack.rsplit(";", 1)[-1] + own[frame] = own.get(frame, 0.0) + float(us) + for line in calls or (): + stack, _, n_calls = line.strip().rpartition(" ") + frame = stack.rsplit(";", 1)[-1] + count[frame] = count.get(frame, 0) + int(n_calls) + whole = sum(own.values()) or 1.0 + return [(frame, round(us / 1e3 / steps, 1), round(100 * us / whole, 1), + round(count[frame] / steps, 1) if frame in count else None) + for frame, us in sorted(own.items(), key=lambda kv: -kv[1])[:n]] + + +def svg(lines, steps=1, title="", width=1200, row=18): + """A flame graph of folded stack lines (root on top, widest child first); each box's + tooltip gives its time per step, its share of all time and its own (self) time.""" + from html import escape + + root = _tree(lines) + whole = root["total"] or 1.0 + boxes = [] + + def place(node, x, depth): + boxes.append((node, x, depth, node["total"] / whole * width)) + for kid in sorted(node["kids"].values(), key=lambda k: -k["total"]): + place(kid, x, depth + 1) + x += kid["total"] / whole * width + place(root, 0.0, 0) + height = (max(d for _, _, d, _ in boxes) + 1) * row + 30 + out = [f'', + f'{escape(title)}: ' + f'{whole / 1e3 / steps:.1f} ms per step over {steps} steps (hover for times)'] + for node, x, depth, w in boxes: + if w < 0.5: + continue + name, y = node["name"], 26 + depth * row + color = next((c for key, c in _COLORS if key in name.lower()), "#d8b890") + ms = node["total"] / 1e3 / steps + text = f"{name} {ms:.1f} ms" if ms >= 0.1 else name # readable without hovering + label = text if len(text) * 6.5 < w - 4 else text[:max(int((w - 4) / 6.5) - 1, 0)] + "…" + out.append( + f'{escape(name)}: {ms:.1f} ms per step ' + f'({100 * node["total"] / whole:.1f}%), self {node["self"] / 1e3 / steps:.1f} ms' + f'' + + (f'{escape(label)}' if w > 14 else "") + + "") + return "\n".join(out + [""]) + + +if __name__ == "__main__": # python -m prain.profile DIR: DIR/KIND.svg and a summary per kind + import sys + + directory = sys.argv[1] + with open(os.path.join(directory, "steps.json")) as f: + counts = json.load(f)["steps"] + for kind, steps in counts.items(): + with open(os.path.join(directory, f"{kind}.folded")) as f: + lines = f.read().splitlines() + calls_path = os.path.join(directory, f"{kind}.calls") + calls = open(calls_path).read().splitlines() if os.path.exists(calls_path) else None + with open(os.path.join(directory, f"{kind}.svg"), "w") as f: + f.write(svg(lines, steps, title=kind)) + print(f"== {kind}: {steps} steps, {_tree(lines)['total'] / 1e3 / steps:.1f} ms per step" + f" (self time per step, share, calls per step, ms per call)") + for frame, ms, share, n in top(lines, steps, calls=calls): + per_call = f"{n:8.1f} x {ms / n:7.3f} ms" if n else "" + print(f"{ms:9.1f} ms {share:5.1f}% {per_call} {frame}") diff --git a/prain/tiered.py b/prain/tiered.py index 3386d40..c316267 100644 --- a/prain/tiered.py +++ b/prain/tiered.py @@ -3,6 +3,8 @@ import torch import torch.nn.functional as F +from prain import profile + HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2} @@ -119,21 +121,23 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight" raw = None skip = miss is not None and miss.min_weight < 1 if draft is not None or skip: - if on_stats is not None: # the trace records the routing, not what was skipped - raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) - rows = draft if draft is not None else torch.zeros(topk_ids.shape[0], dtype=torch.bool, - device=topk_ids.device) - need = min_weight - if skip: # every row may skip; draft rows by the smaller share - need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None] - rows = torch.ones_like(rows) - available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots), - cache.num_experts) - _, kept = degrade(topk_weights, topk_ids, rows, available, need) - # experts some row copies anyway serve every row for free - available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0 - topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need) - stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) # host sync + with profile.span("degrade"): + if on_stats is not None: # the trace records the routing, not what was skipped + raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) + rows = draft if draft is not None else torch.zeros( + topk_ids.shape[0], dtype=torch.bool, device=topk_ids.device) + need = min_weight + if skip: # every row may skip; draft rows by the smaller share + need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None] + rows = torch.ones_like(rows) + available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots), + cache.num_experts) + _, kept = degrade(topk_weights, topk_ids, rows, available, need) + # experts some row copies anyway serve every row for free + available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0 + topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need) + with profile.span("route"): + stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life) # sync tokens = int(topk_ids.shape[0]) if on_stats is not None: on_stats(raw if raw is not None else stats, tokens) @@ -150,14 +154,20 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight" nonlocal total total = out if total is None else total.to(acc).add_(out) # .to: a no-op once fp32 - for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu): + with profile.span("plan"): + rounds = cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu) + for slot_of, loads in rounds: for expert, slot in loads: - load(expert, slot) + with profile.span("copy"): + load(expert, slot) if slot_of: # each round computes only its own routes - table = slot_table(slot_of, cache.num_experts, x.device) - add(kernel(x, topk_weights, ids, table, list(slot_of.values()))) + with profile.span("round"): + table = slot_table(slot_of, cache.num_experts, x.device) + add(kernel(x, topk_weights, ids, table, list(slot_of.values()))) if after_enqueue is not None: - after_enqueue(cpu, stats) + with profile.span("prefetch"): + after_enqueue(cpu, stats) if cpu: # runs while the GPU works through its rounds - add(cpu_experts(*host, cpu, cpu_weights).to(x.device)) + with profile.span("cpu experts"): + add(cpu_experts(*host, cpu, cpu_weights).to(x.device)) return torch.zeros_like(x) if total is None else total.to(x.dtype) diff --git a/prain/web.py b/prain/web.py index 4e3040e..0ea9a9b 100644 --- a/prain/web.py +++ b/prain/web.py @@ -3,8 +3,8 @@ A small stdlib proxy in front of vLLM. GET / serves the page. POST /warm takes {"session", "messages", "draft"} and returns at once. POST /chat takes {"session", "messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" | -"refine" | "error"} events of its own. Binds to 127.0.0.1; use an SSH tunnel from other -machines. +"refine" | "error"} events of its own. Binds to 127.0.0.1 unless --host says otherwise +(0.0.0.0 serves every interface, without a login). Several browsers at once ([serve] and [draft] tables of the placement file): - At most draft.full_requests answers run with all their experts. An answer that arrives @@ -267,8 +267,9 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser return server -def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft()): - server = make_server(base_url, port, warmup=warmup, priority=priority, serve=serve, draft=draft) - print(f"prain web: http://127.0.0.1:{server.server_port} (vLLM at {base_url}; " +def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1"): + server = make_server(base_url, port, host=host, warmup=warmup, priority=priority, serve=serve, + draft=draft) + print(f"prain web: http://{host}:{server.server_port} (vLLM at {base_url}; " f"{draft.full_requests} full of {serve.max_requests} parallel answers)") server.serve_forever() diff --git a/tests/test_cli.py b/tests/test_cli.py index aa2bec7..a286c4c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -61,6 +61,13 @@ class Cli(TmpCase): main(["serve", config, "/models/m", "--", "--max-num-seqs=4"]) execvp.assert_not_called() + def test_web_binds_the_host_it_is_given(self): + with mock.patch("prain.web.serve") as web: + main(["web", "--port", "8081"]) + self.assertEqual(web.call_args.kwargs["host"], "127.0.0.1") + main(["web", "--host", "0.0.0.0"]) + self.assertEqual(web.call_args.kwargs["host"], "0.0.0.0") + def test_synth_sim_and_tune(self): hf = self.write("config.json", json.dumps(GEMMA4)) trace = os.path.join(self.dir, "trace.jsonl.gz") diff --git a/tests/test_profile.py b/tests/test_profile.py new file mode 100644 index 0000000..a274a63 --- /dev/null +++ b/tests/test_profile.py @@ -0,0 +1,100 @@ +import json +import os +import threading +import time + +import torch + +from prain import profile +from tests.helpers import TmpCase + + +class Inner(torch.nn.Module): + def forward(self, x): + time.sleep(0.002) + return x + + +class Outer(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([Inner(), Inner()]) + self.head = Inner() + + def forward(self, x): + for layer in self.layers: + x = layer(x) + with profile.span("copy"): + time.sleep(0.001) + return self.head(x) + + +class TestProfile(TmpCase): + def setUp(self): + super().setUp() + self.syncs = [] + saved = profile.DIR, profile.SYNC[0] + profile.DIR, profile.SYNC[0] = self.dir, lambda: self.syncs.append(1) + self.addCleanup(lambda: (setattr(profile, "DIR", saved[0]), + profile.SYNC.__setitem__(0, saved[1]))) + + def steps(self, model, tokens): + for n in tokens: + profile.step_begin() + with profile.span("execute_model"): + profile.note_step(model, n) + model(torch.zeros(1)) + profile.step_end() + + def folded(self, kind): + with open(os.path.join(self.dir, f"{kind}.folded")) as f: + return {stack: int(us) for stack, us in (l.rsplit(" ", 1) for l in f.read().splitlines())} + + def test_records_the_requested_steps_as_folded_stacks(self): + model = Outer() + self.steps(model, [1]) # no start file: nothing recorded, nothing synced + self.assertEqual(self.syncs, []) + self.write("start", "3") + self.steps(model, [1, 1, 300, 1]) # records 2 decode steps and 1 prefill, not the 4th + self.assertFalse(os.path.exists(os.path.join(self.dir, "start"))) + with open(os.path.join(self.dir, "steps.json")) as f: + self.assertEqual(json.load(f)["steps"], {"decode": 2, "prefill": 1}) + decode = self.folded("decode") + self.assertEqual(set(decode), {"step", "step;execute_model", "step;execute_model;Outer", + "step;execute_model;Outer;Inner", # both layers, merged + "step;execute_model;Outer;copy", + "step;execute_model;Outer;head:Inner"}) + for stack, ms in (("Outer;Inner", 8), ("Outer;head:Inner", 4), ("Outer;copy", 2)): + self.assertGreaterEqual(decode[f"step;execute_model;{stack}"], ms * 1000) + self.assertLess(decode[f"step;execute_model;{stack}"], ms * 1000 * 3) + self.assertGreater(self.folded("prefill")["step;execute_model;Outer;Inner"], 4000) + with open(os.path.join(self.dir, "decode.calls")) as f: + calls = dict(line.rsplit(" ", 1) for line in f.read().splitlines()) + self.assertEqual(calls["step;execute_model;Outer;Inner"], "4") # 2 layers x 2 steps + self.assertEqual(calls["step"], "2") + recorded = len(self.syncs) + self.steps(model, [1]) # finished: the module hooks are gone + self.assertEqual(len(self.syncs), recorded) + self.assertFalse(model.head._forward_hooks or model.head._forward_pre_hooks) + + def test_hooks_leave_other_threads_inputs_and_outputs_alone(self): + model = Outer() + self.write("start", "1") + profile.step_begin() + profile.note_step(model, 1) # installs the hooks + out = [] + thread = threading.Thread(target=lambda: out.append(model(torch.ones(1)))) + thread.start() + thread.join() + self.assertTrue(torch.equal(out[0], torch.ones(1))) + profile.step_end() + + def test_svg_shows_time_per_step(self): + lines = ["step 500", "step;attn 3000", "step;attn;copy 1000"] + svg = profile.svg(lines, steps=2, title="decode") + self.assertTrue(svg.startswith("