5. Architektura LLM
Reading time: 16 minutes
Architektura LLM
tip
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 z powrotem.
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:
- Wejście (Tokenizowany tekst): Proces zaczyna się od tokenizowanego tekstu, który jest przekształcany w reprezentacje numeryczne.
- Warstwa osadzenia tokenów i warstwa osadzenia pozycyjnego: Tokenizowany tekst przechodzi przez warstwę osadzenia tokenów i warstwę osadzenia pozycyjnego, które uchwycają pozycję tokenów w sekwencji, co jest kluczowe dla zrozumienia kolejności słów.
- Bloki transformatorowe: Model zawiera 12 bloków transformatorowych, z których każdy ma wiele warstw. Te bloki powtarzają następującą sekwencję:
- Maskowana wielogłowa uwaga: Pozwala modelowi skupić się na różnych częściach tekstu wejściowego jednocześnie.
- Normalizacja warstwy: Krok normalizacji w celu stabilizacji i poprawy treningu.
- Warstwa feed forward: Odpowiedzialna za przetwarzanie informacji z warstwy uwagi i dokonywanie prognoz dotyczących następnego tokenu.
- Warstwy dropout: Te warstwy zapobiegają przeuczeniu, losowo eliminując jednostki podczas treningu.
- Ostateczna warstwa wyjściowa: 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.
- Cel: 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.
Reprezentacja kodu
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, niezerowe wartości dla ujemnych wejść.
- Definicja Matematyczna:
note
Celem użycia tej funkcji po warstwach liniowych wewnątrz warstwy FeedForward jest zmiana danych liniowych na 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
do4 * emb_dim
. - Aktywacja GELU: Zastosowuje nieliniowość.
- Druga warstwa liniowa: Zmniejsza wymiarowość z powrotem do
emb_dim
.
note
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 wytrenowania w 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 tokenu.
- 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
), z każdą o zmniejszonej 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.
note
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 tokenu, 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 (wymiarów 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
ishift
: Uczące się parametry (nn.Parameter
), które pozwalają modelowi na skalowanie i przesuwanie znormalizowanego wyjścia. Są inicjowane odpowiednio do jedynek i zer.- Proces Normalizacji:
- Obliczanie Średniej (
mean
): Oblicza średnią wejściax
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. Parametrunbiased=False
zapewnia, że wariancja jest obliczana przy użyciu obciążonego estymatora (dzieląc przezN
zamiastN-1
), co jest odpowiednie przy normalizacji wzdłuż cech, a nie próbek. - Normalizacja (
norm_x
): Odejmuje średnią odx
i dzieli przez pierwiastek kwadratowy z wariancji pluseps
. - Skalowanie i Przesunięcie: Zastosowuje uczące się parametry
scale
ishift
do znormalizowanego wyjścia.
note
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łowową 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. - Normalizacja Warstw (
norm1
): Normalizuj wejście. - Wielogłowowa 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. - Normalizacja Warstw (
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.
note
Blok transformera grupuje wszystkie sieci razem i stosuje pewne normalizacje i dropouty w celu poprawy stabilności treningu i wyników.
Zauważ, jak 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 transformatorowe (
trf_blocks
): Stosn_layers
bloków transformatorowych 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.
note
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 użyje tylu bloków transformatorowych, ile wskazano i że każdy blok transformatorowy korzysta z jednej sieci z wieloma głowicami uwagi, jednej sieci feed forward i kilku normalizacji. Więc jeśli użyto 12 bloków transformatorowych, 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ż, że 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łkowite parametry 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:
- Dwie instancje
LayerNorm
na blok. - Każda
LayerNorm
ma2 * 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 poprzedni, 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]))