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:

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