Project: Tiny GPT from Scratch

কী বানাবে

একটা GPT। ছোট, কিন্তু আসল: token আর position embedding, causal multi-head self-attention, MLP, residual connection, LayerNorm আর একটা tied output head সহ একটা decoder-only transformer, পুরো একটা বইয়ের উপর পরের token predict করে train করা (The Adventures of Sherlock Holmes, public domain), laptop এর CPU তে, তারপর সেই বইয়ের ঢংয়ে নতুন text লেখা।

এই part এর প্রতিটা block এখানে একসাথে চলে। Training loop টা হুবহু সেটাই যেটা production এর LLM train করে; পার্থক্য শুধু সাইজে: কোটি কোটি parameter আর কয়েক মাসের বদলে 8 লাখ parameter আর 10 মিনিট।

যা আগে জানা লাগবে

• Self-Attention আর Transformer Architecture।
• Pretraining and Fine-tuning: পরের token predict করা।
• Inference, Decoding and Context Windows: temperature, top-k।
• PyTorch সহ Python, আর একটা plain-text বই (যেকোনো ভাষার; এটা Project Gutenberg থেকে)।

Architecture

Play

Setting	মান	কেন
Tokenizer	প্রতি অক্ষরে একটা token, 88 টা symbol	সবচেয়ে সহজ; BPE tokenizer (আগের project) sequence ছোট করত
Model	d = 128, 4 head, 4 block, context 128	কয়েক মিনিটের CPU বাজেটে আঁটে
Training	1,500 step × 32 জানালা, AdamW, one-cycle LR, gradient clipping	standard recipe, ছোট করে
Data ভাগ	প্রথম 90% train, শেষ 10% validation	শেষের গল্পগুলো model training এ কখনো দেখে না

ধাপে ধাপে বানাও

Step 1: model। দুইটা class। Block pre-norm attention করে (একটা linear layer একসাথে Q, K আর V বানায়, 4 টা head এ ভাগ, তারপর PyTorch এর fused causal attention) আর MLP, প্রতিটার সাথে residual। GPT embedding যোগ করে, block চাপায়, শেষ LayerNorm দেয়, আর output head কে token embedding এর সাথে tie করে। দুইটা detail জরুরি:

• ছোট শুরুর weight (std 0.02)। PyTorch এর default N(0, 1) embedding output layer এর সাথে tie করলে, এই project এর প্রথম run শুরু হয়েছিল validation loss 79.8 দিয়ে, প্রায় ln 88 = 4.48 এর বদলে: বিশাল random logits, আত্মবিশ্বাসের সাথে ভুল। Step 0 তে loss ln(vocab) এর অনেক উপরে মানেই খারাপ init।
• generate context window পর্যন্ত কেটে নেয় আর শুধু শেষ position এর logits ব্যবহার করে: decoding lesson এর generation loop।

"""A small decoder-only transformer (GPT) for characters."""
import torch
import torch.nn as nn
import torch.nn.functional as F

class Block(nn.Module):
    def __init__(self, d, heads, dropout):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d), nn.Dropout(dropout))
        self.heads, self.dropout = heads, dropout

    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(D, dim=2)
        q, k, v = (t.view(B, T, self.heads, D // self.heads).transpose(1, 2) for t in (q, k, v))
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, D))
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, vocab, d=128, heads=4, layers=4, ctx=128, dropout=0.1):
        super().__init__()
        self.ctx = ctx
        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(ctx, d)
        self.blocks = nn.Sequential(*[Block(d, heads, dropout) for _ in range(layers)])
        self.ln, self.head = nn.LayerNorm(d), nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok.weight
        self.apply(self._init)

    @staticmethod
    def _init(m):
        # GPT-2 style: small random weights. Default N(0, 1) embeddings, tied to the output layer,
        # would start with huge logits and a loss far above ln(vocab).
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.zeros_(m.bias)

    def forward(self, ids, targets=None):
        x = self.tok(ids) + self.pos(torch.arange(ids.shape[1]))
        logits = self.head(self.ln(self.blocks(x)))
        loss = None if targets is None else F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, ids, n, temperature=0.8, top_k=20, generator=None):
        for _ in range(n):
            logits, _ = self(ids[:, -self.ctx:])              # crop to the context window
            logits = logits[:, -1, :] / temperature            # only the last position predicts the next char
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = -float("inf")       # top-k
            nxt = torch.multinomial(F.softmax(logits, -1), 1, generator=generator)
            ids = torch.cat([ids, nxt], dim=1)
        return ids

Step 2: data আর training। বইকে character id হিসেবে encode করো, শেষ 10% আলাদা রাখো, আর random জানালা তোলো। Target হলো একই জানালা এক অক্ষর সরানো, তাই প্রতিটা জানালা একসাথে 128 টা training example দেয়। Validation জানালায় evaluate করো আর training চলতে চলতে একটা sample print করো।

import math
import torch
from model import GPT

torch.manual_seed(0)
torch.set_num_threads(4)
text = open("/data/sherlock.txt").read()
chars = sorted(set(text))
stoi, itos = {c: i for i, c in enumerate(chars)}, dict(enumerate(chars))
data = torch.tensor([stoi[c] for c in text])
split = int(0.9 * len(data))
train, val = data[:split], data[split:]                     # the last 10% of the book is held out
print(f"book: {len(text):,} characters, vocabulary {len(chars)}, train {len(train):,}, val {len(val):,}")

CTX, BATCH = 128, 32
def batch(src, g):
    i = torch.randint(len(src) - CTX - 1, (BATCH,), generator=g)
    x = torch.stack([src[j:j + CTX] for j in i])
    return x, torch.stack([src[j + 1:j + CTX + 1] for j in i])   # targets = inputs shifted by one

model = GPT(len(chars), ctx=CTX)
print(f"parameters: {sum(p.numel() for p in model.parameters()):,}")
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
STEPS = 1500
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=3e-3, total_steps=STEPS, pct_start=0.1)
gtrain, gval = torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)

@torch.no_grad()
def val_loss():
    model.eval()
    g = torch.Generator().manual_seed(3)
    loss = sum(model(*batch(val, g))[1].item() for _ in range(10)) / 10
    model.train()
    return loss

def sample(prompt, n=200, seed=4):
    model.eval()
    ids = torch.tensor([[stoi[c] for c in prompt]])
    out = model.generate(ids, n, generator=torch.Generator().manual_seed(seed))
    model.train()
    return "".join(itos[i] for i in out[0].tolist())

print(f"step     0  val loss {val_loss():.3f}  (uniform guess ln {len(chars)} = {math.log(len(chars)):.3f})")
print("sample:", repr(sample("Holmes ", 60)))
for step in range(1, STEPS + 1):
    _, loss = model(*batch(train, gtrain))
    opt.zero_grad(); loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step()
    if step in (250, 750, 1500):
        print(f"step {step:>5}  train loss {loss.item():.3f}  val loss {val_loss():.3f}")
        print("sample:", repr(sample("Holmes ", 60)))

torch.save({"model": model.state_dict(), "chars": chars}, "tiny_gpt.pt")
print("\n--- 400 characters from the trained model, prompt 'Watson,' ---")
print(sample("Watson,", 400))

Output (4 টা CPU thread এ 10 মিনিট):

book: 562,202 characters, vocabulary 88, train 505,981, val 56,221
parameters: 820,992
step     0  val loss 4.510  (uniform guess ln 88 = 4.477)
sample: 'Holmes ENb,8KxN—àfT,—æè‘uGrHGl&!TGlG0uHtVur‘T!!-æ)èNHPk‘æææF””ys1H½'
step   250  train loss 2.217  val loss 2.209
sample: 'Holmes had yom. Of asew wars. Hotelmerly crowned, inink weringesh, '
step   750  train loss 1.660  val loss 1.611
sample: 'Holmes had him. He as with he, and then, tronged his before passe,\n'
step  1500  train loss 1.465  val loss 1.449
sample: 'Holmes had been if a good came in\nin the crowded. And I went that, '

--- 400 characters from the trained model, prompt 'Watson,' ---
Watson, and he his as we seculiar
in his crowded his bride in the had eight only and who had at was very shall
discant. On he monothinst such to the right and at evhair of manage.

“He was grey not by as I have seen to saw if you not in a feel papers and
wedding and at the are of carth
the enought of money.”

“He stood in so.”

“The maid, where it far the bring of the business. My starty and a cafe you,

কী দেখলাম:

• Step 0: validation loss ঠিক ln 88 এ, ছোট init এ যেমন হওয়ার কথা, আর sample শুধু noise।
• Step 250: model space, দাঁড়ির পরে বড় হাতের অক্ষর আর ছোট common শব্দ ("had", "crowned") শিখেছে, সাথে বানানো শব্দও আছে ("Hotelmerly")।
• Step 750: প্রায় প্রতিটা শব্দ আসল আর বাক্যাংশগুলো বইয়ের মতো শোনায় ("Holmes had him. He as with he, and then"), কিন্তু ব্যাকরণ ভেঙে পড়ে।
• Step 1,500: validation loss নেমেছে 1.449 এ। শেষের 400 অক্ষর পড়ো: লাইনের দৈর্ঘ্য, বাঁকা উদ্ধৃতি চিহ্নে সংলাপ, একজন কাজের মেয়ে, কাগজপত্র, টাকা, বিয়ে, Victorian গদ্যের ছন্দ সব আছে, বেশিরভাগ শব্দও আসল; বাক্যগুলোর কোনো মানে নেই। 10 মিনিট ধরে একটা বইয়ে train করা 8 লাখ parameter এটুকুই শিখতে পারে: ভাষার আকার, তার বিষয়বস্তু না। একই code কে parameter, data আর compute এ 10,000 গুণ বড় করলে সেই আকার অর্থ বহন করতে শুরু করে।
• Train আর validation loss কাছাকাছি থাকে, তাই model এখনো overfit করছে না; বরং এখনো underfit। আরো step বা বড় model এ উন্নতি চলতে থাকত।

Test করো

import torch
from model import GPT

torch.manual_seed(0)
m = GPT(vocab=20, d=32, heads=4, layers=2, ctx=16, dropout=0.0)

def test_initial_loss_is_near_uniform():
    x = torch.randint(0, 20, (4, 16))
    _, loss = m(x, torch.randint(0, 20, (4, 16)))          # random targets: nothing to learn yet
    assert abs(loss.item() - torch.log(torch.tensor(20.0)).item()) < 0.1

def test_causal_future_does_not_change_the_past():
    x = torch.randint(0, 20, (1, 10))
    y = x.clone(); y[0, 7:] = (y[0, 7:] + 1) % 20          # change only tokens 7..9
    m.eval()
    assert torch.allclose(m(x)[0][0, :7], m(y)[0][0, :7], atol=1e-6)

def test_generate_respects_context_window():
    out = m.generate(torch.zeros(1, 1, dtype=torch.long), 40, generator=torch.Generator().manual_seed(0))
    assert out.shape == (1, 41)                           # 40 new tokens even though ctx is only 16

def test_can_memorise_a_short_sequence():
    mm = GPT(vocab=20, d=32, heads=4, layers=2, ctx=16, dropout=0.0)
    x = torch.arange(16).unsqueeze(0) % 20
    opt = torch.optim.Adam(mm.parameters(), lr=1e-2)
    for _ in range(200):
        _, loss = mm(x[:, :-1], x[:, 1:]); opt.zero_grad(); loss.backward(); opt.step()
    assert loss.item() < 0.05

if __name__ == "__main__":
    for name, fn in list(globals().items()):
        if name.startswith("test_"):
            fn()
            print("pass", name)

pass test_initial_loss_is_near_uniform
pass test_causal_future_does_not_change_the_past
pass test_generate_respects_context_window
pass test_can_memorise_a_short_sequence

• Random target এ শুরুর loss ≈ ln(vocab): খারাপ init (উপরের 79.8 এর মতো) ধরে ফেলে, 10 মিনিট training এ খরচ করার আগেই।
• ভবিষ্যৎ অতীত বদলায় না: token 7 থেকে 9 বদলালে position 0 থেকে 6 এর output বদলানো চলবে না। এটাই causal mask এর test; fail করলে training চুপচাপ চুরি করত।
• Generation context window মানে: 16 token context এর model থেকে 40 টা token চাইলে কেটে নিয়ে কাজ করতে হবে, position 17 এ crash না করে।
• ছোট একটা sequence মুখস্থ করতে পারে: Part 3 এর এক-batch-এ-overfit check, transformer এর জন্য।

Trade-offs

Choice	ভালো দিক	খরচ
Character token	train করার মতো tokenizer নেই; যেকোনো text চলে	লম্বা sequence: 128 অক্ষর মানে মাত্র প্রায় 25 শব্দের context
8 লাখ parameter, 1,500 step	CPU তে 10 মিনিট	বানান আর ঢং শেখে, অর্থ না; থামানোর সময় loss তখনো কমছে
Weight tying	output layer এর parameter বাঁচায়; প্রায়ই মান বাড়ায়	init এ সাবধান থাকতে হয় (79.8 এর শিক্ষা)
Fused scaled_dot_product_attention	দ্রুত, কম memory, ঠিক causal masking	self-attention lesson এর হাতে লেখা version এর চেয়ে কম স্বচ্ছ
একটা বই	দ্রুত, ঢং এক রকম	model শুধু একজন লেখকের জগৎ আর শব্দভাণ্ডার জানে

এখান থেকে একটা কাজের model এ যেতে ক্রমানুসারে বদলাতে: একটা BPE tokenizer, অনেক বেশি আর বৈচিত্র্যময় text, একটা GPU, বড় model, আর অনেক বেশি step। তারপর উপরে instruction tuning আর preference training (এই part এর lesson 5)।

আরো এগিয়ে নাও

1. আগের project এর BPE tokenizer এ যাও (বইয়ের উপর আবার train করে), আর একই সংখ্যক step এ sample মেলাও।
2. d দ্বিগুণ করে 256 করো আর 3,000 step train করো। Validation loss কত নিচে নামে, আর train আর validation কখন আলাদা হতে শুরু করে?
3. Train করা checkpoint থেকে temperature 0.3 আর 1.2 দিয়ে generate করো, আর text কীভাবে বদলায় বর্ণনা করো।
4. একটা বাংলা public-domain বইয়ে train করো (public archive থেকে রবীন্দ্রনাথ বা শরৎচন্দ্র), আর দেখো character-level model বাংলা লিপি থেকে কী বানায়।