Source: examples/training/torchtitan

TorchTitan: Large-Scale LLM Training with SkyPilot#

TorchTitan is a PyTorch native platform for large-scale LLM training, featuring multi-dimensional parallelisms (FSDP2, Tensor/Pipeline/Context Parallel), distributed checkpointing, torch.compile, and Float8 support.

This example demonstrates how to run TorchTitan on your Kubernetes clusters, or any hypersclaers, neoclouds using SkyPilot, in addition to the instructions for runnning on Slurm.

Quick start#

Here is how to finetune Llama 3.1 on 2 nodes with 8 H100 (or 8 H200):

# Install SkyPilot (if not already installed)
# More cloud setup instructions in: https://docs.skypilot.co/en/latest/getting-started/installation.html
pip install "skypilot[kubernetes,aws]"  # or your cloud: [gcp], [azure], etc.

# Launch a cluster and start training
export HF_TOKEN=... # if using a gated model from the HF Hub
sky launch -c torchtitan-multinode torchtitan.yaml --env HF_TOKEN

# Tail logs
sky logs torchtitan-multinode

# Terminate the cluster when done
sky down torchtitan-multinode

Configuration#

The provided torchtitan.yaml configuration:

  • Sets up a 2-node cluster with 8 H100 (or H200) GPUs per node

  • Installs PyTorch nightly and TorchTitan requirements

  • Downloads the Llama 3.1 tokenizer

Available model configurations#

TorchTitan includes pre-configured training recipes for:

  • Llama 3.1 8B: llama3_8b.toml

  • Llama 3.1 70B: llama3_70b.toml

  • Llama 3.1 405B: llama3_405b.toml

Each configuration file specifies model architecture, parallelism strategies, and training hyperparameters optimized for different scales.

To use a specific training recipe, you can set it through the CONFIG_FILE env var:

sky launch -c torchtitan-multinode torchtitan.yaml \
   --env HF_TOKEN \
   --env CONFIG_FILE=./torchtitan/models/llama3/train_configs/llama3_70b.toml # relative to the torchtitan's repo

Scaling Up#

To scale up your training, you can increase the number of nodes or try larger models:

# Scale to more nodes
sky launch -c torchtitan-8node torchtitan.yaml --num-nodes 8

# Try different model sizes (update CONFIG_FILE in torchtitan.yaml)
sky launch -c torchtitan-llama3-70b torchtitan.yaml --env CONFIG_FILE=./torchtitan/models/llama3/train_configs/llama3_70b.toml

Why SkyPilot for Distributed Training?#

  • Simple multi-node setup: SkyPilot automatically provides environment variables (SKYPILOT_NODE_RANK, SKYPILOT_NODE_IPS, etc.) that integrate seamlessly with PyTorch distributed training - no manual networking configuration needed.

  • Auto-recovery: Built-in fault tolerance automatically recovers from node failures and spot preemptions, resuming from checkpoints.

  • Easily run on Kubernetes or clouds without code changes: SkyPilot offers a simple interface to run TorchTitan on any infrastructure: sky launch --infra k8s torchtitan.yaml

  • Launch distributed training with a single command: SkyPilot automatically provides environment variables(SKYPILOT_NODE_RANK, SKYPILOT_NODE_IPS, etc.) that integrate seamlessly with PyTorch distributed training - no manual networking configuration needed.

  • Auto-recovery: Built-in fault tolerance automatically recovers from node failures and spot preemptions, resuming from checkpoints.

Multi-node training details#

The configuration automatically:

  • Detects the head node IP and sets it as the master address

  • Configures the correct node rank for each node

  • Sets up the distributed environment for PyTorch’s torchrun with key settings:

    • --nnodes: Uses $SKYPILOT_NUM_NODES to specify total number of nodes

    • --nproc_per_node: Uses $SKYPILOT_NUM_GPUS_PER_NODE for GPUs per node

    • --node_rank: Uses $SKYPILOT_NODE_RANK to identify each node’s position

    • --master_addr: Extracts head node IP from $SKYPILOT_NODE_IPS

    • --master_port: Sets communication port to 8008 for distributed coordination

Included files#

torchtitan.yaml

# SkyPilot configuration for TorchTitan multi-node training
# This configuration reproduces the functionality of multinode_trainer.slurm
#
# To launch:
#   sky launch -c torchtitan-cluster sky.yaml
#
# To stop:
#   sky down torchtitan-cluster
#
# To monitor:
#   sky status --refresh

name: torchtitan-multinode

resources:
  accelerators: {H100:8, H200:8}
  disk_size: 1024GB

num_nodes: 2

workdir: .

envs:
  CONFIG_FILE: "./torchtitan/models/llama3/train_configs/llama3_8b.toml"
  HF_TOKEN: ""

setup: |
  git clone https://github.com/pytorch/torchtitan.git
  cd torchtitan
  pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126 --force-reinstall
  pip install -r requirements.txt
  python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=$HF_TOKEN

run: |
  # Get head node IP (first node in the list)
  HEAD_NODE_IP=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
  echo "Head node IP: $HEAD_NODE_IP"
  
  # SKYPILOT_NODE_RANK is automatically set by SkyPilot
  torchrun \
    --nnodes $SKYPILOT_NUM_NODES \
    --nproc_per_node $SKYPILOT_NUM_GPUS_PER_NODE \
    --node_rank $SKYPILOT_NODE_RANK \
    --master_addr=$HEAD_NODE_IP \
    --master_port=8008 \
    -m torchtitan.train \
    --job.config_file $CONFIG_FILE \
    --training.dataset c4_test