NanoGPT Speedrun Worklog

I saw some really awesome GPT-2 speedrun results from people like Keller Jordan, Fern, Braden Koszarsky, and others. I got a little inspired and decided to see how fast I could train GPT-2 on my own hardware.

Technically, the NanoGPT speedrun is to train a neural network to 3.28 validation loss on FineWeb as fast as possible on an 8xH100 node. Keller Jordan maintains a leaderboard here. When I started this experiment on January 16, 2025, the record was 3.14 minutes (!).

I had access to 2xRTX 4090 GPUs, so I followed the same rules on my own hardware. Over six iterations, I reduced the training time from 8.13 hours to 2.55 hours. This worklog records the changes that produced that result; the code and run logs are available in the project repository.

Results

1. Initial setup and baseline

Part of the goal of this project was to learn as I went, so I started at the beginning—with Andrej Karpathy's PyTorch GPT-2 trainer from llm.c. This is the script that Keller Jordan used for his initial baseline. The trainer is very similar to NanoGPT with some minor modifications and simplifications, such as removing dropout.

I upstreamed some QOL improvements and basic tweaks to the training script from Keller's fork, but did not change any of the core training or modeling logic. Specifically:

1. Implemented gradient accumulation so that my 2x24GB GPUs simulate the training experience of an 8xH100 machine.
2. Increased learning rate to 0.0015 and halved the batch size (total batch size is 262144 - that is bs of 32/device * 2 devices * 1024 sequence length * 4 gradient accum steps).
3. Improved learning rate schedule (linear warmup then linear decay).
4. Removed all affine scale/bias parameters and switched to RMSNorm.
5. Padded the vocab size from 50257 to 50304 to make it a multiple of 128 (for better tensor core utilization).
6. Used PyTorch 2.5.1 (the switch from 2.4 to 2.5 gave ~9% speedup on the 8xH100 leaderboard).

Additionally, I added wandb logging to make it easier to track the training runs.

Commit with the initial setup is here: b3c32f8.

The baseline run time on my 2xRTX 4090 setup was 8.13 hours.

2. Implementing major improvements from the 8xH100 leaderboard

Waiting 8 hours for a result was too slow for effective experimentation, so I began by implementing some of the notable improvements from the 8xH100 leaderboard. I started with the most impactful and easiest changes:

1. Architectural changes and training tweaks
2. Muon optimizer
3. Dataloading tweaks
4. Logit Softcapping

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:

1. RoPE (Rotary Positional Embeddings). There are many good explanations of RoPE out there so I won't go into detail here.
2. ReLU^2 ActivationReLU^2 activation function. . Many activations that are better than GeLU have been proposed since GPT-2. ReLU^2 is a simple one that has been shown to be effective in decreasing training time required to reach a certain validation loss.
3. No gradient clipping. Gradient clipping can help stabilize training, but it also slows down training. Since this was a speedrun, I removed it. This also eliminated a hyperparameter that needed to be tuned.
4. 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 shown to match the performance of cosine schedules.

In addition, I tuned the learning rate and batch size.

Once again, many of these changes were downstreamed from the modded-nanogpt repository / 8xH100 speedrun. It wasn't efficient to reinvent the wheel, and I wanted to reduce training time quickly before doing more targeted experiments.

After implementing these changes (commit b7bb93f), the new run time was 7.51 hours. This run was more data-efficient than the baseline, requiring only 5.07B tokens. However, the tokens/second decreased, likely due to the larger batch size (more gradient accumulation steps tends to translate to lower throughput) and architectural changes such as the inclusion of RoPE. The shorter run time made further experimentation more practical.

2.2 Muon Optimizer

The Muon Optimizer was developed with and for the NanoGPT speedrun by Jordan et al. It is a variant of SGD with Momentum that applies a postprocessing step to the gradient updates to approximately orthogonalize each update matrix. Muon has some connections to approximate second-order optimizersBut are these approximate second-order methods actually second-order? New research suggests that methods like Shampoo and Adam can be viewed as variants of steepest descent under specific norms, and thus are actually first-order methods. like Shampoo.

I highly recommend reading the original Muon blog post for more details, as well as checking out the optimizer comparison for GPT-2 speedrunning that Keller Jordan put together here. For those interested in a more step-by-step walkthrough of Muon, check out this excellent post by Jeremy Bernstein.

Muon is designed to work on Linear layers, so it is not quite a drop-in replacement for AdamW (e.g. it isn't meant to optimize Embedding layers). However, it can be used to optimize all of the hidden layers of our GPT-2 model. The output lm_head layer and token embeddings were still optimized with AdamW.

Just like on the 8xH100 leaderboard, I observed a massive speedup after switching to Muon. The new run time was 4.53 hours, requiring only 3.04B tokens. The tokens/second remained very similar to the previous run, indicating that the optimizer change did not sacrifice throughput.

2.3 Dataloading Tweaks

The architecture and optimizer changes improved data efficiency, but training throughput dropped from 221k tokens/second to 187k tokens/second—a decrease of roughly 15%. Recovering that throughput offered another path to a shorter run time, so I turned to the dataloading and gradient accumulation logic.

Up to this point, I had loaded a full batch of data on each device and then split it into smaller chunks (micro-batches) for each gradient accumulation step. I changed the logic to load only the next micro-batch and advance the dataloader for each accumulation step.

I also upgraded PyTorch from 2.5 to 2.6, which had recently been released, and removed torch._inductor.config.coordinate_descent_tuning in accordance with the official rules introduced on February 1, 2025.

These tweaks brought throughput back up to 216k tokens/second. To make runs more consistently hit the 3.28 validation loss targetThere is some variance in how long a speedrun candidate takes to reach the target. An official record must attain a mean validation loss below 3.28. I treated this somewhat loosely in the early experiments because the differences between runs were much larger than the observed variance., I also slightly increased the total number of training steps, bringing consumption to 3.31B tokens. The new run time was 4.26 hours, and the changes can be found at d59944d.

At this point, the training time was almost half the baseline.

2.4 Logit Soft-capping

Logit soft-capping is a technique popularized by Gemma 2 and initially used to improve the NanoGPT speedrun by @Grad62304977.

Soft-capping is essentially a smooth and differentiable version of clipping⊕Soft-capping vs Clipping at ±5: : \[ \text{softcap(x, cap)} = \text{cap} \cdot \tanh\left(\frac{\text{x}}{\text{cap}}\right) \]

Logit soft-capping prevents logits from growing excessively large by scaling them to a fixed range, which seems to help improve training dynamics. One could argue that this imposes an inductive bias—and in this relatively small-model, low-data regime, that bias appeared to be helpful.

After implementing logit soft-capping with a cap of 30 (and doing some learning-rate tuning), the new run time was 4.01 hours, requiring 3.15B tokens (commit 12eab44). Throughput remained steady at ~218k tokens/second.

3. Longer Training and Evaluation Sequence Length

Up to this point, I had trained and evaluated on sequences of 1024 tokens without being particularly clever about how those sequences were constructed. At each step, the dataloader simply placed the next 1024 tokens into an element of the batch without regard for document boundaries. That meant frequently starting in the middle of a document, cutting it off before its end, and attending to tokens across documents through a simple causal mask.

Cutting off documents in the middle is an especially large issue. See this plot of average loss vs sequence position:

Notice how the first twenty-five or so positions have a much higher average loss than later positions. At the beginning of a sequence, the model has much less context with which to predict the next token. Avoiding needless sequence restarts offered a way to reduce this loss penalty.

A natural question to ask at this point is: how long are sequences in our dataset, on average?

The data revealed that approximately 20% of documents exceeded the 1024-token sequence length. Increasing the sequence length to >=8192 tokens would accommodate virtually all documents in the dataset without truncation.

To address these issues, I made two key changes. First, I extended the sequence length to minimize document splitting across sequence boundaries. Taking this approach to its logical conclusion, I eliminated the traditional batch dimension and effectively used a "batch size" of 1 containing multiple concatenated documents. Second, I added attention masking that prevented cross-document attention while retaining the computational efficiency of sparse attention patterns.

Fortunately, FlexAttention provides an elegant solution that maintains the performance benefits of FlashAttention while enabling these improvements. One of FlexAttention's primary strengths is its ability to efficiently handle sparse, custom attention masks, making it ideal for our use case.

To implement FlexAttention, I defined a mask that handled the specific requirements of the dataset:

def make_attn_mask(idx, eot_token, window_size=1024):
    # Create a causal mask (only attend to past tokens)
    def causal_mask(b, h, q_idx, kv_idx):
        return q_idx >= kv_idx

    # Track document boundaries using end-of-text tokens
    documents = (idx == eot_token).cumsum(dim=1)

    # Only allow attention within the same document
    def document_mask(b, h, q_idx, kv_idx):
        return documents[b, q_idx] == documents[b, kv_idx]

    # Limit attention to an N-token window for efficiency
    def sliding_window_mask(b, h, q_idx, kv_idx):
        return q_idx - kv_idx <= window_size

    return and_masks(document_mask, causal_mask, sliding_window_mask)

Let's break down each mask:

1. Causal Mask: Standard in autoregressive language modeling. Ensures that tokens can only attend to previous tokens in the sequence, preventing information leakage from future tokens.

2. Document Mask: This restricts attention to tokens within the same document. By tracking document boundaries using end-of-text tokens, we prevent tokens from attending across different documents, which helps the model maintain coherent context within a single document.

3. Sliding Window Mask: This limits attention to a fixed window of tokens before the current position. This approach balances efficiency with context retention with a clear tradeoff: smaller windows are more efficient but may miss long-range dependencies, while larger windows capture more context at the expense of resources.

The individual component masks are visualized below:

When combined with the and_masks function, these three masksNote that the causal mask is actually redundant to the sliding window mask, as the sliding window mask already ensures that tokens can only attend to previous tokens in the sequence. The causal mask is included here for clarity. work together to create an efficient attention pattern that respects document boundaries, maintains causality, and limits computational overhead for long sequences.

After incorporating FlexAttention with these masks and increasing the sequence length to 32768 tokens, I observed another massive speedupThis speedup is a bit of a hack against the target metric. Supporting longer sequences is a straightforward way to drop the loss on the validation set, but is unlikely to provide a meaningful improvement to the overall performance of the model on practical benchmarks.. The final run time was 2.55 hours, requiring only 1.88B tokens—a large data-efficiency improvement. Throughput dropped slightly to ~205k tokens/second. See commit d982ed5 for the full details.

Summary

The final 2.55-hour run was 3.2x faster than the 8.13-hour baseline and used 1.88B training tokens instead of 6.44B. Muon produced the largest improvement that was independent of the evaluation setup. The final reduction came from longer, document-aware sequences and was more specific to the validation-loss target.

This post records the completed 2025 round of experiments, with 2.55 hours as the last measured result. I may revisit the speedrun in the future, either on the same GPUs or on different hardware. The numbers here are specific to the two-RTX-4090 setup and are not directly comparable to the 8xH100 leaderboard times.

References