TPU Training Scripts

October 24, 2025 ยท View on GitHub

This directory contains scripts for running training/inference jobs on TPU VMs.

Prerequisites

  • TPU VM must be set up using the env-setup scripts

Usage

Running scripts on TPU

./run_on_tpu.sh --tpu-name <tpu-name> -z <zone>  -p <project> --commit <commit-hash> --script <script-path> --config <config-path>

Arguments

  • --tpu-name: Name of the TPU to run training on
  • -z, --zone: GCP zone where the TPU is located
  • -p, --project: GCP project ID
  • --commit: Git commit hash to check out
  • --script: Path to python script to run relative to the root of the repository
  • --config: Path to the training config file

What it does

  1. Clones the repository
  2. Checks out the specified commit
  3. Runs the python script in the background (ctrl+C will not kill the job after it started)
  4. Logs output to ~/cc_git/tmp_script.log on the remote machine

Available scripts

Here is a partial list of the available scripts in this repo that can be used to train models on TPUs.

Every script is accompanied by a structured config file that can be used to run the script. You can find a script's config class in the script's file. Examples of config files can be found in the config_examples directory.

Teacher training

  • Running a teacher training job: context_compression/teacher_training/train.py
  • Resuming a teacher training job: context_compression/teacher_training/resume_training.py

Compressor training

  • Running a training job: context_compression/training_jax/train.py
  • Resuming a training job: context_compression/training_jax/resume_training.py

Inference (sampling)

  • Sampling from a compressor model: context_compression/evaluation/sampling/compressor/compressor_sampling.py
  • Sampling from a teacher model: context_compression/evaluation/sampling/teacher/teacher_sampling.py

Example

./run_on_tpu.sh --tpu-name my-tpu -z us-central1-a -p my-gcloud-proj --commit abc123 --script context_compression/training_jax/train.py --config_path config_examples/train_qwen3-0.6b_multi-ratio.yaml

Monitoring

To view logs:

gcloud compute tpus tpu-vm ssh <tpu-name> --zone <zone> --project <project> --command="tail -f ~/cc_git/tmp_script.log"

Or, if it's a training job, just view on wandb.