Train and chat with a GPT-2-level language model on a $100 budget
Imagine chatting with a model you just trained in your terminal. This repository is a minimal experimental tool that handles the entire LLM training process, from tokenization to inference, on a single GPU node. It takes about two hours on eight H100 GPUs to create a GPT-2-level model, at a cost of approximately $48 (or about $15 with spot instances).
Model structure determined by a single layer count
You only need to specify the number of transformer layers via the --depth value. Based on this value, all hyperparameters such as model width, number of heads, learning rate, and training duration are automatically calculated and set to optimal states. Users only need to focus on adjusting the model size, while the remaining detailed settings are handled internally by the code. Currently, models that achieve GPT-2-level performance are generated in the range of depth 24 to 26. This design allows for experimentation based on readable and easy-to-modify code, without the need to deal with complex configuration objects or model factories.
GPT-2 arrival time leaderboard
Development focus is on optimizing the pre-training stage, which consumes the most computing resources. To this end, a leaderboard is maintained that measures the actual time taken to achieve GPT-2-grade performance. The evaluation metric is the DCLM CORE score, and the baseline GPT-2 CORE score is 0.256525. Unlike the approximately $43,000 cost for GPT-2 training in 2019, the same performance can now be achieved in about 1.5 hours on an 8XH100 node, at a cost of approximately $48 (or about $15 with spot instances). The runs/speedrun.sh script includes the full pipeline for training a model that meets this benchmark and chatting with it.
Installation and execution
Dependency management uses uv, and installation in CUDA environments is done with the uv sync --extra gpu command. For CPU or Apple Silicon (MPS) environments, use uv sync --extra cpu. If GPU memory is less than 80GB, you must reduce the --device-batch-size value to avoid out-of-memory errors. Execution is possible on a single GPU, but the time required increases eightfold.
Tools for researchers and experimenters
For research purposes, the runs/scaling_laws.sh and runs/miniseries.sh scripts are useful. An example of training a 12-layer model (GPT-1 size) for quick experiments is also provided, which completes in about 5 minutes. You can monitor validation loss (val_bpb), CORE score, VRAM usage, and model FLOPS utilization (MFU) via wandb. Precision (dtype) is automatically determined based on hardware, but it can also be manually specified using the NANOCHAT_DTYPE environment variable. bfloat16 is the default for CUDA SM 80 and above, while float32 is the default for lower versions.
Things to check before starting experiments
This repository aims to improve micro-models that work end-to-end on a budget of less than $1,000. It targets a minimal codebase that is easy to fork and modify, rather than a framework offering extensive configuration options. Therefore, it is suitable for those who want to understand the basic principles of LLM training or conduct research on increasing pre-training speed, rather than for adding specific features or building complex pipelines. If you do not have a GPU environment, execution is possible on CPU or Apple Silicon, but the model size is significantly reduced, making it difficult to obtain meaningful results. Additionally, there is a policy requiring disclosure when using AI for code contributions, so keep this in mind if you are considering contributing.
