Optimizing a Tiny Transformer

Motivation

Most transformer tutorials stop once the loss goes down. I wanted to know what happens when you actually tune the boring stuff: batch size, memory usage, precision, kernel choice. Does throughput go up? Does val loss get worse? Do the fancy optimizations even matter on a tiny model?

tiny-transformer is a small GPT-style decoder trained on character-level Shakespeare. It's ~10M params, so experiments run fast, but it's big enough that GPU memory and batching behavior show up for real. Each run changes one thing at a time so the comparisons stay fair.

Model Architecture

Same general setup as nanoGPT / Karpathy's build:

• 6 decoder blocks with causal self-attention + a 4× MLP each
• 384-dim embeddings, 6 heads (64 dims per head)
• 256-char context window
• Character-level vocab from the unique chars in input.txt
• Pre-norm residuals (LayerNorm before attention/FFN, then add the skip connection back)

Attention has two implementations behind a flag. Manual mode builds the full (T, T) attention matrix, applies a causal mask, softmaxes, multiplies by V. Flash mode just calls F.scaled_dot_product_attention and lets PyTorch pick a fused kernel if it can.

tok_emb + pos_emb → [DecoderBlock × 6] → LayerNorm → lm_head → logits

~10.7M params total. At batch 64 and seq length 256, each step sees 16,384 tokens.

Training Setup

Random contiguous chunks from the train/val split. Targets are just the input shifted by one token, normal LM setup.

All experiments ran on a RunPod GPU pod with an RTX 4090. I picked RunPod because it's cheap for short benchmark runs — spin up a pod, clone the repo, run bash run_experiments.sh, tear it down. The 4090's 24 GB was plenty for this model (peak usage was ~14 GB at batch 256); most runs stayed under 4 GB.

Experimental Methodology

Configs are YAML files under configs/, grouped by what they test (batch/, attention/, mp/, etc.):

python train.py --config batch/batch128
bash run_experiments.sh   # runs all configs sequentially
python summarize.py       # aggregates outputs/*/train_log.csv

Each run gets its own folder under outputs/ with:

• The config file copied in
• train_log.csv with loss, lr, step time, tok/s, and memory stats
• best.pt / latest.pt checkpoints

Instrumentation Details

Throughput is a rolling average over the last 20 steps. Token count includes grad accum: batch_size × block_size × grad_accum_steps. I skip step 0 since warmup skews the numbers.

Memory comes from torch.cuda.max_memory_allocated(), i.e. peak tensor memory, not total reserved pool. That's the number you care about for OOMs.

Grad accum divides loss by grad_accum_steps, backprops each micro-batch, then does one optimizer.step(). Same effective batch size, less activation memory per step.

Baseline Results

Batch 64, fp32, manual attention, no compile:

• 291k tokens/sec
• 3.6 GB peak memory
• 1.6111 val loss after 5k steps
• 56.9 ms per step

Everything below is compared to this.

Batch Size Sweep

This was the most interesting part. Memory grows roughly linearly with batch size (activations dominate), but throughput definitely doesn't. Batch 128 is only ~5% faster than 64 even though you're doing 2× the work per step. Batch 256 is actually slower than 64. Looks like memory bandwidth is the bottleneck, not raw compute.

Why Smaller Batches Generalize Better Here

The dataset is tiny (~1 MB of Shakespeare). Batch 32 gives noisier gradients, which actually helps on a small corpus. Batch 256 averages over way more examples per step, so the gradients are smoother but the model lands in sharper minima and val loss gets worse. I didn't scale the learning rate for the large batch runs, which probably doesn't help either.

The Linear LR Scaling Rule (Pending)

Common rule of thumb: double the batch, double the lr. I have configs set up for that (batch128_scaled_lr at 6e-4, batch256_scaled_lr at 1.2e-3) but haven't run them yet. Curious whether that closes the val loss gap or if the small-data noise effect is doing most of the work.

Gradient Accumulation

Grad accum lets you keep a big effective batch while using a smaller micro-batch for memory. Both of these target effective batch 64:

32 × 2 is the one I'd actually use: ~84% of baseline speed, half the memory, val loss basically the same (1.6044 vs 1.6111).

16 × 4 gets you down to 1.1 GB but throughput gets cut in half. Fine if you're memory-starved, but painful otherwise. Val loss is still fine though, which confirms it's the effective batch that matters, not the micro-batch.

Mixed Precision (BF16)

Turned on torch.autocast with bf16. Results matched fp32 exactly:

• 291k tok/s
• 3.6 GB peak memory
• 1.6111 val loss

BF16 has the same exponent range as fp32 so you usually don't need loss scaling (unlike fp16). On newer NVIDIA GPUs it can hit tensor cores too. Didn't see a speedup at this scale, probably because memory bandwidth and Python overhead dominate. But no accuracy hit, so might as well leave it on.

Flash Attention

Swapping manual attention for F.scaled_dot_product_attention: 291k → 293k tok/s. Basically noise.

Flash attention avoids materializing the full attention matrix in memory. At T=256 and 6 heads, that matrix is small enough that the naive implementation is already fine. I'd expect this to matter more at longer context lengths (2k+) where the O(T²) memory actually hurts.

torch.compile

Wrapped the model in torch.compile. Still 291k tok/s.

Compile fuses ops and cuts Python dispatch overhead, but on a 10M param model with short sequences, most of the step time is stuff like data loading, optimizer updates, and kernel launch latency. Plus you pay a compilation warmup cost on a 5k-step run. Probably shows up more on bigger models.

Full Results Table

Takeaways

1. On a model this size, batch size and memory matter way more than flash attention or compile. Profile first.
2. Doubling batch size doesn't double throughput. Sometimes it makes things slower.
3. Smaller batches can give better val loss on small datasets, even when they're slower.
4. Grad accum (2×) is a good way to save memory without tanking throughput or loss.
5. BF16 cost nothing here. Just use it if your GPU supports it.
6. Logging tok/s and peak memory to CSV made this whole comparison possible. Would've been guessing otherwise.