baseGPT

A minimal, single-file PyTorch implementation of GPT-2 / GPT-3 style pretraining, ported from karpathy/llm.c's train_gpt2.cu. No custom CUDA kernels, no multi-GPU/ZeRO, no cuDNN attention — just plain torch.nn modules, F.scaled_dot_product_attention, and standard PyTorch autograd, so it's easy to read end to end and hack on.

• GPT-2/GPT-3 model definitions selectable by a short descriptor string (d12, d48, gpt3:c768, ...), weight-tied embeddings, matching weight init.
• AdamW with selective weight decay (2D+ tensors only), cosine LR schedule with linear warmup.
• Gradient accumulation, gradient clipping, bfloat16 autocast on CUDA.
• A parquet-shard data loader for the ClimbMix dataset, tokenized on the fly with tiktoken.
• Checkpointing/resume, periodic validation loss, and autoregressive text sampling.
• Console diagnostics matching train_gpt2.cu's output: a parameter summary table, z-score outlier detection on loss/grad-norm, and MFU% (matrix FLOPS utilization) estimates for common GPUs.
• A metrics CSV logger (--log-dir) plus a Jupyter notebook to plot loss, grad norm, z-scores, learning rate, throughput, and MFU% after a run.

• Python 3.10+
• torch (CUDA build recommended for real training; CPU works for smoke tests)
• tiktoken, pandas, pyarrow, numpy, requests, tqdm

pip install torch tiktoken pandas pyarrow numpy requests tqdm

Downloads parquet shards of the ClimbMix dataset from HuggingFace into data/:

python train_dataset.py --num-shards 170 --data-dir data

python train_base.py --model d12 --input-dir data --max-steps 1000 --log-dir logs

python train_base.py --model d12 --input-dir data --output-dir checkpoints --max-steps 32000 --warmup-iters 300 --final-lr-frac 0.1 --val-loss-every 500 --val-max-steps 40 --sample-every 500 --checkpoint-every 2000 --log-dir logs

Model size is picked via --model:

Run python train_base.py --help for the full list of flags (batch size, sequence length, learning rate schedule, checkpointing, sampling, etc).

After a run with --log-dir logs, open visualize_training.ipynb and run all cells to plot loss curves, gradient norm / z-score diagnostics, the learning rate schedule, and throughput/MFU%.

train_base.py           # model, optimizer, training loop, diagnostics
train_dataset.py        # ClimbMix parquet shard downloader
visualize_training.ipynb  # post-training matplotlib visualization
data/                    # downloaded parquet shards (not tracked in git)
logs/                    # metrics.csv written by --log-dir (not tracked in git)

This project is a Python/PyTorch port of the training loop in karpathy/llm.c, and uses the karpathy/climbmix-400b-shuffle dataset for pretraining.

MIT