Run the nanoGPT speedrun

modded-nanogpt is a community speedrun: train a GPT-2-sized language model to a validation loss of 3.28 on FineWeb as fast as possible on 8 NVIDIA H100s. The 2024 baseline took 45 minutes. The current record trains in under a minute. In this example you rent an H100 SUNK cluster on Andromeda and run the current record on one of its nodes. SUNK is Slurm running on the cluster’s Kubernetes: you get a login node over SSH, Slurm with Pyxis for containers, and a Ceph filesystem shared by every node. The code and data live in your home directory on that filesystem, the training runs in a container, and the job is submitted with sbatch so it keeps running when your SSH session ends.

What you need

An Andromeda account in an organization, with permission to place orders and the clusters.read and clusters.ssh permissions. See Permissions.
The andromeda CLI, jq, and OpenSSH, set up as in the quickstart.

Rent a SUNK cluster

In the console, open On demand and pick an H100 offer. On Configure on-demand capacity:

Setting	Choose	Why
Nodes	2	SUNK needs at least two nodes. Every node has 8 GPUs; the speedrun uses one node.
Scheduler	SUNK (Slurm on Kubernetes)	Gives you Slurm, a login node, and Pyxis.
Shared storage	Ceph	SUNK needs shared storage. It holds `/home` and `/shared` on every node.

H200 offers work too. The speedrun’s attention kernel is built for Hopper, so other GPU types do not run it. On-demand clusters are billed until you end them; see When you’re done.

Wait until Slurm and storage are ready

Find the new cluster and save its ID in CLUSTER_ID, as in the quickstart. The cluster shows provisioning, then active:

andromeda call platform GET "/clusters?businessOrgId=$ORG_ID" |
  jq -r '.items[] | [.id, .spec.name, .status.lifecycleState] | @tsv'

active means the nodes are assigned. Slurm and Ceph come up after that. Check both:

andromeda call platform GET "/clusters/$CLUSTER_ID/options?businessOrgId=$ORG_ID" |
  jq -c '{scheduler: .scheduler.configured.value, slurm: .scheduler.ready.status,
          ceph: .storage.configured.ceph, storage: .storage.ready.status}'

{"scheduler":"sunk","slurm":"True","ceph":"configured","storage":"True"}

Connect and check the nodes

Register your SSH key and add the training host alias from Connect over SSH. Use your own cluster ID and username in the alias.

Your first key restarts the cluster’s compute nodes, and job steps fail on a node until it has restarted. On a new cluster nothing is running yet, so register the key now, before you submit the job.

List the partitions with their GPUs and node features:

ssh training "sinfo -o '%P %D %t %G %f'"

PARTITION NODES STATE GRES AVAIL_FEATURES
compute-h100 2 idle gpu:nvidia_h100:8 compute-h100,gpu,cu128
all* 2 idle gpu:nvidia_h100:8 compute-h100,gpu,cu128

compute-h100 is the partition for the cluster’s H100 nodes, and the job uses it. On an H200 cluster it is compute-h200. all is the default partition, with every node. cu128 means the nodes provide CUDA 12.8.

Get the code

Clone the repository into your home directory on the cluster:

ssh training 'mkdir -p nanogpt && git -C nanogpt clone https://github.com/KellerJordan/modded-nanogpt.git'

/home/<username> is on the cluster’s Ceph filesystem, so the login node slurm-login-0 and the compute node that runs the job see the same checkout.

The batch script

Resources

#SBATCH --job-name=nanogpt-speedrun
#SBATCH --partition=compute-h100
#SBATCH --nodes=1
#SBATCH --gres=gpu:8
#SBATCH --exclusive
#SBATCH --mem=0
#SBATCH --time=01:00:00
#SBATCH --output=%x-%j.out

The job takes one whole H100 node: all 8 GPUs, every CPU (--exclusive), and all memory (--mem=0). The speedrun itself is short, but the first run also builds a container and compiles the model, so the time limit leaves room for that. The log goes to nanogpt-speedrun-<job ID>.out in the directory you submit from. On an H200 cluster, set --partition=compute-h200.

The container

BASE_IMAGE=docker://nvidia/cuda:13.1.1-cudnn-devel-ubuntu24.04

if [ ! -f "$IMAGE" ]; then
  srun --container-image="$BASE_IMAGE" --container-save="$IMAGE" bash -c "
    set -euo pipefail
    apt-get update
    apt-get install -y --no-install-recommends build-essential git python3-dev python3-venv
    python3 -m venv /opt/venv
    /opt/venv/bin/pip install -r $REPO/requirements.txt
    /opt/venv/bin/pip install torch==2.10 --index-url https://download.pytorch.org/whl/cu128
  "
fi

The speedrun pins PyTorch 2.10 and an attention kernel that needs the CUDA 13 runtime. SUNK compute nodes run CUDA 12.8 on NVIDIA driver 580, which also runs CUDA 13 programs. So the job brings CUDA 13 in a container, and the node’s driver runs it. SUNK compute nodes have Pyxis and enroot installed, so srun --container-image works without setup. The image comes straight from Docker Hub. Inside the container your home directory is mounted and you are root, so apt-get works. --container-save writes the finished container to ~/nanogpt/nanogpt.sqsh on the Ceph filesystem, where every node sees it. Later runs start from that file and skip this step. The build runs inside the job, so on the first run the whole 8-GPU node is held, with its GPUs idle, while packages install. Delete the file to rebuild it, for example after the repository changes its requirements.

The data

srun --container-image="$IMAGE" --container-workdir="$REPO" \
  /opt/venv/bin/python data/cached_fineweb10B.py 9

This downloads the FineWeb validation set and the first 9 training chunks, 900M tokens of pre-tokenized text, from Hugging Face into modded-nanogpt/data/fineweb10B. Later runs skip files already downloaded.

Training

srun --container-image="$IMAGE" --container-workdir="$REPO" \
  --ntasks=1 --gres=gpu:8 --cpus-per-task="$SLURM_CPUS_ON_NODE" \
  bash -c 'unset NCCL_NET; exec /opt/venv/bin/torchrun --standalone --nproc_per_node=8 train_gpt.py'

torchrun starts one training process per GPU. Slurm sees one task, and SUNK confines each task to the CPUs Slurm gives it, so --cpus-per-task gives that task every CPU on the node. Without it, all eight processes would share one CPU. The ranks talk to each other over NVLink inside the node. On clusters with a RoCE network, SUNK sets NCCL_NET=IB for multi-node jobs. unset NCCL_NET keeps NCCL from looking for that network on a one-node run. On InfiniBand clusters the variable is not set, and unsetting it changes nothing.

The full script

Save this as speedrun.sbatch on your machine:

speedrun.sbatch

#!/bin/bash
# Train modded-nanogpt to 3.28 validation loss on FineWeb, on one 8-GPU H100 node of an Andromeda SUNK cluster.
# Submit from ~/nanogpt, next to the modded-nanogpt checkout: sbatch speedrun.sbatch
#SBATCH --job-name=nanogpt-speedrun
#SBATCH --partition=compute-h100
#SBATCH --nodes=1
#SBATCH --gres=gpu:8
#SBATCH --exclusive
#SBATCH --mem=0
#SBATCH --time=01:00:00
#SBATCH --output=%x-%j.out

set -euo pipefail

WORK=$HOME/nanogpt
REPO=$WORK/modded-nanogpt
IMAGE=$WORK/nanogpt.sqsh
BASE_IMAGE=docker://nvidia/cuda:13.1.1-cudnn-devel-ubuntu24.04

# The nodes run CUDA 12.8 and the speedrun's attention kernel links the CUDA 13 runtime, so it runs in a CUDA 13
# container through Pyxis. Build the container once, with Python and the pinned PyTorch installed, and save it on
# the shared filesystem next to the code for later runs.
if [ ! -f "$IMAGE" ]; then
  srun --container-image="$BASE_IMAGE" --container-save="$IMAGE" bash -c "
    set -euo pipefail
    apt-get update
    apt-get install -y --no-install-recommends build-essential git python3-dev python3-venv
    python3 -m venv /opt/venv
    /opt/venv/bin/pip install -r $REPO/requirements.txt
    /opt/venv/bin/pip install torch==2.10 --index-url https://download.pytorch.org/whl/cu128
  "
fi

# Download the validation set and the first 900M training tokens. Files already present are skipped.
srun --container-image="$IMAGE" --container-workdir="$REPO" \
  /opt/venv/bin/python data/cached_fineweb10B.py 9

# Train. One task starts torchrun, which starts one process per GPU, so give the task every CPU on the node.
# NCCL_NET is unset so that NCCL uses NVLink inside the node, even on RoCE clusters where SUNK sets NCCL_NET=IB.
srun --container-image="$IMAGE" --container-workdir="$REPO" \
  --ntasks=1 --gres=gpu:8 --cpus-per-task="$SLURM_CPUS_ON_NODE" \
  bash -c 'unset NCCL_NET; exec /opt/venv/bin/torchrun --standalone --nproc_per_node=8 train_gpt.py'

Run it

Copy the script to the cluster and submit it:

scp speedrun.sbatch training:nanogpt/speedrun.sbatch
ssh training 'cd nanogpt && sbatch speedrun.sbatch'

Submitted batch job 1042

The job now runs on a compute node, independent of your SSH session. Check its state with ssh training 'squeue --me'. PD means it is waiting for a free node, and R means it is running.

Watch it

ssh training 'tail -f nanogpt/nanogpt-speedrun-1042.out'

The SSH gateway can close this session at any time, for example when someone changes the cluster’s configuration. That does not affect the job. Run the command again to keep watching. On the first run the log starts with the container build and the data download. Then the model compiles:

Compiling model, warming up kernels and capturing CUDA graphs (~7 minutes on first execution)

After that, training logs its progress. The last validation line reports the result:

step:<total>/<total> val_loss:<loss> train_time:<milliseconds>ms step_avg:<milliseconds>ms

The run reaches the target when val_loss is at most 3.28. train_time counts only the timed training loop. Records are timed on H100s, so expect a different time on H200s. The job’s total time is longer: the first run adds the container build, the download, and about seven minutes of compilation. The trainer also writes a full log, including the source code it ran, under modded-nanogpt/logs/.

When you’re done

On-demand clusters are billed until they end. To stop paying, open the cluster in the console and choose End cluster. Copy anything you want to keep out of ~/nanogpt first, for example with scp; see Transfer files.

Troubleshooting

Symptom	Cause and fix
The job stays `PD` with reason `Resources`	No H100 node is free. It starts when one is.
`sbatch` rejects the partition	The cluster’s GPUs are not H100s. Run `sinfo` and use the `compute-` partition it lists.
A job step fails right after you register your first SSH key	The compute nodes are restarting to add your account. Wait until `sinfo` shows them `idle`, then submit again. See Access changes restart nodes.
`ssh training` takes minutes to connect	The cluster uses `PerUser` login mode, and your login node is starting. That takes up to 3 minutes. See Login modes.
The container build fails pulling the image	Pyxis pulls straight from Docker Hub, which rate-limits anonymous pulls. Submit again later.
The kernel download from Hugging Face fails	Hugging Face limits anonymous downloads. Create a token at huggingface.co, then submit with `HF_TOKEN=<token> sbatch speedrun.sbatch`; the job passes your environment to the container.
An error mentioning `libcudart.so.13` or an unsupported GPU	The job ran outside the container, or on GPUs other than H100 or H200.
The job ends at the time limit	The first run took longer than an hour. Raise `--time` and submit again; finished steps are skipped.

Run it again. With the container and data in place, the job goes straight to compiling and training.
Read train_gpt.py and the track_1_short/ package. The repository README explains what each record changed.
Configure Slurm, for example to add a prolog that checks GPU health before every job.

Andromeda Reference

Guides

What you need

The batch script

Resources

The container

The data

Training

The full script

Run it

Watch it

When you’re done

Troubleshooting

Next

​What you need

​The batch script

​Resources

​The container

​The data

​Training

​The full script

​Run it

​Watch it

​When you’re done

​Troubleshooting

​Next

What you need

The batch script

Resources

The container

The data

Training

The full script

Run it

Watch it

When you’re done

Troubleshooting

Next