How to train LearnGPT
A copy-ready path from a clean clone to a validated FineWeb-Edu dataset, a complete 17.7M-parameter training run, and seeded text generation.
ResultLearnGPT controlled 18M - 17.7M
- parameters
- 256
- context tokens
- 6
- Transformer blocks
- 8,192
- tokens per update
- 45,000
- optimizer updates
Reproducible 1 GiB subset of the canonical 10 GiB FineWeb-Edu preparation
Apple Silicon · MPSmacOS on an M-series Mac. This is the profile used by the verified 45,000-step reference run.
-
git clone https://github.com/ferdinandobons/learn-gpt.git
cd learn-gpt
python3 -m venv .venv
.venv/bin/python -m pip install --upgrade pip setuptools wheel
.venv/bin/python -m pip install -r final_project/requirements.txt
Continue whenThe environment interpreter imports PyTorch without an error.
-
Verify MPS and project gatesbash .venv/bin/python -c "import torch; print(torch.__version__); print(torch.backends.mps.is_built(), torch.backends.mps.is_available()); print(torch.ones(1, device='mps'))"
.venv/bin/python -B tools/validate_learngpt.py
.venv/bin/python -B -m unittest discover -s tests -v
Continue whenThe backend is available and every repository test passes before any large download starts.
-
Prepare the canonical corpusbash .venv/bin/python -B -m final_project.prepare_data \
--target-gb 10 \
--output-dir data/processed/fineweb_edu
Continue whendata/processed/fineweb_edu contains train.bin, val.bin, and meta.json with complete=true.
-
Create and validate the subsetbash .venv/bin/python -B -m final_project.prepare_subset \
--source-data-dir data/processed/fineweb_edu \
--output-dir data/processed/fineweb_edu_experiment_1g \
--target-gb 1 \
--validation-ratio 0.01 \
--seed 1337 \
--chunk-tokens 65536
.venv/bin/python -B tools/validate_learngpt.py \
--training-data-dir data/processed/fineweb_edu_experiment_1g
Continue whenValidation passes and the subset metadata records seed 1337, chunk size 65536, and a 1% validation split.
-
One-step disposable MPS smoke testbash .venv/bin/python -B -m final_project.training \
--device mps \
--data-dir data/processed/fineweb_edu_experiment_1g \
--checkpoint-path /tmp/learngpt-mps-smoke.pt \
--overwrite-checkpoints \
--context-size 8 \
--embedding-size 16 \
--num-heads 4 \
--num-transformer-blocks 1 \
--batch-size 1 \
--gradient-accumulation-steps 1 \
--training-steps 1 \
--eval-interval 1 \
--eval-batches 1 \
--base-learning-rate 1e-4 \
--min-learning-rate 1e-5 \
--warmup-steps 0 \
--decay-steps 1
Continue whenLoss is finite, the first loss is near 10.82, and best/latest checkpoint files are created.
-
Fresh 45,000-step MPS runbash caffeinate -i .venv/bin/python -B -m final_project.training \
--device mps \
--data-dir data/processed/fineweb_edu_experiment_1g \
--checkpoint-path checkpoints/learngpt-mps-18m-stable-1g-v2.pt \
--encoding-name gpt2 \
--seed 1337 \
--context-size 256 \
--embedding-size 256 \
--num-heads 4 \
--num-transformer-blocks 6 \
--dropout 0.0 \
--use-scaled-dot-product-attention \
--output-chunk-size 32768 \
--batch-size 4 \
--gradient-accumulation-steps 8 \
--training-steps 45000 \
--eval-interval 250 \
--eval-batches 20 \
--base-learning-rate 3e-4 \
--min-learning-rate 3e-5 \
--warmup-steps 1000 \
--decay-steps 45000 \
--weight-decay 0.05 \
--gradient-clip 1.0 \
--max-grad-norm-before-clip 100 \
--gradient-retry-attempts 3 \
--context-sensitivity-contexts 32
Continue whenValidation loss trends downward, retries normally remain zero, and the latest checkpoint advances every evaluation interval.
-
Generate three seeded MPS samplesbash .venv/bin/python -B -m final_project.generate \
--device mps \
--checkpoint-path checkpoints/learngpt-mps-18m-stable-1g-v2.pt \
--prompt "The purpose of education is" \
--max-new-tokens 120 \
--temperature 0.8 \
--top-k 40 \
--num-samples 3 \
--seed 1337
Continue whenContinuations are plausible, react to different prompts, and are evaluated as base-model completions rather than chat answers.
If training stops, resume it.
--training-steps 45000 is the total target, not 45,000 additional steps. Resume from -latest.pt; generate from the best file without that suffix.
Resume an interrupted MPS runbash caffeinate -i .venv/bin/python -B -m final_project.training \
--device mps \
--data-dir data/processed/fineweb_edu_experiment_1g \
--checkpoint-path checkpoints/learngpt-mps-18m-stable-1g-v2.pt \
--resume-checkpoint-path checkpoints/learngpt-mps-18m-stable-1g-v2-latest.pt \
--training-steps 45000
Model shape6 parameters+
--context-size 256 Number of tokens visible at once. Larger contexts cost substantially more memory.
--embedding-size 256 Width of every token representation; the main capacity and memory control.
--num-heads 4 Splits the 256-wide state into four attention views.
--num-transformer-blocks 6 Adds depth while keeping the model near 17.7M parameters.
--dropout 0.0 The controlled subset and run use no stochastic dropout, improving reproducibility.
--encoding-name gpt2 Selects the 50,257-token GPT-2 BPE vocabulary used by the dataset.
Batch and duration4 parameters+
--batch-size 4 Micro-batches processed together; lower this first if the GPU runs out of memory.
--gradient-accumulation-steps 8 Combines eight micro-batches before AdamW, producing 8,192 tokens per update.
--training-steps 45000 Total optimizer-update target, including completed steps after resume.
--seed 1337 Makes initialization, batches, subset selection, and samples reproducible.
Learning schedule5 parameters+
--base-learning-rate 3e-4 Peak AdamW learning rate after warm-up.
--min-learning-rate 3e-5 Cosine-decay floor used near the end of training.
--warmup-steps 1000 Raises the rate gradually so random initial weights do not receive an abrupt full update.
--decay-steps 45000 Spreads cosine decay across the entire controlled run.
--weight-decay 0.05 Regularizes matrix weights through AdamW without applying decay to biases and normalization scales.
Evaluation and safety10 parameters+
--eval-interval 250 Measures train and validation loss every 250 updates without evaluating constantly.
--eval-batches 20 Averages enough batches to reduce noise while keeping evaluation affordable.
--gradient-clip 1.0 Limits the applied update after raw-gradient integrity checks pass.
--max-grad-norm-before-clip 100 Rejects a suspicious raw gradient instead of hiding it with clipping.
--gradient-retry-attempts 3 MPS / 0 CUDA MPS can retry an invalid raw gradient; CUDA handles transient FP16 overflow separately.
--context-sensitivity-contexts 32 Checks whether correct context eventually beats shuffled context.
--output-chunk-size 32768 MPS / 0 CUDA MPS chunks the large vocabulary projection to preserve correct gradients; CUDA uses one projection.
--fused-attention CUDA Combines all heads into one QKV projection and one batched SDPA call per block; legacy checkpoints retain separate heads.
--log-interval 50 CUDA / optional MPS Prints lightweight progress without running validation, context diagnostics, or checkpoint writes.
--mixed-precision CUDA only CUDA uses FP16 autocast and a resumable GradScaler; the controlled MPS profile intentionally disables it.
Read the run without guessing.
✓ Healthy run signals
- The first loss is close to 10.82, the random baseline for a 50,257-token vocabulary.
- MPS reports a successful repeated-gradient and CPU-parity self-check before step 1.
- grad_norm is finite and grad_retries normally remains 0.
- On CUDA, amp_retries normally remains 0; a small stable amp_overflows count means transient FP16 overflow was retried safely.
- Validation loss improves as a trend even when individual evaluations fluctuate.
- context_loss_gain becomes positive as the correct context becomes more useful than shuffled context.
- Different prompts produce meaningfully different seeded continuations.
! Stop and investigate
- A non-finite loss or a non-AMP gradient appears.
- The MPS startup parity check fails.
- The raw gradient guard exhausts all retries.
- CUDA exhausts eight same-step AMP retries or amp_overflows keeps growing persistently.
- Validation loss persistently worsens while training loss keeps falling.
- The operating system reports sustained thermal or memory pressure.
- A resume command reports a dataset fingerprint mismatch.
Open the detailed references only when needed.