NanoGPT Speedrun Living Worklog
#
+Description
Record time
Training Tokens
-Description
+Tokens/Second
Date
Commit
Log
@@ -88,14 +92,25 @@ NanoGPT Speedrun Living Worklog
-1
-8.13 hours
-6.44e+09
+1
Initial baseline
-2025–01–16
+8.13 hours
+6.44B
+220.7k
+2025/01/16
b3c32f8
here
+
+2.1
+Architectural changes
+7.51 hours
+5.07B
+187.7k
+2025/01/18
+b7bb93f
+here
+
1. Initial setup and baseline
Part of the goal of this project is for me to learn as I go, so I am going to start at the beginning - with with Andrej Karpathy’s PyTorch GPT-2 trainer from llm.c. This is the script that Keller Jordan used for his initial baseline. This trainer is very similar to the NanoGPT trainer with some minor modifications / simplifications (such as no dropout).
I have upstreamed some QOL improvements and basic tweaks to the training script from Keller’s fork, but have not changed any of the core training / modeling logic. Specifically:
@@ -104,15 +119,77 @@ NanoGPT Speedrun Living Worklog
- Improved learning rate schedule (linear warmup then linear decay).
- Removed all affine scale/bias parameters and switched to RMSNorm.
- Padded the vocab size from 50257 to 50304 to make it a multiple of 128 (for better tensor core utilization).
-
Additionally, I added wandb logging for easy tracking of training progress - optimistically I may need to remove this one day as it slightly increases step time.
Commit with the initial setup is here: b3c32f8.
The baseline run time on my 2xRTX 4090 setup is 8.13 hours.
NanoGPT Speedrun Living Worklog
1. Initial setup and baseline
Part of the goal of this project is for me to learn as I go, so I am going to start at the beginning - with with Andrej Karpathy’s PyTorch GPT-2 trainer from llm.c. This is the script that Keller Jordan used for his initial baseline. This trainer is very similar to the NanoGPT trainer with some minor modifications / simplifications (such as no dropout).
I have upstreamed some QOL improvements and basic tweaks to the training script from Keller’s fork, but have not changed any of the core training / modeling logic. Specifically:
-
@@ -104,15 +119,77 @@
- Improved learning rate schedule (linear warmup then linear decay).
- Removed all affine scale/bias parameters and switched to RMSNorm.
- Padded the vocab size from 50257 to 50304 to make it a multiple of 128 (for better tensor core utilization). -
NanoGPT Speedrun Living Worklog
Additionally, I added wandb logging for easy tracking of training progress - optimistically I may need to remove this one day as it slightly increases step time.
Commit with the initial setup is here: b3c32f8.
The baseline run time on my 2xRTX 4090 setup is 8.13 hours.
Additionally, I added wandb logging for easy tracking of training progress - optimistically I may need to remove this one day as it slightly increases step time.
Commit with the initial setup is here: b3c32f8.
The baseline run time on my 2xRTX 4090 setup is 8.13 hours.
+2. Implementing major improvements from the 8xH100 leaderboard
Waiting 8 hours for a result, so I’m going to begin by implementing some of the notable improvements from the 8xH100 leaderboard. I’ll start with the most impactful/easiest changes first:
-
+
-
+
Architectural changes (31.8% speedup, then 24% speedup)
+ + +
+
2.1 Architectural changes and training tweaks
There are some basic architectural changes and modernizations that can be made to the model that will speed up training. These changes are general improvements to the transformer decoder architecture that have been generally adopted since the original GPT-2 paper. The changes are:
-
+
- RoPE (Rotary Positional Embeddings). There are many good explanations of RoPE out there so I won’t go into detail here. +
- ReLU^2 ActivationReLU^2 activation function.
+ - No gradient clipping. Gradient clipping can help stabilize training but it also slows down training. Since we are speed-running, we will remove gradient clipping. This also eliminates a hyperparameter that needs to be tuned. +
- Trapezoidal learning rate schedule. While cosine learning rate schedules are the de-facto standard, they can be difficult to work with since changing the number of training steps changes the entire schedule. Trapezoidal learning rate schedules are often easier to reason about / tune around, and they have been show to match the performance of cosine schedules. +
In addition, learning rate and batch size have been tuned.
Once again, many of these changes are downstreamed from the modded-nanogpt repository / 8xH100 speedrun. Its not efficient to reinvent the wheel, and I want to get training time down as fast as possible in the beginning.
After implementing these changes (commit b7bb93f), the new run time is 7.51 hours. This run was more data-efficient than the baseline, requiring only 5.07B tokens. However, the tokens/second increased, likely due to the larger batch size (more gradient accumulation steps which tends to translate to lower throughput) and the architectural changes, such as the inclusion of RoPE. Once I have a shorter run time, I will be able to tune more effectively and see if I can remove gradient accumulation.
