5. LLM Architecture

LLM Architecture

Мета цього п'ятого етапу дуже проста: Розробити архітектуру повного LLM. З'єднайте все разом, застосуйте всі шари та створіть усі функції для генерації тексту або перетворення тексту в ID та назад.

Ця архітектура буде використовуватися як для навчання, так і для прогнозування тексту після його навчання.

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

A high level representation can be observed in:

Input (Tokenized Text): Процес починається з токенізованого тексту, який перетворюється на числові представлення.
Token Embedding and Positional Embedding Layer: Токенізований текст проходить через шар векторизації токенів та шар позиційної векторизації, який захоплює позицію токенів у послідовності, що є критично важливим для розуміння порядку слів.
Transformer Blocks: Модель містить 12 блоків трансформера, кожен з яких має кілька шарів. Ці блоки повторюють наступну послідовність:

Masked Multi-Head Attention: Дозволяє моделі зосереджуватися на різних частинах вхідного тексту одночасно.
Layer Normalization: Крок нормалізації для стабілізації та покращення навчання.
Feed Forward Layer: Відповідає за обробку інформації з шару уваги та прогнозування наступного токена.
Dropout Layers: Ці шари запобігають перенавчанню, випадковим чином відключаючи одиниці під час навчання.

Final Output Layer: Модель виводить тензор розміром 4x50,257, де 50,257 представляє розмір словника. Кожен рядок у цьому тензорі відповідає вектору, який модель використовує для прогнозування наступного слова в послідовності.
Goal: Мета полягає в тому, щоб взяти ці векторизації та перетворити їх назад у текст. Конкретно, останній рядок виходу використовується для генерації наступного слова, представленого як "вперед" у цій діаграмі.

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)

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

Мета та Функціональність

GELU (Гаусова Помилка Лінійний Одиниця): Активаційна функція, яка вводить нелінійність у модель.
Плавна Активація: На відміну від ReLU, яка обнуляє негативні вхідні дані, GELU плавно відображає вхідні дані на виходи, дозволяючи невеликі, ненульові значення для негативних вхідних даних.
Математичне Визначення:

Мета використання цієї функції після лінійних шарів всередині шару FeedForward полягає в тому, щоб змінити лінійні дані на нелінійні, щоб дозволити моделі вивчати складні, нелінійні зв'язки.

FeedForward Нейронна Мережа

Форми були додані як коментарі для кращого розуміння форм матриць:

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

Мета та Функціональність

Позиційна FeedForward мережа: Застосовує двошарову повністю з'єднану мережу до кожної позиції окремо та ідентично.
Деталі Шарів:
Перший Лінійний Шар: Розширює розмірність з emb_dim до 4 * emb_dim.
Активація GELU: Застосовує нелінійність.
Другий Лінійний Шар: Зменшує розмірність назад до emb_dim.

Як ви можете бачити, мережа Feed Forward використовує 3 шари. Перший - це лінійний шар, який помножить розміри на 4, використовуючи лінійні ваги (параметри для навчання всередині моделі). Потім функція GELU використовується у всіх цих вимірах, щоб застосувати нелінійні варіації для захоплення багатших представлень, і нарешті, ще один лінійний шар використовується для повернення до початкового розміру вимірів.

Механізм Багатоголової Уваги

Це вже було пояснено в попередньому розділі.

Мета та Функціональність

Багатоголова Самоувага: Дозволяє моделі зосереджуватися на різних позиціях у вхідній послідовності під час кодування токена.
Ключові Компоненти:
Запити, Ключі, Значення: Лінійні проекції вхідних даних, які використовуються для обчислення оцінок уваги.
Голови: Кілька механізмів уваги, що працюють паралельно (num_heads), кожен з зменшеною розмірністю (head_dim).
Оцінки Уваги: Обчислюються як скалярний добуток запитів і ключів, масштабовані та замасковані.
Маскування: Застосовується каузальна маска, щоб запобігти моделі звертатися до майбутніх токенів (важливо для авторегресивних моделей, таких як GPT).
Ваги Уваги: Softmax замаскованих і масштабованих оцінок уваги.
Контекстний Вектор: Вагова сума значень відповідно до ваг уваги.
Вихідна Проекція: Лінійний шар для об'єднання виходів усіх голів.

Мета цієї мережі - знайти відносини між токенами в одному контексті. Більше того, токени діляться на різні голови, щоб запобігти перенавчанню, хоча фінальні відносини, знайдені для кожної голови, об'єднуються в кінці цієї мережі.

Крім того, під час навчання застосовується каузальна маска, щоб пізні токени не враховувалися при пошуку специфічних відносин до токена, і також застосовується dropout для запобігання перенавчанню.

Нормалізація Шарів

# 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

Мета та Функціональність

Layer Normalization: Техніка, що використовується для нормалізації вхідних даних по ознаках (вимірах вбудовування) для кожного окремого прикладу в партії.
Компоненти:
eps: Маленька константа (1e-5), додана до дисперсії, щоб запобігти діленню на нуль під час нормалізації.
scale та shift: Навчальні параметри (nn.Parameter), які дозволяють моделі масштабувати та зміщувати нормалізований вихід. Вони ініціалізуються одиницями та нулями відповідно.
Процес Нормалізації:
Обчислення Середнього (mean): Обчислює середнє значення вхідного x по виміру вбудовування (dim=-1), зберігаючи вимір для трансляції (keepdim=True).
Обчислення Дисперсії (var): Обчислює дисперсію x по виміру вбудовування, також зберігаючи вимір. Параметр unbiased=False забезпечує, що дисперсія обчислюється за допомогою упередженого оцінювача (ділення на N замість N-1), що є доречним при нормалізації по ознакам, а не зразкам.
Нормалізація (norm_x): Віднімає середнє від x і ділить на квадратний корінь з дисперсії плюс eps.
Масштабування та Зміщення: Застосовує навчальні параметри scale та shift до нормалізованого виходу.

Мета полягає в забезпеченні середнього значення 0 з дисперсією 1 по всіх вимірах одного й того ж токена. Мета цього - стабілізувати навчання глибоких нейронних мереж шляхом зменшення внутрішнього зсуву коваріат, що відноситься до зміни розподілу активацій мережі через оновлення параметрів під час навчання.

Transformer Block

Форми були додані як коментарі для кращого розуміння форм матриць:

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

Мета та Функціональність

Складення Шарів: Поєднує багатоголову увагу, мережу прямого зв'язку, нормалізацію шару та залишкові з'єднання.
Нормалізація Шару: Застосовується перед шарами уваги та прямого зв'язку для стабільного навчання.
Залишкові З'єднання (Швидкі З'єднання): Додають вхід шару до його виходу для покращення потоку градієнтів та можливості навчання глибоких мереж.
Випадкове Випадання: Застосовується після шарів уваги та прямого зв'язку для регуляризації.

Покрокова Функціональність

Перший Залишковий Шлях (Самоувага):

Вхід (shortcut): Зберегти оригінальний вхід для залишкового з'єднання.
Нормалізація Шару (norm1): Нормалізувати вхід.
Багатоголова Увага (att): Застосувати самоувагу.
Випадкове Випадання (drop_shortcut): Застосувати випадкове випадання для регуляризації.
Додати Залишковий (x + shortcut): Об'єднати з оригінальним входом.

Другий Залишковий Шлях (Прямий Зв'язок):

Вхід (shortcut): Зберегти оновлений вхід для наступного залишкового з'єднання.
Нормалізація Шару (norm2): Нормалізувати вхід.
Мережа Прямого Зв'язку (ff): Застосувати перетворення прямого зв'язку.
Випадкове Випадання (drop_shortcut): Застосувати випадкове випадання.
Додати Залишковий (x + shortcut): Об'єднати з входом з першого залишкового шляху.

Блок трансформера об'єднує всі мережі разом і застосовує деяку нормалізацію та випадкове випадання для покращення стабільності навчання та результатів. Зверніть увагу, що випадкове випадання застосовується після використання кожної мережі, тоді як нормалізація застосовується перед.

Крім того, він також використовує швидкі з'єднання, які полягають у додаванні виходу мережі до її входу. Це допомагає запобігти проблемі зникнення градієнта, забезпечуючи, щоб початкові шари вносили "стільки ж", скільки останні.

GPTModel

Форми були додані як коментарі для кращого розуміння форм матриць:

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

Мета та Функціональність

Embedding Layers:
Token Embeddings (tok_emb): Перетворює індекси токенів на embeddings. Нагадаємо, це ваги, які надаються кожному виміру кожного токена в словнику.
Positional Embeddings (pos_emb): Додає позиційну інформацію до embeddings, щоб зафіксувати порядок токенів. Нагадаємо, це ваги, які надаються токену відповідно до його позиції в тексті.
Dropout (drop_emb): Застосовується до embeddings для регуляризації.
Transformer Blocks (trf_blocks): Стек з n_layers трансформерних блоків для обробки embeddings.
Final Normalization (final_norm): Нормалізація шару перед вихідним шаром.
Output Layer (out_head): Проектує фінальні приховані стани на розмір словника для отримання логітів для прогнозування.

Мета цього класу полягає в тому, щоб використовувати всі інші згадані мережі для прогнозування наступного токена в послідовності, що є основоположним для завдань, таких як генерація тексту.

Зверніть увагу, як він використовуватиме стільки трансформерних блоків, скільки вказано, і що кожен трансформерний блок використовує одну мережу з багатоголовим увагою, одну мережу прямого проходження та кілька нормалізацій. Тож, якщо використовується 12 трансформерних блоків, помножте це на 12.

Крім того, шар нормалізації додається перед виходом, а фінальний лінійний шар застосовується в кінці, щоб отримати результати з правильними розмірами. Зверніть увагу, що кожен фінальний вектор має розмір використаного словника. Це тому, що він намагається отримати ймовірність для кожного можливого токена в словнику.

Кількість параметрів для навчання

Маючи визначену структуру GPT, можна дізнатися кількість параметрів для навчання:

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

Покроковий Розрахунок

1. Вкладки Векторів: Векторне Вкладення та Позиційне Вкладення

Шар: nn.Embedding(vocab_size, emb_dim)
Параметри: vocab_size * emb_dim

token_embedding_params = 50257 * 768 = 38,597,376

Шар: nn.Embedding(context_length, emb_dim)
Параметри: context_length * emb_dim

position_embedding_params = 1024 * 768 = 786,432

Загальна кількість параметрів вбудовування

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

2. Transformer Blocks

Є 12 блоків трансформера, тому ми розрахуємо параметри для одного блоку, а потім помножимо на 12.

Параметри на один блок трансформера

a. Багатоголове увага

Компоненти:
Лінійний шар запиту (W_query): nn.Linear(emb_dim, emb_dim, bias=False)
Лінійний шар ключа (W_key): nn.Linear(emb_dim, emb_dim, bias=False)
Лінійний шар значення (W_value): nn.Linear(emb_dim, emb_dim, bias=False)
Вихідна проекція (out_proj): nn.Linear(emb_dim, emb_dim)
Розрахунки:
Кожен з W_query, W_key, W_value:

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

Оскільки є три такі шари:

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

Вихідна проекція (out_proj):

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

Загальна кількість параметрів багатоголової уваги:

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

b. Мережа зворотного зв'язку

Компоненти:
Перший лінійний шар: nn.Linear(emb_dim, 4 * emb_dim)
Другий лінійний шар: nn.Linear(4 * emb_dim, emb_dim)
Розрахунки:
Перший лінійний шар:

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

Другий лінійний шар:

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

Загальна кількість параметрів зворотного зв'язку:

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

c. Нормалізації шару

Компоненти:
Два екземпляри LayerNorm на блок.
Кожен LayerNorm має 2 * emb_dim параметрів (масштаб і зсув).
Розрахунки:

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

d. Загальна кількість параметрів на один блок трансформера

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

Загальна кількість параметрів для всіх блоків трансформера

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

3. Остаточні шари

a. Нормалізація остаточного шару

Параметри: 2 * emb_dim (масштаб і зсув)

pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536

b. Вихідний проекційний шар (out_head)

Шар: nn.Linear(emb_dim, vocab_size, bias=False)
Параметри: emb_dim * vocab_size

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

4. Підсумування всіх параметрів

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

Генерація тексту

Маючи модель, яка передбачає наступний токен, як і попередній, потрібно просто взяти останні значення токенів з виходу (оскільки вони будуть значеннями передбаченого токена), які будуть значенням на запис у словнику, а потім використати функцію softmax, щоб нормалізувати виміри в ймовірності, що в сумі дорівнюють 1, а потім отримати індекс найбільшого запису, який буде індексом слова в словнику.

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

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

Last updated 1 day ago