File size: 5,109 Bytes
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 | #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""ttnn-visualizer memory + graph report of ONE eager forward of the served graph (default precision policy, ETH
dispatch, the device grid). Never use it for timing: graph capture and detailed buffer tracing synchronise after
every op.
ROOT=/home/ubuntu/experiments/tt-models; cd $ROOT/bundles/transfusion-p150
$ROOT/bin/devrun -t 2400 -- env PYTHONPATH=$PWD/code:$PYTHONPATH \\
python code/scripts/visualizer_report.py $ROOT/generated/ttnn_visualizer/transfusion_baseline
ttnn-visualizer --profiler-path $ROOT/generated/ttnn_visualizer/transfusion_baseline \\
--performance-path $ROOT/generated/profiler/transfusion_baseline/reports/<ts>
What the pytest fixture of ttnn does (ttnn/graph_report.py ``_run_pytest_graph_report_fixture``), in a script:
warm-up forward (JIT, weight upload), then the second forward inside ``ttnn.graph.full_graph_capture`` (Python I/O,
stack traces, detailed buffer tracing), ``ttnn.graph_report.import_report`` -> ``db.sqlite``, and ``config.json``
(the settings the capture ran with). The intermediate graph JSON is deleted after the import.
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
SAMPLE = Path(__file__).resolve().parents[1] / "tt_transfusion" / "samples" / "test_pcd.npz"
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
ap.add_argument("out_dir")
ap.add_argument("--variant", default=None)
ap.add_argument("--input", default=str(SAMPLE))
ap.add_argument("--weights-dir", default=None)
a = ap.parse_args()
out_dir = Path(a.out_dir).resolve()
out_dir.mkdir(parents=True, exist_ok=True)
import ttnn
from tt_transfusion.api import TransFusion
from tt_transfusion.device import close_device, describe_device, open_device
from tt_transfusion.host.preprocess import TransFusionPreprocessor
from tt_transfusion.reference.config import DEFAULT_VARIANT, TransFusionConfig
from tt_transfusion.reference.weights import load_weights
from tt_transfusion.tt.model import TtTransFusion, device_inputs
from tt_transfusion.ttaw.tensors import to_host_tensor
weights_path = TransFusion.resolve_weights(weights_dir=a.weights_dir)
config = TransFusionConfig.from_weights(weights_path, a.variant or DEFAULT_VARIANT)
prepared = TransFusionPreprocessor(config)(a.input)
dev = open_device(allow_fallback=False)
try:
print("device:", json.dumps(describe_device(dev, 0)), flush=True)
graph = TtTransFusion(dev, load_weights(weights_path, config), config) # the default knobs / policy
graph.query_head.prepare(dev)
host = device_inputs(prepared, graph.num_pillars)
feats = ttnn.allocate_tensor_on_device(ttnn.Shape(list(graph.features_shape)), ttnn.bfloat16,
ttnn.ROW_MAJOR_LAYOUT, dev, ttnn.DRAM_MEMORY_CONFIG)
index = ttnn.allocate_tensor_on_device(ttnn.Shape(list(graph.index_shape)), ttnn.uint32,
ttnn.ROW_MAJOR_LAYOUT, dev, ttnn.DRAM_MEMORY_CONFIG)
ttnn.copy_host_to_device_tensor(to_host_tensor(host["features"], "bfloat16", ttnn.ROW_MAJOR_LAYOUT), feats)
ttnn.copy_host_to_device_tensor(to_host_tensor(host["canvas_index"], "uint32", ttnn.ROW_MAJOR_LAYOUT), index)
def forward():
out = graph.forward(feats, index).tensor # ttaw.trace.Packed: the one output row
ttnn.synchronize_device(dev)
if out.is_allocated():
ttnn.deallocate(out)
t0 = time.perf_counter()
forward() # warm-up: JIT + weight upload (host writes)
print(f"warm-up forward {time.perf_counter() - t0:.1f} s", flush=True)
json_path = out_dir / "graph_capture.json"
t0 = time.perf_counter()
with ttnn.graph.full_graph_capture(str(json_path)):
forward()
print(f"captured forward {time.perf_counter() - t0:.1f} s -> {json_path} "
f"({json_path.stat().st_size / 1e6:.1f} MB)", flush=True)
ttnn.deallocate(feats)
ttnn.deallocate(index)
graph.release()
finally:
close_device(dev)
t0 = time.perf_counter()
db = ttnn.graph_report.import_report(json_path, out_dir)
ttnn.save_config_to_json_file(out_dir / "config.json")
cfg = json.loads((out_dir / "config.json").read_text()) # record what the capture ran with
cfg.update({"enable_fast_runtime_mode": False, "enable_graph_report": True, "enable_graph_python_stack_traces": True,
"enable_detailed_buffer_report": True, "root_report_path": str(out_dir.parent),
"report_name": out_dir.name})
(out_dir / "config.json").write_text(json.dumps(cfg, indent=4) + "\n")
print(f"imported in {time.perf_counter() - t0:.1f} s -> {db} ({Path(db).stat().st_size / 1e6:.1f} MB)", flush=True)
json_path.unlink(missing_ok=True)
if __name__ == "__main__":
main()
|