Visit TT-Forge for our latest compiler project

Pytorch 2.0 TT-NN is no longer maintained, please consider using TT-Forge instead.

The PyTorch 2.0 TT-NN Compiler enables seamless execution of PyTorch models on Tenstorrent AI accelerators. By leveraging the TT-NN backend, you can achieve significant performance improvements while maintaining PyTorch's familiar API.

🚀 Quick Start

Installation

Install from the repo:

pip install git+https://bitbucket.org/tenstorrent/pytorch2.0_ttnn

or as an editable package from source:

git clone https://github.com/tenstorrent/pytorch2.0_ttnn.git
cd pytorch2.0_ttnn
pip install -e .

✨ Basic Usage

Option 1: Eager Mode: get your model running by switching to a TT device

import torch
import torch_ttnn

model = YourModel()

device = ttnn.open_device(device_id=0)
model.to(torch_ttnn.ttnn_device_as_torch_device(device))

output = model(input_data)

Option 2: Compilation Mode (Recommended): get more perf with a JIT compiler

import torch
import torch_ttnn

model = YourModel()

device = ttnn.open_mesh_device(ttnn.MeshShape(1, 2))  # 1x2 device grid
option = torch_ttnn.TorchTtnnOption(device=device, data_parallel=2)

model = torch.compile(model, backend=torch_ttnn.backend, options=option)
output = model(input_data)

📊 Model Support

We've extensively tested the compiler across a diverse range of model architectures. Here's a summary of our validation results:

Model	Status	Batch	Compiled First Run (ms)	Original Throughput (Inferences Per Second)	Compiled Throughput (Inferences Per Second)	Accuracy (%)	Torch Ops Before (Unique Ops)	Torch Ops Remain (Unique Ops)	To/From Device Ops
Autoencoder (linear)	✅	1	311.22	0.469125	431.0344827586207	100.0	22 (3)	0 (0)	0
BERT	✅	8	46304.71	0.115219	24.906600249066006	99.67	987 (19)	0 (0)	0
DPR	✅	1	17436.8	0.538315	40.30632809351068	99.35	518 (20)	0 (0)	1
HardNet	✅	1	67585.85	0.202277	19.029495718363464	98.42	245 (10)	0 (0)	124
Mnist	✅	1	9216.93	29.1971	327.8688524590164	99.42	14 (8)	0 (0)	1
MobileNetV2	✅	1	130241.5	1.08704	28.457598178713713	98.62	154 (9)	0 (0)	0
OpenPose V2	✅	1	23585.55	0.277855	33.145508783559826	91.08	155 (7)	0 (0)	6
Perceiver IO	✅	1	56269.31	0.0951014	18.821757952192733	99.94	1531 (20)	0 (0)	1
ResNet18	✅	1	24086.48	0.457216	67.6132521974307	98.78	70 (9)	0 (0)	1
ResNet50	✅	4	90657.23	0.758689	56.274620146314014	98.43	176 (9)	0 (0)	1
RoBERTa	✅	1	38704.76	0.246022	20.30456852791878	13.29	517 (20)	0 (0)	2
U-Net	✅	1	38981.53	0.0158235	55.52470849528039	100.0	68 (6)	0 (0)	12
Unet-brain	✅	1	37742.61	0.0170157	54.76451259583789	N/A	68 (6)	0 (0)	12
Unet-carvana	✅	1	102409.37	0.0122212	35.149384885764505	99.43	67 (5)	0 (0)	12
albert/albert-base-v2	✅	1	37930.94	1.03405	33.288948069241016	98.36	591 (19)	0 (0)	2
albert/albert-base-v2-classification	✅	1	8825.88	1.33477	36.64345914254306	99.97	579 (19)	0 (0)	2
albert/albert-large-v2	✅	1	17674.75	0.874439	19.29384526336099	97.09	1143 (19)	0 (0)	2
albert/albert-xlarge-v2	✅	1	44929.23	0.567077	15.95150741745095	96.85	1143 (19)	0 (0)	2
densenet121	✅	1	257260.77	0.269953	12.256403971074887	99.54	432 (10)	0 (0)	597
densenet161	✅	1	262930.54	0.100993	8.577064928381507	99.22	572 (10)	0 (0)	1147
densenet169	✅	1	77540.64	0.228431	8.143985666585225	99.45	600 (10)	0 (0)	1241
densenet201	✅	1	407171.93	0.207018	6.69478476266988	99.01	712 (10)	0 (0)	1905
distilbert-base-uncased	✅	1	31281.24	1.08541	78.92659826361484	99.9	250 (16)	0 (0)	1
dla34.in1k	✅	1	127568.5	0.270627	37.147102526002975	99.32	135 (9)	0 (0)	23
ese_vovnet19b_dw.ra_in1k	✅	1	100735.87	0.516273	41.56275976724855	99.18	111 (12)	0 (0)	19
ghostnet_100.in1k	✅	1	229980.19	1.53139	17.26817475392851	99.32	515 (14)	0 (0)	64
mobilenet_v2	✅	1	104479.28	1.04519	31.03662321539417	98.62	154 (9)	0 (0)	0
mobilenet_v3_large	✅	1	203098.39	1.26883	28.01905295601009	99.27	188 (11)	0 (0)	0
mobilenet_v3_small	✅	1	144514.17	1.94341	32.8515111695138	99.0	158 (11)	0 (0)	0
mobilenetv1_100.ra4_e3600_r224_in1k	✅	1	96973.84	0.770048	47.08097928436912	95.43	85 (7)	0 (0)	0
regnet_x_16gf	✅	1	71873.68	0.0681221	14.744913005013272	99.32	235 (8)	0 (0)	0
regnet_x_1_6gf	✅	1	72529.83	0.539904	26.116479498563596	99.33	195 (8)	0 (0)	0
regnet_x_32gf	✅	1	84967.4	0.0335761	9.022827754218174	98.76	245 (8)	0 (0)	0
regnet_x_3_2gf	✅	1	50834.01	0.293084	20.712510356255176	99.15	265 (8)	0 (0)	0
regnet_x_400mf	✅	1	70402.6	1.08918	24.582104228121928	99.28	235 (8)	0 (0)	0
regnet_x_800mf	✅	1	105997.17	0.85068	31.71582619727244	99.19	175 (8)	0 (0)	0
regnet_x_8gf	✅	1	111156.08	0.131345	21.55636990730761	98.52	245 (8)	0 (0)	0
regnet_y_16gf	✅	1	73889.04	0.0659471	9.965122072745391	99.5	303 (10)	0 (0)	0
regnet_y_1_6gf	✅	1	78875.58	0.323209	13.968431345159939	99.33	447 (10)	0 (0)	0
regnet_y_32gf	✅	1	184387.66	0.03089	8.207485226526591	99.17	335 (10)	0 (0)	0
regnet_y_3_2gf	✅	1	162774.45	0.29579	17.59943681802182	99.64	351 (10)	0 (0)	0
regnet_y_400mf	✅	1	152323.33	1.00028	23.353573096683792	99.26	271 (10)	0 (0)	0
regnet_y_800mf	✅	1	97060.58	0.737485	24.95632642874969	99.46	239 (10)	0 (0)	0
regnet_y_8gf	✅	1	114078.86	0.120364	16.181229773462785	99.72	287 (10)	0 (0)	0
resnet101	✅	1	12076.63	0.132859	15.865460891638904	98.96	346 (9)	0 (0)	1
resnet152	✅	1	18440.05	0.0923735	10.912265386294195	97.94	516 (9)	0 (0)	1
resnet18	✅	1	36033.95	0.402478	68.91798759476224	99.44	70 (9)	0 (0)	1
resnet34	✅	1	4727.9	0.229456	40.65040650406504	98.51	126 (9)	0 (0)	1
resnet50	✅	1	110324.24	0.230154	31.665611146295124	98.49	176 (9)	0 (0)	1
resnext101_32x8d	✅	1	45034.76	0.0648038	8.13140348024069	99.31	346 (9)	0 (0)	1
resnext101_64x4d	✅	1	19511.79	0.0672185	8.165264962848045	99.45	346 (9)	0 (0)	1
resnext50_32x4d	✅	1	107909.87	0.223939	29.949086552860138	98.92	176 (9)	0 (0)	1
textattack/albert-base-v2-imdb	✅	1	43008.15	0.952381	38.299502106472616	100.0	582 (20)	0 (0)	2
tf_efficientnet_lite0.in1k	✅	1	163827.81	0.832016	26.98327037236913	98.94	149 (9)	0 (0)	5
tf_efficientnet_lite1.in1k	✅	1	103670.21	0.758374	20.024028834601523	99.15	194 (9)	0 (0)	5
tf_efficientnet_lite2.in1k	✅	1	143822.51	0.612329	18.178512997636794	98.74	194 (9)	0 (0)	5
twmkn9/albert-base-v2-squad2	✅	1	23010.63	0.815082	34.08316291751875	99.88	583 (21)	0 (0)	2
vgg11	✅	1	73151.64	0.0843261	99.8003992015968	99.5	33 (8)	0 (0)	5
vgg11_bn	✅	1	7317.34	0.105333	88.49557522123894	98.81	41 (9)	0 (0)	5
vgg13	✅	1	6149.23	0.0541703	82.50825082508251	99.29	37 (8)	0 (0)	5
vgg13_bn	✅	1	83132.14	0.0524954	73.6377025036819	97.29	47 (9)	0 (0)	5
vgg16	✅	1	2179.87	0.0389061	74.62686567164178	99.54	43 (8)	0 (0)	5
vgg16_bn	✅	1	3153.12	0.039009	64.8508430609598	98.17	56 (9)	0 (0)	5
vgg19	✅	1	77033.66	0.0295933	61.957868649318456	99.2	49 (8)	0 (0)	5
vgg19_bn	✅	1	8779.61	0.029195	58.07200929152149	95.73	65 (9)	0 (0)	5
wide_resnet101_2	✅	1	13378.94	0.0441978	16.05136436597111	98.82	346 (9)	0 (0)	1
wide_resnet50_2	✅	1	103416.34	0.0812497	27.816411682892905	98.48	176 (9)	0 (0)	1
xception71.tf_in1k	✅	1	148902.58	0.0480272	7.127583749109052	98.63	393 (9)	0 (0)	0
Autoencoder (conv)	🚧	1	6269.5	0.628366	265.2519893899204	100.0	9 (3)	1 (1)	1
Autoencoder (conv)-train	🚧	1	24091.86	0.289166	135.1351351351351	100.0	30 (7)	11 (4)	0
Autoencoder (linear)-train	🚧	1	17176.17	0.474886	73.74631268436578	100.0	116 (8)	14 (2)	0
Bloom	🚧	1	40360.68	0.179315	10.982976386600768	98.74	1512 (30)	1 (1)	0
CLIP	🚧	1	69209.04	0.276222	6.0698027314112295	99.55	1161 (29)	7 (6)	2
CLIP-train	🚧	1	80908.66	0.082211	0.7799886121662625	100.0	3352 (42)	291 (17)	5
DETR	🚧	1	192804.27	0.00969074	0.1869347560314499	89.95	1672 (41)	9 (6)	3
DINOv2	🚧	1	32883.37	0.259654	18.01801801801802	98.83	696 (24)	20 (2)	2
GLPN-KITTI	🚧	1	294036.88	0.0100672	0.01693750872281699	99.72	2949 (26)	22 (2)	6
GPT-2	🚧	1	34722.02	0.331807	33.145508783559826	99.97	565 (29)	3 (3)	2
GaussianSplatting	🚧	1	48105.78	0.132922	0.0929294618919509	49.71	489 (35)	39 (4)	4
GaussianSplatting-train	🚧	1	61072.96	0.0765136	0.024651933192274975	43.7	1685 (53)	207 (13)	8
HardNet-train	🚧	1	195041.98	0.150316	0.4870208931963181	100.0	800 (21)	278 (8)	120
Mnist-train	🚧	1	28705.96	0.454548	37.622272385252074	100.0	54 (15)	10 (6)	0
MobileNetSSD	🚧	1	332164.54	1.62715	0.7590997077466126	45.92	520 (30)	7 (4)	32
OpenPose V2-train	🚧	1	90531.26	0.211543	0.721807984639926	100.0	490 (14)	180 (6)	6
ResNet18-train	🚧	1	59470.42	0.353551	1.2256403971074887	100.0	221 (19)	81 (8)	0
ResNet50-train	🚧	1	102246.89	0.156891	0.5537865152983525	100.0	563 (19)	212 (8)	0
SegFormer	🚧	1	30064.14	0.0241822	3.5042225882188034	99.87	680 (21)	16 (1)	4
SegFormer-train	🚧	1	190721.05	0.0211994	0.2452783909737552	100.0	1829 (37)	155 (11)	4
Stable Diffusion V2	🚧	1	434509.73	0.00078495	0.0040985143500258305	99.59	1738 (26)	71 (3)	28
U-Net-train	🚧	1	79629.48	0.0144585	0.1526083824732325	100.0	220 (15)	86 (7)	8
Unet-brain-train	🚧	1	78777.06	0.0126975	0.1598782367349027	100.0	220 (15)	86 (7)	8
Unet-carvana-train	🚧	1	80383.19	0.0101184	0.0594745306270043	100.0	214 (13)	85 (6)	8
YOLOS	🚧	1	45630.26	0.497787	3.7851546235663727	98.22	721 (27)	21 (3)	2
YOLOv3	🚧	1	92599.28	0.00502589	17.596339961288052	98.23	250 (7)	2 (1)	4
albert/albert-xxlarge-v2	🚧	1	17928.08	0.309049	8.399126490844951	98.26	591 (19)	24 (1)	2
dla34.in1k-train	🚧	1	92803.86	0.22342	0.7613594834937264	100.0	432 (18)	156 (7)	17
ese_vovnet19b_dw.ra_in1k-train	🚧	1	96678.65	0.405785	1.2328632014991616	100.0	360 (25)	130 (9)	16
facebook/deit-base-patch16-224	🚧	1	24527.93	0.338954	9.260974254491572	98.37	541 (15)	1 (1)	1
facebook/deit-base-patch16-224-train	🚧	1	29755.52	0.0303366	0.9905795881170072	100.0	1518 (24)	139 (9)	2
ghostnet_100.in1k-train	🚧	1	260290.62	0.801892	0.7104038645970234	100.0	1439 (29)	419 (10)	64
ghostnetv2_100.in1k	🚧	1	343874.91	0.838195	8.410428931875526	99.41	683 (18)	24 (2)	68
ghostnetv2_100.in1k-train	🚧	1	95658.66	0.584994	0.3513024538476401	100.0	1957 (36)	625 (16)	68
googlenet	🚧	1	139609.89	0.522253	22.73760800363802	99.4	214 (15)	1 (1)	51
hrnet_w18.ms_aug_in1k	🚧	1	159902.02	0.190664	4.422430567840085	99.56	1209 (11)	31 (1)	0
hrnet_w18.ms_aug_in1k-train	🚧	1	161973.76	0.15064	0.3487747542881856	100.0	3757 (21)	1323 (8)	0
inception_v4.tf_in1k	🚧	1	168391.54	0.0724696	5.946010227137591	98.67	495 (11)	14 (1)	84
inception_v4.tf_in1k-train	🚧	1	158333.71	0.0555337	0.20988517182249591	100.0	1702 (24)	634 (10)	80
mixer_b16_224.goog_in21k	🚧	1	18465.33	0.263587	7.621370322383964	0.32	392 (12)	1 (1)	0
mixer_b16_224.goog_in21k-train	🚧	1	40335.93	0.0370841	1.1090655014085133	100.0	959 (19)	101 (6)	0
mobilenetv1_100.ra4_e3600_r224_in1k-train	🚧	1	81606.62	0.76365	1.0999164063531173	100.0	231 (15)	110 (6)	0
regnet_y_128gf	🚧	1	298255.02	0.00185385	0.01669818170159815	98.68	447 (10)	3 (1)	0
ssd300_vgg16	🚧	1	255629.89	0.276845	1.3041040153362633	N/A	330 (29)	8 (5)	37
ssdlite320_mobilenet_v3_large	🚧	1	272755.24	1.62251	0.7410096998169706	47.01	520 (30)	7 (4)	32
swin_b	🚧	1	110297.89	0.226322	4.1769349651225935	99.38	2492 (32)	110 (2)	479
swin_s	🚧	1	19245.31	0.369505	4.457917261055635	99.58	2492 (32)	110 (2)	479
swin_t	🚧	1	192988.43	0.589352	8.39137366786943	99.6	1238 (32)	50 (2)	227
swin_v2_b	🚧	1	107336.84	0.198254	3.4696922382984634	99.9	3164 (40)	158 (3)	473
swin_v2_s	🚧	1	22575.38	0.295119	3.552902721523485	99.65	3164 (40)	158 (3)	473
swin_v2_t	🚧	1	187456.62	0.472916	6.518904823989569	99.46	1574 (40)	74 (3)	221
tf_efficientnet_lite0.in1k-train	🚧	1	147764.65	0.553131	0.22743865410902045	100.0	403 (17)	187 (7)	5
tf_efficientnet_lite1.in1k-train	🚧	1	101632.13	0.475971	0.2175445313655705	100.0	523 (17)	242 (7)	5
tf_efficientnet_lite2.in1k-train	🚧	1	161603.3	0.384084	0.131343220850999	100.0	523 (17)	242 (7)	5
tf_efficientnet_lite3.in1k	🚧	1	157894.51	0.378362	3.732318142798492	98.96	221 (9)	5 (1)	5
tf_efficientnet_lite3.in1k-train	🚧	1	128574.91	0.263609	0.06615694942367373	100.0	595 (17)	280 (8)	5
tf_efficientnet_lite4.in1k	🚧	1	148212.38	0.201735	2.2092612230470134	99.09	275 (9)	6 (1)	5
tf_efficientnet_lite4.in1k-train	🚧	1	142411.39	0.137126	0.06980447069713026	100.0	739 (17)	347 (8)	5
vit_b_16	🚧	1	18571.82	0.203005	1.0204914686913218	99.51	216 (14)	13 (2)	1
vit_b_32	🚧	1	17291.17	0.4411	3.5888601780074647	98.39	216 (14)	13 (2)	1
vit_h_14	🚧	1	382279.74	0.00709972	0.011202390679790192	95.81	556 (14)	33 (2)	1
vit_l_16	🚧	1	25280.47	0.0995329	0.3911781503532339	99.51	420 (14)	25 (2)	1
vit_l_32	🚧	1	12820.28	0.149968	1.5353200374618088	98.84	420 (14)	25 (2)	1
xception71.tf_in1k-train	🚧	1	155176.6	0.0412191	0.06015638253203027	100.0	1276 (18)	514 (6)	0
FLAN-T5	❌	N/A	N/A	0.390116	N/A	N/A	4182 (42)	N/A	N/A
Falcon-7B	❌	N/A	N/A	0.0316554	N/A	N/A	2623 (31)	N/A	N/A
GPTNeo	❌	N/A	N/A	0.262519	N/A	N/A	2895 (44)	N/A	N/A
Llama	❌	N/A	N/A	0.0219249	N/A	N/A	2704 (29)	N/A	N/A
OPT	❌	N/A	N/A	0.205786	N/A	N/A	3419 (34)	N/A	N/A
ViLT	❌	N/A	N/A	0.22532	N/A	N/A	766 (29)	N/A	N/A
Whisper	❌	N/A	N/A	0.028717	N/A	N/A	4842 (43)	N/A	N/A
YOLOv5	❌	N/A	N/A	0.0533503	N/A	N/A	228 (13)	N/A	N/A
codegen	❌	N/A	N/A	0.249193	N/A	N/A	9279 (42)	N/A	N/A
speecht5-tts	❌	N/A	N/A	0.0524124	N/A	N/A	2725 (38)	N/A	N/A
t5-base	❌	N/A	N/A	0.197262	N/A	N/A	5569 (44)	N/A	N/A
t5-large	❌	N/A	N/A	0.157036	N/A	N/A	10837 (44)	N/A	N/A
t5-small	❌	N/A	N/A	0.501487	N/A	N/A	2935 (44)	N/A	N/A

Explanation of Metrics

Model: Name of the model.
Status: Indicates whether the model is:

• ✅ End-to-end on device: All PyTorch operations have been converted to TT-NN operations.
• 🚧 Compiled: The converted model runs but some operations still fallback to PyTorch. This may be due to an unsupported operation or configuration.
• ❌ Traced: The model does not run but its PyTorch operations are traced for future development. This may indicate a temporary incompatibility with a compiler pass.
  Batch: Batch size used for inference
  Compiled First Run (ms): Time until the first compiled run finishes (ms), including compilation time and warming caches.
  Original Throughput (Inferences Per Second): Execution throughput (in inferences per second) of the model before conversion.
  Compiled Throughput (Inferences Per Second): Execution throughput (in inferences per second) of the model after conversion, once caches are warm.
  Accuracy (%): Model accuracy on a predefined test dataset after conversion.
  Torch Ops Before (Unique Ops): The total number of operations used by the model in the original Torch implementation. The number in parentheses represents the total unique ops.
  Torch Ops Remain (Unique Ops): The total number of operations used after conversion to TT-NN. The number in parentheses represents the total unique ops.
  To/From Device Ops: The number of to/from_device operations (data transfer to/from the device).

---

Contributing

Whether you are new to Tenstorrent hardware or an experienced developer, there are many ways to contribute.

Getting Started

Start with our high level Contribution guide. You can find more information here:

• Discussions
• Operations Report
• Lowering TT-NN Operation to PyTorch
• Native Device Integration Extension
• Build with Metal from Source
• Known Issues
• Problem Solving

We encourage contributions and offer 🤑 Bounties for some issues.

Development Environment

To get started with development, you'll need a Wormhole or Blackhole Tenstorrent accelerator card, which:

• can be ordered on the Tenstorrent website
• can be requested on Koyeb

Install the development dependencies:

pip install -r requirements-dev.txt
pip install -e .

You can build the wheel file with

python -m build

Project Structure

• torch_ttnn/: Main package directory containing the core implementation
• tests/: Test files for the project including model suites. We use pytest as our testing framework.
• tools/: Development and utility scripts
• docs/: Project documentation and reports
• demo/: Example code and usage demonstrations

Questions and Support

If you have questions or need help getting started, please:

1. Review the existing documentation
2. Ask PyTorch TT-NN DeepWiki or TT-Metal DeepWiki
3. Ask on Discord
4. Open an issue on GitHub