I trained a 124M-parameter large language model from scratch on FineWeb-10B. This is the work-log where I document every observation, every bottleneck, and every optimization.

Here’s what the model says after training:

prompt: What is the answer to life, the universe, and everything?
generated_text: The answer is yes. What can I do for my life, my family or even a little piece of me with all that I have lost in this time? How could I have changed that and made it work again on my own? How could I have made some changes to my life over the years so that it has been taken into account when it comes back from this loss? How could I make myself more productive? In addition, what would be my priority if I had no other words to say to me now? Would I choose to follow up any of these things and give up? What would happen if someone else was involved in my
  1. Baseline
  2. Making training fast
  3. Architectural Changes
  4. Going All In

Baseline

I started simple. Plain PyTorch, no tricks. Just the standard GPT-2 architecture so I have something that works before making it fast.

Model

config:

@dataclass
class GPTConfig:
    vocab_size: int = 50257
    context_length: int = 1024
    n_embed: int = 768
    num_layers: int = 12
    n_head: int = 12

Architecture:

  • Token embeddings + position embeddings (learned)
  • 12 transformer blocks (attention + MLP)
  • Pre-norm (LayerNorm before attention/MLP)
  • Residual connections
  • GELU activation
class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embed),
            wpe = nn.Embedding(config.context_length, config.n_embed),
            h = nn.ModuleList([Block(config) for _ in range(config.num_layers)]),
            ln_f = nn.LayerNorm(config.n_embed)
        ))
        self.lm_head = nn.Linear(config.n_embed, config.vocab_size, bias=False)
        # Weight sharing: saves ~40M parameters
        self.transformer.wte.weight = self.lm_head.weight

Weight initialization

def _init_weights(self, module):
    if isinstance(module, nn.Linear):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            torch.nn.init.zeros_(module.bias)
    if isinstance(module, nn.Embedding):
        torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

I use std=0.02. GPT-2 paper used this. It is close to Xavier scaling 1/√768 ≈ 0.036 but 0.02 works fine. I also scale the attention output projection by 1/√(2 * num_layers) to keep the residual stream stable as the network gets deeper. This trick comes from the modded-nanogpt speedrun codebase.

Data

Dataset: FineWeb 10B (sample-10BT)

  • Development: 20M-token subset
  • Full training later: 10B

Tokenizer: GPT-2 BPE (tiktoken).

Loader:

class TokenDataLoader:
    def __init__(self, data_root, B, T):
        self.B = B
        self.T = T
        self.shards = sorted([...])
    
    def next_batch(self):
        buf = self.tokens[self.current_position : self.current_position + needed]
        x = buf[:-1].view(B, T)  
        y = buf[1:].view(B, T)    
        return x, y

Training Setup

Hyperparameters:

total_batch_size = 524288       
batch_size = 32               
sequence_length = 1024        
grad_accum_steps = 16         
learning_rate = 3e-4

Training loop:

for step in range(max_steps):
    optimizer.zero_grad()
    loss_accum = 0.0
    for micro_step in range(grad_accum_steps):
        x, y = train_loader.next_batch()
        x, y = x.to(device), y.to(device)   
        _, loss = model(x, y)
        loss = loss / grad_accum_steps 
        loss_accum += loss.detach()
        loss.backward()
    optimizer.step()

So here is our starting point: it took 20 minutes to process 20M tokens on an A100 GPU. We hit 16,700 tokens/sec with a final loss of 7.04. The code works, but it’s so slow. Let’s make it faster. Check the log

Making training fast

The math is correct, but it takes too long. PyTorch has a few built-in tricks to make the GPU work smarter.

Mixed precision

32-bit floats give you the best accuracy, but training is slow for two reasons: standard 32-bit math runs on slower CUDA cores instead of the ultra-fast Tensor Cores, and moving all those heavy 32-bit numbers around eats up your memory bandwidth. The idea behind mixed precision is simple: drop to lower precision where we can get away with it, keep 32-bit where it matters, and let the tensor cores do the heavy work.

Precision comparison

TensorFloat-32 (TF32): TF32 is not a storage format. Your data stays in FP32 in memory, all 32 bits. But when a matrix multiply hits the tensor cores, the hardware internally truncates each input’s mantissa from 23 bits down to 10 bits, does the multiply in this trimmed format, then accumulates the result back in full FP32. So the math inside the tensor cores is faster, but the data still moves around in 32 bits. You get faster compute, but no memory bandwidth savings. One line turns it on:

torch.set_float32_matmul_precision('high')

On the A100, TF32 tensor core throughput is 156 TFLOPS, compared to 19.5 TFLOPS for standard FP32 on CUDA cores. That’s an 8x speedup on the math alone, with no code changes.

True mixed precision: BF16 goes further. It’s an actual 16-bit storage format. Massive activation tensors are stored and moved as 16-bit values. That cuts memory bandwidth in half and doubles the amount of data that fits in cache. Your master weights safely stay in 32-bit to prevent precision loss. On top of that, the A100’s tensor cores run BF16 at 312 TFLOPS.

PyTorch’s torch.autocast handles the precision decisions automatically: operations like matmuls and convolutions drop to BF16, while sensitive operations like loss computation and reductions stay in FP32.

You have two choices for lower precision: Float16 or Bfloat16.

Float16 is fast and runs anywhere, but it has a small exponent (5 bits), which gives it a narrow numeric range. Small gradients can underflow to zero, which breaks training. To fix that, use a GradScaler to scale up the loss during the backward pass, then scale it back down before the optimizer steps.

Float16 recipe:

scaler = torch.cuda.amp.GradScaler() 
with torch.autocast(device_type='cuda', dtype=torch.float16):
    _, loss = model(x, y)
    loss = loss / grad_accum_steps

scaler.scale(loss).backward() # Scale gradients
scaler.step(optimizer)        # Unscale & Step
scaler.update()               # Update scale factor

Bfloat16 has the same 8-bit exponent as FP32, so it covers the same numeric range. Gradients don’t underflow. No scaler needed. The tradeoff is less decimal precision (7-bit mantissa vs FP16’s 10-bit), but in practice this doesn’t hurt training.

Bfloat16 recipe:

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    _, loss = model(x, y)
    loss = loss / grad_accum_steps
    loss_accum += loss.detach()

# Backward pass (No scaler needed)
loss.backward()

I set torch.set_float32_matmul_precision('high') globally and wrap the forward pass in torch.autocast(dtype=torch.bfloat16). Autocast runs most operations in BF16 to save memory bandwidth and maximize speed. For the sensitive operations it leaves in FP32, the TF32 setting catches them and forces them onto the fast tensor cores too.

torch.compile

One line that makes the whole model faster.

model = torch.compile(model, dynamic=True)

Normally, PyTorch is eager. It launches one GPU kernel per operation. Each time, the GPU loads data from its main memory (HBM), runs the math, and writes it back.

torch.compile() does two things:

  1. Kernel fusion: It merges multiple operations into one kernel. The data stays in on-chip SRAM instead of bouncing back and forth to HBM.
  2. Graph optimization: It removes redundant operations, reorders the math, and picks the right CUDA kernels for the hardware.

Why dynamic=True? PyTorch typically compiles kernels expecting exact, fixed tensor shapes. If your input dimensions vary even slightly, the compiler halts training for a CPU recompilation. The dynamic=True flag forces it to generate flexible kernels that gracefully handle varying dimensions.

The first run takes 1–2 minutes to build these custom kernels. After that, the training steps just fly. If you want a deep dive into this compilation magic, the official PyTorch compiler tutorial is an excellent resource.

Flash Attention

Self-attention is mathematically simple: just two matrix multiplies (Q @ K^T, then @ V) and a softmax. But the memory round trips are what kills performance. Here is what happens in baseline code:

  1. You compute Q @ K^T and write that N×N matrix to the GPU’s slow main memory (HBM).
  2. You read it back from HBM to apply softmax, then write the new N×N probabilities back to HBM.
  3. You read it from HBM one last time to multiply it by the V matrix.

During the backward pass, the same thing repeats. You either keep that giant matrix in HBM the whole time, or you throw it away and recompute it, which means doing all those HBM round-trips again. Either way, the tensor cores end up sitting idle, waiting for data.

The bottleneck isn’t the math. It’s the memory bandwidth. Flash Attention fixes this with one idea: stay in SRAM.

It splits the matrices into small blocks that fit inside the on-chip SRAM, and fuses the entire attention computation into a single kernel:

  1. Move a tiny block from slow HBM to fast SRAM.
  2. Multiply it.
  3. Normalize it on the fly.
  4. Multiply by the V block.
  5. Write only the final, small output back to HBM.

That’s it. The full attention matrix never materializes in HBM. During backprop, it recomputes everything block-by-block in SRAM. The extra math is cheap compared to the cost of reading from HBM.

PyTorch ≥ 2.0 has it built-in. One line:

attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)

The Results. Just by adding these three things, the same model on the same 20M tokens on an A100 went from 16,700 tok/sec to 200,000 tokens/sec. Training time dropped from 20 minutes to 2.66 minutes. The loss stayed at 7.03. Same hardware, same math, just faster. Check the log

Architectural Changes

The memory bottlenecks are gone, but the model itself is still standard 2019-era GPT-2. We can swap out a few components for newer alternatives that are faster and train better.

Rotary Positional Embeddings

So far, the model implicitly encoded position like the original GPT-2: it learned a giant 1024x768 lookup table for the token positions, fetching one vector per index and adding it directly to the token embedding.

wpe = nn.Embedding(config.context_length, config.n_embed)

This works, but the model just memorizes absolute positions as vectors. It has no way to know that the distance between token 5 and 6 is the same as between 505 and 506. We want the attention mechanism to understand relative distance.

Rotary Positional Embeddings (RoPE). Instead of adding a learned position vector, we rotate the Query and Key vectors by an angle that corresponds to their position.

If you want a good deep-dive into this, I recommend Zachary Huang’s Give me 30 min, I will make RoPE click forever video.

RMSNorm

The standard GPT-2 uses LayerNorm. LayerNorm computes the mean, subtracts it to center the values, and then applies a learned scale and bias.

I replaced it with an RMSNorm (Root Mean Square Normalization). It drops the mean-centering step, the learned bias, and the learned scale weights. It just scales by the root mean square. Simpler, cheaper, and doesn’t hurt performance.

def norm(x: torch.Tensor):
    return F.rms_norm(x, (x.size(-1),))

Squared ReLU

I also swapped out the standard GELU activation function inside the MLP for Squared ReLU (ReLU²).

GELU is mathematically squishy. It relies on tanh and expensive exponential functions under the hood to smooth out its curve, which are slow to evaluate. ReLU² is literally just max(0, x)^2.

Vocabulary Padding

The official GPT-2 tokenizer vocabulary size is 50,257.

This is an awkward number for GPUs. Tensor Cores work best with dimensions that are clean multiples. When the embedding matrix doesn’t line up, you get misalignment and wasted cycles.

So I padded the vocabulary size to 50,304 (50,257 + 47). This makes the embedding and lm-head dimensions nice round numbers. Small change, easy speedup.

The Results. With these architectural changes on the same 20M-token slice, throughput went up to 205,000 tokens/sec. Training finished in 1.82 minutes and the final loss dropped to 6.93. Check the log

Going All In

The first three sections were just experiments on 20M tokens. Now for the real thing: training on the full FineWeb-10B dataset. The target is to hit around val_loss ≤ 3.28 in about two hours on a single A100.

To hit this target, I made some upgrades in three areas: the model architecture, the data loader, and the optimizer.

Model Upgrades

QK-Norm. At high learning rates, the Query and Key vectors can grow too large. Their dot products overflow, NaNs appear, and the run dies.

To fix this, I applied RMSNorm to the Queries and Keys right after their linear projections (before RoPE).

q = F.rms_norm(q, (H,))
k = F.rms_norm(k, (H,))

Without these two lines, the attention scores explode at Muon’s aggressive 0.02 learning rate.

U-Net Skips. Vanilla transformer layers only see out to the layer directly below them. Taking inspiration from U-Nets, I split the stack. Layers 1-6 stash their residuals in a list. Layers 7-12 pop those residuals off and add skips from their symmetrically mirrored shallow layer (Layer 7 pulls from 6, Layer 8 pulls from 5, up to Layer 12 pulls from 1).

skips = []
for i, block in enumerate(self.blocks):
    if i < self.half_n:  
        skips.append(x)
    else:        
        x = x + skips.pop()
    x = block(x, cu_seqlens, pos, max_seqlen)

Shallow signals (like fundamental syntax and frequency patterns) reach deep into the network without going through twelve sequential transforms.

Logit Soft-Capping. High entropy is good, but unbounded logits can let a single class dominate the softmax distribution. The model collapses, entropy vanishes, and it stops learning the long tail of the vocabulary.

Logit soft-capping, popularized by Gemma 2 and the modded-nanogpt is essentially a smooth, differentiable version of clipping. Instead of a hard cut-off, it uses a scaled tanh to gently squash logits into a fixed range (usually ±30):

logits = self.config.logit_softcap * torch.tanh(logits / self.config.logit_softcap)

Training Upgrades

Document Packing. The basic data loader blindly cuts one long stream of text into chunks. This is bad because the model ends up reading half of one document and half of a completely unrelated document in the same sequence.

The easy fix is to just pack your documents perfectly back-to-back, separating them with a <|endoftext|> token and give it to the model as a single long sequence. But how do you prevent the model from looking across those packed boundaries?

The easiest way to handle this is by using FlashAttention library’s varlen function. Simply give it a list of your document lengths (cu_seqlens). It automatically resets the attention mask at every boundary, guaranteeing the model never reads across different documents. No complicated custom masking required:

y = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
    dropout_p=0.0, causal=True,
)

Muon Optimizer. AdamW is the standard default, but we can do a lot better.

Muon (Momentum Orthogonalized by Newton–Schulz) is an optimizer designed for the NanoGPT speedrun by Jordan et al. See the original Muon blog post for the deep mathematical details.

Muon relies on 2D matrix properties of internal linear layers, so it’s not a drop-in replacement for everything. I split the workload: Muon handles the transformer blocks, AdamW handles the token embeddings and the language model head.

muon_optimizer = Muon(model.blocks.parameters(), lr=0.02, weight_decay=0.01)
adam_optimizer = torch.optim.AdamW(
    model.wte.parameters(), lr=0.0036, betas=(0.9, 0.95), weight_decay=0.1, fused=True,
)

Learning-Rate Schedule. I used a simple trapezoidal schedule. It has three easy phases: a fast warmup, a long flat hold at maximum speed, and a gentle cooldown.

3,000 Steps on 2× A100

With all of the above combined, a single A100 run for 1,500 steps took 69 minutes. Loss started at 11.0, dropped below 5.0 by step 125, and below 4.0 by step 400. Training loss bottomed out at 3.41, best validation loss at 3.46. The model processed 786 million tokens at 200,000 tok/s and 37.8% MFU, reaching 50.04 perplexity on WikiText-2.

Logs: nanoGPT-1500 · Raw log: nanogpt-1500.txt · Model: ashrafs1/nanogpt-1500s

To go further, I added a second A100. Each GPU gets a full copy of the model and processes different data. Gradients are averaged across both. This doubles the tokens per step.

3,000 steps took ~70 minutes. The model processed ~1.57 billion tokens at 375,000 tok/s and 70% MFU. Validation loss landed at 3.3147, dropping perplexity to 40.50 on WikiText-2. OpenAI’s GPT-2 reaches 25.17 on ~40 billion tokens. We’re not close yet, but it’s learning fast on far less data.

Logs: nanogpt-3000 · Raw log: shehab-ashraf/nanoGPT · Model: ashrafs1/nanogpt-3000

Resources

Architecture and training recipe heavily inspired by Keller Jordan’s modded-nanogpt speedrun and Tyler Romero’s contributions. Built on top of Andrej Karpathy’s nanoGPT. Sebastian Raschka’s LLMs-from-scratch was a great learning resource.