Setup

nanochat uses uv for dependency management. To install:

uv sync --extra gpu    # Use for CUDA (A100/H100/etc.)
uv sync --extra cpu    # (or) Use for CPU-only / MPS
source .venv/bin/activate

For development (adds pytest, matplotlib, ipykernel, transformers, etc.):

uv sync --extra gpu --group dev

Reproduce and talk to GPT-2

The most fun you can have is to train your own GPT-2 and talk to it. The entire pipeline to do so is contained in the single file runs/speedrun.sh, which is designed to be run on an 8XH100 GPU node. Boot up a new 8XH100 GPU box from your favorite provider (e.g. I use and like Lambda), and kick off the training script:

bash runs/speedrun.sh

You may wish to do so in a screen session as this will take ~1.5 hours to run. Once it's done, you can talk to your model over the CLI. Make sure again that your local uv virtual environment is active (run source .venv/bin/activate), and chat:

python -m scripts.chat_cli

Get it to write stories or poems. Ask it to tell you who you are to see a hallucination. Ask it why the sky is blue. Or why it's green. The speedrun is a 4e19 FLOPs capability model so it's a bit like talking to a kindergartener :). An example conversation with a speedrun model:

User: Hello whats up?

Assistant: Hello! How can I help you today?

User: Why is the sky blue?

Assistant: The sky is blue due to an optical illusion called the Rayleigh Scattering of Light, which occurs when sunlight passes through tiny particles in the air. These particles are made up of tiny blue and violet particles that cause the light to bend in a particular way.

User: That's cool. Write a poem about it.

Assistant: In the twilight's glow, the sky is blue,<br> A canvas of shimmering blue, where skies unfold,<br> A world of blue, where blue and more blue,<br> Are woven together in a tapestry so fine,<br> Where every hue seems to sing a story.<br> ...

A few more notes:

• The code will run just fine on the Ampere 8XA100 GPU node as well, but a bit slower.
• All code will run just fine on even a single GPU by omitting torchrun, and will produce ~identical results (code will automatically switch to gradient accumulation), but you'll have to wait 8 times longer.
• If your GPU(s) have less than 80GB, you'll have to tune some of the hyperparameters or you will OOM / run out of VRAM. Look for --device-batch-size in the scripts and reduce it until things fit. E.g. from 32 (default) to 16, 8, 4, 2, or even 1. Less than that you'll have to know a bit more what you're doing and get more creative.
• Most of the code is fairly vanilla PyTorch so it should run on anything that supports that - xpu, mps, or etc, but I haven't personally exercised all of these code paths so there might be sharp edges.