# Genesis 1B, Run 2: Nearly 3× Throughput, Same Hardware - Kroonen AI

> Redesigning Genesis 1B from 20 to 32 layers. Same param count, same GPUs, 3× training throughput.

Canonical: https://www.kroonen.ai/blog/genesis-architecture-v2/

[Lab notebook](https://www.kroonen.ai/blog/) / [Genesis 1B](https://www.kroonen.ai/blog/#genesis) / Entry 005

Engineering note Run 2 Entry 005

# Genesis 1B, Run 2: Nearly 3× Throughput, Same Hardware

Why Run 2 moved from 20 to 32 layers, changed its RoPE implementation, and rebuilt the training path around the same two GPUs.

Author

Robin Kroonen

Published

Mar 24, 2026

Reading time

5 min read

Entry 005

Program

Genesis 1B

Format

Engineering note

Phase

Run 2

Subjects

Genesis / Run 2 / architecture / torch.compile

[Program index →](https://www.kroonen.ai/blog/#genesis)

Genesis 1B, Run 2 is a full architecture redesign. Same ~1B parameters, same 2x RTX 4090 setup, but 32 layers instead of 20, real-valued RoPE, `torch.compile`, batch size 4, and proper LR scheduling. Result: ~19k tok/s (up from 6,500), ~6 days to 20k steps instead of 13.

## Architecture Comparison

|  | Run 1 | Run 2 |
| --- | --- | --- |
| Parameters | 1,003M | 1,000M |
| Layers | 20 | 32 |
| Hidden dim | 2048 | 1536 |
| Attention heads | 16 | 12 |
| KV heads (GQA) | 4 | 6 |
| FFN dim | 5632 | 4736 |
| Seq length | 2048 | 2048 |
| Batch size | 1 | 4 |
| torch.compile | ✗ | ✓ |
| Activation ckpt | ✗ | ✓ |
| Throughput | 6,500 tok/s | ~19,000 tok/s |
| Time/step | ~41s | ~21s |
| Est. 20k steps | ~13 days | ~6 days |

## Why Deeper, Not Wider

Run 1 was wide: 20 layers at dim 2048. Run 2 trades width for depth: 32 layers at dim 1536. Same parameter budget, fundamentally different compute graph.

More layers means more sequential transformations, more chances for the model to build compositional representations. For reference, Llama 3.2 1B uses only 16 layers. Genesis 1B, Run 2 has 32. Twice the depth at the same parameter count is a bet on reasoning over memorization.

The narrower hidden dimension (1536 vs 2048) also plays better with `torch.compile`: smaller per-layer tensors mean less memory pressure and better kernel fusion.

## Where the Nearly 3× Speedup Came From

Two changes account for nearly all of the throughput gain:

-   **`torch.compile`**: Fuses operations, eliminates Python overhead, generates optimized CUDA kernels. This alone was a ~40% speedup with zero code changes to the model.
-   **Batch size 1 → 4**: Activation checkpointing freed enough VRAM to quadruple the batch. Combined with `torch.compile` and real-valued RoPE (avoiding complex64 graph breaks), throughput jumped to ~19k tok/s.

Same hardware. Same parameter count. ~3× throughput. No tricks, just using PyTorch properly.

## LR Schedule Fix

Run 1 had a bug: pure cosine decay from step 0. No linear warmup. The learning rate started high and the first few hundred steps were essentially random noise.

Run 2 uses proper linear warmup over 1,000 steps followed by cosine decay to 10% of peak LR. Standard practice, but it was missing before.

## Checkpoint Infrastructure

Run 2 introduces DCP checkpoint versioning with full architecture metadata embedded in every checkpoint. Each save includes the complete model config (layers, dimensions, head counts, LR schedule parameters) so any checkpoint is self-describing.

Auto rotation keeps the last 5 checkpoints and prunes older ones. Try the latest checkpoint in the [live playground on HuggingFace](https://huggingface.co/spaces/rob-x-ai/genesis-1b-run2-playground).

## What's Next

Run 2 finished at the 20,000-step target, validating the deeper architecture and improved training stack. For final metrics and results, see the [training results post](https://www.kroonen.ai/blog/genesis-training-progress/) or [try the live playground](https://huggingface.co/spaces/rob-x-ai/genesis-1b-run2-playground).

Notebook / Reading path

## Continue in Genesis 1B

[Complete index →](https://www.kroonen.ai/blog/)

[

Next entry 006 / 8 min read

### Genesis 1B: Run 2 Complete at 80,000 Steps

Read entry →](https://www.kroonen.ai/blog/genesis-training-progress/)[

Previous entry 004 / 8 min read

### The Optimizer State Bug: A Silent Failure in DCP Resume

Read entry →](https://www.kroonen.ai/blog/genesis-optimizer-state-bug/)

[← Lab notebook](https://www.kroonen.ai/blog/) [Back to top ↑](#main-content)