5. LLM Architecture

LLM Arhitektura

Cilj ove pete faze je vrlo jednostavan: Razviti arhitekturu celog LLM-a. Spojiti sve zajedno, primeniti sve slojeve i kreirati sve funkcije za generisanje teksta ili transformaciju teksta u ID-ove i obrnuto.

Ova arhitektura će se koristiti za obuku i predikciju teksta nakon što je obučena.

Primer LLM arhitekture sa https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

Visok nivo reprezentacije može se posmatrati u:

https://camo.githubusercontent.com/6c8c392f72d5b9e86c94aeb9470beab435b888d24135926f1746eb88e0cc18fb/68747470733a2f2f73656261737469616e72617363686b612e636f6d2f696d616765732f4c4c4d732d66726f6d2d736372617463682d696d616765732f636830345f636f6d707265737365642f31332e776562703f31

1. Ulaz (Tokenizovani Tekst): Proces počinje sa tokenizovanim tekstom, koji se konvertuje u numeričke reprezentacije.

2. Sloj Token Embedding i Sloj Pozicionog Embedding-a: Tokenizovani tekst prolazi kroz sloj token embedding-a i sloj pozicionog embedding-a, koji hvata poziciju tokena u sekvenci, što je ključno za razumevanje reda reči.

3. Transformer Blokovi: Model sadrži 12 transformer blokova, svaki sa više slojeva. Ovi blokovi ponavljaju sledeću sekvencu:

• Maskirana Multi-Head Pažnja: Omogućava modelu da se fokusira na različite delove ulaznog teksta odjednom.

• Normalizacija Slojeva: Korak normalizacije za stabilizaciju i poboljšanje obuke.

• Sloj Feed Forward: Odgovoran za obradu informacija iz sloja pažnje i pravljenje predikcija o sledećem tokenu.

• Dropout Slojevi: Ovi slojevi sprečavaju prekomerno prilagođavanje tako što nasumično isključuju jedinice tokom obuke.

1. Završni Izlazni Sloj: Model izlazi sa 4x50,257-dimenzionalnim tenzorom, gde 50,257 predstavlja veličinu rečnika. Svaki red u ovom tenzoru odgovara vektoru koji model koristi za predikciju sledeće reči u sekvenci.

2. Cilj: Cilj je uzeti ove embedding-e i konvertovati ih nazad u tekst. Konkretno, poslednji red izlaza se koristi za generisanje sledeće reči, predstavljene kao "napred" u ovoj dijagramu.

Reprezentacija Koda

import torch
import torch.nn as nn
import tiktoken

class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

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

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
self.dropout = nn.Dropout(dropout)
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

def forward(self, x):
b, num_tokens, d_in = x.shape

keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)

# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)

# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)

# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)

# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection

return context_vec

class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"])
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# Shortcut connection for attention block
shortcut = x
x = self.norm1(x)
x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

# Shortcut connection for feed forward block
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

return x


class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(
cfg["emb_dim"], cfg["vocab_size"], bias=False
)

def forward(self, in_idx):
batch_size, seq_len = in_idx.shape
tok_embeds = self.tok_emb(in_idx)
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_emb(x)
x = self.trf_blocks(x)
x = self.final_norm(x)
logits = self.out_head(x)
return logits

GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
out = model(batch)
print("Input batch:\n", batch)
print("\nOutput shape:", out.shape)
print(out)

GELU Aktivacijska Funkcija

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))

Svrha i Funkcionalnost

• GELU (Gaussian Error Linear Unit): Aktivacijska funkcija koja uvodi nelinearnost u model.

• Gladka Aktivacija: Za razliku od ReLU, koja nula negativne ulaze, GELU glatko mapira ulaze na izlaze, omogućavajući male, nenulte vrednosti za negativne ulaze.

• Matematička Definicija:

Cilj korišćenja ove funkcije nakon linearnih slojeva unutar FeedForward sloja je da se promeni linearni podaci u nelinearne kako bi se modelu omogućilo učenje složenih, nelinearnih odnosa.

FeedForward Neuronska Mreža

Oblici su dodati kao komentari kako bi se bolje razumele forme matrica:

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

x = self.layers[0](x)# x shape: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[1](x) # x shape remains: (batch_size, seq_len, 4 * emb_dim)
x = self.layers[2](x) # x shape: (batch_size, seq_len, emb_dim)
return x  # Output shape: (batch_size, seq_len, emb_dim)

Svrha i Funkcionalnost

• Mreža sa Prolazom po Pozicijama: Primena dvoslojne potpuno povezane mreže na svaku poziciju odvojeno i identično.

• Detalji Slojeva:

• Prvi Linearni Sloj: Proširuje dimenzionalnost sa emb_dim na 4 * emb_dim.

• GELU Aktivacija: Primena nelinearnosti.

• Drugi Linearni Sloj: Smanjuje dimenzionalnost nazad na emb_dim.

Kao što možete videti, mreža sa prolazom koristi 3 sloja. Prvi je linearni sloj koji će pomnožiti dimenzije sa 4 koristeći linearne težine (parametre za obučavanje unutar modela). Zatim, GELU funkcija se koristi u svim tim dimenzijama da primeni nelinearne varijacije kako bi uhvatila bogatije reprezentacije, a na kraju se koristi još jedan linearni sloj da se vrati na originalnu veličinu dimenzija.

Mehanizam Višestruke Glave Pažnje

Ovo je već objašnjeno u ranijem odeljku.

Svrha i Funkcionalnost

• Višestruka Samo-Pažnja: Omogućava modelu da se fokusira na različite pozicije unutar ulazne sekvence prilikom kodiranja tokena.

• Ključne Komponente:

• Upiti, Ključevi, Vrednosti: Linearne projekcije ulaza, korišćene za izračunavanje ocena pažnje.

• Glave: Više mehanizama pažnje koji rade paralelno (num_heads), svaki sa smanjenom dimenzijom (head_dim).

• Ocene Pažnje: Izračunate kao skalarni proizvod upita i ključeva, skalirane i maskirane.

• Maskiranje: Primena uzročnog maskiranja kako bi se sprečilo da model obraća pažnju na buduće tokene (važno za autoregresivne modele poput GPT).

• Težine Pažnje: Softmax maskiranih i skaliranih ocena pažnje.

• Vektor Konteksta: Teženi zbir vrednosti, prema težinama pažnje.

• Projekcija Izlaza: Linearni sloj za kombinovanje izlaza svih glava.

Cilj ove mreže je da pronađe odnose između tokena u istom kontekstu. Štaviše, tokeni su podeljeni u različite glave kako bi se sprečilo prekomerno prilagođavanje, iako se konačni odnosi pronađeni po glavi kombinuju na kraju ove mreže.

Takođe, tokom obuke se primenjuje uzročno maskiranje kako kasniji tokeni ne bi bili uzeti u obzir prilikom gledanja specifičnih odnosa sa tokenom, a takođe se primenjuje i neki dropout da se spreči prekomerno prilagođavanje.

Normalizacija Sloja

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5 # Prevent division by zero during normalization.
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift

Svrha i Funkcionalnost

• Normalizacija slojeva: Tehnika koja se koristi za normalizaciju ulaza preko karakteristika (dimenzije ugradnje) za svaki pojedinačni primer u seriji.

• Komponente:

• eps: Mala konstanta (1e-5) koja se dodaje varijansi kako bi se sprečila deljenje sa nulom tokom normalizacije.

• scale i shift: Parametri koji se mogu učiti (nn.Parameter) koji omogućavaju modelu da skalira i pomera normalizovani izlaz. Inicijalizovani su na jedinice i nule, redom.

• Proces normalizacije:

• Izračunaj Srednju Vrednost (mean): Izračunava srednju vrednost ulaza x preko dimenzije ugradnje (dim=-1), zadržavajući dimenziju za emitovanje (keepdim=True).

• Izračunaj Varijansu (var): Izračunava varijansu x preko dimenzije ugradnje, takođe zadržavajući dimenziju. Parametar unbiased=False osigurava da se varijansa izračunava koristeći pristrasnog estimatora (deljenje sa N umesto N-1), što je prikladno kada se normalizuje preko karakteristika umesto uzoraka.

• Normalizuj (norm_x): Oduzima srednju vrednost od x i deli sa kvadratnim korenom varijanse plus eps.

• Skaliraj i Pomeri: Primena parametara scale i shift koji se mogu učiti na normalizovani izlaz.

Cilj je osigurati srednju vrednost 0 sa varijansom 1 preko svih dimenzija istog tokena. Cilj ovoga je da stabilizuje obuku dubokih neuronskih mreža smanjenjem unutrašnjeg pomeranja kovarijate, što se odnosi na promenu u distribuciji aktivacija mreže zbog ažuriranja parametara tokom obuke.

Transformer Blok

Oblici su dodati kao komentari kako bi se bolje razumele forme matrica:

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04

class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"]
)
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# x shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for attention block
shortcut = x  # shape: (batch_size, seq_len, emb_dim)
x = self.norm1(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.att(x)    # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

# Shortcut connection for feedforward block
shortcut = x       # shape: (batch_size, seq_len, emb_dim)
x = self.norm2(x)  # shape remains (batch_size, seq_len, emb_dim)
x = self.ff(x)     # shape: (batch_size, seq_len, emb_dim)
x = self.drop_shortcut(x)  # shape remains (batch_size, seq_len, emb_dim)
x = x + shortcut   # shape: (batch_size, seq_len, emb_dim)

return x  # Output shape: (batch_size, seq_len, emb_dim)

Svrha i Funkcionalnost

• Sastav Slojeva: Kombinuje višekratnu pažnju, feedforward mrežu, normalizaciju slojeva i rezidualne veze.

• Normalizacija Slojeva: Primena pre slojeva pažnje i feedforward za stabilno treniranje.

• Rezidualne Veze (Prečice): Dodaju ulaz sloja njegovom izlazu kako bi poboljšale protok gradijenata i omogućile treniranje dubokih mreža.

• Dropout: Primena nakon slojeva pažnje i feedforward za regularizaciju.

Funkcionalnost Korak po Korak

1. Prvi Rezidualni Put (Samo-Pažnja):

• Ulaz (shortcut): Sačuvaj originalni ulaz za rezidualnu vezu.

• Layer Norm (norm1): Normalizuj ulaz.

• Višekratna Pažnja (att): Primeni samo-pažnju.

• Dropout (drop_shortcut): Primeni dropout za regularizaciju.

• Dodaj Rezidual (x + shortcut): Kombinuj sa originalnim ulazom.

1. Drugi Rezidualni Put (FeedForward):

• Ulaz (shortcut): Sačuvaj ažurirani ulaz za sledeću rezidualnu vezu.

• Layer Norm (norm2): Normalizuj ulaz.

• FeedForward Mreža (ff): Primeni feedforward transformaciju.

• Dropout (drop_shortcut): Primeni dropout.

• Dodaj Rezidual (x + shortcut): Kombinuj sa ulazom iz prvog rezidualnog puta.

Transformator blok grupiše sve mreže zajedno i primenjuje neku normalizaciju i dropout kako bi poboljšao stabilnost treniranja i rezultate. Obratite pažnju kako se dropout primenjuje nakon korišćenja svake mreže dok se normalizacija primenjuje pre.

Pored toga, koristi i prečice koje se sastoje od dodavanja izlaza mreže sa njenim ulazom. Ovo pomaže u sprečavanju problema nestajućeg gradijenta osiguravajući da inicijalni slojevi doprinose "onoliko" koliko i poslednji.

GPTModel

Oblici su dodati kao komentari kako bi se bolje razumele forme matrica:

# From https://github.com/rasbt/LLMs-from-scratch/tree/main/ch04
class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
# shape: (vocab_size, emb_dim)

self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
# shape: (context_length, emb_dim)

self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
)
# Stack of TransformerBlocks

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
# shape: (emb_dim, vocab_size)

def forward(self, in_idx):
# in_idx shape: (batch_size, seq_len)
batch_size, seq_len = in_idx.shape

# Token embeddings
tok_embeds = self.tok_emb(in_idx)
# shape: (batch_size, seq_len, emb_dim)

# Positional embeddings
pos_indices = torch.arange(seq_len, device=in_idx.device)
# shape: (seq_len,)
pos_embeds = self.pos_emb(pos_indices)
# shape: (seq_len, emb_dim)

# Add token and positional embeddings
x = tok_embeds + pos_embeds  # Broadcasting over batch dimension
# x shape: (batch_size, seq_len, emb_dim)

x = self.drop_emb(x)  # Dropout applied
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.trf_blocks(x)  # Pass through Transformer blocks
# x shape remains: (batch_size, seq_len, emb_dim)

x = self.final_norm(x)  # Final LayerNorm
# x shape remains: (batch_size, seq_len, emb_dim)

logits = self.out_head(x)  # Project to vocabulary size
# logits shape: (batch_size, seq_len, vocab_size)

return logits  # Output shape: (batch_size, seq_len, vocab_size)

Svrha i Funkcionalnost

• Ugrađeni slojevi:

• Token Ugrađivanja (tok_emb): Pretvara indekse tokena u ugrađivanja. Kao podsetnik, ovo su težine date svakoj dimenziji svakog tokena u rečniku.

• Pozicijska Ugrađivanja (pos_emb): Dodaje pozicijske informacije u ugrađivanja kako bi se uhvatio redosled tokena. Kao podsetnik, ovo su težine date tokenu prema njegovoj poziciji u tekstu.

• Dropout (drop_emb): Primena na ugrađivanja za regularizaciju.

• Transformer Blokovi (trf_blocks): Stek od n_layers transformer blokova za obradu ugrađivanja.

• Finalna Normalizacija (final_norm): Normalizacija sloja pre izlaznog sloja.

• Izlazni Sloj (out_head): Projektuje konačne skrivene state-ove na veličinu rečnika kako bi proizveo logite za predikciju.

Cilj ove klase je da koristi sve ostale pomenute mreže da predvidi sledeći token u sekvenci, što je fundamentalno za zadatke poput generisanja teksta.

Napomena kako će koristiti onoliko transformer blokova koliko je naznačeno i da svaki transformer blok koristi jednu mrežu sa više glava, jednu mrežu za unapred i nekoliko normalizacija. Dakle, ako se koristi 12 transformer blokova, pomnožite ovo sa 12.

Štaviše, normalizacija sloj se dodaje pre izlaza i konačni linearni sloj se primenjuje na kraju kako bi se dobili rezultati sa odgovarajućim dimenzijama. Napomena kako svaki konačni vektor ima veličinu korišćenog rečnika. To je zato što pokušava da dobije verovatnoću po mogućem tokenu unutar rečnika.

Broj parametara za obuku

Imajući definisanu GPT strukturu, moguće je saznati broj parametara za obuku:

GPT_CONFIG_124M = {
"vocab_size": 50257,    # Vocabulary size
"context_length": 1024, # Context length
"emb_dim": 768,         # Embedding dimension
"n_heads": 12,          # Number of attention heads
"n_layers": 12,         # Number of layers
"drop_rate": 0.1,       # Dropout rate
"qkv_bias": False       # Query-Key-Value bias
}

model = GPTModel(GPT_CONFIG_124M)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")
# Total number of parameters: 163,009,536

Korak-po-korak Proračun

1. Ugradne Slojeve: Ugradnja Tokena & Ugradnja Pozicije

• Sloj: nn.Embedding(vocab_size, emb_dim)

• Parametri: vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376

• Sloj: nn.Embedding(context_length, emb_dim)

• Parametri: context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

Ukupni parametri ugrađivanja

embedding_params = token_embedding_params + position_embedding_params
embedding_params = 38,597,376 + 786,432 = 39,383,808

2. Transformer Blokovi

Postoji 12 transformer blokova, pa ćemo izračunati parametre za jedan blok, a zatim pomnožiti sa 12.

Parametri po Transformer Bloku

a. Multi-Head Attention

• Komponente:

• Query Linear Layer (W_query): nn.Linear(emb_dim, emb_dim, bias=False)

• Key Linear Layer (W_key): nn.Linear(emb_dim, emb_dim, bias=False)

• Value Linear Layer (W_value): nn.Linear(emb_dim, emb_dim, bias=False)

• Output Projection (out_proj): nn.Linear(emb_dim, emb_dim)

• Izračunavanja:

• Svaki od W_query, W_key, W_value:

qkv_params = emb_dim * emb_dim = 768 * 768 = 589,824

Pošto postoje tri takva sloja:

total_qkv_params = 3 * qkv_params = 3 * 589,824 = 1,769,472

• Output Projection (out_proj):

out_proj_params = (emb_dim * emb_dim) + emb_dim = (768 * 768) + 768 = 589,824 + 768 = 590,592

• Ukupni Multi-Head Attention Parametri:

mha_params = total_qkv_params + out_proj_params
mha_params = 1,769,472 + 590,592 = 2,360,064

b. FeedForward Mreža

• Komponente:

• Prvi Linear Layer: nn.Linear(emb_dim, 4 * emb_dim)

• Drugi Linear Layer: nn.Linear(4 * emb_dim, emb_dim)

• Izračunavanja:

• Prvi Linear Layer:

ff_first_layer_params = (emb_dim * 4 * emb_dim) + (4 * emb_dim)
ff_first_layer_params = (768 * 3072) + 3072 = 2,359,296 + 3,072 = 2,362,368

• Drugi Linear Layer:

ff_second_layer_params = (4 * emb_dim * emb_dim) + emb_dim
ff_second_layer_params = (3072 * 768) + 768 = 2,359,296 + 768 = 2,360,064

• Ukupni FeedForward Parametri:

ff_params = ff_first_layer_params + ff_second_layer_params
ff_params = 2,362,368 + 2,360,064 = 4,722,432

c. Normalizacije Slojeva

• Komponente:

• Dva LayerNorm instance po bloku.

• Svaki LayerNorm ima 2 * emb_dim parametara (skala i pomeraj).

• Izračunavanja:

layer_norm_params_per_block = 2 * (2 * emb_dim) = 2 * 768 * 2 = 3,072

d. Ukupni Parametri po Transformer Bloku

pythonCopy codeparams_per_block = mha_params + ff_params + layer_norm_params_per_block
params_per_block = 2,360,064 + 4,722,432 + 3,072 = 7,085,568

Ukupni parametri za sve Transformer blokove

pythonCopy codetotal_transformer_blocks_params = params_per_block * n_layers
total_transformer_blocks_params = 7,085,568 * 12 = 85,026,816

3. Final Layers

a. Final Layer Normalization

• Parameters: 2 * emb_dim (skala i pomeraj)

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. Sloj izlaza (out_head)

• Sloj: nn.Linear(emb_dim, vocab_size, bias=False)

• Parametri: emb_dim * vocab_size

pythonCopy codeoutput_projection_params = 768 * 50257 = 38,597,376

4. Sumiranje svih parametara

pythonCopy codetotal_params = (
embedding_params +
total_transformer_blocks_params +
final_layer_norm_params +
output_projection_params
)
total_params = (
39,383,808 +
85,026,816 +
1,536 +
38,597,376
)
total_params = 163,009,536

Generate Text

Imajući model koji predviđa sledeći token kao prethodni, potrebno je uzeti poslednje vrednosti tokena iz izlaza (jer će to biti vrednosti predviđenog tokena), što će biti vrednost po unosu u rečniku i zatim koristiti softmax funkciju da normalizuje dimenzije u verovatnoće koje se sabiraju na 1 i zatim dobiti indeks najvećeg unosa, koji će biti indeks reči unutar rečnika.

Code from https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

def generate_text_simple(model, idx, max_new_tokens, context_size):
# idx is (batch, n_tokens) array of indices in the current context
for _ in range(max_new_tokens):

# Crop current context if it exceeds the supported context size
# E.g., if LLM supports only 5 tokens, and the context size is 10
# then only the last 5 tokens are used as context
idx_cond = idx[:, -context_size:]

# Get the predictions
with torch.no_grad():
logits = model(idx_cond)

# Focus only on the last time step
# (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
logits = logits[:, -1, :]

# Apply softmax to get probabilities
probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)

# Get the idx of the vocab entry with the highest probability value
idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)

# Append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

return idx


start_context = "Hello, I am"

encoded = tokenizer.encode(start_context)
print("encoded:", encoded)

encoded_tensor = torch.tensor(encoded).unsqueeze(0)
print("encoded_tensor.shape:", encoded_tensor.shape)

model.eval() # disable dropout

out = generate_text_simple(
model=model,
idx=encoded_tensor,
max_new_tokens=6,
context_size=GPT_CONFIG_124M["context_length"]
)

print("Output:", out)
print("Output length:", len(out[0]))

References

• https://www.manning.com/books/build-a-large-language-model-from-scratch