PyTorch The LLM Training Loop

One pass through the step rhythm moves a tiny layer once. The loss on the same small batch prints lower after the step than before it.

import torch
from torch import nn
from torch.nn import functional as F

torch.manual_seed(0)
layer = nn.Linear(4, 3)
optimizer = torch.optim.AdamW(layer.parameters(), lr=3e-3)
x = torch.randn(2, 4)
y = torch.tensor([0, 2])
before = F.cross_entropy(layer(x), y)
optimizer.zero_grad()
before.backward()
optimizer.step()
after = F.cross_entropy(layer(x), y)
print("before:", f"{before.item():.4f}")
print("after:", f"{after.item():.4f}")

The modulo check picks the steps where the loop pauses to measure. A 20-step run with a check every 10 steps measures at the three printed steps.

STEPS = 20
EVAL_EVERY = 10
marks = [s for s in range(STEPS + 1) if s % EVAL_EVERY == 0]
print(marks)

estimate_loss runs the model over several batches without training and averages the losses. The two printed averages stand for the train split and the validation split.

train_losses = [2.41, 2.38, 2.45]
val_losses = [2.52, 2.49, 2.55]
print("train estimate:", f"{sum(train_losses) / len(train_losses):.2f}")
print("val estimate:", f"{sum(val_losses) / len(val_losses):.2f}")

import torch
from torch import nn
from torch.nn import functional as F

CORPUS = """The sun came up over the hill. Moss the cat woke in the warm house. Nan put milk in a blue bowl. Moss drank the milk and licked her lips. Then Moss went out to the garden.

Pip the sparrow sat on the wall. Good morning, Moss, said Pip. Good morning, Pip, said Moss. The sun is warm, said Pip. The sun is warm and the sky is clear, said Moss.

Bram the dog ran down the lane. Bram had a stick in his mouth. Drop the stick, Bram, said Nan. Bram dropped the stick. Good dog, said Nan. Bram wagged his tail and barked at the sky.

The three friends walked to the river. The river ran cold and clear over the white stones. A silver fish jumped in the river. Splash! The fish jumped again. Moss watched the fish. Pip watched the fish. Bram put one paw in the water. The water was cold. Bram shook his wet paw. The friends laughed in the sun.

At noon, Nan called the friends home. Bread and cheese waited on the table. The friends ate in the garden. The garden smelled of rain and roses. Bees hummed in the warm air. It was a good day, and the day was not over yet.

The rain came in the night. Rain tapped on the roof. Rain ran down the lane. Rain filled the blue bowl in the garden. Moss sat at the window and watched the rain. Pip hid under the wall. Bram slept by the fire and dreamed of sticks.

In the morning, the rain was gone. The sun came up over the hill. The lane was wet and the garden was green. A snail crossed the step, slow and silver. Good morning, snail, said Moss. The snail said nothing. The snail had far to go.

Nan gave Moss a fish for lunch. Nan gave Bram a bone for lunch. Nan gave Pip a crust of bread. Pip pecked the crust. Bram chewed the bone. Moss ate the fish, bite by bite. Then all three slept in the sun, full and warm and glad.

One day, a red ball rolled into the garden. The ball rolled past Moss. The ball rolled past Bram. The ball rolled under the wall and stopped by the river. Whose ball is this? said Pip. Tom came down the lane. Tom was a small boy with red boots. My ball! said Tom. Thank you, friends!

Tom threw the ball. Bram ran and caught the ball. Tom threw the ball again. Pip flew over the ball. Moss hid in the grass and jumped out. Boo! said Moss. Tom laughed and laughed. The friends played until the sun went down.

When the sun went down, the sky turned red and gold. The friends walked home in the soft light. Nan had soup on the fire. The soup was hot and good. Eat, said Nan, and rest. Tomorrow is a new day.

The moon came up over the hill. The stars came out, one by one. Moss curled on the step. Pip slept in the tree. Bram slept by the door, one ear up. The house was quiet. The river ran on in the dark, cold and clear. Sleep well, little friends, said the moon. And they did.

In the spring, the garden woke up. Green shoots came out of the dark earth. A bee came to the garden. The bee hummed from rose to rose. Pip sang in the tree. Bram dug a hole by the wall. Moss rolled in the warm grass. Spring is here, said Moss. Spring is here at last.

Nan planted seeds in the garden. Small seeds, brown and round. Moss watched the seeds go into the earth. Pip watched for worms. Bram watched for Tom. Rain came, soft and warm. Sun came, warm and gold. The seeds slept in the dark. Then, one morning, a green shoot! And then two! And then ten!

The shoots grew tall in the summer sun. Red roses opened on the wall. The garden was full of bees and birds. Tom picked peas in the garden. Nan picked beans. Moss picked a warm spot and slept in it. It was a good summer, a long, warm summer.

One hot day, the friends went up the hill. The hill was high and the grass was dry. From the top of the hill, the friends saw the river. They saw the lane and the house and the garden. They saw the whole world, small and bright. I can see my house, said Tom. I can see my river, said Moss. I can see everything, said Pip, and Pip flew up, up, up into the blue.

At the top of the hill, the wind blew. The wind blew the grass. The wind blew the tree. The wind took the song of Pip and carried it over the river. Sing again, Pip, said Bram. Pip sang again, loud and sweet. The sun shone. The wind blew. It was good to be alive.

When autumn came, the leaves turned gold and red. Leaves fell from the tree. Leaves covered the lane. Bram jumped in the leaves. Moss hid in the leaves. Pip rode on a leaf down to the river. Splash! Pip shook the water off and laughed. Autumn is fun, said Pip. Autumn is cold, said Moss, and Moss went home to the fire.

Nan made a pie in the autumn. The pie was full of apples. The house smelled of apples and spice. Tom ate two slices of pie. Bram ate the crumbs. Moss drank warm milk. Pip pecked at a seed cake. Outside, the leaves fell. Inside, the fire burned, warm and bright.

One cold morning, frost covered the garden. The grass was white. The bowl was ice. The river ran slow under a skin of ice. Moss put one paw on the ice. The ice was cold and hard. Winter is coming, said Moss. Winter is coming soon.

Snow fell in the night, soft and deep. In the morning, the world was white. The hill was white. The lane was white. The tree was white, and the wall was white. Bram ran in the snow. Tom rolled in the snow. Pip sat in the tree and shivered. Come down, Pip, said Nan. Come in by the fire.

By the fire, the friends were warm. Nan read a story out loud. The story was about a cat, a dog, and a bird. That is us! said Tom. That is us, said Moss, and Moss purred. The snow fell outside. The fire burned inside. It was the best place in the whole white world.

And when the snow melted, the river ran fast and full. The sun came up over the hill, warm as ever. Moss sat on the step. Pip sang in the tree. Bram ran down the lane with a stick. Spring had come again, and the friends were glad. The end."""

batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4

chars = sorted(set(CORPUS))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}
data = torch.tensor([stoi[ch] for ch in CORPUS], dtype=torch.long)
split = int(0.9 * len(data))
train_data = data[:split]
val_data = data[split:]


def get_batch(split_name):
    source = train_data if split_name == "train" else val_data
    ix = torch.randint(len(source) - block_size, (batch_size,))
    x = torch.stack([source[i:i + block_size] for i in ix])
    y = torch.stack([source[i + 1:i + block_size + 1] for i in ix])
    return x, y


class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        scores = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
        scores = scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        return weights @ v


class MultiHeadAttention(nn.Module):
    def __init__(self, n_head, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(n_head)])
        self.proj = nn.Linear(n_embd, n_embd)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        return self.proj(out)


class FeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x):
        return self.net(x)


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        head_size = n_embd // n_head
        self.attn = MultiHeadAttention(n_head, head_size)
        self.ff = FeedForward()
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[Block() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)
        self.lm_head.weight = self.token_embedding.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.token_embedding(idx) + self.position_embedding(torch.arange(T))
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(B * T, vocab_size), targets.view(B * T))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new):
        for _ in range(max_new):
            window = idx[:, -block_size:]
            logits, _ = self(window)
            probs = F.softmax(logits[:, -1, :], dim=-1)
            nxt = torch.multinomial(probs, 1)
            idx = torch.cat([idx, nxt], dim=1)
        return idx


STEPS = 20
EVAL_EVERY = 10
EVAL_BATCHES = 5


@torch.no_grad()
def estimate_loss(model, split_name):
    model.eval()
    total = 0.0
    for _ in range(EVAL_BATCHES):
        xb, yb = get_batch(split_name)
        _, loss = model(xb, yb)
        total += loss.item()
    model.train()
    return total / EVAL_BATCHES


torch.manual_seed(0)
model = GPT()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)

print("step  train   val")
for step in range(STEPS + 1):
    if step % EVAL_EVERY == 0:
        train_loss = estimate_loss(model, "train")
        val_loss = estimate_loss(model, "val")
        print(f"{step:>4}  {train_loss:.2f}  {val_loss:.2f}")
    if step == STEPS:
        break
    xb, yb = get_batch("train")
    _, loss = model(xb, yb)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()