Transformer Architecture

জিনিসটা কী, কেন দরকার

সব টুকরো এখন হাতে: token position সহ embedding হয়, self-attention token দের তথ্য ভাগাভাগি করতে দেয়, MLP একটা ছোট neural network। Transformer হলো এই টুকরোগুলো একটা সহজ, বারবার আসা pattern এ সাজানো, আর সেই একই pattern 12 বার বা 120 বার চাপালেই GPT, Llama, Claude, Gemini আর ব্যবহারে থাকা প্রায় সব LLM।

আকারটা জানলে আসল system নিয়ে ভাবতে পারবে: "7B parameter" এর model কেন প্রায় 14 GB GPU memory চায়, সেই parameter গুলো আসলে কোথায় বসে, fine-tuning এর সময় কোন layer গুলো সস্তায় বদলানো যায় (Part 9), আর model card এ "32 layer এর model" পড়লে তার মানে কী।

ভেতরে কীভাবে চলে

একটা decoder-only transformer (GPT পরিবার) token id এর একটা sequence নিয়ে এটা করে:

1. Embed: token embedding + position (এই part এর lesson 2)।
2. N টা block বারবার। প্রতিটা block এ দুইটা sub-layer:
   • Attention: LayerNorm, তারপর causal multi-head self-attention, তারপর result টা input এর সাথে যোগ করো। Token রা একে অপরের থেকে এখানেই পড়ে।
   • MLP (feed-forward): LayerNorm, তারপর দুই layer এর একটা network, সাধারণত মাঝে 4 গুণ চওড়া আর GELU, তারপর result আবার যোগ। এটা প্রতিটা token এর উপর আলাদা আলাদা কাজ করে; model এর জমা "জ্ঞান" এর অনেকটাই থাকে এই weight গুলোতে।
3. শেষ LayerNorm, তারপর একটা linear head যেটা প্রতিটা token এর vector কে vocabulary র প্রতিটা token এর জন্য একটা করে score বানায়: পরের token এর logits।

Residual connection হলো ওই "result আবার যোগ করো" অংশগুলো। একটা residual stream কল্পনা করো: প্রতিটা token এর একটা vector, model এর নিচ থেকে উপরে বয়ে যাচ্ছে। প্রতিটা sub-layer একটা কপি পড়ে, কিছু হিসাব করে, আর নিজের অবদান যোগ করে; stream কে কেউ বদলে দেয় না। এতে loss থেকে প্রথম layer পর্যন্ত gradient এর একটা সোজা রাস্তা থাকে, আর এজন্যই 100 layer এর model train করা যায়।

LayerNorm প্রতিটা sub-layer এর আগে প্রতিটা token এর vector কে mean 0 আর spread 1 এ আনে (শেখা scale আর shift সহ), যাতে stack যত গভীরই হোক সংখ্যাগুলো স্থির range এ থাকে। এখনকার model এটা প্রতিটা sub-layer এর আগে রাখে (pre-norm); অনেকে সস্তা RMSNorm ব্যবহার করে।

Parameter কোথায়। প্রতি block এ: attention এ d × d এর 4 টা matrix (Q, K, V, output) = 4d²; MLP তে d × 4d এর দুইটা = 8d²। তাই একটা block মোটামুটি 12d², তার দুই তৃতীয়াংশ MLP তে। সাথে embedding table (vocab × d), প্রায়ই output head এর সাথে ভাগ করা (weight tying)।

Decoder-only বনাম অন্যরা। আসল transformer এ ছিল একটা encoder (পুরো input দেখে, mask নেই) আর একটা decoder (causal), translation এর জন্য। BERT ধরনের model শুধু encoder, classification আর search embedding এ ভালো। আজকের chat LLM গুলো decoder-only: একটা stack, causal mask, পরের token predict।

আগে অনুমান করো

একটা model card বলছে: d = 4,096, 32 layer, vocabulary 32,000। প্রতি block এ মোটামুটি 12·d² ধরে, এর parameter মোটামুটি কত?

প্রায় 6.6 billionপ্রায় 20 কোটিপ্রায় 70 billionপ্রায় 13 কোটি

ছবিতে দেখো

Play

Code এ পুরো একটা GPT

PyTorch এ মোটামুটি 30 লাইনে পুরো একটা GPT ধরনের model, GPT-2 small এর হুবহু আকারে বানানো যাতে parameter কোথায় থাকে গোনা যায়, তারপর ছোট একটা version দিয়ে input আর output এর shape দেখা।

import torch
import torch.nn as nn

class Block(nn.Module):
    """One decoder block: pre-norm attention and MLP, each wrapped in a residual connection."""
    def __init__(self, d, heads):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        n = x.shape[1]
        mask = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)   # causal
        h = self.ln1(x)
        a = self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        x = x + a                                    # residual: add, don't replace (tokens exchange information)
        m = self.mlp(self.ln2(x))
        return x + m                                 # residual again (each token thinks on its own)

class TinyGPT(nn.Module):
    def __init__(self, vocab, d, heads, layers, ctx):
        super().__init__()
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(ctx, d)
        self.blocks = nn.Sequential(*[Block(d, heads) for _ in range(layers)])
        self.ln = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok.weight           # weight tying: output layer reuses the embedding table

    def forward(self, ids):
        x = self.tok(ids) + self.pos(torch.arange(ids.shape[1]))
        return self.head(self.ln(self.blocks(x)))   # (batch, tokens, vocab) scores for the next token

# Where do the parameters live? GPT-2 small's shape: 50,257 tokens, d=768, 12 heads, 12 layers, 1,024 context.
gpt2 = TinyGPT(50257, 768, 12, 12, 1024)
groups = {"token embeddings (shared with output)": gpt2.tok, "position embeddings": gpt2.pos,
          "attention (12 layers)": [b.attn for b in gpt2.blocks], "MLP (12 layers)": [b.mlp for b in gpt2.blocks],
          "layer norms": [gpt2.ln] + [m for b in gpt2.blocks for m in (b.ln1, b.ln2)]}
count = lambda ms: sum(p.numel() for m in (ms if isinstance(ms, list) else [ms]) for p in m.parameters())
total = sum(p.numel() for p in gpt2.parameters())
for name, ms in groups.items():
    print(f"{name:<40} {count(ms):>11,}  {count(ms) / total:5.1%}")
print(f"{'total':<40} {total:>11,}")

# Shapes through a small model.
torch.manual_seed(0)
tiny = TinyGPT(vocab=100, d=32, heads=4, layers=2, ctx=16)
ids = torch.randint(0, 100, (1, 6))
print("\nids", tuple(ids.shape), "-> logits", tuple(tiny(ids).shape), " (one score per vocab token, at every position)")

# What the parameter count means in practice.
print(f"\nweights in float32: {total * 4 / 1e6:,.0f} MB, in float16: {total * 2 / 1e6:,.0f} MB")
print(f"compute per generated token: about 2 x params = {2 * total / 1e9:.2f} GFLOPs")

Output:

token embeddings (shared with output)     38,597,376  31.0%
position embeddings                          786,432   0.6%
attention (12 layers)                     28,348,416  22.8%
MLP (12 layers)                           56,669,184  45.5%
layer norms                                   38,400   0.0%
total                                    124,439,808

ids (1, 6) -> logits (1, 6, 100)  (one score per vocab token, at every position)

weights in float32: 498 MB, in float16: 249 MB
compute per generated token: about 2 x params = 0.25 GFLOPs

কী দেখলাম:

• 124,439,808 টা parameter, GPT-2 small এর পরিচিত "124M", 30 লাইন থেকে। Architecture টা সত্যিই এত ছোট; সাইজ আসে d, layer আর vocabulary থেকে।
• MLP তে 45.5%, attention এ 22.8%: সেই 8d² বনাম 4d² এর ভাগ। Token embedding নেয় 31%, ছোট model এ বড় ভাগ কারণ vocabulary বড়; 70B model এ একই table প্রায় 1%।
• Position embedding (শেখা, 1,024 × 768) মাত্র 0.6%। Layer norm প্রায় কিছুই না।
• Shape: 6 টা token id ঢোকে, বের হয় (1, 6, 100) এর tensor: প্রতিটা position এ, vocabulary র 100 টা token এর প্রতিটার জন্য একটা score। Training এ ছয়টা position একসাথে score হয় (causal mask এর কারণে প্রতিটা তার পরের token predict করে); generation এ শুধু শেষেরটা ব্যবহার হয়।
• বাস্তবে সাইজ: float32 এ 498 MB, float16 এ 249 MB, আর প্রতি generate করা token এ প্রায় 0.25 GFLOP (প্রতি parameter এ 2: একটা গুণ, একটা যোগ)। 7B model এর জন্য এগুলোকে 56 দিয়ে গুণ করো: float16 এ প্রায় 14 GB আর প্রতি token এ 14 GFLOP।

Production এ যা মাথায় রাখবে

• Memory = parameter × প্রতি parameter এর byte, সাথে KV cache আর activation। 7B model 16-bit এ ~14 GB, 4-bit এ ~4 GB (Part 9)। কোন GPU লাগবে, অন্য সবকিছুর আগে এই হিসাবটাই ঠিক করে।
• প্রতিটা generate করা token প্রতিটা weight ছোঁয়। Generation এর speed সাধারণত আটকায় GPU কত দ্রুত memory থেকে সব weight পড়তে পারে তাতে, অঙ্কে না। এজন্যই quantization জিনিস দ্রুত করে।
• Context length এর খরচ আলাদা। Context এর সাথে parameter এর সংখ্যা বদলায় না, কিন্তু attention এর কাজ আর KV cache বাড়ে।
• Model card এ যে architecture variant গুলো দেখবে: শেখা position এর বদলে RoPE, LayerNorm এর বদলে RMSNorm, MLP তে GELU এর বদলে SwiGLU, grouped-query attention, mixture-of-experts MLP (প্রতি token এ শুধু কিছু MLP চলে)। কঙ্কালটা একই থাকে।

এক কথায়

Decoder-only transformer token embed করে, তারপর একটা residual stream এর উপর [LayerNorm → attention → যোগ] আর [LayerNorm → MLP → যোগ] এর block বারবার চালায়, আর শেষে একটা head পরের token score করে; প্রতিটা block এ মোটামুটি 12·d² parameter, দুই তৃতীয়াংশ MLP তে।

নিজে করে দেখো

Code এ GPT-2 এর config বদলে করো d = 1,024, 16 head আর 24 layer (GPT-2 medium এর আকার), আবার গোনো। 12·d²·layers + embedding দিয়ে result মেলাও, আর বের করো weight এর জন্য float16 এ কত memory লাগবে।

Answer

Block গুলোতে 12 × 1,024² × 24 ≈ 30.2 কোটি, সাথে 50,257 × 1,024 ≈ 5.1 কোটি token embedding আর প্রায় 10 লাখ position embedding: মোটামুটি 35.4 কোটি, GPT-2 medium এর প্রকাশিত 355M এর কাছে (code bias ও গোনে)। Float16 এ weight এর জন্য প্রায় 710 MB।