This repository contains a high-performance training framework developed in C++ designed to efficiently leverage the computational capabilities of Tenstorrent hardware. The framework is optimized to accelerate model training tasks, providing a seamless interface for training deep learning models on Tenstorrent's advanced hardware architecture.
tt-train is built as part of tt-metal. Before building, ensure you have:
- Cloned the tt-metal repository with submodules:
git submodule update --init --recursive- Followed the tt-metal setup instructions in the main repository README.
tt-train is built using the build_metal.sh script from the tt-metal root directory.
# Release build (default)
./build_metal.sh --build-tt-train
# Debug build
./build_metal.sh --build-tt-train --debug
# Development build (RelWithDebInfo)
./build_metal.sh --build-tt-train --development- Open the tt-metal root directory in VSCode
- Install the CMake extension
- Set
BUILD_TT_TRAIN=ONin your CMake configure settings - Build all targets using the CMake extension
# Navigate to the root directory of the repository
./build/sources/examples/mnist_mlp/mnist_mlp --model_path mnist_mlp.msgpack --num_epochs 10
# Navigate to the root directory of the repository
./build/sources/examples/mnist_mlp/mnist_mlp --model_path mnist_mlp.msgpack -e 1
# Navigate to the root directory of the repository
TT_LOGGER_LEVEL=FATAL ./build/sources/examples/nano_gpt/nano_gpt
Training loss example from wandb project:
For more information on training with the C++ version, please see the NanoGPT Example README.
For both training and evaluation, we encourage you to use the Python version of NanoGPT, which can be used via Jupyter Notebook found in the NanoGPT example directory.
More information on available configuration options can be found in the configs directory.
If CI fails, but local tests pass as expected, please consider changing the
is_nightly_tt_train_tests_enabled in the nano_gpt_test.cpp
TT-Train nightly tests are all tests with "NIGHTLY_" in the name.
To run it in github please search for Nightly tt-metal L2 tests.
If you don't have an account to wandb (or don't want to use it), use -w 0 argument or run wandb offline beforehand (creates wandb/settings file)
This repository can be used to compare speed and accuracy to GPU implementation of nanoGPT and gpt2s
Use of profiler requires additional setup. Follow instructions here.
- Create a new branch.
- Make your changes and commit them.
- Add new tests and run existing ones
- Open a pull request (PR).
- Ensure the PR is approved by at least one code owner before merging.