Training Your GPT
Teaching it to write.
Everything comes together now. We give our network some of the most famous lines ever written, and we teach it to predict the next character. At the start, it will produce nothing but noise. A few minutes later, on a laptop, it will start to write.
- Famous lines as training data
- Predict the next character
- From noise to writing
The training run
The data is our famous text, turned into numbers. Each batch takes random chunks, and the target is simply the same chunk, shifted one place to the right. That's the key to language models. Every position in the text is a little exercise in next character prediction. The first loss should be near the log of our vocabulary size, the sanity check that you learned earlier. Then we train, and we ask it to write. Run it. The loss falls steeply. And look at the output! It's writing words, spaces, and phrases, all learned from nothing but raw characters. Nobody taught it spelling.
import math
import torch
from torch import nn
class Block(nn.Module):
def __init__(self, dim, heads):
super().__init__()
self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(),
nn.Linear(4 * dim, dim))
self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
def forward(self, x, mask):
y = self.norm1(x)
x = x + self.attn(y, y, y, attn_mask=mask, need_weights=False)[0]
return x + self.mlp(self.norm2(x))
class TinyGPT(nn.Module):
def __init__(self, vocab, block_size=48, dim=64, heads=4, layers=3):
super().__init__()
self.block_size = block_size
self.tok = nn.Embedding(vocab, dim)
self.pos = nn.Embedding(block_size, dim)
self.blocks = nn.ModuleList(Block(dim, heads) for _ in range(layers))
self.norm = nn.LayerNorm(dim)
self.head = nn.Linear(dim, vocab)
def forward(self, ids):
T = ids.size(1)
x = self.tok(ids) + self.pos(torch.arange(T))
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), 1)
for b in self.blocks:
x = b(x, mask)
return self.head(self.norm(x))
@torch.no_grad()
def generate(self, ids, n, temperature=1.0):
for _ in range(n):
logits = self(ids[:, -self.block_size:])[:, -1] / temperature
next_id = torch.multinomial(logits.softmax(-1), 1)
ids = torch.cat([ids, next_id], dim=1)
return ids
text = ("to be or not to be that is the question. "
"the slings and arrows of outrageous fortune. "
"to sleep perchance to dream. " * 30 +
"all the world is a stage and all the men and women merely players. "
"the fault is not in our stars but in ourselves. " * 30)
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
torch.manual_seed(0)
model = TinyGPT(len(chars))
opt = torch.optim.AdamW(model.parameters(), lr=2e-3)
print(f"Vocabulary {len(chars)}, first loss should be near {math.log(len(chars)):.2f}")
for step in range(301):
i = torch.randint(0, len(data) - 49, (32,))
x = torch.stack([data[j:j + 48] for j in i])
y = torch.stack([data[j + 1:j + 49] for j in i])
loss = nn.functional.cross_entropy(model(x).transpose(1, 2), y)
opt.zero_grad(); loss.backward(); opt.step()
if step % 100 == 0:
print(f"step {step} loss {loss.item():.3f}")
out = model.generate(torch.tensor([[stoi["t"]]]), 90, temperature=0.8)
print("".join(chars[i] for i in out[0].tolist()))Vocabulary 24, first loss should be near 3.18 step 0 loss 3.373 step 100 loss 0.463 step 200 loss 0.101 step 300 loss 0.084 t is the question. the slings and arrows of outrageous fortune. to sleep perchance to dream
How it picks the next word
How does generation work? The model reads what's been written so far, and gives a chance for every possible next character. We sample one, add it to the text, and go round again. That's all. One token at a time. Chatbots do exactly this, using word pieces instead of characters, which is why you see their answers appearing bit by bit.
Temperature
One small dial changes the personality of the model. Temperature divides the scores before softmax. A low temperature makes the top choice dominate, so it's safe, but repetitive. A high one flattens the chances, so it gets creative, and then chaotic. We train a model, then generate with three temperatures. Run it. Look at the three samples. Point three repeats familiar phrases. Point eight is the sweet spot. And one point eight falls apart. Every chatbot has this dial, under the hood.
import math
import torch
from torch import nn
class Block(nn.Module):
def __init__(self, dim, heads):
super().__init__()
self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(),
nn.Linear(4 * dim, dim))
self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
def forward(self, x, mask):
y = self.norm1(x)
x = x + self.attn(y, y, y, attn_mask=mask, need_weights=False)[0]
return x + self.mlp(self.norm2(x))
class TinyGPT(nn.Module):
def __init__(self, vocab, block_size=48, dim=64, heads=4, layers=3):
super().__init__()
self.block_size = block_size
self.tok = nn.Embedding(vocab, dim)
self.pos = nn.Embedding(block_size, dim)
self.blocks = nn.ModuleList(Block(dim, heads) for _ in range(layers))
self.norm = nn.LayerNorm(dim)
self.head = nn.Linear(dim, vocab)
def forward(self, ids):
T = ids.size(1)
x = self.tok(ids) + self.pos(torch.arange(T))
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), 1)
for b in self.blocks:
x = b(x, mask)
return self.head(self.norm(x))
@torch.no_grad()
def generate(self, ids, n, temperature=1.0):
for _ in range(n):
logits = self(ids[:, -self.block_size:])[:, -1] / temperature
next_id = torch.multinomial(logits.softmax(-1), 1)
ids = torch.cat([ids, next_id], dim=1)
return ids
text = ("to be or not to be that is the question. "
"the slings and arrows of outrageous fortune. "
"to sleep perchance to dream. " * 30 +
"all the world is a stage and all the men and women merely players. "
"the fault is not in our stars but in ourselves. " * 30)
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
torch.manual_seed(0)
model = TinyGPT(len(chars))
opt = torch.optim.AdamW(model.parameters(), lr=2e-3)
for step in range(220):
i = torch.randint(0, len(data) - 49, (32,))
x = torch.stack([data[j:j + 48] for j in i])
y = torch.stack([data[j + 1:j + 49] for j in i])
loss = nn.functional.cross_entropy(model(x).transpose(1, 2), y)
opt.zero_grad(); loss.backward(); opt.step()
for t in [0.3, 0.8, 1.8]:
torch.manual_seed(1)
out = model.generate(torch.tensor([[stoi["t"]]]), 60, temperature=t)
print(f"temperature {t}:")
print(" ", "".join(chars[i] for i in out[0].tolist()))temperature 0.3: t is not in our stars but in ourselves. all the world is a st temperature 0.8: t is not in our stars but in ourselves. all the world is a st temperature 1.8: t all the mefan world en hen womeousln mene. at wmerlmelyerlt
Summary
You've built and trained a real G P T. It learned to write, from characters alone, in minutes, on a laptop. What separates this from the giants? Mostly scale. Next, what it really takes to go from tiny to enormous.
Key takeaways- Targets are the text shifted by one
- Generate one token at a time
- Temperature sets creativity