tt-torch

tt-torch is a PyTorch2.0 and torch-mlir based front-end for tt-mlir.

tt-torch uses venv to keep track of all dependencies. After compiling you can activate the venv by running from the project root directory:

source env/activate

The currently supported models can be found here. There is a brief demo showing how to use the compiler in demos/resnet/resnet50_demo.py

The general compile flow is:

1. Pytorch model -> torch.compile which creates an fx graph
2. Several compiler passes on the fx graph including consteval and dead code removal
3. Conversion to torch-mlir -> torch-backend-mlir -> stableHLO through torch-mlir
4. Conversion to TTIR -> TTNN -> flatbuffer through tt-mlir
5. Creating executor with flatbuffer and passing back to user
6. Copying inputs to device and executing flatbuffer through tt-mlir on each user invocation

In order to speed up model bring-up, users have the option of compiling models op-by-op. This allows in-parallel testing of the model since compilation does not stop at the first error. If enabled, see Controlling Compilation, after step 2, compilation stops and the fx graph is passed to the executor which is returned to the user. Upon execution, whenever a new, unique op is seen (based on op-type and shape on inputs), a new fx graph is created with just one operation, inputs and outputs. This small graph then proceeds through steps 3-4 and is executed in place.

Results of each unique op execution are stored in a json file to be later parsed into either a spreadsheet, or uploaded to a database.

Op-by-op execution is currently performed on the pytorch fx graph, we'll be adding support for op-by-op on the stableHLO graph soon to allow op-by-op bringup of onnx models.

The repository uses pre-commit, read more about it here.

Getting Started

System Dependencies

tt-torch requires the python 3.11 dev package, as well as the venv package. If not already installed, please run the following:

sudo apt-get install python3.11-dev python3.11-venv

Creating a Virtual Environment (skip if you already have one)

Create a virtual environment if you do not already have one in your project:

python3.11 -m venv myvenv

This will create a virtual environemnt in the folder myvenv in the current directory.

Activate the environemnt:

source myvenv/bin/activate

Installing tt-torch

Installation Notes

• tt-torch requires a pytorch installation that ships with their ABI.
  • The tt-torch wheel lists the following version of torch as an installation requirement: torch@https://download.pytorch.org/whl/cpu-cxx11-abi/torch-2.5.0%2Bcpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
  • This will be installed by pip upon installing the tt-torch wheel
• The tt-torch wheel contains a fork of torch-mlir. Please ensure that torch-mlir has not been installed in your venv before installing the tt-torch wheel.

Torchvision Install (Required if you need to install torchvision)

If you intend to use torchvision in your project then this step must be done before installing the tt-torch wheel

You will need to build the torchvision wheel yourself with certain build flags. This is because torchvision does not publish a wheel which uses the PyTorch CXX11 ABI.

To install torchvision:

git clone https://github.com/pytorch/vision.git
cd vision
git checkout v0.20.0 # tt-torch requires PyTorch 2.5.0. torchvision 0.20 is the latest version of torchvision that is compatible with PyTorch 2.5.0
pip uninstall -y torchvision # Ensure torchvision is not in your virtual environment
pip install wheel
pip install torch@https://download.pytorch.org/whl/cpu-cxx11-abi/torch-2.5.0%2Bcpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
TORCHVISION_USE_VIDEO_CODEC=0 TORCHVISION_USE_FFMPEG=0 _GLIBCXX_USE_CXX11_ABI=1 USE_CUDA=OFF python setup.py bdist_wheel
pip install dist/torchvision*.whl --force-reinstall

If the install was successful then there's no need to keep the torchvision source around:

cd ..
rm -rf vision

Installing the tt-torch wheel

Download a tt-torch wheel from here

Install the wheel:

pip install <PATH_TO_TT_TORCH_WHEEL>.whl

Updating PYTHONPATH

In addition to the tt-torch python library that gets installed in <YOUR_ENV_ROOT>/lib/python3.x/site-packages, some binaries will be installed in <YOUR_ENV_ROOT>/lib, and some files from tt-metal will be installed under <YOUR_ENV_ROOT>/tt-metal. Python needs to see these installations and so you should update your PYTHONPATH environment variable to include them:

export PYTHONPATH=$PYTHONPATH:<YOUR_ENV_ROOT>:<YOUR_ENV_ROOT>/lib

Compiling and Running a Model

Once you have your torch.nn.Module compile the model:

from tt_torch.dynamo.backend import backend
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        ...

    def foward(self, ...):
        ...

model = MyModel()

model = torch.compile(model, backend=backend)

inputs = ...

outputs = model(inputs)

Example - Add Two Tensors

Here is an exampe of a small model which adds its inputs running through tt-torch. Try it out!

from tt_torch.dynamo.backend import backend
import torch

class AddTensors(torch.nn.Module):
  def forward(self, x, y):
    return x + y


model = AddTensors()
tt_model = torch.compile(model, backend=backend)

x = torch.ones(5, 5)
y = torch.ones(5, 5)
print(tt_model(x, y))

Prerequisites:

Main project dependencies are:

• clang 17
• Ninja
• CMake >= 3.30
• git LFS
• python 3.11

On Ubuntu 22.04 systems these can be installed using the following commands:

# Update package list
sudo apt update -y
sudo apt upgrade -y

# Install Clang
sudo apt install clang-17

# Install Ninja
sudo apt install ninja-build

# Install Git LFS
sudo apt install git-lfs

# Install CMake
sudo apt remove cmake -y
pip3 install cmake --upgrade

Ensure cmake can by found in this path pip installed it to. E.g. by adding PATH=$PATH:$HOME/.local/bin to your .bashrc file, and verify installation:

cmake --version

This project requires the GCC 11 toolchain. To check which GCC toolchain is currently in use, run:

clang -v

Look for the line that starts with: Selected GCC installation:. If it is something other than GCC 11, please uninstall that and install GCC 11 using:

sudo apt-get install gcc-11 lib32stdc++-11-dev lib32gcc-11-dev

The project also requires a toolchain build. By default, the toolchain is built in /opt/ttmlir-toolchain. This path is controlled by the TTMLIR_TOOLCHAIN_DIR environment variable.

The toolchain installation only needs to be done once, by running the following commands:

# Create toolchain dir
sudo mkdir -p /opt/ttmlir-toolchain
sudo chown -R $USER /opt/ttmlir-toolchain


# Build environment
cd third_party
export TTMLIR_TOOLCHAIN_DIR=/opt/ttmlir-toolchain/
cmake -B toolchain -DBUILD_TOOLCHAIN=ON
cd -

For more information see tt-mlir build steps.

Compile Steps:

Run the following commands to compile:

source env/activate
cmake -G Ninja -B build
cmake --build build
cmake --install build

Run a basic test to verify:

pytest tests/torch/test_basic.py

tt-torch uses pytest for all unit and model tests.

Tests are organized into unit tests for pytorch (tests/torch), unit tests for onnx (test/onns) and models (tests/models). They can be run locally by running:

source env/activate
pytest -svv tests/torch

Model tests (tests/models) have the option to run op-by-op, see overview. This allows for faster model bring-up as it allows users to find any potential issues in parallel. This is controlled by the --op_by_op_torch or --op_by_op_stablehlo flags. Example:

pytest -svv tests/models/albert --op_by_op_torch

Controlling Compiler Behaviour

You can use the following environment variables to override default behaviour:

Controlling Compiler Behaviour Programatically

Instead of using the above environment variables, compiler behaviour can be configured programatically as well.

Here is an example of enabling consteval:

from tt_torch.dynamo.backend import backend
from tt_torch.tools.utils import CompilerConfig
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        ...

    def foward(self, ...):
        ...

model = MyModel()

cc = CompilerConfig()
cc.enable_consteval = True
cc.consteval_parameters = True # This will enable constant folding on the parameters in addition to any constants

model = torch.compile(model, backend=backend, options=cc)

inputs = ...

outputs = model(inputs)

Pre-Commit

Pre-Commit applies a Git hook to the local repository, ensuring linting is checked and applied on every git commit action. Install it from the root of the repository using:

source env/activate
pre-commit install

If you have already made commits before installing the pre-commit hooks, you can run the following to “catch up”:

pre-commit run --all-files

For more information visit pre-commit

Supported Models

The following models can be currently run through tt-torch as of Feb 3rd, 2025. Please note, there is a known bug causing incorrect output for some models. The PCC is displayed at the end of each test below. This issue will be addressed soon.

Ops Documentation

This section contains documentation for Ops operations.

Stablehlo Documentation

This section contains documentation for Stablehlo operations.

arith.constant

stablehlo.add::ttnn.add

stablehlo.and::ttnn.and

stablehlo.broadcast_in_dim