TensorZero Recipe: Supervised Fine-Tuning with torchtune

The torchtune.ipynb notebook provides a step-by-step recipe to perform supervised fine-tuning of models using torchtune based on data collected by the TensorZero Gateway.

You will need to set a few environment variables in the shell your notebook will run in.

Set TENSORZERO_CLICKHOUSE_URL=http://chuser:chpassword@localhost:8123/tensorzero.
Set HF_TOKEN=<your-hf-token> to your huggingface token to use models like Llama and Gemma.
Set CHECKPOINT_HOME=</path/to/store/large/models> as the directory to save models downloaded from huggingface.
Install the CLI tool firectl on your machine and sign in with firectl signin. You can test that this all worked with firectl whoami. We use firectl for deployment to Fireworks in this example but you can serve the model however you prefer.

Setup

We recommend using uv.

bash

uv sync

Optional: Dev Container

We have provided a Dev Container config in .devcontainer to help users of VS Code who want to run the notebook on a remote server. To use our container, follow the VS Code Instructions, then proceed with the "Using uv" instructions below.