Pre-training a 95M Parameter LLM on a Consumer GPU
Can you train a language model from scratch on a single RTX 3070 Ti with only 8GB VRAM? Spoiler: yes, you can. In this post, I'll walk through how I pre-trained TinyGPT2, a 95 million parameter transformer model, on consumer hardware—and the engineering tricks that made it possible.
This isn't about building the next Claude or GPT. It's about understanding what happens under the hood when you train a language model, and proving that you don't need a cluster of A100s to learn these fundamentals hands-on.
The Challenge: 8GB Isn't Much
Let's do some back-of-napkin math. A 95M parameter model in float32 needs about 380MB just for the weights. But during training, you also need:
- Gradients: Another 380MB
- Optimizer states: AdamW stores two momentum terms per parameter, so ~760MB more
- Activations: These scale with batch size and sequence length—easily several GB
Add it up and you're looking at 5-6GB minimum before you've even loaded a single batch of data. With an 8GB card, there's not much room for error.
The Architecture: TinyGPT2
TinyGPT2 is a decoder-only transformer with some modern architectural improvements over the original GPT-2:
1@dataclass
2class TinyGPT2Config:
3 vocab_size: int = 50304 # GPT-2 vocab + padding token
4 block_size: int = 512 # Context window
5 n_embd: int = 768 # Embedding dimension
6 n_head: int = 12 # Attention heads
7 n_layer: int = 12 # Transformer blocks
8 gqa_kv_head: int = 4 # KV heads for GQA
9 hidden_size: int = 2048 # FFN hidden dimension
10 dropout: float = 0.1Key architectural choices:
- Grouped Query Attention (GQA): Instead of having 12 key-value heads matching our 12 query heads, we use only 4 KV heads. This reduces memory usage during attention computation and enables efficient KV caching during inference.
- Rotary Position Embeddings (RoPE): No learned position embeddings. RoPE encodes position information directly into the attention computation, which generalizes better to different sequence lengths.
- RMSNorm: Replaces LayerNorm. It's computationally cheaper (no mean subtraction) and often more stable during training.
- Weight Tying: The input embedding matrix and output projection share weights, cutting ~38M parameters.
Grouped Query Attention: A Closer Look
Standard multi-head attention has separate K, V projections for each head. With 12 heads and 768 embedding dimensions, that's significant memory. GQA groups multiple query heads to share the same K, V heads:
1class GroupedQueryAttention(nn.Module):
2 def __init__(self, config):
3 super().__init__()
4 self.n_head = config.n_head # 12 query heads
5 self.n_kv_head = config.gqa_kv_head # 4 KV heads
6 self.head_dim = config.n_embd // config.n_head
7
8 # Q has full heads, K and V have fewer
9 self.q_proj = nn.Linear(config.n_embd, config.n_embd)
10 self.k_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim)
11 self.v_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim)
12 self.o_proj = nn.Linear(config.n_embd, config.n_embd)
13
14 def forward(self, x, freqs_cis, mask=None):
15 B, T, C = x.shape
16
17 q = self.q_proj(x).view(B, T, self.n_head, self.head_dim)
18 k = self.k_proj(x).view(B, T, self.n_kv_head, self.head_dim)
19 v = self.v_proj(x).view(B, T, self.n_kv_head, self.head_dim)
20
21 # Apply RoPE to queries and keys
22 q, k = apply_rotary_emb(q, k, freqs_cis)
23
24 # Repeat KV heads to match query heads (12 / 4 = 3 repeats)
25 k = k.repeat_interleave(self.n_head // self.n_kv_head, dim=2)
26 v = v.repeat_interleave(self.n_head // self.n_kv_head, dim=2)
27
28 # Standard attention from here
29 ...This gives us 3× memory savings on K and V projections and caches, with minimal quality loss.
Training Strategy: Making 8GB Work
The key insight: we can trade compute for memory. Here's how:
1. Mixed Precision Training (AMP)
Instead of float32 everywhere, we use bfloat16 for forward and backward passes. This halves memory for activations and speeds up computation on tensor cores:
1# Training loop with automatic mixed precision
2dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
3
4for batch in dataloader:
5 inputs = batch['input'].to(device)
6 targets = batch['target'].to(device)
7
8 # Forward pass in reduced precision
9 with torch.autocast(device_type='cuda', dtype=dtype):
10 logits, loss = model(inputs, targets)
11
12 # Scale loss for gradient accumulation
13 loss = loss / gradient_accumulation_steps
14 loss.backward()
15
16 if step % gradient_accumulation_steps == 0:
17 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
18 optimizer.step()
19 optimizer.zero_grad(set_to_none=True)Why bfloat16 over float16? bfloat16 has the same exponent range as float32, just less precision. This means no gradient scaling needed—one less thing to worry about.
2. Gradient Accumulation
We can't fit a large batch in memory, but we can simulate one. Instead of batch_size=512, we use batch_size=8 and accumulate gradients over 64 steps:
1# Config
2batch_size = 8
3gradient_accumulation_steps = 64
4block_size = 512
5
6# Effective batch size: 8 × 64 × 512 = 262,144 tokens per optimizer step
7effective_batch_tokens = batch_size * gradient_accumulation_steps * block_sizeThis means we process 262K tokens before each weight update. Large batch training is crucial for stable language model training—the gradient estimates are much less noisy.
3. torch.compile
PyTorch 2.0's compiler fuses operations and generates optimized CUDA kernels:
1model = TinyGPT2(config).to(device)
2model = torch.compile(model) # That's it. Free speedup.In practice, this gave me a 20-30% speedup with no code changes. The compiler identifies patterns like "LayerNorm followed by Linear" and fuses them into single CUDA kernels.
4. Streaming Data
OpenWebText is 40GB+ on disk. Loading it all into RAM isn't practical. HuggingFace datasets with streaming mode to the rescue:
1import datasets
2
3# Stream data, never load full dataset into memory
4ds = datasets.load_dataset(
5 "Skylion007/openwebtext",
6 split="train",
7 streaming=True,
8 trust_remote_code=False
9)
10
11def collate_fn(batch):
12 """Tokenize and prepare batch on-the-fly"""
13 texts = [item['text'] for item in batch]
14 all_tokens = []
15 for text in texts:
16 tokens = tokenizer.encode(text)
17 all_tokens.extend(tokens + [tokenizer.eos_id])
18
19 # Create input-target pairs with sliding window
20 inputs, targets = [], []
21 for i in range(0, len(all_tokens) - block_size, block_size):
22 inputs.append(all_tokens[i:i + block_size])
23 targets.append(all_tokens[i + 1:i + block_size + 1])
24
25 return {
26 'input': torch.tensor(inputs, dtype=torch.long),
27 'target': torch.tensor(targets, dtype=torch.long)
28 }The Training Run
With everything set up, here's what the actual training looked like:
- Dataset: OpenWebText (~6.7B tokens)
- Hardware: Single NVIDIA RTX 3070 Ti (8GB VRAM)
- Training time: ~5 days for 25,000 optimizer steps
- Tokens processed: 6.5B+ tokens (multiple epochs)
- Throughput: ~15,000-18,000 tokens/second
Optimizer Configuration
1optimizer = torch.optim.AdamW(
2 model.parameters(),
3 lr=3e-4,
4 weight_decay=0.1,
5 betas=(0.9, 0.95),
6 fused=True # Use CUDA-fused implementation
7)
8
9# OneCycleLR: ramp up, then decay
10scheduler = lr_scheduler.OneCycleLR(
11 optimizer,
12 max_lr=8e-4,
13 total_steps=50000,
14 pct_start=0.05 # 5% warmup
15)The learning rate schedule is important. We warm up for the first 5% of training (prevents early instability), hit peak learning rate, then gradually decay. This is standard practice for transformer training.
Loss Curves
Here's roughly what training looked like:
- Step 0: Loss ~10.0 (random initialization, predicting ~50K classes)
- Step 1,000: Loss ~4.5 (learning common tokens)
- Step 5,000: Loss ~3.2 (grammatical structure emerging)
- Step 15,000: Loss ~2.8 (coherent sentences)
- Step 25,000: Loss ~2.5 (good fluency, reasonable coherence)
Validation loss tracked training loss closely, indicating we weren't overfitting. The model saw each token only 1-2 times on average over the full training run.
Supervised Fine-Tuning (SFT)
A pre-trained model is good at predicting next tokens, but not at following instructions. Enter SFT: we fine-tune on instruction-response pairs.
1# Alpaca-style prompt template
2def format_prompt(instruction, input_text=""):
3 if input_text:
4 return f"""### Instruction:
5{instruction}
6
7### Input:
8{input_text}
9
10### Response:
11"""
12 return f"""### Instruction:
13{instruction}
14
15### Response:
16"""
17
18# Only compute loss on the response, not the prompt
19labels = [-100] * len(prompt_tokens) + response_tokensKey SFT details:
- Dataset: Stanford Alpaca (52K instruction-response pairs)
- Learning rate: 2e-5 (10× lower than pre-training to avoid forgetting)
- Epochs: 3 (~85 minutes total)
- Loss masking: Only train on response tokens, ignore prompt
Results: What Can 95M Parameters Do?
Let's be real: this isn't GPT-4. But it's surprisingly capable for its size:
1Prompt: What is the capital of France?
2Response: The capital of France is Paris. It is the largest city in France
3and serves as the country's political, economic, and cultural center.
4
5Prompt: Write a short poem about the moon.
6Response: The moon shines bright in the night,
7A silver orb of gentle light.
8It watches over land and sea,
9A faithful friend for you and me.
10
11Prompt: Explain what machine learning is in simple terms.
12Response: Machine learning is a type of artificial intelligence that allows
13computers to learn from data without being explicitly programmed. Instead of
14following strict rules, the computer finds patterns in examples and uses those
15patterns to make predictions or decisions.It handles factual questions, creative writing, and explanations reasonably well. It struggles with complex reasoning, long-form consistency, and nuanced tasks—but that's expected at this scale.
Lessons Learned
After spending a week watching loss curves, here's what I took away:
- Memory is the bottleneck. Most of my optimization effort went into fitting more into 8GB. The actual training code is straightforward once you have memory under control.
- Large effective batch sizes matter. Early experiments with smaller batches (32K tokens) showed much noisier training. 262K tokens per update made training smooth and stable.
- torch.compile is magic. Zero code changes, significant speedup. No reason not to use it.
- Streaming data works. I was skeptical about not having data in RAM, but disk I/O was never the bottleneck. The GPU is always the limiting factor.
- Checkpointing is essential. Training crashed twice (power outages). Saving checkpoints every 500 steps saved days of compute.
Try It Yourself
Everything is open source:
- GitHub Repository - Full training code, model architecture, and scripts
- HuggingFace - Pre-trained weights you can download and use
- Live Demo - Try the model in your browser
To train your own:
1# Clone the repo
2git clone https://github.com/NotShrirang/tinygpt.git
3cd tinygpt
4
5# Install dependencies
6pip install -r requirements.txt
7
8# Start pre-training (will take several days on consumer hardware)
9python train_liger.py
10
11# Or resume from checkpoint
12python train_liger.py --resume
13
14# Fine-tune on Alpaca after pre-training
15python sft.py --checkpoint checkpoints/ckpt_step25000.pthWhat's Next?
There's always more to explore:
- Flash Attention: Even more memory-efficient attention that would allow larger context windows
- LoRA fine-tuning: Train adapters instead of full weights for domain adaptation
- Quantization: INT8 or INT4 inference for deployment on even smaller devices
- RLHF: Human feedback training to improve response quality
The democratization of AI training is real. You don't need a datacenter to understand how language models work—or to train your own. A gaming GPU, some patience, and curiosity are enough to get started.
Resources
- Andrej Karpathy's nanoGPT - The inspiration for this project
- GQA Paper - Grouped Query Attention explained
- RoFormer Paper - Rotary Position Embeddings
- PyTorch AMP Documentation
If you found this useful, give the repo a star. And if you train your own model, I'd love to hear about it!
If this resonated, pass it along.