-
-
Training: validation loss over 5B tokens on one RTX 5060 Ti (29.9 h)
-
Zero-shot benchmarks tracked every 100M tokens during training (500-example subsets)
-
Architecture: 3 unique Transformer blocks applied twice, 42,070,080 parameters
Inspiration
How much language ability can you squeeze out of a model that fits in 50M parameters, trained on one consumer GPU? Instead of copying a recipe, we wanted to earn every architectural choice: start from a plain GPT-2-style model and only keep a change if it measurably lowered validation loss under the parameter cap.
What it does
Looped GPT-42M is a decoder-only language model trained completely from scratch (no pretrained weights, no fine-tuning, no distillation). Its key idea is weight sharing in depth: 3 unique Transformer blocks are applied twice, so the network computes with an effective depth of 6 while only paying parameters for 3. The parameters saved are spent on width (1,024-dimensional hidden state), which mattered more at this scale.
| Trainable parameters | 42,070,080 (token embeddings and tied output head included) |
| Architecture | 3 unique pre-norm blocks × 2 loops (effective depth 6), width 1,024, 8 heads |
| Attention | Causal MHA with RoPE, fused scaled_dot_product_attention (Flash/cuDNN kernels) |
| Feed-forward | SwiGLU, 1,024 → 1,376 → 1,024 |
| Tokenizer | SentencePiece BPE, 16,384 tokens (vocab shrunk from GPT-2's 50,257 to free ~35M parameters) |
| Context | 128 tokens |
How we built it
We followed a strict one-change-at-a-time ablation ladder, every run tracked in Weights & Biases:
- GPT-2-style baseline (learned positions, GELU, 50K vocab)
- Cosine LR schedule with warmup
- Xavier initialization
- SwiGLU feed-forward
- More depth vs. more width
- Looped Transformer (shared blocks) — slightly lower validation loss than the unshared SwiGLU baseline, while reusing weights
- Better batching / gradient accumulation
- RoPE instead of learned positions
- 16K SentencePiece vocabulary — the biggest single parameter saving
- Also explored: Muon optimizer, Grouped-Query Attention vs. MHA, fused attention kernels (with equivalence tests), tied embeddings, and an agent-driven "autoresearch" loop that proposes and tests one training-script change at a time under a fixed 1M-token budget.
Final training run (SmolLM-Corpus, cosmopedia-v2 subset, streamed from Hugging Face):
| Phase | Tokens | Learning rate | Time |
|---|---|---|---|
| 1 | 0 → 2B | AdamW, 3e-4, 5% warmup, cosine decay to 3e-5 | 11.96 h |
| 2 | 2B → 5B | AdamW, constant 3e-5 | 17.9 h |
| Total | 5B tokens | weight decay 0.05, grad-clip 1.0, bf16 autocast | 29.9 h |
- Hardware: a single NVIDIA GeForce RTX 5060 Ti (16 GB) for the final model, provided by collaborators; ablations ran on a mix of collaborator GPUs (RTX 3060/3080/3090/4070 SUPER, TITAN RTX, GTX 1660 Ti) and rented cloud GPUs (Tesla V100).
- Efficiency: ~46K tokens/s average throughput, peak GPU memory under 2 GB, batch of 8,192 tokens per update (16 × 128 × 4 accumulation steps).
- Compute: ≈ 2.0 × 10¹⁸ FLOPs (6·N·D estimate, counting both loop passes and the output head).
Results
Official evaluation with lm-evaluation-harness 0.4.9.1, zero-shot, full evaluation splits (scripts/eval_lm_harness.py, receipt in results/eval_lm_harness_5B.json):
| Benchmark | acc | acc_norm | Random chance |
|---|---|---|---|
| HellaSwag | 27.5% | 28.4% | 25% |
| ARC-Easy | 39.9% | 37.6% | 25% |
| PIQA | 59.3% | 59.0% | 50% |
| WinoGrande | 50.2% | — | 50% |
WikiText-103 test perplexity (held-out, 128-token windows): 129.9 per token (16K SentencePiece), 2105 per word, 2.07 bits per byte. Validation loss on the training distribution: 2.045 nats/token.
ARC-Easy and PIQA are clearly above chance, HellaSwag modestly so; WinoGrande (pronoun resolution) stays at chance, as is typical at this size.
Loops at inference. Because the blocks are shared, we can change how many times they are applied when running the model. On held-out WikiText-103 text: 1 loop → perplexity 254, 2 loops (as trained) → 130, 3 → 135, 4 → 158. The second pass halves perplexity; looping more than in training does not help.
Challenges we ran into
- Fitting the 50M cap honestly: with a GPT-2 vocabulary, embeddings alone ate most of the budget. Switching to a 16K tokenizer and tied embeddings freed room for a wider model.
- Stability in mixed precision: our custom LayerNorm overflowed in FP16; we moved its reduction to FP32.
- Long runs on borrowed hardware: we built resumable checkpoints that store optimizer, scheduler, RNG and data-stream state, periodic Hugging Face uploads and a "promotion" mechanism to continue a finished 2B run to 5B without restarting.
- Diminishing returns: after ~2B tokens the loss curve flattens; phase 2 bought only a small improvement (2.071 → 2.045), which taught us where the compute is best spent.
What we learned
- At 50M parameters, where parameters go (vocabulary vs. width vs. depth) matters more than adding exotic components.
- Weight-shared depth is a cheap way to get extra computation per token without breaking the parameter cap.
- Measuring everything (loss, benchmarks, per-loop activation telemetry) saved us from "improvements" that were really noise.
What's next
Supervised fine-tuning for instruction following, a data-mixture study during mid-training (already prototyped), longer context, and training with more loops (extra loops at inference do not help a model trained with two).
Built with AI assistance
Following the rules, we disclose that AI coding assistants were used: OpenAI Codex (most of the code refactors, training scripts and notebooks) and Claude / Claude Code (debugging, the evaluation adapter, documentation, submission material and the demo video). The demo video was animated with manim and narrated with a synthetic voice (Google Cloud Text-to-Speech). All design decisions, experiments and training runs were directed and verified by the team.
Log in or sign up for Devpost to join the conversation.