Train a Chat Model with NanoChat

Set up Docker permissions

To manage containers without sudo, add your user to the docker group. Open a terminal and test Docker access:

docker ps

If you see a permission-denied error, add your user to the docker group (skip if it already works):

sudo usermod -aG docker $USER
newgrp docker

Set environment variables

Nanochat uses Weights & Biases for training visualization and a Hugging Face token for evaluation datasets. Export both in your shell:

export WANDB_API_KEY=<YOUR_WANDB_API_KEY>
export HF_TOKEN=<YOUR_HF_TOKEN>
export WANDB_RUN=speedrun   # optional run name

Create a W&B account and a Hugging Face token if you do not have them. The launch script exits immediately if either key is unset.

Clone and set up

Clone the playbook repository and navigate to the assets directory:

git clone https://github.com/NVIDIA/dgx-spark-playbooks
cd dgx-spark-playbooks/nvidia/playbook-nanochat/assets

Run the single-node setup script. It clones nanochat, checks out the supported commit, copies speedrun_single.sh, and builds the nanochat Docker image:

chmod +x setup.sh launch.sh
./setup.sh

Confirm the image built:

docker images | grep nanochat

Your directory structure after setup should look like:

assets/
├── Dockerfile
├── launch.sh
├── setup.sh
├── speedrun_single.sh
└── nanochat/

Launch training

Ensure your API keys are still exported, then launch:

./launch.sh

GPU_DEVICE='"device=N"' ./launch.sh

The training container runs the full pipeline automatically:

1. Tokenizer — downloads pretraining text, trains a BPE tokenizer
2. Base pretraining — pretrains the default d24 model (~1B params) with FP8
3. SFT — fine-tunes for chat with identity conversations + SmolTalk
4. Report — writes metrics and samples to report.md

A full d24 run is on the order of 12+ hours. Keep the terminal open or use tmux / screen.

Monitor training

Track progress at wandb.ai under the nanochat project (the run URL appears in the training logs). Key metrics:

• Training loss
• Validation BPB
• Throughput (tokens/sec)

Inference

After training, checkpoints are under nanochat_cache/. Run inference from a container with the same GPU selection as training.

Web UI (recommended):

docker run --rm --gpus all --net=host \
    -v $(pwd)/nanochat:/workspace/nanochat \
    -v $(pwd)/nanochat_cache:/root/.cache/nanochat \
    -w /workspace/nanochat \
    nanochat \
    python -m scripts.chat_web

Open a browser to http://<HOST_IP>:8000. If you are on a remote SSH session, forward the port:

ssh -L 8000:localhost:8000 username@<HOST_IP>

CLI:

docker run --rm -it --gpus all \
    -v $(pwd)/nanochat:/workspace/nanochat \
    -v $(pwd)/nanochat_cache:/root/.cache/nanochat \
    -w /workspace/nanochat \
    nanochat \
    python -m scripts.chat_cli -p "Why is the sky blue?"

Cleanup

To stop training early:

# If launch.sh is running: press Ctrl+C

# Or stop the container directly
docker stop $(docker ps -q --filter ancestor=nanochat)

To free disk space (cache dirs are often root-owned because the container runs as root):

sudo rm -rf ./nanochat_cache ./hf_cache
docker rmi nanochat

Customization

Smaller / faster run: Edit speedrun_single.sh before setup to reduce data and model size:

# Fewer data shards
python -m nanochat.dataset -n 10 &

# Smaller model (d4 instead of d24), smaller batch size
python -m scripts.base_train --depth=4 --device-batch-size=32

Then re-run ./setup.sh to rebuild with the changes.

Batch size: The default --device-batch-size=64 is tuned for high-memory hardware platforms. Lower it (32, 16, 8) if you hit OOM.

Next steps

• Try sample prompts with the trained model (web UI or CLI)
• Experiment with larger or smaller depths in the speedrun script
• Customize model personality via identity conversations — see the nanochat customization guide