3 billion tokens on one GPU.

Once you separate those two questions, "3 billion tokens on one GPU" stops sounding impossible and starts sounding like a scheduling problem. The rest of this post is the arithmetic behind both halves.

CorX1.5 has 157.8 million parameters. Everything below follows from that one number. In fp32, each parameter costs 4 bytes — and the optimiser needs its own copies.

Here is where the 3 billion actually lives. Before training starts, the entire corpus — 25 open datasets spanning code, maths, reasoning, conversation and general text — gets tokenised once, offline, and written into a single flat binary file of raw token IDs. No JSON, no per-example records, no tokenising during training.

CorX1.5's vocabulary is 32,768 tokens. That number matters more than it looks:

import numpy as np, torch

# 3e9 tokens, uint16, sitting on disk. Nothing is loaded here.
data = np.memmap("train.bin", dtype=np.uint16, mode="r")

def get_batch(batch_size, block_size):
    # Sample random windows straight out of the mapped file.
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([torch.from_numpy(data[i     : i+block_size  ].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(data[i + 1 : i+block_size+1].astype(np.int64)) for i in ix])
    return x.pin_memory().to("cuda", non_blocking=True), \
           y.pin_memory().to("cuda", non_blocking=True)

Because the corpus is one continuous stream of token IDs, documents are concatenated end to end with a separator and then sliced at exactly 1,024 tokens. Every position in every batch carries a real token.

The naive alternative — one document per row, padded out to the block size — quietly throws away a large share of every batch. If your documents average 300 tokens and you pad them to 1,024, roughly 70% of the compute in every step is spent on padding tokens that contribute nothing. On a cluster that is money. On one GPU it is the difference between finishing and not finishing.

The related trick is on the other side of the loss function. CorX1.5 uses assistant-only loss masking: during instruction tuning the model reads the user's tokens as context but only ever gets a gradient for the tokens it is supposed to produce. Same forward pass, same cost, and none of the capacity is spent learning to predict the user's half of the conversation — which is not a skill the model needs.

Language models train best with large batches — hundreds of thousands of tokens per optimiser step. A single consumer GPU cannot hold that. These two facts look like a contradiction and are not, because batch size means two separate things:

• The micro-batch — how many sequences are in flight at once. This is a memory constraint.
• The effective batch — how many sequences the optimiser sees before it takes a step. This is a learning-dynamics constraint.

Gradient accumulation decouples them. Run the forward and backward pass on a small micro-batch, keep the gradients instead of stepping, repeat, and only call optimizer.step() once you have accumulated enough. The gradients sum exactly as if you had run one big batch — the memory just never went up.

accum_steps = 64          # micro-batches per optimiser step
optimizer.zero_grad(set_to_none=True)

for micro in range(accum_steps):
    x, y = get_batch(batch_size=8, block_size=1024)
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(x, y) / accum_steps      # scale so the sum is a mean
    loss.backward()                            # gradients accumulate in place

torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
scheduler.step()

Three settings do a disproportionate amount of the work.

bfloat16 autocast. Running the forward and backward pass in bf16 roughly halves activation memory and lets the GPU use its tensor cores. bf16 keeps fp32's exponent range, so unlike fp16 it does not need loss scaling and does not silently produce infinities halfway through a long run. If the card supports it, there is no reason not to.

Fused attention. A naive attention implementation materialises the full scores matrix — shape batch × heads × 1024 × 1024. At a micro-batch of 8 with, say, 16 heads, that single tensor is about 268 MB in bf16, per layer. A fused kernel (F.scaled_dot_product_attention, or FlashAttention) never writes it to memory at all, computing the softmax in tiles instead. This is usually the largest single activation saving available.

Gradient checkpointing. The escape hatch. Discard intermediate activations during the forward pass and recompute them during the backward pass — roughly 30% slower per step in exchange for a large drop in activation memory. At 10 layers and a 1,024-token window CorX1.5 is small enough that this is a last resort rather than a default, but it is what you reach for when nothing else will fit.

One more, almost free: torch.set_float32_matmul_precision("high") to enable TF32 on the remaining fp32 matmuls.

This was a budget decision, not a technical limit. Attention cost grows with the square of sequence length while everything else grows linearly, so doubling the context window to 2,048 more than doubles the memory and compute per sequence.

Given a fixed compute budget, that money can be spent two ways: fewer, longer sequences, or more, shorter ones. For a model this small, more tokens wins. A 158M model has nowhere near enough capacity to exploit long-range structure across 4,000 tokens; it has plenty of room to get better at ordinary sentences. Spending the budget on breadth rather than reach was the right trade — and it is honestly documented as a limitation on the model page rather than dressed up as a feature.

Once everything fits, the only remaining question is time — and it is pure division. Throughput depends on the card, the kernels and how well the data loader keeps up, so here is the shape of it rather than one number:

This is the part that separates a bedroom setup from a lab, and it has nothing to do with machine learning. A run measured in days on a single machine will be interrupted. Power, a driver update, a full disk, an out-of-memory error at step 4,000.

So the checkpoint has to contain everything needed to continue as if nothing happened — not just the weights:

• Model state_dict
• Optimiser state_dict — those Adam moments are hard-won; losing them costs real progress
• Learning-rate scheduler state and the step counter
• RNG state, so the data sampling continues rather than restarts
• The config used, so a checkpoint is never ambiguous about what produced it

Write them on a fixed step interval, keep the last few plus the best-by-validation, and never overwrite in place — write to a temporary file and rename, so a crash mid-write cannot destroy the only good checkpoint you have. I have learned every item on that list the expensive way.

A held-out validation split is carved out before training starts and never touched by the training sampler. Validation loss is evaluated on a fixed schedule, and training stops when it stops improving — early stopping.

The reason this matters more for a small model than a large one: with 158M parameters and billions of tokens, training loss will keep drifting down long after the model has stopped getting genuinely better. Without a held-out split you cannot tell the difference between learning and memorising, and you will happily burn a day of compute making the model worse.

The training loss tells you the optimiser is working. Only the validation loss tells you the model is.

• Instrument tokens per second from step one. I spent time tuning things that did not matter before I was measuring the thing that did.
• Validate the tokenised file before training, not after. A bug in a preprocessing script is invisible until it has cost you a day.
• Write the resume path first. Building checkpoint-and-resume before the first long run, instead of after the first crash, would have saved more time than any optimisation on this list.
• More data cleaning, less architecture fiddling. At this scale the corpus moved validation loss more than any hyperparameter I touched.

None of this is exotic. Every technique here is published, and most of it is a few lines of PyTorch. The reason it is worth writing down is that when you are starting out, the hard part is not the trick — it is not knowing which question you are actually asking. "How do I fit 3 billion tokens on one GPU" has no answer. "How do I fit 158 million parameters plus one batch on one GPU, and how long does streaming 3 billion tokens past it take" has two, and both of them are arithmetic.