Lab: Train a 10M to 125M Parameter Model in an Afternoon

Validated on: written from the documentation and source cited above; not yet validated on hardware on any track. The commit, torch build and wall clock each track was run with will be recorded here when the validation pass is done.

Objective

By the end of this lab you will have made a language model that did not exist before. Concretely, you will have a tokeniser you trained and can inspect, a corpus subset you chose, a checkpoint on disk, an evaluation of that checkpoint in bits per byte and on the one benchmark the codebase ships, a set of samples you read rather than skimmed, and one line in the lab notebook that records all of it with the machine and the settings that produced it.

You will also have done one thing that matters more than the model: predicted the training throughput from the matrix-multiply rate you measured in Part 5, then measured it, and accounted for the gap. A training run you cannot predict is a training run you cannot budget.

The codebase is nanochat, nanochat main · verified 2026-09-08, which the course pins at main. The training run itself is calibrated to finish in about twenty-five minutes on whatever machine you have, because the script measures your throughput first and then chooses the number of steps.

⌂Single-machine path

This is a single-machine lab on every track. nanochat supports multiple GPUs through torchrun, and the README notes that a single GPU produces “~identical results” because the code switches to gradient accumulation automatically. If you own more than one GPU, run this lab on one of them and record which; the multi-device work starts in Part 18.

Requirements

Every track needs the lab notebook from Part 1, the Part 5 entry recording your machine’s BF16 matrix-multiply rate, uv and git, and about 12 GB of free disk: roughly 800 MB for eight corpus shards, 1 to 6 GB for the virtual environment depending on whether the wheel brings CUDA or ROCm libraries with it, around 1 GB for checkpoints, and a few hundred megabytes for the benchmark bundle that downloads on first use.

Time. About seventy-five minutes of attended work. The training run is twenty-five of those by construction. The environment sync and the corpus download are unattended and depend on your connection; start them before you read the rest of the page.

The model is small on every track. The default depth of six, with the head dimension nanochat’s own CPU recipe uses, gives a width of 384 and six heads, and a parameter count in the tens of millions of which most sits in the embedding tables rather than in the layers: two tables of 32,768 by 384 for the token embedding and the output projection, plus the value-embedding tables gpt.py adds on alternating layers. The training script prints the exact breakdown at startup and you will read it.

Where the memory goes during training, at the lab defaults on an 8 GB machine (estimate)

Weights, FP32 master copy

0.3 GB

Gradients

0.3 GB

Optimiser state (AdamW moments, Muon momentum)

0.6 GB

Activations and the logits tensor

1.8 GB

Free

5 GB

Total

8 GB

An estimate from the architecture, not a measurement. The dominant term is the logits: nanochat's forward pass materialises a tensor of batch x sequence length x vocabulary and casts it to FP32 for the loss, so at a device batch of 8, a sequence length of 512 and a vocabulary of 32,768 that single tensor is already about half a gigabyte before the softcap makes another copy. This is why the lab lowers the device batch size from the 32 nanochat's own CPU recipe uses. Replace every figure here with the run's printed peak memory.

Track S — NVIDIA DGX SparkPartial

nanochat pins torch 2.9.1 from the CUDA 12.8 wheel index. That index publishes aarch64 wheels, but whether that build carries kernels for the Spark's GB10 has not been confirmed on hardware; the verification step in the script tells you within a minute, and the fallback is one environment variable.

A DGX Spark with 128 GB of unified memory has room for a much larger run than the defaults, and the aarch64 dependency question is the only thing to check first. Two of the three answers are already known: rustbpe, the tokeniser trainer, publishes manylinux_2_17_aarch64 wheels for every supported Python, so it installs without a Rust toolchain; and torch==2.9.1+cu128, the version the project pins, publishes manylinux_2_28_aarch64 wheels. Both were read from their indexes on 2026-09-09.

What is not confirmed is whether the CUDA 12.8 build includes kernels for this GPU. Run the preparation script and read its verification line. If it prints a cuda device and names the GPU, continue. If it prints no accelerator visible, or the first training step fails with a kernel error, point uv at the CUDA 13.0 index instead, which publishes torch 2.9.1+cu130 for aarch64:

RunnableTrack S · DGX Spark

use the CUDA 13.0 wheels instead

TORCH_INDEX=https://download.pytorch.org/whl/cu130 TORCH_VERSION=2.9.1 \
    TRACK=spark bash prepare-data.sh

The script does the sync as published and then installs torch again from the index you named, so nothing in the project’s files is edited and the change is undone by rerunning without the variable. Record which index worked in the notebook; it is the most useful thing this track can contribute to the course.

Track X — AMD Ryzen AI Max+ 395Partial

nanochat pins torch 2.9.1, which the PyTorch wheel index publishes for ROCm 6.4 but not for any ROCm 7.x index, whose earliest listings are 2.10.0 or newer (read 2026-09-09). The CPU path finishes the lab on any machine and is stated with its arithmetic.

A Ryzen AI Max+ 395 machine with the ROCm user-space packages installed. The preparation script installs the CPU dependency set first and then replaces torch with a ROCm build, because nanochat’s pyproject.toml only knows about a CPU index and a CUDA index.

The version arithmetic is worth knowing before you start. Reading the PyTorch wheel index on 2026-09-09, torch 2.9.1+rocm6.4 exists; none of the ROCm 7.x indexes lists 2.9.1 at all, and their earliest listings are 2.10.0 or newer. So you either take ROCm 6.4 with the pinned version, which is what the script does by default, or you take a newer ROCm and let torch move a minor version, checking the index page for the earliest version it actually publishes:

RunnableTrack X · Ryzen AI Max+

a newer ROCm, with torch moved to match

TORCH_INDEX=https://download.pytorch.org/whl/rocm7.1 TORCH_VERSION=2.10.0 \
    TRACK=strix bash prepare-data.sh

Separately, the ROCm compatibility matrix lists gfx1151 without a support-tier qualifier (checked 2026-09-09), so treat a working GPU run as a good result to record rather than as something the vendor has committed to.

If the GPU path does not work, the CPU path finishes the lab. Do the arithmetic before you commit to it: the six-operations-per-parameter rule with a CPU matrix rate rather than a GPU one typically gives a throughput one to two orders of magnitude lower, which the calibration step turns into a much smaller step count for the same twenty-five minutes. You get a worse model in the same time, and every other part of the lab is unchanged. Run it with TRACK=cpu so the notebook records what happened.

Track M — Apple siliconPartial

nanochat has no MLX port. It runs on PyTorch's MPS backend, which the README lists as a target while noting the author 'hasn't personally exercised all of these code paths so there might be sharp edges', and the project's own CPU/MPS script says plainly that you will not get strong results.

An Apple silicon Mac on macOS 14 or later. nanochat’s pyproject.toml routes the cpu extra to the CPU wheel index, and that is the wheel that carries the MPS backend, so the installation is uv sync --extra cpu and nothing else. nanochat/common.py autodetects the device and will report mps.

There is no MLX port of nanochat, in its README or in mlx-examples, as of 2026-09-09. The closest MLX equivalent is the transformer_lm example in mlx-examples, described by its README as “an example of a decoder-only Transformer LM” whose “only dependency is MLX”, defaulting to the Penn Treebank corpus. That is a different and much smaller exercise; it is worth running afterwards if you want to see MLX’s training loop, but it is not this lab and its results are not comparable.

Two Mac-specific settings. The README’s precision table gives the compute dtype on MPS as float32 by default and notes that “On recent macOS, MPS also runs NANOCHAT_DTYPE=bfloat16 fine (~25% less memory, similar speed)”. Try it, because the logits tensor is the memory bottleneck and halving its element size helps directly:

RunnableTrack M · Apple silicon

train with bfloat16 compute on MPS

NANOCHAT_DTYPE=bfloat16 TRACK=mac bash train-small.sh

And close everything else. The GPU and the desktop share one memory pool, and a browser with forty tabs is competing with your training run for it.

Track N — NVIDIA desktop or laptop

A desktop or laptop with an NVIDIA GPU, on Linux or in WSL2 on Windows. This is the path nanochat is written for, minus the multi-GPU part: uv sync --extra gpu pulls the CUDA 12.8 build of the pinned torch, and the training script is invoked without torchrun, which the README says produces “~identical results” via gradient accumulation.

On a card with 8 or 12 GB of VRAM the defaults are chosen for you and should fit. On 16 GB and above, raise DEVICE_BATCH_SIZE to 16 or 32 after the first run and watch the peak memory the run prints; a larger device batch reduces the number of gradient-accumulation micro-steps and usually raises throughput.

Flash Attention 3 is unlikely to be available on a consumer card, and the training script says so loudly when it is not. Leave the window pattern at L, which is what the script does, because the same warning says that without Flash Attention 3 “SDPA has no support for sliding window attention” and that utilisation will suffer.

Tasks

Working directory and terminal roles

Prepare the course execution workspace once before this procedure. It includes this part's scripts, data and shared Python helpers. In the client or training terminal, select this directory:

RunnableAll tracks

select this part’s execution directory

export LABS_ROOT="${LABS_ROOT:-$HOME/llm-course/labs}"
export LAB_DIR="$LABS_ROOT/part-12-pretraining-from-scratch"
cd "$LAB_DIR"
pwd
test -f "prepare-data.sh"

Expected result: pwd ends in part-12-pretraining-from-scratch and the file check returns successfully. If it does not, finish workspace preparation before continuing. Activate the environment in the requirements for your track. Bare script and data filenames below are relative to this directory; paths to earlier experiments must point at the artefacts you actually retained.

Keep each foreground server in a separate terminal and send requests from this terminal. Reapply lesson-specific environment variables in each new shell. Stop at the first failed checkpoint and retain its output; the execution guide explains how to distinguish missing files, endpoint failures and capacity problems.

1. Get the code and record what you got

RunnableAll tracks

clone and pin

cd "$LAB_DIR"
git clone https://github.com/karpathy/nanochat.git ~/nanochat
git -C ~/nanochat rev-parse --short HEAD

The checkout command follows nanochat’s default branch; that branch is not an immutable pin. The commit you just printed identifies your source version. Preserve it and use that commit when reproducing this run. Put it in the notebook’s Environment section now. Every number this lab produces is a number about that commit.

2. Prepare the environment, the corpus and the tokeniser

One script does the three setup stages, because they are the three that have no decisions in them once you have chosen a track. Read it before running it: the interesting part is the case statement that selects the wheel index.

RunnableAll tracks

prepare-data.sh

#!/usr/bin/env bash
# Purpose: prepare a from-scratch pretraining run: build the nanochat virtual
#          environment for this platform track, download a chosen number of corpus
#          shards, train a byte-pair tokeniser on them, and evaluate its compression
#          against the GPT-2 and GPT-4 vocabularies
# Platform: all (CUDA on Tracks S and N, ROCm on Track X, MPS on Track M, CPU anywhere)
# Minimum memory: 8 GB
# Assumes: nanochat has been cloned to $NANOCHAT (default ~/nanochat); uv, git and
#          curl are on PATH; there is about 150 MB of free disk per corpus shard
#          under $NANOCHAT_BASE_DIR (default ~/.cache/nanochat); the machine can
#          reach pypi.org, download.pytorch.org and huggingface.co
set -euo pipefail

# ---------------------------------------------------------------- settings ---
TRACK="${TRACK:-}"                 # spark | strix | mac | nvidia | cpu
NANOCHAT="${NANOCHAT:-$HOME/nanochat}"
SHARDS="${SHARDS:-8}"              # corpus shards to download; ~250M characters each
VOCAB_SIZE="${VOCAB_SIZE:-32768}"  # nanochat's default, 2^15
MAX_CHARS="${MAX_CHARS:-2000000000}"   # characters the tokeniser is trained on
ROCM_INDEX="${ROCM_INDEX:-https://download.pytorch.org/whl/rocm6.4}"
# Optional: replace torch after the sync with a build from another wheel index.
# Track X always does this because nanochat's pyproject.toml knows only a CPU index
# and a CUDA index. Any track can do it, which is how Track S moves to a different
# CUDA build without editing the project's files.
TORCH_INDEX="${TORCH_INDEX:-}"
TORCH_VERSION="${TORCH_VERSION:-2.9.1}"   # the version nanochat pins
LOG_DIR="${LOG_DIR:-$PWD/part-12-logs}"

export NANOCHAT_BASE_DIR="${NANOCHAT_BASE_DIR:-$HOME/.cache/nanochat}"

usage() {
    cat <<'USAGE'
Usage: TRACK=<spark|strix|mac|nvidia|cpu> bash prepare-data.sh

Environment variables (all optional except TRACK):
  NANOCHAT           path to the cloned repository        (default ~/nanochat)
  SHARDS             corpus shards to download            (default 8)
  VOCAB_SIZE         tokeniser vocabulary size            (default 32768)
  MAX_CHARS          characters to train the tokeniser on (default 2000000000)
  NANOCHAT_BASE_DIR  where shards and checkpoints live    (default ~/.cache/nanochat)
  ROCM_INDEX         PyTorch ROCm wheel index, Track X   (default rocm6.4)
  TORCH_INDEX        replace torch from this wheel index after the sync
  TORCH_VERSION      version to pull from TORCH_INDEX     (default 2.9.1)
  LOG_DIR            where this script writes its logs    (default ./part-12-logs)
USAGE
}

case "$TRACK" in
    spark|strix|mac|nvidia|cpu) ;;
    *) usage; echo; echo "ERROR: set TRACK to one of spark, strix, mac, nvidia, cpu." >&2; exit 2 ;;
esac

# ------------------------------------------------------------- preflight ----
need() {
    command -v "$1" >/dev/null 2>&1 || { echo "ERROR: '$1' is not on PATH. $2" >&2; exit 1; }
}
need uv "Install it from https://docs.astral.sh/uv/ as in Part 1."
need git "Install your distribution's git package."

if [ ! -d "$NANOCHAT/.git" ]; then
    echo "ERROR: no git repository at $NANOCHAT." >&2
    echo "       git clone https://github.com/karpathy/nanochat.git \"$NANOCHAT\"" >&2
    exit 1
fi

mkdir -p "$LOG_DIR" "$NANOCHAT_BASE_DIR"
cd "$NANOCHAT"

COMMIT="$(git rev-parse --short HEAD)"
echo "==> nanochat at $NANOCHAT, commit $COMMIT"
echo "==> artefacts in $NANOCHAT_BASE_DIR, logs in $LOG_DIR"

# --------------------------------------------------------- the environment --
# nanochat pins torch in pyproject.toml and selects the wheel index with an extra:
# "gpu" points at the CUDA index and "cpu" at the CPU index, which is also the one
# that carries the Metal-capable macOS wheels. Track X installs the CPU set first
# and then replaces torch with a ROCm build from AMD's index.
echo "==> creating the virtual environment"
[ -d .venv ] || uv venv

case "$TRACK" in
    spark|nvidia) uv sync --extra gpu ;;
    mac|cpu)      uv sync --extra cpu ;;
    strix)
        # The CPU extra brings every dependency except an accelerated torch.
        uv sync --extra cpu
        TORCH_INDEX="${TORCH_INDEX:-$ROCM_INDEX}"
        ;;
esac

if [ -n "$TORCH_INDEX" ]; then
    echo "==> replacing torch $TORCH_VERSION from $TORCH_INDEX"
    # --index-url makes that index authoritative for torch itself and the wheels it
    # ships alongside; --extra-index-url keeps PyPI available for the ordinary
    # dependencies (filelock, sympy, networkx and friends).
    uv pip install \
        --index-url "$TORCH_INDEX" \
        --extra-index-url https://pypi.org/simple \
        "torch==$TORCH_VERSION"
fi

# shellcheck disable=SC1091
source .venv/bin/activate

# ------------------------------------------------------------ verification --
# Print what we actually got. A wrong wheel here is the single most common reason
# a run is mysteriously slow, and it is silent unless you look.
echo "==> environment"
python - <<'PY'
import platform
import torch

print(f"    python  {platform.python_version()} on {platform.system()} {platform.machine()}")
print(f"    torch   {torch.__version__}")
if torch.cuda.is_available():
    print(f"    device  cuda -> {torch.cuda.get_device_name(0)}")
    print(f"    hip     {getattr(torch.version, 'hip', None)}")
    print(f"    cuda    {getattr(torch.version, 'cuda', None)}")
elif getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
    print("    device  mps -> Apple silicon GPU")
else:
    print("    device  cpu -> no accelerator visible; the run will be slow")
PY

# ------------------------------------------------------------------ corpus --
# Shards are parquet files with a "text" column. The download always adds the
# final shard as well, which nanochat holds out as the validation split.
echo "==> downloading $SHARDS corpus shard(s) plus the validation shard"
python -m nanochat.dataset -n "$SHARDS" 2>&1 | tee "$LOG_DIR/dataset.log"

SHARD_DIR="$NANOCHAT_BASE_DIR/base_data_climbmix"
if [ -d "$SHARD_DIR" ]; then
    echo "==> on disk: $(du -sh "$SHARD_DIR" | cut -f1) in $(find "$SHARD_DIR" -name '*.parquet' | wc -l) file(s)"
fi

# --------------------------------------------------------------- tokeniser --
echo "==> training the tokeniser: vocabulary $VOCAB_SIZE, from $MAX_CHARS characters"
python -m scripts.tok_train \
    --max-chars="$MAX_CHARS" \
    --vocab-size="$VOCAB_SIZE" 2>&1 | tee "$LOG_DIR/tok-train.log"

echo "==> evaluating the tokeniser against the GPT-2 and GPT-4 vocabularies"
python -m scripts.tok_eval 2>&1 | tee "$LOG_DIR/tok-eval.log"

# ----------------------------------------------------------------- summary --
cat <<SUMMARY

==> done
    commit         $COMMIT
    track          $TRACK
    shards         $SHARDS (+1 validation shard)
    vocabulary     $VOCAB_SIZE
    tokeniser      $NANOCHAT_BASE_DIR/tokenizer
    logs           $LOG_DIR

    Record in the lab notebook: the commit, the torch version and device the
    verification step printed, the shard count, and the bytes-per-token ratio
    your tokeniser achieved on the corpus rows of the tok_eval table, next to
    what GPT-2 and GPT-4 achieved on the same text.

    Next: TRACK=$TRACK bash train-small.sh
SUMMARY

Download prepare-data.sh162 lines

RunnableAll tracks

prepare everything (substitute your track)

TRACK=nvidia bash prepare-data.sh

Use spark, strix, mac, nvidia or cpu. It prints what it found, downloads eight corpus shards plus the validation shard, trains a byte-pair tokeniser with a vocabulary of 32,768 on two billion characters, and evaluates it.

Output — what you should see

==> environment
    python  3.12.x on Linux x86_64
    torch   2.9.1+cu128
    device  cuda -> <your GPU>
    cuda    12.8
==> downloading 8 corpus shard(s) plus the validation shard
...
==> on disk: <size> in 9 file(s)
==> training the tokeniser: vocabulary 32768, from 2000000000 characters
max_chars: 2,000,000,000
doc_cap: 10,000
vocab_size: 32,768
Training time: <seconds>s

=Expected result

The verification block must name an accelerator. If it says no accelerator visible on Track S, X or N, stop and fix that before training: everything downstream will work and take a very long time, which is the worst failure mode because it looks like success. Track M should print mps.

3. Read the tokeniser evaluation before you train anything

The last thing the script printed is the most interesting artefact of this stage. tok_eval encodes the same texts with the GPT-2 vocabulary, the GPT-4 vocabulary and yours, and prints bytes, tokens and bytes per token for each.

Output — what you should see

Vocab sizes:
GPT-2: 50257
GPT-4: 100277
Ours: 32768

Comparison with GPT-2:
Text Type  Bytes    GPT-2           Ours            Relative     Better
                    Tokens  Ratio   Tokens  Ratio   Diff %
news       <n>      <n>     <r>     <n>     <r>     <+/-%>       <which>
korean     <n>      <n>     <r>     <n>     <r>     <+/-%>       <which>
code       <n>      <n>     <r>     <n>     <r>     <+/-%>       <which>
...

Read two rows in particular. On the rows drawn from the training corpus your vocabulary should compress better than both baselines despite being a third the size of GPT-2’s and a quarter of GPT-4’s, because it was trained on exactly this text. On the Korean row it should do considerably worse, because English web text does not teach a byte-pair merger anything about Korean.

That contrast is the whole tokeniser argument from the first lesson, on your own machine, in half a minute. Copy the ratios into the notebook.

4. Find out what compute-optimal would have cost

Before you train the model you can afford, look at the model you cannot. Start the training script with no iteration override so that it derives the horizon from its default tokens-per-parameter ratio, read four lines, and stop it.

RunnableAll tracks

read the compute-optimal horizon, then interrupt

cd ~/nanochat && source "$HOME/llm-course/.venv/bin/activate"
python -m scripts.base_train --depth=6 --head-dim=64 --window-pattern=L \
    --max-seq-len=512 --device-batch-size=8 --total-batch-size=16384

Output — what you should see

Parameter counts:
wte                     : <n>
value_embeds            : <n>
lm_head                 : <n>
transformer_matrices    : <n>
scalars                 : <n>
total                   : <n>
Estimated FLOPs per token: <x>e+07
Calculated number of iterations from target data:param ratio: <steps>
Total number of training tokens: <tokens>
Tokens : Scaling params ratio: 12.00
Total training FLOPs estimate: <x>e+16

Press Ctrl-C once the step lines start. Write down the parameter breakdown, the compute-optimal step count and the total tokens. Two things should be visible immediately: the embedding groups are most of the parameter count, and the compute-optimal horizon is far longer than an afternoon.

✦Tip

Multiply the FLOPs per token by the total tokens and check it against the FLOPs estimate the script printed. Then divide by your Part 5 matrix rate times a utilisation fraction and you have the compute-optimal wall clock for this model on this machine. That single division is the most useful piece of arithmetic in this part.

5. Predict your throughput before you measure it

From the scaling-laws lesson: tokens per second ≈ (matrix rate × utilisation) ÷ FLOPs per token, with the FLOPs per token taken from the line the script just printed. Guess a utilisation between a fifth and a half. Write the prediction in the notebook before the next step, because a prediction written afterwards is not a prediction.

6. Calibrate, then train

The training script measures rather than assumes. It runs forty steps at exactly the settings of the real run, averages the last ten throughput figures, divides your time budget by them to get a step count, deletes the calibration checkpoint, and then trains.

RunnableAll tracks

train-small.sh

#!/usr/bin/env bash
# Purpose: pretrain a small nanochat base model inside a stated wall-clock budget:
#          calibrate the throughput with a short run, turn the budget into an
#          iteration count, train with intermediate checkpoints, then evaluate the
#          result and write the configuration out for the recording script
# Platform: all (CUDA on Tracks S and N, ROCm on Track X, MPS on Track M, CPU anywhere)
# Minimum memory: 8 GB
# Assumes: prepare-data.sh has already been run for this track, so $NANOCHAT holds
#          a synced .venv, $NANOCHAT_BASE_DIR holds corpus shards and a trained
#          tokeniser, and about 2 GB of free disk is available for checkpoints
set -euo pipefail

# ---------------------------------------------------------------- settings ---
TRACK="${TRACK:-}"                 # spark | strix | mac | nvidia | cpu
NANOCHAT="${NANOCHAT:-$HOME/nanochat}"
TARGET_MINUTES="${TARGET_MINUTES:-25}"   # wall clock the real run should take
DEPTH="${DEPTH:-6}"                # the one size dial; width and heads follow
MAX_SEQ_LEN="${MAX_SEQ_LEN:-512}"
# The logits tensor is device_batch_size x max_seq_len x vocab_size, materialised in
# fp32 for the loss, so it dominates peak memory: at 8 x 512 x 32768 x 4 bytes it is
# already about half a gigabyte before the copies the softcap makes. 8 is chosen for
# the 8 GB memory floor; raise it while the peak memory the run prints leaves room.
DEVICE_BATCH_SIZE="${DEVICE_BATCH_SIZE:-8}"
TOTAL_BATCH_SIZE="${TOTAL_BATCH_SIZE:-16384}"  # must be a multiple of batch x seq len
HEAD_DIM="${HEAD_DIM:-64}"
WINDOW_PATTERN="${WINDOW_PATTERN:-L}"   # full context on every layer; see the page
CALIBRATE_STEPS="${CALIBRATE_STEPS:-40}"
MIN_ITERATIONS="${MIN_ITERATIONS:-200}"
EVAL_EVERY="${EVAL_EVERY:-100}"
EVAL_TOKENS="${EVAL_TOKENS:-524288}"
SAMPLE_EVERY="${SAMPLE_EVERY:-100}"
EVAL_SPLIT_TOKENS="${EVAL_SPLIT_TOKENS:-16384}"
EVAL_MAX_PER_TASK="${EVAL_MAX_PER_TASK:-16}"
MODEL_TAG="${MODEL_TAG:-d${DEPTH}-lab}"
ITERATIONS="${ITERATIONS:-0}"      # set this to skip calibration entirely
LOG_DIR="${LOG_DIR:-$PWD/part-12-logs}"
FIELDS="${FIELDS:-$PWD/train-fields.json}"

export NANOCHAT_BASE_DIR="${NANOCHAT_BASE_DIR:-$HOME/.cache/nanochat}"
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"

CAL_TAG="calibration-tmp"

usage() {
    cat <<'USAGE'
Usage: TRACK=<spark|strix|mac|nvidia|cpu> bash train-small.sh

Environment variables (all optional except TRACK):
  TARGET_MINUTES      wall clock the training run should take   (default 25)
  DEPTH               transformer layers; the only size dial    (default 6)
  MAX_SEQ_LEN         context length used for training          (default 512)
  DEVICE_BATCH_SIZE   rows per forward/backward                 (default 8)
  TOTAL_BATCH_SIZE    tokens per optimiser step                 (default 16384)
  ITERATIONS          skip calibration and use this step count  (default 0 = calibrate)
  CALIBRATE_STEPS     steps in the calibration run              (default 40)
  NANOCHAT_DTYPE      float32 | bfloat16 | float16, passed through to nanochat
  LOG_DIR             where logs are written                    (default ./part-12-logs)
  FIELDS              where the run configuration is written    (default ./train-fields.json)
USAGE
}

case "$TRACK" in
    spark|strix|mac|nvidia|cpu) ;;
    *) usage; echo; echo "ERROR: set TRACK to one of spark, strix, mac, nvidia, cpu." >&2; exit 2 ;;
esac

# ------------------------------------------------------------- preflight ----
[ -d "$NANOCHAT/.git" ] || { echo "ERROR: no nanochat repository at $NANOCHAT." >&2; exit 1; }
[ -d "$NANOCHAT/.venv" ] || { echo "ERROR: no .venv in $NANOCHAT; run prepare-data.sh first." >&2; exit 1; }
[ -d "$NANOCHAT_BASE_DIR/tokenizer" ] || { echo "ERROR: no tokeniser in $NANOCHAT_BASE_DIR; run prepare-data.sh first." >&2; exit 1; }

# total_batch_size must divide evenly into whole forward/backward passes
TOKENS_PER_FWDBWD=$(( DEVICE_BATCH_SIZE * MAX_SEQ_LEN ))
if [ $(( TOTAL_BATCH_SIZE % TOKENS_PER_FWDBWD )) -ne 0 ]; then
    echo "ERROR: TOTAL_BATCH_SIZE ($TOTAL_BATCH_SIZE) must be a multiple of" >&2
    echo "       DEVICE_BATCH_SIZE x MAX_SEQ_LEN ($TOKENS_PER_FWDBWD). Keep both powers of two." >&2
    exit 2
fi

mkdir -p "$LOG_DIR"
cd "$NANOCHAT"
# shellcheck disable=SC1091
source .venv/bin/activate
COMMIT="$(git rev-parse --short HEAD)"

COMMON_ARGS=(
    --depth="$DEPTH"
    --head-dim="$HEAD_DIM"
    --window-pattern="$WINDOW_PATTERN"
    --max-seq-len="$MAX_SEQ_LEN"
    --device-batch-size="$DEVICE_BATCH_SIZE"
    --total-batch-size="$TOTAL_BATCH_SIZE"
)

# ----------------------------------------------------------- calibration ----
# A short run at exactly the real run's settings, so the throughput it reports is
# the throughput the real run will get. Everything that is not the training step
# is switched off, and the checkpoint it leaves behind is deleted afterwards.
RATE=""
if [ "$ITERATIONS" -le 0 ]; then
    CAL_LOG="$LOG_DIR/calibrate-$MODEL_TAG.log"
    echo "==> calibrating with $CALIBRATE_STEPS steps (this pays the compile cost once)"
    python -m scripts.base_train \
        "${COMMON_ARGS[@]}" \
        --num-iterations="$CALIBRATE_STEPS" \
        --eval-every=-1 \
        --core-metric-every=-1 \
        --sample-every=-1 \
        --model-tag="$CAL_TAG" 2>&1 | tee "$CAL_LOG"

    # Average the last ten step rates: the first steps include warm-up and compile.
    RATE="$(grep -oE 'tok/sec: [0-9,]+' "$CAL_LOG" \
        | tr -d ',' \
        | awk '{print $2}' \
        | tail -n 10 \
        | awk '{s+=$1; n+=1} END {if (n > 0) printf "%d", s/n}')"

    if [ -z "$RATE" ] || [ "$RATE" -le 0 ]; then
        echo "ERROR: could not read a tok/sec figure from $CAL_LOG." >&2
        echo "       Read the log, then rerun with ITERATIONS=<count> to skip calibration." >&2
        exit 1
    fi

    ITERATIONS="$(awk -v m="$TARGET_MINUTES" -v r="$RATE" -v b="$TOTAL_BATCH_SIZE" \
        'BEGIN { printf "%d", (m * 60 * r) / b }')"
    if [ "$ITERATIONS" -lt "$MIN_ITERATIONS" ]; then
        echo "==> calibration suggests $ITERATIONS steps; raising to the $MIN_ITERATIONS-step floor"
        echo "    (the learning-rate warm-up alone is 40 steps, so shorter runs are not meaningful)"
        ITERATIONS="$MIN_ITERATIONS"
    fi

    # The calibration checkpoint has served its purpose; reclaim the disk.
    CAL_DIR="$NANOCHAT_BASE_DIR/base_checkpoints/$CAL_TAG"
    if [ -d "$CAL_DIR" ]; then
        rm -rf "$CAL_DIR"
        echo "==> removed the calibration checkpoint at $CAL_DIR"
    fi

    echo "==> measured throughput and the budget give $ITERATIONS steps"
    echo "    ${TARGET_MINUTES} min x 60 s x ${RATE} tokens/s / ${TOTAL_BATCH_SIZE} tokens per step"
fi

SAVE_EVERY=$(( ITERATIONS / 4 ))
[ "$SAVE_EVERY" -lt 100 ] && SAVE_EVERY=100

# -------------------------------------------------------------- training ----
TRAIN_LOG="$LOG_DIR/train-$MODEL_TAG.log"
echo "==> training $MODEL_TAG for $ITERATIONS steps, checkpointing every $SAVE_EVERY"
python -m scripts.base_train \
    "${COMMON_ARGS[@]}" \
    --num-iterations="$ITERATIONS" \
    --eval-every="$EVAL_EVERY" \
    --eval-tokens="$EVAL_TOKENS" \
    --core-metric-every=-1 \
    --sample-every="$SAMPLE_EVERY" \
    --save-every="$SAVE_EVERY" \
    --model-tag="$MODEL_TAG" 2>&1 | tee "$TRAIN_LOG"

# ------------------------------------------------------------ evaluation ----
# bpb on both splits, the CORE benchmark on a small sample of each task, and the
# fixed sample prompts. The CORE bundle is downloaded on first use.
EVAL_LOG="$LOG_DIR/eval-$MODEL_TAG.log"
echo "==> evaluating $MODEL_TAG"
python -m scripts.base_eval \
    --eval=core,bpb,sample \
    --model-tag="$MODEL_TAG" \
    --device-batch-size=1 \
    --split-tokens="$EVAL_SPLIT_TOKENS" \
    --max-per-task="$EVAL_MAX_PER_TASK" 2>&1 | tee "$EVAL_LOG"

# ------------------------------------------------ configuration for the log --
cat > "$FIELDS" <<JSON
{
  "lab": "part-12/train-a-small-model",
  "track": "$TRACK",
  "commit": "$COMMIT",
  "model_tag": "$MODEL_TAG",
  "dataset": "nanochat default shards in $NANOCHAT_BASE_DIR",
  "depth": $DEPTH,
  "head_dim": $HEAD_DIM,
  "window_pattern": "$WINDOW_PATTERN",
  "max_seq_len": $MAX_SEQ_LEN,
  "device_batch_size": $DEVICE_BATCH_SIZE,
  "total_batch_size": $TOTAL_BATCH_SIZE,
  "num_iterations": $ITERATIONS,
  "save_every": $SAVE_EVERY,
  "target_minutes": $TARGET_MINUTES,
  "calibrated_tokens_per_second": ${RATE:-null},
  "dtype_override": "${NANOCHAT_DTYPE:-auto}",
  "train_log": "$TRAIN_LOG",
  "eval_log": "$EVAL_LOG"
}
JSON

cat <<SUMMARY

==> done
    model tag      $MODEL_TAG
    steps          $ITERATIONS
    checkpoints    $NANOCHAT_BASE_DIR/base_checkpoints/$MODEL_TAG
    training log   $TRAIN_LOG
    eval log       $EVAL_LOG
    configuration  $FIELDS

    Read these lines out of the training log before you go on:
      "Parameter counts:"                  where the parameters actually are
      "Estimated FLOPs per token:"         the six-per-parameter rule, applied
      "Tokens : Scaling params ratio:"     compare with the compute-optimal 12
      "Total training time:"               against your TARGET_MINUTES
      "Minimum validation bpb:"            the number worth comparing between runs

    Next: python sample-and-record.py --model-tag $MODEL_TAG --fields $FIELDS \\
              --labbook labbook.md
SUMMARY

Download train-small.sh214 lines

⚠Warning

The script removes the calibration checkpoint directory with rm -rf when it is finished with it. It only ever removes $NANOCHAT_BASE_DIR/base_checkpoints/calibration-tmp, a directory it created moments earlier, and it checks that the directory exists first. If you have set NANOCHAT_BASE_DIR to somewhere with data you care about, read those lines before running the script.

RunnableAll tracks

the training run

cd "$LAB_DIR"
TRACK=nvidia TARGET_MINUTES=25 bash train-small.sh

Output — what you should see

==> calibrating with 40 steps (this pays the compile cost once)
...
==> measured throughput and the budget give <steps> steps
    25 min x 60 s x <rate> tokens/s / 16384 tokens per step
==> removed the calibration checkpoint at ...
==> training d6-lab for <steps> steps, checkpointing every <n>
step 00000/0<steps> (0.00%) | loss: <x> | lrm: 0.03 | dt: <x>ms | tok/sec: <n> | bf16_mfu: <x> | epoch: 1 ... | eta: <x>m

Compare the calibrated rate with your prediction from step 5 now, while the run is going. If they differ by a factor rather than a fraction, the utilisation guess was wrong; if they differ by ten or more, something is on the CPU that should not be.

7. Watch three fields, not the whole log

While it runs, three things are worth your attention.

epoch must stay at 1. If it reaches 2 the loader has been round the corpus and is re-reading text. Stop, download more shards with SHARDS=32 bash prepare-data.sh, and start again.

bf16_mfu is the utilisation you guessed. Write the real one down.

The sample blocks, printed every hundred steps. The first ones are token soup. Somewhere in the middle the grammar arrives. By the end you should have plausible sentences with the wrong content in them. Copy the sample block from an early step and a late step into the notebook side by side; that pair is the most memorable artefact of this lab and it disappears when the terminal scrolls.

8. Read the evaluation

The script runs base_eval when training finishes, with the three evaluations the codebase provides: bits per byte on the training and validation splits, the CORE metric, and the fixed sample prompts.

Output — what you should see

train bpb: <x>
val bpb: <x>
...
CORE metric: <x>

CORE is the benchmark nanochat ships, taken from the DataComp-LM work; the leaderboard in the README is scored on it. Two cautions about your number. It is computed here on a small sample of each task, sixteen examples by default, so it is noisy. And a model of this size trained for twenty-five minutes will score close to what random guessing gives, which is the correct result and is worth recording as such rather than hiding.

Bits per byte is the number to actually compare. It is what you will compare between your own runs, and it is the only one of the three that is meaningful across different tokenisers.

9. Sample from the checkpoint and record everything

RunnableAll tracks

sample-and-record.py

"""Sample from a pretrained nanochat checkpoint and append one run-log line.

Purpose: close the loop on a from-scratch pretraining run: load the checkpoint you
         just trained, sample from it with a fixed prompt set so that two runs can
         be compared, read the losses and scores out of the training and evaluation
         logs, read the accelerator's power draw where the platform reports it, and
         append one JSON line to the lab notebook in the run-log format Part 11
         introduces.
Platform: all (cuda on Tracks S and N and on Track X with ROCm, mps on Track M,
          cpu anywhere as a fallback); the device is autodetected and recorded.
Minimum memory: 8 GB
Assumes: nanochat is cloned at --nanochat (default $NANOCHAT or ~/nanochat) with its
         virtual environment synced, a checkpoint exists under
         $NANOCHAT_BASE_DIR/base_checkpoints/<model-tag>, and this script is run with
         that environment's interpreter, for example
         ~/nanochat/.venv/bin/python sample-and-record.py ...

Usage: python sample-and-record.py --model-tag d6-lab [--step N]
                                   [--fields train-fields.json]
                                   [--labbook labbook.md]
                                   [--max-tokens 48] [--temperature 0.8]
                                   [--print-only]

The run log is one JSON object per line with the fields Part 11's runlog.py writes:
run_id, lab, date, config_commit, model, dataset, hyperparameters, seed, hardware,
versions, losses, scores and notes. A field that cannot be filled in is written as
null rather than omitted, so that a reader can always tell "not recorded" from "not
applicable".
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import re
import secrets
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

# Prompts that probe different things: a fact, a lookup, a small piece of
# reasoning, an antonym, a list, an open continuation and arithmetic. They are the
# set nanochat samples with during training, kept identical here so that a sample
# taken now can be compared with the ones in the training log.
FACT_PROMPTS = [
    "The capital of France is",
    "The chemical symbol of gold is",
    "If yesterday was Friday, then tomorrow will be",
    "The opposite of hot is",
    "The planets of the solar system are:",
    "My favorite color is",
    "If 5*x + 3 = 13, then x is",
]

# Longer, sampled rather than greedy, to show what the model does when it is not
# being led. This is where an undertrained model is most obvious.
FREE_PROMPTS = [
    "The history of the printing press begins",
    "In order to bake bread you will need",
]

REQUIRED_FIELDS = (
    "run_id", "lab", "date", "config_commit", "model", "dataset",
    "hyperparameters", "seed", "hardware", "versions", "losses", "scores", "notes",
)


# --------------------------------------------------------------------------- #
# Reading the logs                                                             #
# --------------------------------------------------------------------------- #

def read_text(path: str | None) -> str:
    """Return a log file's contents, or an empty string if it is missing."""
    if not path:
        return ""
    try:
        return Path(path).read_text(encoding="utf-8", errors="replace")
    except OSError:
        return ""


def last_float(pattern: str, text: str) -> float | None:
    """The last match of a single-group numeric pattern, as a float, or None.

    MULTILINE is on because several of the lines below are matched from the start
    of a line, and the logs are whole files rather than single lines.
    """
    matches = re.findall(pattern, text, flags=re.MULTILINE)
    if not matches:
        return None
    try:
        return float(str(matches[-1]).replace(",", ""))
    except ValueError:
        return None


def parse_train_log(text: str) -> dict:
    """Pull the figures nanochat prints during and after training out of its log."""
    step_rates = [float(m.replace(",", "")) for m in re.findall(r"tok/sec: ([0-9,]+)", text)]
    return {
        "min_val_bpb": last_float(r"Minimum validation bpb: ([0-9.]+)", text),
        "last_val_bpb": last_float(r"Validation bpb: ([0-9.]+)", text),
        "last_train_loss": last_float(r"\| loss: ([0-9.]+)", text),
        "total_training_minutes": last_float(r"Total training time: ([0-9.]+)m", text),
        "peak_memory_mib": last_float(r"Peak memory usage: ([0-9.]+)MiB", text),
        "parameters_total": last_float(r"^total\s+: ([0-9,]+)", text),
        "flops_per_token": last_float(r"Estimated FLOPs per token: ([0-9.e+]+)", text),
        "tokens_per_scaling_param": last_float(r"Tokens : Scaling params ratio: ([0-9.]+)", text),
        "training_tokens": last_float(r"Total number of training tokens: ([0-9,]+)", text),
        "median_tokens_per_second": (
            sorted(step_rates)[len(step_rates) // 2] if step_rates else None
        ),
        "last_mfu_percent": last_float(r"bf16_mfu: ([0-9.]+)", text),
    }


def parse_eval_log(text: str) -> dict:
    """Pull the evaluation figures out of the base_eval log."""
    return {
        "train_bpb": last_float(r"^train bpb: ([0-9.]+)", text),
        "val_bpb": last_float(r"^val bpb: ([0-9.]+)", text),
        "core_metric": last_float(r"CORE metric: ([0-9.]+)", text),
    }


# --------------------------------------------------------------------------- #
# The machine                                                                  #
# --------------------------------------------------------------------------- #

def run_command(args: list[str]) -> str | None:
    """Run a command and return its stripped output, or None if it is unavailable."""
    if not shutil.which(args[0]):
        return None
    try:
        out = subprocess.run(args, capture_output=True, text=True, timeout=20, check=True)
    except (subprocess.SubprocessError, OSError):
        return None
    return out.stdout.strip() or None


def read_power() -> dict:
    """Accelerator power draw, where the platform has a documented way to report it.

    NVIDIA and AMD both ship a query tool. macOS has no equivalent one-line query
    that this course has been able to cite, so on Track M the power field is null
    and the notebook records wall-clock time instead.
    """
    nvidia = run_command(
        ["nvidia-smi", "--query-gpu=power.draw", "--format=csv,noheader,nounits"]
    )
    if nvidia:
        first = nvidia.splitlines()[0].strip()
        try:
            return {"watts": float(first), "source": "nvidia-smi --query-gpu=power.draw"}
        except ValueError:
            return {"watts": None, "source": f"nvidia-smi returned {first!r}"}

    amd = run_command(["rocm-smi", "--showpower"])
    if amd:
        match = re.search(r"([0-9]+\.?[0-9]*)\s*W", amd)
        return {
            "watts": float(match.group(1)) if match else None,
            "source": "rocm-smi --showpower",
        }

    return {"watts": None, "source": "not reported on this platform"}


def describe_hardware(device) -> dict:
    import torch

    accelerator = "cpu"
    name = platform.processor() or platform.machine()
    if device.type == "cuda":
        accelerator = "cuda"
        name = torch.cuda.get_device_name(0)
    elif device.type == "mps":
        accelerator = "mps"
        name = "Apple silicon GPU (Metal Performance Shaders)"
    return {
        "os": f"{platform.system()} {platform.release()}",
        "arch": platform.machine(),
        "accelerator": accelerator,
        "device_name": name,
        "power": read_power(),
    }


def describe_versions() -> dict:
    import torch

    versions = {
        "python": platform.python_version(),
        "torch": torch.__version__,
        "cuda": getattr(torch.version, "cuda", None),
        "hip": getattr(torch.version, "hip", None),
    }
    for name in ("tiktoken", "rustbpe"):
        try:
            module = __import__(name)
        except ImportError:
            versions[name] = None
        else:
            versions[name] = getattr(module, "__version__", "installed")
    return versions


def new_run_id() -> str:
    """A short identifier that sorts by time and does not collide between runs."""
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    return f"{stamp}-{secrets.token_hex(3)}"


# --------------------------------------------------------------------------- #
# Sampling                                                                     #
# --------------------------------------------------------------------------- #

def sample_from_model(engine, tokenizer, prompts, *, max_tokens, temperature, top_k, seed):
    """Return one continuation per prompt, decoded back to text."""
    out = []
    for prompt in prompts:
        tokens = tokenizer(prompt, prepend="<|bos|>")
        kwargs = {"num_samples": 1, "max_tokens": max_tokens, "temperature": temperature, "seed": seed}
        if temperature > 0 and top_k:
            kwargs["top_k"] = top_k
        completion, _ = engine.generate_batch(tokens, **kwargs)
        out.append({"prompt": prompt, "completion": tokenizer.decode(completion[0])})
    return out


# --------------------------------------------------------------------------- #

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model-tag", required=True, help="checkpoint directory name, e.g. d6-lab")
    parser.add_argument("--step", type=int, default=None, help="checkpoint step (default: the last one)")
    parser.add_argument("--nanochat", default=os.environ.get("NANOCHAT", str(Path.home() / "nanochat")))
    parser.add_argument("--fields", default=None, help="train-fields.json written by train-small.sh")
    parser.add_argument("--labbook", default=None, help="append one JSON line to this file")
    parser.add_argument("--max-tokens", type=int, default=48)
    parser.add_argument("--temperature", type=float, default=0.8, help="for the free-continuation prompts")
    parser.add_argument("--top-k", type=int, default=50)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device-type", default="", help="cuda|mps|cpu (empty = autodetect)")
    parser.add_argument("--print-only", action="store_true", help="show the record without writing it")
    args = parser.parse_args()

    # nanochat is imported from the clone rather than installed, so put it on the path.
    repo = Path(args.nanochat).expanduser()
    if not (repo / "nanochat").is_dir():
        sys.exit(f"No nanochat package under {repo}. Pass --nanochat or set NANOCHAT.")
    sys.path.insert(0, str(repo))

    try:
        from nanochat.checkpoint_manager import load_model
        from nanochat.common import autodetect_device_type, compute_cleanup, compute_init
        from nanochat.engine import Engine
    except ImportError as exc:
        sys.exit(
            f"Could not import nanochat ({exc}). Run this script with the repository's own\n"
            f"interpreter, for example {repo}/.venv/bin/python sample-and-record.py ..."
        )

    device_type = autodetect_device_type() if args.device_type == "" else args.device_type
    _, _, _, _, device = compute_init(device_type)
    model, tokenizer, meta = load_model(
        "base", device, phase="eval", model_tag=args.model_tag, step=args.step
    )
    engine = Engine(model, tokenizer)

    # Greedy on the fact prompts so the answer is the model's most likely one, and
    # sampled on the free prompts so the text is representative rather than flat.
    greedy = sample_from_model(
        engine, tokenizer, FACT_PROMPTS,
        max_tokens=16, temperature=0.0, top_k=None, seed=args.seed,
    )
    free = sample_from_model(
        engine, tokenizer, FREE_PROMPTS,
        max_tokens=args.max_tokens, temperature=args.temperature, top_k=args.top_k, seed=args.seed,
    )

    print("\nGreedy continuations (temperature 0)")
    for item in greedy:
        print("-" * 78)
        print(item["completion"])
    print("\nSampled continuations (temperature %.2f, top-k %d)" % (args.temperature, args.top_k))
    for item in free:
        print("-" * 78)
        print(item["completion"])

    fields = {}
    if args.fields:
        try:
            fields = json.loads(Path(args.fields).read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            print(f"\nwarning: could not read {args.fields} ({exc}); recording without it")

    train_stats = parse_train_log(read_text(fields.get("train_log")))
    eval_stats = parse_eval_log(read_text(fields.get("eval_log")))

    record = {
        "run_id": new_run_id(),
        "lab": fields.get("lab", "part-12/train-a-small-model"),
        "date": datetime.now(timezone.utc).strftime("%Y-%m-%d"),
        "config_commit": fields.get("commit"),
        "model": {
            "name": f"nanochat base, tag {args.model_tag}",
            "trained_from": "random initialisation",
            "step": meta.get("step"),
            "config": meta.get("model_config"),
            "parameters_total": train_stats["parameters_total"],
        },
        "dataset": {
            "name": fields.get("dataset"),
            "training_tokens": train_stats["training_tokens"],
            "tokens_per_scaling_param": train_stats["tokens_per_scaling_param"],
        },
        "hyperparameters": {
            key: fields.get(key)
            for key in (
                "depth", "head_dim", "window_pattern", "max_seq_len",
                "device_batch_size", "total_batch_size", "num_iterations",
                "save_every", "target_minutes", "dtype_override",
            )
        },
        "seed": args.seed,
        "hardware": describe_hardware(device),
        "versions": describe_versions(),
        "losses": {
            "final_train_loss": train_stats["last_train_loss"],
            "min_val_bpb": train_stats["min_val_bpb"],
            "train_bpb": eval_stats["train_bpb"],
            "val_bpb": eval_stats["val_bpb"],
        },
        "scores": {
            "core_metric": eval_stats["core_metric"],
            "median_tokens_per_second": train_stats["median_tokens_per_second"],
            "last_mfu_percent": train_stats["last_mfu_percent"],
            "flops_per_token": train_stats["flops_per_token"],
            "total_training_minutes": train_stats["total_training_minutes"],
            "peak_memory_mib": train_stats["peak_memory_mib"],
        },
        "notes": {
            "greedy_samples": greedy,
            "sampled_completions": free,
            "track": fields.get("track"),
            "calibrated_tokens_per_second": fields.get("calibrated_tokens_per_second"),
        },
    }
    for key in REQUIRED_FIELDS:
        record.setdefault(key, None)

    line = json.dumps(record, ensure_ascii=False, sort_keys=False)
    if args.print_only or not args.labbook:
        print("\n" + line)
    else:
        notebook = Path(args.labbook)
        if not notebook.exists():
            notebook.write_text("# Lab notebook\n\n## Results\n\n", encoding="utf-8")
        with notebook.open("a", encoding="utf-8") as handle:
            handle.write(line + "\n")
        print(f"\nrecorded in {notebook}")

    compute_cleanup()


if __name__ == "__main__":
    main()

Download sample-and-record.py371 lines

Run it with nanochat’s own interpreter, so that it can import the package without activating anything:

RunnableAll tracks

sample, then append one run-log line

~/nanochat/.venv/bin/python sample-and-record.py \
    --model-tag d6-lab \
    --fields train-fields.json \
    --labbook labbook.md

It loads the last checkpoint, generates greedily from the seven fact prompts and with sampling from two open ones, reads the losses and scores out of the two logs, asks the platform for the accelerator’s power draw where there is a documented way to ask, and appends one JSON line to the notebook in the run-log format Part 11 introduces.

Power is where the tracks differ and the page says so rather than pretending otherwise. On Tracks S and N the script queries nvidia-smi with --query-gpu=power.draw --format=csv, both of which are documented selective-query options. On Track X it parses rocm-smi --showpower; AMD’s documentation notes that AMD SMI is the successor to ROCm SMI, so on a newer stack you may need its replacement. On Track M the field is written as null: macOS has no one-line power query this course has been able to cite, so record the wall clock and note that power was not measured.

10. Fill in the comparison sheet

Pending validationYour pretraining run — the sheet the notebook line summarises

Quantity	Where it comes from	Value
Commit	git rev-parse --short HEAD	—
Parameters, total	the Parameter counts block	—
Parameters in transformer matrices	the same block	—
FLOPs per token	Estimated FLOPs per token	—
Compute-optimal steps	task 4, before you interrupted it	—
Steps actually run	the calibrated count	—
Tokens : scaling params ratio	Tokens : Scaling params ratio	—
Predicted throughput	your arithmetic in task 5	—
Measured throughput	the calibration line	—
Utilisation	bf16_mfu	—
Wall clock	Total training time	—
Peak memory	Peak memory usage	—
Validation bits per byte	Minimum validation bpb	—
CORE metric	the eval log	—
Accelerator power	nvidia-smi or rocm-smi, or "not measured"	—

your machine: track, chip and memory, your operating system and version · nanochat the commit you cloned · depth 6, head dimension 64, as printed at startup, BF16 or FP32 compute per the platform default, FP32 master weights · 512 tokens of context · the date you ran it

Empty on purpose: this is the shape to fill in from your own run. The recording script writes the same fields as one JSON line, and this table is the human-readable version the capstone in Part 28 is written from. The two rows worth arguing about are the predicted and measured throughput; if they agree within a few tens of per cent your mental model of the machine is sound.

11. Go further: change exactly one thing

Now that you have a baseline, change one variable and rerun with the same time budget. Good choices, in increasing order of interest:

RunnableAll tracks

a deeper model in the same twenty-five minutes

TRACK=nvidia DEPTH=8 MODEL_TAG=d8-lab bash train-small.sh

A deeper model in the same wall clock means fewer tokens. That is the compute-optimal trade in one experiment, and the bits-per-byte comparison between the two runs is your own small scaling-law result.

RunnableAll tracks

the same model on twice the time

TRACK=nvidia TARGET_MINUTES=50 MODEL_TAG=d6-long bash train-small.sh

Doubling the tokens at a fixed size tells you whether the first run was on the steep part of the curve or the flat part.

The one thing not to do is change several at once. The whole reason nanochat exposes one size dial is that a codebase with twenty knobs produces results nobody can attribute.

Keep source, training environment and results paths distinct

This lab has two working locations: the course part directory holds the wrapper scripts and your evidence, while ~/nanochat holds the external training project and its environment. Whenever a task changes directory, verify pwd. Return to LAB_DIR before invoking a course script by its bare filename. Record the nanochat commit: a moving branch name is not an immutable pin.

Before the main run, confirm corpus shards, tokeniser files and the held-out split exist. Inspect a decoded token sample and the tokeniser evaluation. Use the calibration run to measure actual training throughput, then calculate the permitted main-run duration from your budget. Treat warm-up and compilation separately from sustained steps.

Watch loss, processed tokens and checkpoint progress. If loss becomes non-finite, stop and retain the first failing log rather than spending the remaining budget. When the time budget ends, evaluate the saved checkpoint on the unchanged held-out data and generate the fixed sample prompts. Report a budget-limited result if training had not converged. Archive configuration, tokeniser, corpus identity, checkpoint and evaluation outputs before removing downloads. A model file without its tokeniser and configuration is an incomplete pretraining artefact.

Validation

You are done when all of the following are true:

• ~/nanochat exists at a commit you recorded, and the preparation script’s verification block named an accelerator or you recorded that it did not and why;
• the tokeniser directory exists under $NANOCHAT_BASE_DIR, and the tokeniser evaluation ratios for your vocabulary against GPT-2 and GPT-4 are in the notebook;
• you recorded the compute-optimal step count from task 4 and the calibrated step count from task 6, and they are different numbers;
• the training log shows epoch: 1 throughout, a validation bits-per-byte line at least twice, and a Total training time close to your TARGET_MINUTES;
• checkpoints exist under $NANOCHAT_BASE_DIR/base_checkpoints/d6-lab, and one of them is not the final step;
• the evaluation log has train bpb, val bpb and CORE metric lines;
• labbook.md has one new JSON line whose losses, scores, hardware and hyperparameters objects are populated, with notes.greedy_samples containing seven continuations;
• your predicted and measured throughput are both written down, with a sentence explaining the gap.

Expected outcome

A model of a few tens of millions of parameters that produces grammatical, confidently wrong English after twenty-five minutes of training on your own machine, and a notebook entry that lets somebody else reproduce the run or compare theirs with it. Both halves matter and the second one is the one that survives.

=Expected result

Greedy continuations of “The capital of France is” from a run this short may well be wrong, and frequently entertainingly so. A correct answer means the corpus contained that fact often enough for a model this size to have kept it. Neither result is a bug. Record what you got, because comparing this list with the same list after the deeper or longer run in task 11 is a cleaner experiment than any benchmark score at this scale.

Troubleshooting

The verification step prints no accelerator visible on Track S or N. Run nvidia-smi first. If that fails the problem is the driver, not the wheel. If it works, the wheel is built for a CUDA the driver does not support or for architectures that do not include your GPU; on the Spark, retry with the CUDA 13.0 index as shown in the Track S requirements.

Track X: torch installs but reports no device. Check rocm-smi sees the GPU and that your user is in the render and video groups, as Part 1 and Part 5 covered. If the ROCm 6.4 build does not work on your stack, try the newer index with the matching torch version from the Track X requirements. If neither works, finish with TRACK=cpu and record that; a CPU result with an explanation is a result.

Out of memory during the first training step. The logits tensor is the cause nine times in ten. Halve DEVICE_BATCH_SIZE and rerun; the script keeps the total batch size the same by doing more gradient-accumulation micro-steps, so the run is equivalent and only the speed changes. The README gives the same advice for the same reason: reduce --device-batch-size “from 32 (default) to 16, 8, 4, 2, or even 1”.

TOTAL_BATCH_SIZE must be a multiple of .... You changed the device batch size or the sequence length to something that does not divide the total batch size. Keep all three powers of two.

The calibration step reports a rate but the real run is much slower. Something else started using the machine, or the run is thermally limited. On a laptop, plug it in. Compare the dt field early and late in the log; a steadily rising one is a thermal story.

The loss goes to nan in the first few steps. Most likely you overrode TOTAL_BATCH_SIZE without letting the script rescale the learning rate. Go back to the default and change one thing at a time. On Track M, try NANOCHAT_DTYPE=float32 to rule out a precision issue in an MPS kernel; the README’s own precision table gives float32 as the MPS default for exactly this kind of caution.

epoch: 2 appears partway through. The corpus ran out. Rerun the preparation script with a larger SHARDS and train again. The training tokens the run needs are the step count times the total batch size, and the corpus you need is more than that because packing crops documents.

The recording script cannot import nanochat. Run it with ~/nanochat/.venv/bin/python rather than a system Python, or pass --nanochat if you cloned somewhere else.

The CORE metric is close to random. That is the expected result for a model this small trained this briefly. Record it. It becomes interesting only when you compare two of your own runs.

Cleanup

The checkpoints and corpus shards are the largest artefacts. Everything under NANOCHAT_BASE_DIR can be deleted and regenerated by rerunning the two scripts, at the cost of the download and the training time.

RunnableAll tracks

see what it cost in disk before deciding

du -sh ~/.cache/nanochat/* 2>/dev/null

Keep the tokeniser directory and at least one checkpoint if you intend to do the project in this part, which reuses the pipeline on a corpus of your own. Keep the logs and train-fields.json regardless: they are small, and they are what the notebook line points at. The virtual environment in ~/nanochat/.venv is the expensive thing to rebuild and the cheapest to keep.

What you learned

• The pipeline is six scripts and each leaves a file. Environment, shards, tokeniser, training, evaluation, samples. When something looks wrong at the end, the file from an earlier stage is where the answer is.
• The tokeniser is a measurable object. Yours compresses your corpus better than two much larger published vocabularies and compresses Korean much worse, and you have the table to prove both.
• A time budget is an arithmetic problem, not a guess. Measure the throughput, divide, choose the steps. The same division sizes every training run you will ever plan.
• Undertrained is a description you can quantify. You know the compute-optimal step count and the count you ran, and the ratio between them is the size of your compromise.
• The logits tensor decides your batch size. Batch times sequence length times vocabulary, in FP32, is the term that sets peak memory on a small model, and it is why the course lowered a default that the upstream recipe sets higher.
• Fluency is not evidence. You watched grammar arrive long before content did, in a model you can account for entirely.

Record in the notebook: the commit and torch build; the tokeniser ratios; the compute-optimal and actual step counts; predicted and measured throughput with the gap explained; the utilisation, wall clock and peak memory; the validation bits per byte and the CORE metric; the power draw or a note that it was not measured; and the early and late sample blocks side by side.

Check your understanding

Question 1. The calibration step reports a throughput five times lower than your prediction from the Part 5 matrix rate. What is the most likely cause?

• The prediction formula is wrong
• The utilisation you assumed was too high; a training step is many small operations interleaved with attention, optimiser work and data movement, so it reaches a fraction of the peak matrix rate
• The tokeniser is inefficient
• The corpus is too small

Show the answer and why

Answer: The utilisation you assumed was too high; a training step is many small operations interleaved with attention, optimiser work and data movement, so it reaches a fraction of the peak matrix rate

The peak matrix rate is the ceiling for one enormous multiplication. Real utilisation on a small model with a short sequence is often well under a fifth, and the training loop prints it as bf16_mfu so you can replace your guess with the truth and redo the arithmetic.

Question 2. You halve DEVICE_BATCH_SIZE because the run ran out of memory. What else changes?

• The total number of tokens the model trains on
• The learning rate schedule
• Nothing about what the run computes: the script keeps the total batch size the same by doing twice as many gradient-accumulation micro-steps, so only the speed and the peak memory change
• The vocabulary size

Show the answer and why

Answer: Nothing about what the run computes: the script keeps the total batch size the same by doing twice as many gradient-accumulation micro-steps, so only the speed and the peak memory change

Gradient accumulation exists precisely so that the optimiser step is unaffected by how the work was split. That is why the total batch size, not the device batch size, is the number the learning rate and weight decay are scaled against.

Question 3. Why does the lab have you start the training script once with no iteration override, and then interrupt it?

• To warm up the disk cache
• To read the compute-optimal horizon the script derives from its tokens-per-parameter ratio, so that you know how far below it your time-budgeted run sits
• To compile the model
• To download the benchmark bundle

Show the answer and why

Answer: To read the compute-optimal horizon the script derives from its tokens-per-parameter ratio, so that you know how far below it your time-budgeted run sits

With no override the script prints the derived step count, the total tokens, the achieved ratio and the total FLOPs. Those four lines are the size of the compromise you are about to make, and recording them turns "undertrained" from a vague worry into a number.

Question 4. Your CORE metric comes out close to the random-guessing baseline. What should you do?

• Retrain with a higher learning rate
• Record it as the result: a model of this size trained for twenty-five minutes on a small sample of each task is expected to score near chance, and the number becomes useful only when comparing two of your own runs
• Leave the benchmark out of the notebook
• Increase the vocabulary size

Show the answer and why

Answer: Record it as the result: a model of this size trained for twenty-five minutes on a small sample of each task is expected to score near chance, and the number becomes useful only when comparing two of your own runs

A near-chance score is information, and hiding it is the failure the course exists to prevent. Bits per byte is the sensitive metric at this scale; CORE is there so that you have run the same benchmark the project leaderboard is scored on, at a scale where it barely moves.

Question 5. On Track M, what does the lab actually run?

• An MLX port of nanochat
• nanochat on PyTorch's MPS backend, installed through the CPU wheel extra; no MLX port of nanochat exists in its README or in mlx-examples as of 2026-09-09, and the mlx-examples transformer_lm example is a different, smaller exercise
• The CPU path only
• A container

Show the answer and why

Answer: nanochat on PyTorch's MPS backend, installed through the CPU wheel extra; no MLX port of nanochat exists in its README or in mlx-examples as of 2026-09-09, and the mlx-examples transformer_lm example is a different, smaller exercise

The pyproject routes the cpu extra to the CPU wheel index, and that wheel carries the MPS backend. The track is marked partial because the README says the author has not exercised all of these code paths and that the shrunken CPU and MPS runs will not give strong results.

Check my answersClear