torch-ttnn 0.5.6

PyTorch dynamo backend for Tenstorrent TT-NN framework

PyTorch 2.0 TTNN Compiler

This project allows to run PyTorch code on Tenstorrent hardware.

Supported Models

The table below summarizes the results of running various ML models through our TTNN compiler. For each model, we track whether the run was successful, the number of operations before and after conversion, the number of to_device and from_device operations, performance metrics, and accuracy.

Model	Status	Torch Ops Before (Unique Ops)	Torch Ops Remain (Unique Ops)	To/From Device Ops	Original Run Time (ms)	Compiled Run Time for 5th Iteration (ms)	Accuracy (%)
Autoencoder (linear)	✅	22 (3)	0 (0)	0	2296.78	1.7	100.0
BERT	✅	1393 (21)	0 (0)	0	82498.1	53.45	99.53
DPR	✅	720 (22)	0 (0)	3	2849.21	12.66	99.38
HardNet	✅	245 (10)	0 (0)	124	5182.31	778.68	98.54
Llama	✅	40 (11)	0 (0)	2	163547	130500.09	100.0
MLPMixer	✅	253 (11)	0 (0)	0	5679.45	168.88	99.99
Mnist	✅	14 (8)	0 (0)	1	32.82	3.53	99.42
MobileNetV2	✅	154 (9)	0 (0)	0	1065.61	1521.88	99.53
OpenPose V2	✅	155 (7)	0 (0)	6	3880.72	705.51	91.47
Perceiver IO	✅	1531 (20)	0 (0)	1	43049.1	52.61	99.95
ResNet18	✅	70 (9)	0 (0)	1	2339.23	389.24	98.92
ResNet50	✅	176 (9)	0 (0)	1	4343.79	714.02	99.15
RoBERTa	✅	719 (21)	0 (0)	3	13536	48.42	7.19
SqueezeBERT	✅	16 (9)	0 (0)	3	1162.88	274.24	100.0
U-Net	✅	68 (6)	0 (0)	12	66608.2	249.19	100.0
Unet-brain	✅	68 (6)	0 (0)	12	68139.1	252.4	N/A
Unet-carvana	✅	67 (5)	0 (0)	12	83677.6	1377.0	99.69
YOLOv5	✅	3 (3)	0 (0)	0	23312.6	17824.75	99.98
albert/albert-base-v2	✅	791 (21)	0 (0)	3	1565.82	18.42	98.82
albert/albert-base-v2-classification	✅	779 (21)	0 (0)	3	1265.88	13.7	99.97
albert/albert-large-v2	✅	1547 (21)	0 (0)	3	2966.88	31.72	98.95
albert/albert-xlarge-v2	✅	1547 (21)	0 (0)	3	9902.67	73.9	97.36
densenet121	✅	432 (10)	0 (0)	597	3276.11	295.44	99.84
densenet161	✅	572 (10)	0 (0)	1147	8716	1234.9	99.75
densenet169	✅	600 (10)	0 (0)	1241	4081.23	513.58	99.81
densenet201	✅	712 (10)	0 (0)	1905	4970.58	714.24	99.79
distilbert-base-uncased	✅	361 (16)	0 (0)	2	1855.84	8.23	99.9
dla34.in1k	✅	135 (9)	0 (0)	23	4559.82	562.9	99.57
ese_vovnet19b_dw.ra_in1k	✅	111 (12)	0 (0)	19	1948.98	251.8	99.87
ghostnet_100.in1k	✅	515 (14)	0 (0)	64	1510.79	1200.87	99.67
mobilenet_v2	✅	154 (9)	0 (0)	0	892.12	1492.61	99.53
mobilenet_v3_large	✅	188 (11)	0 (0)	0	1126.06	2244.08	99.42
mobilenet_v3_small	✅	158 (11)	0 (0)	0	855.44	794.13	99.34
mobilenetv1_100.ra4_e3600_r224_in1k	✅	85 (7)	0 (0)	0	1267.47	820.32	96.9
regnet_x_16gf	✅	235 (8)	0 (0)	0	15239.9	9577.32	99.84
regnet_x_1_6gf	✅	195 (8)	0 (0)	0	1969.25	2027.28	99.76
regnet_x_32gf	✅	245 (8)	0 (0)	0	28777.9	15723.12	99.61
regnet_x_3_2gf	✅	265 (8)	0 (0)	0	3879.53	3262.72	99.64
regnet_x_400mf	✅	235 (8)	0 (0)	0	1093.64	1179.31	99.73
regnet_x_800mf	✅	175 (8)	0 (0)	0	1247.91	1138.99	99.8
regnet_x_8gf	✅	245 (8)	0 (0)	0	7888.55	8034.42	99.72
regnet_y_16gf	✅	303 (10)	0 (0)	0	15047.3	14997.29	99.82
regnet_y_1_6gf	✅	447 (10)	0 (0)	0	2463.28	2070.64	99.79
regnet_y_32gf	✅	335 (10)	0 (0)	0	31418.6	22493.05	99.79
regnet_y_3_2gf	✅	351 (10)	0 (0)	0	7851.62	3006.29	99.9
regnet_y_400mf	✅	271 (10)	0 (0)	0	850.57	773.18	99.72
regnet_y_800mf	✅	239 (10)	0 (0)	0	1086.88	1027.7	99.76
regnet_y_8gf	✅	287 (10)	0 (0)	0	11377.4	4947.38	99.93
resnet101	✅	346 (9)	0 (0)	1	7582.59	1054.7	99.48
resnet152	✅	516 (9)	0 (0)	1	11398.8	1586.87	99.49
resnet18	✅	70 (9)	0 (0)	1	2224	328.96	99.61
resnet34	✅	126 (9)	0 (0)	1	4381.48	581.03	99.07
resnet50	✅	176 (9)	0 (0)	1	4489.92	606.06	99.15
resnext101_32x8d	✅	346 (9)	0 (0)	1	15908.9	18850.89	99.91
resnext101_64x4d	✅	346 (9)	0 (0)	1	14271.5	12560.81	99.86
resnext50_32x4d	✅	176 (9)	0 (0)	1	4565.8	2593.35	99.88
textattack/albert-base-v2-imdb	✅	782 (22)	0 (0)	3	1291.04	13.61	100.0
tf_efficientnet_lite0.in1k	✅	149 (9)	0 (0)	5	1070.97	5160.82	99.41
tf_efficientnet_lite1.in1k	✅	194 (9)	0 (0)	5	1213.68	6695.89	99.73
tf_efficientnet_lite2.in1k	✅	194 (9)	0 (0)	5	2839.05	8474.35	99.3
twmkn9/albert-base-v2-squad2	✅	783 (23)	0 (0)	3	4573.01	16.39	99.86
vgg11	✅	33 (8)	0 (0)	5	10032.6	282.34	99.65
vgg11_bn	✅	41 (9)	0 (0)	5	11220.5	273.44	98.99
vgg13	✅	37 (8)	0 (0)	5	17295.2	306.4	99.38
vgg13_bn	✅	47 (9)	0 (0)	5	17640.8	294.08	97.35
vgg16	✅	43 (8)	0 (0)	5	23775.7	578.64	99.44
vgg16_bn	✅	56 (9)	0 (0)	5	24227.5	457.74	98.38
vgg19	✅	49 (8)	0 (0)	5	29814	594.11	99.27
vgg19_bn	✅	65 (9)	0 (0)	5	29991.1	581.39	96.97
wide_resnet101_2	✅	346 (9)	0 (0)	1	22527.8	3423.68	99.47
wide_resnet50_2	✅	176 (9)	0 (0)	1	12377	2385.8	99.31
xception71.tf_in1k	✅	393 (9)	0 (0)	0	21920.2	17733.36	99.74
Autoencoder (conv)	🚧	9 (3)	1 (1)	1	1457.1	4.08	100.0
Autoencoder (conv)-train	🚧	24 (7)	11 (4)	0	2548.74	7.32	100.0
Autoencoder (linear)-train	🚧	104 (8)	14 (2)	0	2580.73	11.73	100.0
Bloom	🚧	1405 (27)	2 (2)	0	27443.9	2759.8	98.86
CLIP	🚧	1397 (30)	7 (6)	5	7873.2	167.91	99.56
CLIP-train	🚧	3944 (44)	265 (16)	5	24051.1	2381.97	100.0
DETR	🚧	1663 (42)	31 (8)	3	111593	6006.83	46.78
DINOv2	🚧	928 (25)	16 (1)	4	19632.3	330.07	98.99
Falcon	🚧	94 (15)	5 (5)	3	64651.4	31227.61	100.0
GLPN-KITTI	🚧	2959 (26)	22 (2)	6	94236.9	53280.47	99.75
GPT-2	🚧	745 (29)	2 (2)	2	3636.86	31.2	99.98
HardNet-train	🚧	867 (21)	412 (9)	120	13062.1	9167.5	100.0
MLPMixer-train	🚧	616 (19)	100 (5)	0	15306.3	8333.47	100.0
Mnist-train	🚧	46 (15)	10 (6)	0	2757.4	36.6	100.0
MobileNetSSD	🚧	447 (28)	5 (2)	32	875.47	2227.24	38.84
OpenPose V2-train	🚧	523 (14)	246 (7)	6	13021.3	9843.26	100.0
ResNet18-train	🚧	241 (19)	121 (9)	0	5310.56	3778.48	100.0
ResNet50-train	🚧	616 (19)	318 (9)	0	14549.6	10941.05	100.0
SegFormer	🚧	676 (22)	16 (1)	4	41649.5	1533.77	99.86
SegFormer-train	🚧	1794 (36)	156 (12)	4	83453.2	36809.23	100.0
Stable Diffusion V2	🚧	29 (14)	11 (5)	4	1.54655e+06	1579326.23	99.97
U-Net-train	🚧	236 (15)	122 (8)	8	78211.8	40213.96	100.0
Unet-brain-train	🚧	236 (15)	122 (8)	8	109491	47281.93	100.0
Unet-carvana-train	🚧	232 (13)	121 (7)	8	172874	88220.53	100.0
ViLT	🚧	42 (16)	5 (4)	3	15801	13922.78	87.94
YOLOS	🚧	952 (27)	17 (2)	6	12880.8	250.03	98.6
YOLOv3	🚧	250 (7)	2 (1)	4	184931	1645.46	98.95
albert/albert-xxlarge-v2	🚧	791 (21)	24 (1)	3	18452.2	147.16	98.54
dla34.in1k-train	🚧	469 (18)	230 (8)	17	9685.39	6379.18	100.0
ese_vovnet19b_dw.ra_in1k-train	🚧	383 (25)	176 (10)	16	4874.93	3216.47	100.0
facebook/deit-base-patch16-224	🚧	685 (17)	1 (1)	2	15378	112.86	98.34
facebook/deit-base-patch16-224-train	🚧	1854 (27)	127 (8)	2	78973.5	2649.03	100.0
ghostnet_100.in1k-train	🚧	1469 (33)	562 (12)	64	1640.83	2301.0	100.0
ghostnetv2_100.in1k	🚧	683 (18)	24 (2)	68	1220.53	2970.08	99.74
ghostnetv2_100.in1k-train	🚧	2001 (39)	852 (17)	68	2308.84	4728.34	100.0
googlenet	🚧	214 (15)	1 (1)	51	2115.51	235.7	99.75
hrnet_w18.ms_aug_in1k	🚧	1209 (11)	31 (1)	0	6337.78	1083.53	99.68
hrnet_w18.ms_aug_in1k-train	🚧	3998 (21)	1973 (9)	0	12575.8	9320.98	100.0
inception_v4.tf_in1k	🚧	495 (11)	14 (1)	84	11889.4	1344.25	99.47
inception_v4.tf_in1k-train	🚧	1851 (24)	932 (11)	80	37762.6	27634.38	100.0
mixer_b16_224.goog_in21k	🚧	356 (11)	1 (1)	0	10931.1	106.64	3.65
mixer_b16_224.goog_in21k-train	🚧	959 (18)	101 (6)	0	55333.3	1575.27	100.0
mobilenetv1_100.ra4_e3600_r224_in1k-train	🚧	258 (16)	164 (7)	0	2681	2820.02	100.0
regnet_y_128gf	🚧	447 (10)	3 (1)	0	516943	161210.57	99.01
speecht5-tts	🚧	864 (23)	2 (1)	2	50498	278957.32	N/A
ssd300_vgg16	🚧	251 (26)	6 (3)	37	3401.03	1262.3	N/A
ssdlite320_mobilenet_v3_large	🚧	447 (28)	5 (2)	32	639.64	1976.3	49.52
swin_b	🚧	2492 (32)	128 (3)	525	14542.4	256.8	90.35
swin_s	🚧	2492 (32)	128 (3)	525	7382.71	214.47	89.91
swin_t	🚧	1238 (32)	68 (3)	249	4246.34	111.15	96.75
swin_v2_b	🚧	3140 (40)	200 (5)	567	15743.9	374.31	63.06
swin_v2_s	🚧	3140 (40)	200 (5)	567	10736.6	373.4	40.28
swin_v2_t	🚧	1562 (40)	104 (5)	267	5922.54	194.97	50.18
tf_efficientnet_lite0.in1k-train	🚧	452 (17)	285 (8)	5	2614.81	11672.95	100.0
tf_efficientnet_lite1.in1k-train	🚧	587 (17)	370 (8)	5	3891.64	8545.87	100.0
tf_efficientnet_lite2.in1k-train	🚧	587 (17)	370 (8)	5	4008.61	11058.79	100.0
tf_efficientnet_lite3.in1k	🚧	221 (9)	5 (1)	5	11758.2	11358.01	99.48
tf_efficientnet_lite3.in1k-train	🚧	668 (17)	426 (9)	5	6323.98	15848.7	100.0
tf_efficientnet_lite4.in1k	🚧	275 (9)	6 (1)	5	15761	17453.54	99.8
tf_efficientnet_lite4.in1k-train	🚧	830 (17)	529 (9)	5	16450.9	55780.68	100.0
vit_b_16	🚧	552 (17)	1 (1)	2	13208.1	147.77	99.54
vit_b_32	🚧	552 (17)	1 (1)	2	4765.49	141.18	98.68
vit_h_14	🚧	1452 (17)	1 (1)	2	765657	2700.56	98.59
vit_l_16	🚧	1092 (17)	1 (1)	2	44060.1	284.95	99.74
vit_l_32	🚧	1092 (17)	1 (1)	2	17469.1	219.1	98.97
xception71.tf_in1k-train	🚧	1378 (18)	806 (7)	0	54065.4	54359.32	100.0
FLAN-T5	❌	20020 (38)	N/A	N/A	2962.34	N/A	N/A
GPTNeo	❌	2733 (35)	N/A	N/A	13438.6	N/A	N/A
OPT	❌	4003 (32)	N/A	N/A	21590.4	N/A	N/A
Whisper	❌	4310 (21)	N/A	N/A	207323	N/A	N/A
codegen	❌	9183 (37)	N/A	N/A	8754.13	N/A	N/A
t5-base	❌	14681 (38)	N/A	N/A	4896.86	N/A	N/A
t5-large	❌	22696 (38)	N/A	N/A	11192.6	N/A	N/A
t5-small	❌	6118 (38)	N/A	N/A	2197.39	N/A	N/A

Explanation of Metrics

Model: Name of the model.
Status: Indicates whether the model is ❌ traced / 🚧 compiled / ✅ E2E on device.
Torch Ops Before (Unique Ops): The total number of operations used by the model in the original Torch implementation. The number in parenthesis represents the total unique ops.
Torch Ops Remain (Unique Ops): The total number of operations used after conversion to TTNN. The number in parenthesis represents the total unique ops.
To/From Device Ops: The number of to/from_device operations (data transfer to/from the device).
Original Run Time (ms): Execution time (in seconds) of the model before conversion.
Compiled Run Time for 5th Iteration (ms): Execution time (in seconds) of the model after conversion for the 5th iteration.
Accuracy (%): Model accuracy on a predefined test dataset after conversion.

---

Quickstart

The torch_ttnn module has a backend function, which can be used with the torch.compile().

import torch
import torch_ttnn

# A torch Module
class FooModule(torch.nn.Module):
    ...
# Create a module
module = FooModule()

# Compile the module, with ttnn backend
device = ttnn.open_device(device_id=0)
option = torch_ttnn.TorchTtnnOption(device=self.device)
ttnn_module = torch.compile(module, backend=torch_ttnn.backend, options=option)

# Running inference / training
ttnn_module(input_data)

Tracer

The tracer dump the information of fx graph such as node's op_name and shape.

For example, you can run this script to parse the information

PYTHONPATH=$(pwd) python3 tools/stat_models.py --trace_orig --backward --profile
ls stat/raw

By default, the raw result will be stored at stat/raw, and you can run this script to generate the report

python3 tools/generate_report.py
ls stat/

Now the stat/ folder have these report

• fw_node_count.csv
• bw_node_count.csv
• fw_total_input_size_dist/
• bw_total_input_size_dist/
• fw_total_output_size_dist/
• bw_total_output_size_dist/
• profile/

The node_count.csv show the node with op_type appear in the fx graph. This report can help analyze the frequency of op type appear in the graph.

The *_total_*_size_dist/ statistics the op_type's input/output_size distribution from all fx graph recored in stat/raw. This report can help analyze the memory footprint durning the calculation of op_type.

• Notice: the default input_shapes in tools/stat_torchvision.py is [1,3,224,224], which has dependency with *_total_*_size_dist/ report.

• Notice: the aten ir interface is in there

The profile/ is the tools provided by pytorch, you can open it by the url: chrome://tracing

For developers

Install torch-ttnn with editable mode

During development, you may want to use the torch-ttnn package for testing. In order to do that, you can install the torch-ttnn package in "editable" mode with

pip install -e .

Now, you can utilize torch_ttnn in your Python code. Any modifications you make to the torch_ttnn package will take effect immediately, eliminating the need for constant reinstallation via pip.

Build wheel file

For developers want to deploy the wheel, you can build the wheel file with

python -m build

Then you can upload the .whl file to the PyPI (Python Package Index).

Run transformer models

To run transformer model with ttnn backend, run:

PYTHONPATH="$TT_METAL_HOME:$(pwd)" python3 tools/run_transformers.py --model "phiyodr/bert-large-finetuned-squad2" --backend torch_ttnn

You can also substitute the backend with torch_stat to run a reference comparison.