This is a structurally honest pretraining loop — the same pieces (optimizer, schedule, accumulation, checkpointing) that appear in a trillion-token run on a thousand GPUs, just at a scale that runs on a single machine so you can actually read and modify every line.
python
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class TinyLM(nn.Module):
def __init__(self, vocab_size, d_model=512, n_layers=8, n_heads=8, max_len=1024):
super().__init__()
self.tok_embed = nn.Embedding(vocab_size, d_model)
self.pos_embed = nn.Embedding(max_len, d_model)
layer = nn.TransformerEncoderLayer(d_model, n_heads, dim_feedforward=4 * d_model,
batch_first=True, norm_first=True)
self.blocks = nn.TransformerEncoder(layer, n_layers)
self.ln_f = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, vocab_size, bias=False)
self.head.weight = self.tok_embed.weight # weight tying: fewer params, better generalization
def forward(self, idx):
B, T = idx.shape
pos = torch.arange(T, device=idx.device)
x = self.tok_embed(idx) + self.pos_embed(pos)[None, :, :]
causal_mask = nn.Transformer.generate_square_subsequent_mask(T).to(idx.device)
x = self.blocks(x, mask=causal_mask, is_causal=True)
return self.head(self.ln_f(x))
def next_token_loss(logits, targets):
# logits: (B, T, V), targets: (B, T) -- shifted by one position by the caller
return F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))