5. LLM Architecture

LLM-Architektur

Das Ziel dieser fünften Phase ist sehr einfach: Entwickeln Sie die Architektur des vollständigen LLM. Fügen Sie alles zusammen, wenden Sie alle Schichten an und erstellen Sie alle Funktionen, um Text zu generieren oder Text in IDs und umgekehrt zu transformieren.

Diese Architektur wird sowohl für das Training als auch für die Vorhersage von Text nach dem Training verwendet.

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

Eine hochrangige Darstellung kann beobachtet werden in:

  1. Eingabe (Tokenisierter Text): Der Prozess beginnt mit tokenisiertem Text, der in numerische Darstellungen umgewandelt wird.

  2. Token-Embedding- und Positions-Embedding-Schicht: Der tokenisierte Text wird durch eine Token-Embedding-Schicht und eine Positions-Embedding-Schicht geleitet, die die Position der Tokens in einer Sequenz erfasst, was entscheidend für das Verständnis der Wortreihenfolge ist.

  3. Transformer-Blöcke: Das Modell enthält 12 Transformer-Blöcke, von denen jeder mehrere Schichten hat. Diese Blöcke wiederholen die folgende Sequenz:

  • Maskierte Multi-Head-Attention: Ermöglicht es dem Modell, sich gleichzeitig auf verschiedene Teile des Eingabetextes zu konzentrieren.

  • Layer-Normalisierung: Ein Normalisierungsschritt zur Stabilisierung und Verbesserung des Trainings.

  • Feed-Forward-Schicht: Verantwortlich für die Verarbeitung der Informationen aus der Attention-Schicht und für die Vorhersage des nächsten Tokens.

  • Dropout-Schichten: Diese Schichten verhindern Overfitting, indem sie während des Trainings zufällig Einheiten fallen lassen.

  1. Endausgabeschicht: Das Modell gibt einen 4x50.257-dimensionalen Tensor aus, wobei 50.257 die Größe des Wortschatzes darstellt. Jede Zeile in diesem Tensor entspricht einem Vektor, den das Modell verwendet, um das nächste Wort in der Sequenz vorherzusagen.

  2. Ziel: Das Ziel ist es, diese Embeddings zu nehmen und sie wieder in Text umzuwandeln. Insbesondere wird die letzte Zeile der Ausgabe verwendet, um das nächste Wort zu generieren, das in diesem Diagramm als "vorwärts" dargestellt ist.

Code-Darstellung

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-Aktivierungsfunktion

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

Zweck und Funktionalität

  • GELU (Gaussian Error Linear Unit): Eine Aktivierungsfunktion, die Nichtlinearität in das Modell einführt.

  • Glatte Aktivierung: Im Gegensatz zu ReLU, das negative Eingaben auf null setzt, mappt GELU Eingaben sanft auf Ausgaben und erlaubt kleine, von null verschiedene Werte für negative Eingaben.

  • Mathematische Definition:

Das Ziel der Verwendung dieser Funktion nach linearen Schichten innerhalb der FeedForward-Schicht ist es, die linearen Daten nicht linear zu machen, um dem Modell zu ermöglichen, komplexe, nicht-lineare Beziehungen zu lernen.

FeedForward-Neuronales Netzwerk

_ Formen wurden als Kommentare hinzugefügt, um die Formen der Matrizen besser zu verstehen:_

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

Zweck und Funktionalität

  • Positionsweise FeedForward-Netzwerk: Wendet ein zweilagiges voll verbundenes Netzwerk auf jede Position separat und identisch an.

  • Schichtdetails:

  • Erste lineare Schicht: Erweitert die Dimensionalität von emb_dim auf 4 * emb_dim.

  • GELU-Aktivierung: Wendet Nichtlinearität an.

  • Zweite lineare Schicht: Reduziert die Dimensionalität zurück auf emb_dim.

Wie Sie sehen können, verwendet das Feed Forward-Netzwerk 3 Schichten. Die erste ist eine lineare Schicht, die die Dimensionen mit 4 multipliziert, indem sie lineare Gewichte (Parameter, die im Modell trainiert werden) verwendet. Dann wird die GELU-Funktion in all diesen Dimensionen verwendet, um nicht-lineare Variationen anzuwenden, um reichhaltigere Darstellungen zu erfassen, und schließlich wird eine weitere lineare Schicht verwendet, um zur ursprünglichen Größe der Dimensionen zurückzukehren.

Multi-Head Attention-Mechanismus

Dies wurde bereits in einem früheren Abschnitt erklärt.

Zweck und Funktionalität

  • Multi-Head Self-Attention: Ermöglicht es dem Modell, sich auf verschiedene Positionen innerhalb der Eingabesequenz zu konzentrieren, wenn es ein Token kodiert.

  • Schlüsselfunktionen:

  • Abfragen, Schlüssel, Werte: Lineare Projektionen der Eingabe, die zur Berechnung der Aufmerksamkeitswerte verwendet werden.

  • Köpfe: Mehrere Aufmerksamkeitsmechanismen, die parallel laufen (num_heads), jeder mit einer reduzierten Dimension (head_dim).

  • Aufmerksamkeitswerte: Berechnet als das Skalarprodukt von Abfragen und Schlüsseln, skaliert und maskiert.

  • Maskierung: Eine kausale Maske wird angewendet, um zu verhindern, dass das Modell zukünftige Tokens berücksichtigt (wichtig für autoregressive Modelle wie GPT).

  • Aufmerksamkeitsgewichte: Softmax der maskierten und skalierten Aufmerksamkeitswerte.

  • Kontextvektor: Gewichtete Summe der Werte, entsprechend den Aufmerksamkeitsgewichten.

  • Ausgabeprojektion: Lineare Schicht zur Kombination der Ausgaben aller Köpfe.

Das Ziel dieses Netzwerks ist es, die Beziehungen zwischen Tokens im gleichen Kontext zu finden. Darüber hinaus werden die Tokens in verschiedene Köpfe unterteilt, um Überanpassung zu verhindern, obwohl die endgültigen Beziehungen, die pro Kopf gefunden werden, am Ende dieses Netzwerks kombiniert werden.

Darüber hinaus wird während des Trainings eine kausale Maske angewendet, sodass spätere Tokens nicht berücksichtigt werden, wenn die spezifischen Beziehungen zu einem Token betrachtet werden, und es wird auch ein Dropout angewendet, um Überanpassung zu verhindern.

Schicht Normalisierung

# 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

Zweck und Funktionalität

  • Layer Normalization: Eine Technik, die verwendet wird, um die Eingaben über die Merkmale (Embedding-Dimensionen) für jedes einzelne Beispiel in einem Batch zu normalisieren.

  • Komponenten:

  • eps: Eine kleine Konstante (1e-5), die zur Varianz hinzugefügt wird, um eine Division durch Null während der Normalisierung zu verhindern.

  • scale und shift: Lernbare Parameter (nn.Parameter), die es dem Modell ermöglichen, die normalisierte Ausgabe zu skalieren und zu verschieben. Sie werden jeweils mit Einsen und Nullen initialisiert.

  • Normalisierungsprozess:

  • Mittelwert berechnen (mean): Berechnet den Mittelwert der Eingabe x über die Embedding-Dimension (dim=-1), wobei die Dimension für Broadcasting beibehalten wird (keepdim=True).

  • Varianz berechnen (var): Berechnet die Varianz von x über die Embedding-Dimension und behält ebenfalls die Dimension bei. Der Parameter unbiased=False stellt sicher, dass die Varianz mit dem verzerrten Schätzer berechnet wird (Division durch N anstelle von N-1), was angemessen ist, wenn über Merkmale anstelle von Proben normalisiert wird.

  • Normalisieren (norm_x): Subtrahiert den Mittelwert von x und dividiert durch die Quadratwurzel der Varianz plus eps.

  • Skalieren und Verschieben: Wendet die lernbaren scale- und shift-Parameter auf die normalisierte Ausgabe an.

Das Ziel ist es, einen Mittelwert von 0 mit einer Varianz von 1 über alle Dimensionen des gleichen Tokens sicherzustellen. Das Ziel davon ist es, das Training von tiefen neuronalen Netzwerken zu stabilisieren, indem der interne Kovariatenverschiebung reduziert wird, die sich auf die Änderung der Verteilung der Netzwerkaktivierungen aufgrund der Aktualisierung der Parameter während des Trainings bezieht.

Transformer Block

Größen wurden als Kommentare hinzugefügt, um die Größen der Matrizen besser zu verstehen:

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

Zweck und Funktionalität

  • Zusammensetzung der Schichten: Kombiniert Multi-Head-Attention, Feedforward-Netzwerk, Layer-Normalisierung und Residualverbindungen.

  • Layer-Normalisierung: Vor der Attention- und Feedforward-Schicht angewendet für stabiles Training.

  • Residualverbindungen (Abkürzungen): Fügen den Eingang einer Schicht zu ihrem Ausgang hinzu, um den Gradientfluss zu verbessern und das Training tiefer Netzwerke zu ermöglichen.

  • Dropout: Nach der Attention- und Feedforward-Schicht zur Regularisierung angewendet.

Schritt-für-Schritt-Funktionalität

  1. Erster Residualpfad (Selbst-Attention):

  • Eingang (shortcut): Speichern Sie den ursprünglichen Eingang für die Residualverbindung.

  • Layer Norm (norm1): Normalisieren Sie den Eingang.

  • Multi-Head-Attention (att): Wenden Sie Selbst-Attention an.

  • Dropout (drop_shortcut): Wenden Sie Dropout zur Regularisierung an.

  • Add Residual (x + shortcut): Kombinieren Sie mit dem ursprünglichen Eingang.

  1. Zweiter Residualpfad (FeedForward):

  • Eingang (shortcut): Speichern Sie den aktualisierten Eingang für die nächste Residualverbindung.

  • Layer Norm (norm2): Normalisieren Sie den Eingang.

  • FeedForward-Netzwerk (ff): Wenden Sie die Feedforward-Transformation an.

  • Dropout (drop_shortcut): Wenden Sie Dropout an.

  • Add Residual (x + shortcut): Kombinieren Sie mit dem Eingang vom ersten Residualpfad.

Der Transformer-Block gruppiert alle Netzwerke zusammen und wendet einige Normalisierungen und Dropouts an, um die Trainingsstabilität und -ergebnisse zu verbessern. Beachten Sie, dass Dropouts nach der Verwendung jedes Netzwerks durchgeführt werden, während die Normalisierung vorher angewendet wird.

Darüber hinaus verwendet er auch Abkürzungen, die darin bestehen, den Ausgang eines Netzwerks mit seinem Eingang zu addieren. Dies hilft, das Problem des verschwindenden Gradienten zu verhindern, indem sichergestellt wird, dass die anfänglichen Schichten "genauso viel" beitragen wie die letzten.

GPTModel

Größen wurden als Kommentare hinzugefügt, um die Größen der Matrizen besser zu verstehen:

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

Zweck und Funktionalität

  • Einbettungsschichten:

  • Token-Einbettungen (tok_emb): Wandelt Token-Indizes in Einbettungen um. Zur Erinnerung, dies sind die Gewichte, die jeder Dimension jedes Tokens im Vokabular zugewiesen werden.

  • Positions-Einbettungen (pos_emb): Fügt den Einbettungen Positionsinformationen hinzu, um die Reihenfolge der Tokens zu erfassen. Zur Erinnerung, dies sind die Gewichte, die Tokens entsprechend ihrer Position im Text zugewiesen werden.

  • Dropout (drop_emb): Wird auf Einbettungen zur Regularisierung angewendet.

  • Transformer-Blöcke (trf_blocks): Stapel von n_layers Transformer-Blöcken zur Verarbeitung von Einbettungen.

  • Endgültige Normalisierung (final_norm): Schichtnormalisierung vor der Ausgabeschicht.

  • Ausgabeschicht (out_head): Projiziert die endgültigen versteckten Zustände auf die Vokabulargröße, um Logits für die Vorhersage zu erzeugen.

Das Ziel dieser Klasse ist es, alle anderen genannten Netzwerke zu nutzen, um das nächste Token in einer Sequenz vorherzusagen, was grundlegend für Aufgaben wie die Textgenerierung ist.

Beachten Sie, wie es so viele Transformer-Blöcke wie angegeben verwenden wird und dass jeder Transformer-Block ein Multi-Head-Attention-Netz, ein Feed-Forward-Netz und mehrere Normalisierungen verwendet. Wenn also 12 Transformer-Blöcke verwendet werden, multiplizieren Sie dies mit 12.

Darüber hinaus wird eine Normalisierungsschicht vor der Ausgabe hinzugefügt und eine endgültige lineare Schicht wird am Ende angewendet, um die Ergebnisse mit den richtigen Dimensionen zu erhalten. Beachten Sie, dass jeder endgültige Vektor die Größe des verwendeten Vokabulars hat. Dies liegt daran, dass versucht wird, eine Wahrscheinlichkeit pro möglichem Token im Vokabular zu erhalten.

Anzahl der zu trainierenden Parameter

Nachdem die GPT-Struktur definiert ist, ist es möglich, die Anzahl der zu trainierenden Parameter zu ermitteln:

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

Schritt-für-Schritt-Berechnung

1. Einbettungsschichten: Token-Einbettung & Positions-Einbettung

  • Schicht: nn.Embedding(vocab_size, emb_dim)

  • Parameter: vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376
  • Schicht: nn.Embedding(context_length, emb_dim)

  • Parameter: context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

Gesamtanzahl der Einbettungsparameter

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

2. Transformer-Blöcke

Es gibt 12 Transformer-Blöcke, daher berechnen wir die Parameter für einen Block und multiplizieren dann mit 12.

Parameter pro Transformer-Block

a. Multi-Head Attention

  • Komponenten:

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

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

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

  • Ausgabeprojektion (out_proj): nn.Linear(emb_dim, emb_dim)

  • Berechnungen:

  • Jede von W_query, W_key, W_value:

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

Da es drei solcher Schichten gibt:

total_qkv_params = 3 * qkv_params = 3 * 589,824 = 1,769,472
  • Ausgabeprojektion (out_proj):

out_proj_params = (emb_dim * emb_dim) + emb_dim = (768 * 768) + 768 = 589,824 + 768 = 590,592
  • Gesamtzahl der Multi-Head Attention-Parameter:

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

b. FeedForward-Netzwerk

  • Komponenten:

  • Erste Linearschicht: nn.Linear(emb_dim, 4 * emb_dim)

  • Zweite Linearschicht: nn.Linear(4 * emb_dim, emb_dim)

  • Berechnungen:

  • Erste Linearschicht:

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
  • Zweite Linearschicht:

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
  • Gesamtzahl der FeedForward-Parameter:

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

c. Layer-Normalisierungen

  • Komponenten:

  • Zwei LayerNorm-Instanzen pro Block.

  • Jede LayerNorm hat 2 * emb_dim Parameter (Skalierung und Verschiebung).

  • Berechnungen:

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

d. Gesamtparameter pro Transformer-Block

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

Gesamtparameter für alle Transformatorblöcke

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

3. Endschichten

a. Endschicht Normalisierung

  • Parameter: 2 * emb_dim (Skalierung und Verschiebung)

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. Ausgabewurfebene (out_head)

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

  • Parameter: emb_dim * vocab_size

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

4. Zusammenfassung aller Parameter

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

Text generieren

Ein Modell, das das nächste Token wie das vorherige vorhersagt, benötigt lediglich die letzten Token-Werte aus der Ausgabe (da dies die Werte des vorhergesagten Tokens sein werden), was ein Wert pro Eintrag im Vokabular sein wird, und dann die softmax-Funktion verwenden, um die Dimensionen in Wahrscheinlichkeiten zu normalisieren, die sich zu 1 summieren, und dann den Index des größten Eintrags zu erhalten, der der Index des Wortes im Vokabular sein wird.

Code von 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]))

Referenzen

Last updated