TPU Environment Setup

October 24, 2025 ยท View on GitHub

This directory contains scripts to set up a TPU VM environment with necessary configurations.

Prerequisites

  • Google Cloud SDK installed and configured
  • Have a GCS bucket to mount to the TPU VM
  • Required environment variables set:
    • WANDB_API_KEY
    • HF_TOKEN
    • GH_TOKEN
    • KAGGLE_KEY
    • KAGGLE_USERNAME

Usage

  1. Run the setup script:

    ./tpu_setup.sh --tpu-name <tpu-name> -z <zone> [-p <project>] [--bucket <bucket-name>]

Arguments

  • --tpu-name: Name of the TPU to set up (required)
  • -z, --zone: GCP zone where the TPU is located (required)
  • -p, --project: GCP project ID (default: cs6741-llm-training)
  • --bucket: GCS bucket to mount (optional, can be specified multiple times)

What it does

  1. Copies necessary setup scripts to the TPU VM
  2. Installs UV package manager
  3. Configures environment variables in ~/.env
  4. Updates ~/.bashrc to source the environment variables
  5. Mounts specified GCS buckets to /mnt/<bucket-name> (if provided)

Bucket Mounting

  • Each bucket is mounted to /mnt/<bucket-name> on all TPU workers
  • Buckets are optional - the script will run successfully without any

Example

export WANDB_API_KEY=your_wandb_key
export HF_TOKEN=your_hf_token
export GH_TOKEN=your_gh_token
export KAGGLE_KEY=your_kaggle_key
export KAGGLE_USERNAME=your_kaggle_username
./tpu_setup.sh --tpu-name my-tpu -z us-central1-a --bucket model-checkpoints