transfusion-p150 / code /scripts /visualizer_report.py

tt-model push transfusion-p150 (container)

68a3286 verified 11 days ago

5.11 kB

#!/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()

	#!/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()