5. LLM Architecture

LLM Architecture

Celem tej piątej fazy jest bardzo prosty: Opracowanie architektury pełnego LLM. Połącz wszystko, zastosuj wszystkie warstwy i stwórz wszystkie funkcje do generowania tekstu lub przekształcania tekstu na identyfikatory i odwrotnie.

Ta architektura będzie używana zarówno do treningu, jak i przewidywania tekstu po jego wytrenowaniu.

Przykład architektury LLM z https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb:

Wysokopoziomowa reprezentacja może być obserwowana w:

Input (Tokenized Text): Proces zaczyna się od tokenizowanego tekstu, który jest przekształcany w reprezentacje numeryczne.
Token Embedding and Positional Embedding Layer: Tokenizowany tekst przechodzi przez warstwę token embedding oraz warstwę positional embedding, która uchwyca pozycję tokenów w sekwencji, co jest kluczowe dla zrozumienia kolejności słów.
Transformer Blocks: Model zawiera 12 bloków transformatorowych, z których każdy ma wiele warstw. Te bloki powtarzają następującą sekwencję:

Masked Multi-Head Attention: Pozwala modelowi skupić się na różnych częściach tekstu wejściowego jednocześnie.
Layer Normalization: Krok normalizacji, aby ustabilizować i poprawić trening.
Feed Forward Layer: Odpowiada za przetwarzanie informacji z warstwy uwagi i dokonywanie prognoz dotyczących następnego tokena.
Dropout Layers: Te warstwy zapobiegają przeuczeniu, losowo eliminując jednostki podczas treningu.

Final Output Layer: Model generuje tensor o wymiarach 4x50,257, gdzie 50,257 reprezentuje rozmiar słownika. Każdy wiersz w tym tensorze odpowiada wektorowi, który model wykorzystuje do przewidywania następnego słowa w sekwencji.
Goal: Celem jest wzięcie tych osadzeń i przekształcenie ich z powrotem w tekst. Konkretnie, ostatni wiersz wyjścia jest używany do generowania następnego słowa, reprezentowanego jako "forward" w tym diagramie.

Code representation

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)

Funkcja aktywacji GELU

# 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))
))

Cel i Funkcjonalność

GELU (Gaussian Error Linear Unit): Funkcja aktywacji, która wprowadza nieliniowość do modelu.
Gładka Aktywacja: W przeciwieństwie do ReLU, która zeruje ujemne wejścia, GELU gładko mapuje wejścia na wyjścia, pozwalając na małe, różne od zera wartości dla ujemnych wejść.
Definicja Matematyczna:

Celem użycia tej funkcji po warstwach liniowych wewnątrz warstwy FeedForward jest przekształcenie danych liniowych w nieliniowe, aby umożliwić modelowi uczenie się złożonych, nieliniowych relacji.

Sieć Neuronowa FeedForward

Kształty zostały dodane jako komentarze, aby lepiej zrozumieć kształty macierzy:

# 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)

Cel i Funkcjonalność

Sieć FeedForward na poziomie pozycji: Zastosowuje dwuwarstwową sieć w pełni połączoną do każdej pozycji osobno i identycznie.
Szczegóły warstwy:
Pierwsza warstwa liniowa: Zwiększa wymiarowość z emb_dim do 4 * emb_dim.
Aktywacja GELU: Zastosowuje nieliniowość.
Druga warstwa liniowa: Redukuje wymiarowość z powrotem do emb_dim.

Jak widać, sieć Feed Forward używa 3 warstw. Pierwsza to warstwa liniowa, która pomnoży wymiary przez 4, używając wag liniowych (parametrów do trenowania wewnątrz modelu). Następnie funkcja GELU jest używana we wszystkich tych wymiarach, aby zastosować nieliniowe wariacje w celu uchwycenia bogatszych reprezentacji, a na końcu używana jest kolejna warstwa liniowa, aby wrócić do oryginalnego rozmiaru wymiarów.

Mechanizm Uwag Wielogłowych

To zostało już wyjaśnione w wcześniejszej sekcji.

Cel i Funkcjonalność

Wielogłowa Uwaga Własna: Pozwala modelowi skupić się na różnych pozycjach w sekwencji wejściowej podczas kodowania tokena.
Kluczowe komponenty:
Zapytania, Klucze, Wartości: Liniowe projekcje wejścia, używane do obliczania wyników uwagi.
Głowy: Wiele mechanizmów uwagi działających równolegle (num_heads), każdy z zredukowaną wymiarowością (head_dim).
Wyniki uwagi: Obliczane jako iloczyn skalarny zapytań i kluczy, skalowane i maskowane.
Maskowanie: Zastosowana jest maska przyczynowa, aby zapobiec modelowi zwracania uwagi na przyszłe tokeny (ważne dla modeli autoregresywnych, takich jak GPT).
Wagi uwagi: Softmax z maskowanych i skalowanych wyników uwagi.
Wektor kontekstu: Ważona suma wartości, zgodnie z wagami uwagi.
Projekcja wyjściowa: Warstwa liniowa do połączenia wyjść wszystkich głów.

Celem tej sieci jest znalezienie relacji między tokenami w tym samym kontekście. Ponadto tokeny są dzielone na różne głowy, aby zapobiec nadmiernemu dopasowaniu, chociaż ostateczne relacje znalezione na głowę są łączone na końcu tej sieci.

Ponadto, podczas treningu stosowana jest maska przyczynowa, aby późniejsze tokeny nie były brane pod uwagę przy poszukiwaniu specyficznych relacji do tokena, a także stosowany jest dropout, aby zapobiec nadmiernemu dopasowaniu.

Normalizacja Warstwy

# 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

Cel i Funkcjonalność

Normalizacja Warstw: Technika używana do normalizacji wejść wzdłuż cech (wymiary osadzenia) dla każdego pojedynczego przykładu w partii.
Składniki:
eps: Mała stała (1e-5) dodawana do wariancji, aby zapobiec dzieleniu przez zero podczas normalizacji.
scale i shift: Uczące się parametry (nn.Parameter), które pozwalają modelowi skalować i przesuwać znormalizowane wyjście. Są inicjowane odpowiednio do jedynek i zer.
Proces Normalizacji:
Obliczanie Średniej (mean): Oblicza średnią z wejścia x wzdłuż wymiaru osadzenia (dim=-1), zachowując wymiar do rozprzestrzeniania (keepdim=True).
Obliczanie Wariancji (var): Oblicza wariancję x wzdłuż wymiaru osadzenia, również zachowując wymiar. Parametr unbiased=False zapewnia, że wariancja jest obliczana przy użyciu obciążonego estymatora (dzieląc przez N zamiast N-1), co jest odpowiednie przy normalizacji wzdłuż cech, a nie próbek.
Normalizacja (norm_x): Odejmuje średnią od x i dzieli przez pierwiastek kwadratowy z wariancji plus eps.
Skalowanie i Przesunięcie: Zastosowuje uczące się parametry scale i shift do znormalizowanego wyjścia.

Celem jest zapewnienie średniej 0 z wariancją 1 we wszystkich wymiarach tego samego tokena. Celem tego jest stabilizacja treningu głębokich sieci neuronowych poprzez redukcję wewnętrznego przesunięcia kowariancji, które odnosi się do zmiany w rozkładzie aktywacji sieci z powodu aktualizacji parametrów podczas treningu.

Blok Transformera

Kształty zostały dodane jako komentarze, aby lepiej zrozumieć kształty macierzy:

# 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)

Cel i Funkcjonalność

Kompozycja Warstw: Łączy wielogłową uwagę, sieć feedforward, normalizację warstw i połączenia resztkowe.
Normalizacja Warstw: Stosowana przed warstwami uwagi i feedforward dla stabilnego treningu.
Połączenia Resztkowe (Skróty): Dodają wejście warstwy do jej wyjścia, aby poprawić przepływ gradientu i umożliwić trening głębokich sieci.
Dropout: Stosowany po warstwach uwagi i feedforward w celu regularyzacji.

Funkcjonalność Krok po Kroku

Pierwsza Ścieżka Resztkowa (Self-Attention):

Wejście (shortcut): Zapisz oryginalne wejście dla połączenia resztkowego.
Layer Norm (norm1): Normalizuj wejście.
Wielogłowa Uwaga (att): Zastosuj self-attention.
Dropout (drop_shortcut): Zastosuj dropout w celu regularyzacji.
Dodaj Resztkę (x + shortcut): Połącz z oryginalnym wejściem.

Druga Ścieżka Resztkowa (FeedForward):

Wejście (shortcut): Zapisz zaktualizowane wejście dla następnego połączenia resztkowego.
Layer Norm (norm2): Normalizuj wejście.
Sieć FeedForward (ff): Zastosuj transformację feedforward.
Dropout (drop_shortcut): Zastosuj dropout.
Dodaj Resztkę (x + shortcut): Połącz z wejściem z pierwszej ścieżki resztkowej.

Blok transformera grupuje wszystkie sieci razem i stosuje pewne normalizacje oraz dropouty w celu poprawy stabilności treningu i wyników. Zauważ, że dropouty są stosowane po użyciu każdej sieci, podczas gdy normalizacja jest stosowana przed.

Ponadto, wykorzystuje również skróty, które polegają na dodawaniu wyjścia sieci do jej wejścia. Pomaga to zapobiegać problemowi znikającego gradientu, zapewniając, że początkowe warstwy przyczyniają się "tak samo" jak ostatnie.

GPTModel

Kształty zostały dodane jako komentarze, aby lepiej zrozumieć kształty macierzy:

# 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)

Cel i Funkcjonalność

Warstwy Osadzeń:
Osadzenia Tokenów (tok_emb): Konwertuje indeksy tokenów na osadzenia. Przypomnienie, są to wagi przypisane do każdego wymiaru każdego tokena w słowniku.
Osadzenia Pozycyjne (pos_emb): Dodaje informacje o pozycji do osadzeń, aby uchwycić kolejność tokenów. Przypomnienie, są to wagi przypisane do tokena zgodnie z jego pozycją w tekście.
Dropout (drop_emb): Stosowane do osadzeń w celu regularyzacji.
Bloki Transformera (trf_blocks): Stos n_layers bloków transformera do przetwarzania osadzeń.
Ostateczna Normalizacja (final_norm): Normalizacja warstwy przed warstwą wyjściową.
Warstwa Wyjściowa (out_head): Projekcja ostatecznych ukrytych stanów do rozmiaru słownika w celu wygenerowania logitów do predykcji.

Celem tej klasy jest wykorzystanie wszystkich innych wspomnianych sieci do przewidywania następnego tokena w sekwencji, co jest fundamentalne dla zadań takich jak generowanie tekstu.

Zauważ, jak będzie używać tylu bloków transformera, ile wskazano i że każdy blok transformera używa jednej sieci z wieloma głowami uwagi, jednej sieci feed forward oraz kilku normalizacji. Więc jeśli używa się 12 bloków transformera, pomnóż to przez 12.

Ponadto, warstwa normalizacji jest dodawana przed wyjściem i na końcu stosowana jest ostateczna warstwa liniowa, aby uzyskać wyniki o odpowiednich wymiarach. Zauważ, jak każdy ostateczny wektor ma rozmiar używanego słownika. Dzieje się tak, ponieważ próbuje uzyskać prawdopodobieństwo dla każdego możliwego tokena w słowniku.

Liczba parametrów do wytrenowania

Mając zdefiniowaną strukturę GPT, możliwe jest ustalenie liczby parametrów do wytrenowania:

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

Krok po Kroku Obliczenia

1. Warstwy Osadzania: Osadzenie Tokenów i Osadzenie Pozycji

Warstwa: nn.Embedding(vocab_size, emb_dim)
Parametry: vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376

Warstwa: nn.Embedding(context_length, emb_dim)
Parametry: context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

Całkowita liczba parametrów osadzenia

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

2. Bloki Transformera

Jest 12 bloków transformera, więc obliczymy parametry dla jednego bloku, a następnie pomnożymy przez 12.

Parametry na blok transformera

a. Uwaga wielogłowa

Składniki:
Warstwa liniowa zapytania (W_query): nn.Linear(emb_dim, emb_dim, bias=False)
Warstwa liniowa klucza (W_key): nn.Linear(emb_dim, emb_dim, bias=False)
Warstwa liniowa wartości (W_value): nn.Linear(emb_dim, emb_dim, bias=False)
Projekcja wyjściowa (out_proj): nn.Linear(emb_dim, emb_dim)
Obliczenia:
Każda z W_query, W_key, W_value:

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

Ponieważ są trzy takie warstwy:

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

Projekcja wyjściowa (out_proj):

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

Całkowite parametry uwagi wielogłowej:

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

b. Sieć FeedForward

Składniki:
Pierwsza warstwa liniowa: nn.Linear(emb_dim, 4 * emb_dim)
Druga warstwa liniowa: nn.Linear(4 * emb_dim, emb_dim)
Obliczenia:
Pierwsza warstwa liniowa:

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

Druga warstwa liniowa:

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

Całkowite parametry FeedForward:

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

c. Normalizacje warstw

Składniki:
Dwa wystąpienia LayerNorm na blok.
Każde LayerNorm ma 2 * emb_dim parametrów (skala i przesunięcie).
Obliczenia:

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

d. Całkowite parametry na blok transformera

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

Całkowita liczba parametrów dla wszystkich bloków transformatorów

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

3. Ostateczne Warstwy

a. Normalizacja Ostatecznej Warstwy

Parametry: 2 * emb_dim (skala i przesunięcie)

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. Warstwa Projekcji Wyjścia (out_head)

Warstwa: nn.Linear(emb_dim, vocab_size, bias=False)
Parametry: emb_dim * vocab_size

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

4. Podsumowanie wszystkich parametrów

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

Generowanie tekstu

Mając model, który przewiduje następny token, jak ten wcześniej, wystarczy wziąć wartości ostatniego tokena z wyjścia (ponieważ będą to wartości przewidywanego tokena), które będą wartością na wpis w słowniku, a następnie użyć funkcji softmax, aby znormalizować wymiary do prawdopodobieństw, które sumują się do 1, a następnie uzyskać indeks największego wpisu, który będzie indeksem słowa w słowniku.

Kod z 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

Previous4. Attention Mechanisms Next6. Pre-training & Loading models

Last updated 6 days ago