File size: 8,361 Bytes
68a3286 8dc9457 68a3286 8dc9457 68a3286 8dc9457 68a3286 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device profile of transfusion-p150: one eager forward and one traced replay of the whole graph between Tracy
signposts.
ROOT=/home/ubuntu/experiments/tt-models
$ROOT/bin/devrun -t 3600 -- python -m tracy -r -p -v --no-web-server \\
-o $ROOT/generated/profiler/transfusion_baseline code/scripts/profile_ops.py
tt-perf-report <ops_perf_results_*.csv> --start-signpost trace --end-signpost trace_end --tracing-mode
python code/scripts/profile_summary.py <the same csv> profile_summary.json --md profile_summary.md
Sections (signposts): ``eager`` .. ``eager_end``, one eager forward after the warm-up (``TraceRunner.run_eager``:
the same graph and programs as the trace, nothing compiled), and ``trace`` .. ``trace_end``, ``--replays`` replays
of the captured ``default`` trace (the served device path; the uploaded inputs are the voxelized ``--input``). The
device profiler buffer is flushed (``ttnn.ReadDeviceProfiler``) before and after each section, and inside the eager
section after every module (``--flush-modules``, default on: a frame is well above the profiler buffer's 1,000
programs). Inside the eager section every module of the graph opens a signpost (``m:pfn``, ``m:scatter``,
``m:backbone.block<i>``, ``m:neck``, ``m:head.shared``, ``m:head.heatmap.hidden``, ``m:head.heatmap``,
``m:head.sigmoid``, ``m:head.local_max``, ``m:query.topk``, ``m:query.decode``, ``m:query.init``, ``m:query.keys``,
``m:query.decoder``, ``m:query.predict``, ``m:query.heat``, ``m:query.pack``) and every backbone conv a layer
signpost (``c:backbone.block<i>.conv<j>``; ``--no-layer-signposts`` turns them off), so the ops of the eager
forward, and through the identical op order those of the replay, can be attributed to modules
(``profile_summary.py``). The trace section is flushed every replay; with ``TT_METAL_PROFILER_MID_RUN_DUMP=1`` the
profiler also dumps mid-trace (needed when one replay exceeds the buffer). The output folder must be absolute and
outside the bundle (PLAN.md 5.1).
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
from typing import Callable, List, Sequence
from tt_transfusion import TransFusion
from tt_transfusion.ttaw.profiling import read_device_profiler, signpost, signposted
SAMPLE = Path(__file__).resolve().parents[1] / "tt_transfusion" / "samples" / "test_pcd.npz"
class _Signposted:
"""A callable stand-in for a module that emits its signposts, then calls the module (attributes forward)."""
def __init__(self, inner, labels: Sequence[str], flush: Callable[[], None] = None):
self.inner, self._labels, self._flush = inner, list(labels), flush
def __call__(self, *args, **kwargs):
if self._flush is not None:
self._flush()
for label in self._labels:
signpost(label)
return self.inner(*args, **kwargs)
def __getattr__(self, name):
return getattr(self.inner, name)
def module_signposts(graph, layers: bool = True, flush: Callable[[], None] = None) -> Callable[[], None]:
"""Wrap the modules of a ``tt.model.TtTransFusion`` (and the module-level ops its forward calls) with
signposting stand-ins; returns the function that restores them. ``flush`` runs before each module signpost
(device profiler read-out: one eager frame holds more programs than the buffer)."""
from tt_transfusion.tt import model as tt_model
undo: List[Callable[[], None]] = []
def swap(obj, attr, labels, fl=flush):
inner = getattr(obj, attr)
setattr(obj, attr, _Signposted(inner, labels, fl))
undo.append(lambda: setattr(obj, attr, inner))
def swap_method(obj, attr, labels, fl=flush):
"""An instance attribute that shadows a bound method (removed again on restore)."""
inner = getattr(obj, attr)
setattr(obj, attr, _Signposted(inner, labels, fl))
undo.append(lambda: delattr(obj, attr))
def swap_item(seq, i, labels, fl=flush):
inner = seq[i]
seq[i] = _Signposted(inner, labels, fl)
undo.append(lambda: seq.__setitem__(i, inner))
def swap_fn(module, name, label):
inner = getattr(module, name)
def wrapped(*args, **kwargs):
if flush is not None:
flush()
signpost(label)
return inner(*args, **kwargs)
setattr(module, name, wrapped)
undo.append(lambda: setattr(module, name, inner))
swap(graph, "pfn", ["m:pfn"])
swap_fn(tt_model, "scatter_rows", "m:scatter")
for i, block in enumerate(graph.backbone.blocks):
for j in range(len(block)):
labels = ([f"m:backbone.block{i}"] if j == 0 else []) + ([f"c:backbone.block{i}.conv{j}"] if layers else [])
swap_item(block, j, labels, flush if j == 0 else None)
swap(graph, "neck", ["m:neck"])
swap(graph, "shared", ["m:head.shared"])
swap(graph, "hm_hidden", ["m:head.heatmap.hidden"])
swap(graph, "hm_out", ["m:head.heatmap"])
swap_fn(tt_model, "sigmoid_heat", "m:head.sigmoid")
swap(graph, "local_max", ["m:head.local_max"])
qh = graph.query_head
swap_method(qh.select, "indices", ["m:query.topk"])
swap_method(qh.select, "decode", ["m:query.decode"])
for attr, label in (("init_queries", "m:query.init"), ("keys", "m:query.keys"), ("decoder", "m:query.decoder"),
("predict", "m:query.predict"), ("query_heat", "m:query.heat"),
("_indices_tile", "m:query.pack")):
swap_method(qh, attr, [label])
def restore() -> None:
for fn in reversed(undo):
fn()
return restore
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
ap.add_argument("--input", default=str(SAMPLE))
ap.add_argument("--dispatch", default=None, choices=["eth", "worker"])
ap.add_argument("--cqs", type=int, default=None, choices=[1, 2])
ap.add_argument("--replays", type=int, default=1, help="replays inside the trace section")
ap.add_argument("--weights-dir", default=None)
ap.add_argument("--no-eager", dest="eager", action="store_false")
ap.add_argument("--no-layer-signposts", dest="layers", action="store_false")
ap.add_argument("--no-flush-modules", dest="flush_modules", action="store_false")
a = ap.parse_args()
import ttnn
from tt_transfusion.tt.model import stage_inputs
t0 = time.perf_counter()
with TransFusion.from_pretrained(dispatch=a.dispatch, num_command_queues=a.cqs,
weights_dir=a.weights_dir) as model:
tr, dev = model.runner, model.device
print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info), flush=True)
prepared = model._prepare(points=a.input)
out = model(a.input) # one served frame: upload + replay + readback
print(f"{out.to_dict()['num_detections']} detections, {prepared.num_pillars} pillars", flush=True)
ttnn.synchronize_device(dev)
read_device_profiler(dev) # warm-up / capture / first frame out of the buffer
inputs = stage_inputs(model.staging, prepared)
variant = getattr(model.staging, "variant", "default") # PFN_BUCKETS: the frame's capacity variant
print(f"variant {variant}", flush=True)
if a.eager:
def flush():
ttnn.synchronize_device(dev)
read_device_profiler(dev)
restore = module_signposts(model.graph, layers=a.layers, flush=flush if a.flush_modules else None)
try:
with signposted("eager"):
tr.run_eager(variant, inputs=inputs)
ttnn.synchronize_device(dev)
finally:
restore()
read_device_profiler(dev)
tr.upload(inputs)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
with signposted("trace"):
for _ in range(a.replays):
tr.replay(variant)
ttnn.synchronize_device(dev)
read_device_profiler(dev)
read_device_profiler(dev)
print("trace:", json.dumps(tr.describe()), flush=True)
if __name__ == "__main__":
main()
|