GPU Throughput Analysis: Nano-GPT2 vs Transformer
Training on 1x H800 GPU, the Nano-GPT2 model achieves ~180,000 tokens/sec while the Transformer (Attention Is All You Need) achieves ~50,000 tokens/sec. This document analyzes why.
Key Differences
| Nano-GPT2 | Transformer | |
|---|---|---|
| Architecture | Decoder-only | Encoder-Decoder |
| Attention | F.scaled_dot_product_attention (Flash Attention) |
Manual matmul + softmax |
| Tokens/micro-batch | 64 x 1024 = 65,536 | ~15,388 |
| Attention ops/layer | 1 (self-attn) | 3 (enc self-attn + dec self-attn + cross-attn) |
| Optimizer | Fused AdamW | Regular Adam |
| Sequence length | 1024 (fixed, dense) | ~50 avg (padded to max in batch) |
Root Causes (biggest to smallest)
-
No Flash Attention -- The biggest factor. The Transformer manually materializes the full (T,T) attention matrix. GPT-2 uses
F.scaled_dot_product_attentionwhich dispatches to Flash Attention on CUDA -- 2-4x faster and far more memory efficient. -
Low GPU utilization -- Only 15K tokens per accumulation step vs GPT-2's 65K per micro-batch. The H800 is starved for work. Short WMT sentences (~30-50 tokens) with padding waste compute.
-
Encoder-Decoder = more ops per token -- Each decoder layer has 3 attention operations vs GPT-2's 1. Roughly 2x more attention compute for the same number of layers.
-
No fused optimizer -- GPT-2 uses
fused=Trueon AdamW, which runs the optimizer step in a single CUDA kernel.
Parameter Count
Nano-GPT2 (Decoder-only)
Config: n_layer=12, n_head=12, n_embd=768, vocab_size=50304
| Component | Calculation | Params |
|---|---|---|
| Token embedding | 50304 x 768 | 38,633,472 |
| Position embedding | 1024 x 768 | 786,432 |
| Per block (x12): | ||
| LN1 | 768 + 768 | 1,536 |
| QKV projection (fused) | 768 x 2304 + 2304 | 1,771,776 |
| Attn output projection | 768 x 768 + 768 | 590,592 |
| LN2 | 768 + 768 | 1,536 |
| MLP up (c_fc) | 768 x 3072 + 3072 | 2,362,368 |
| MLP down (c_proj) | 3072 x 768 + 768 | 2,360,064 |
| Block subtotal | 7,087,872 | |
| 12 blocks total | 7,087,872 x 12 | 85,054,464 |
| Final LN | 768 + 768 | 1,536 |
| lm_head | tied with token embedding | 0 |
| Total | 124,475,904 (~124M) |
Transformer (Encoder-Decoder)
Config: N=6, h=8, d_model=512, d_ff=2048, vocab=37000, shared_vocab=True
Shared embedding: weight tied between src_embed, tgt_embed, and generator.
| Component | Calculation | Params |
|---|---|---|
| Shared embedding | 37000 x 512 | 18,944,000 |
| Positional encoding | sinusoidal (buffer, not learned) | 0 |
| Per encoder layer (x6): | ||
| Self-attn: 4 x Linear(512,512) | 4 x (512 x 512 + 512) | 1,050,624 |
| FFN w1 | 512 x 2048 + 2048 | 1,050,624 |
| FFN w2 | 2048 x 512 + 512 | 1,049,088 |
| 2 x LayerNorm(512) | 2 x (512 + 512) | 2,048 |
| Encoder layer subtotal | 3,152,384 | |
| 6 encoder layers | 3,152,384 x 6 | 18,914,304 |
| Encoder final LN | 1,024 | |
| Per decoder layer (x6): | ||
| Masked self-attn: 4 x Linear | 4 x (512 x 512 + 512) | 1,050,624 |
| Cross-attn: 4 x Linear | 4 x (512 x 512 + 512) | 1,050,624 |
| FFN w1 + w2 | (same as encoder) | 2,099,712 |
| 3 x LayerNorm(512) | 3 x 1024 | 3,072 |
| Decoder layer subtotal | 4,204,032 | |
| 6 decoder layers | 4,204,032 x 6 | 25,224,192 |
| Decoder final LN | 1,024 | |
| Generator | tied with embedding | 0 |
| Total | 63,084,544 (~63M) |
FLOPs Per Token (Forward Pass)
For matrix multiply: FLOPs = 2 x M x N x K. Using d = d_model, T = seq_len.
Nano-GPT2 -- per token, per layer
| Op | FLOPs | d=768, T=1024 |
|---|---|---|
| QKV projection | 2 x d x 3d = 6d^2 | 3,538,944 |
| Attention (QK^T + attn x V) | 2 x 2 x T x d = 4Td | 3,145,728 |
| Attn output projection | 2d^2 | 1,179,648 |
| MLP up + down | 2 x 2 x d x 4d = 16d^2 | 9,437,184 |
| Per layer | 24d^2 + 4Td | 17,301,504 |
- 12 layers: 207,618,048
- lm_head (2 x d x V): 77,266,944
- Forward total: ~285M FLOPs/token
- Training (~3x fwd): ~855M FLOPs/token
Transformer -- per token pair (1 src + 1 tgt), per layer
Assuming T_src ~ T_tgt ~ T ~ 40 (typical WMT sentence length).
Encoder layer (processes src token):
| Op | FLOPs |
|---|---|
| Self-attn QKV + output | 8d^2 |
| Self-attn compute | 4 x T_src x d |
| FFN | 16d^2 |
| Per encoder layer | 24d^2 + 4Td |
Decoder layer (processes tgt token):
| Op | FLOPs |
|---|---|
| Masked self-attn (QKV + out) | 8d^2 |
| Self-attn compute | 4 x T_tgt x d |
| Cross-attn (QKV + out) | 8d^2 |
| Cross-attn compute | 4 x T_src x d |
| FFN | 16d^2 |
| Per decoder layer | 32d^2 + 4d(T_src + T_tgt) |
Per layer pair (1 enc + 1 dec): 56d^2 + 4d(2T_src + T_tgt)
With d=512, T=40:
| Per layer | 6 layers | |
|---|---|---|
| Encoder layer | 24 x 512^2 + 4 x 40 x 512 = 6,373,376 | 38,240,256 |
| Decoder layer | 32 x 512^2 + 4 x 512 x 80 = 8,552,448 | 51,314,688 |
| Generator (2 x d x V) | 37,888,000 | |
| Total forward | 127,442,944 (~127M) | |
| Training (~3x) | ~382M FLOPs/token-pair |
Summary
| Nano-GPT2 | Transformer | Ratio | |
|---|---|---|---|
| Parameters | 124M | 63M | 2.0x |
| Layers | 12 decoder | 6 enc + 6 dec | -- |
| Attention ops/layer | 1 (self) | 1 (enc) + 2 (dec) = 3 | 3x |
| d_model | 768 | 512 | 1.5x |
| Seq length | 1024 | ~40 | 25x |
| FLOPs/token (fwd) | 285M | 127M (per src+tgt pair) | 2.2x |
GPT-2 does ~2.2x more FLOPs per token -- yet it achieves 3.6x higher throughput. This confirms the gap is not from model complexity but from implementation efficiency (manual attention vs Flash Attention, low batch utilization, no fused optimizer).