mirror of
https://github.com/HackTricks-wiki/hacktricks.git
synced 2025-10-10 18:36:50 +00:00
672 lines
35 KiB
Markdown
672 lines
35 KiB
Markdown
# 5. LLM Architecture
|
||
|
||
{{#include ../../banners/hacktricks-training.md}}
|
||
|
||
## LLM Architecture
|
||
|
||
> [!TIP]
|
||
> Мета цього п'ятого етапу дуже проста: **Розробити архітектуру повного LLM**. З'єднайте все разом, застосуйте всі шари та створіть усі функції для генерації тексту або перетворення тексту в ID та назад.
|
||
>
|
||
> Ця архітектура буде використовуватися як для навчання, так і для прогнозування тексту після його навчання.
|
||
|
||
LLM architecture example from [https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb):
|
||
|
||
A high level representation can be observed in:
|
||
|
||
<figure><img src="../../images/image (3) (1) (1) (1).png" alt="" width="563"><figcaption><p><a href="https://camo.githubusercontent.com/6c8c392f72d5b9e86c94aeb9470beab435b888d24135926f1746eb88e0cc18fb/68747470733a2f2f73656261737469616e72617363686b612e636f6d2f696d616765732f4c4c4d732d66726f6d2d736372617463682d696d616765732f636830345f636f6d707265737365642f31332e776562703f31">https://camo.githubusercontent.com/6c8c392f72d5b9e86c94aeb9470beab435b888d24135926f1746eb88e0cc18fb/68747470733a2f2f73656261737469616e72617363686b612e636f6d2f696d616765732f4c4c4d732d66726f6d2d736372617463682d696d616765732f636830345f636f6d707265737365642f31332e776562703f31</a></p></figcaption></figure>
|
||
|
||
1. **Input (Tokenized Text)**: Процес починається з токенізованого тексту, який перетворюється на числові представлення.
|
||
2. **Token Embedding and Positional Embedding Layer**: Токенізований текст проходить через **шар вбудовування токенів** та **шар позиційного вбудовування**, який захоплює позицію токенів у послідовності, що є критично важливим для розуміння порядку слів.
|
||
3. **Transformer Blocks**: Модель містить **12 блоків трансформера**, кожен з яких має кілька шарів. Ці блоки повторюють наступну послідовність:
|
||
- **Masked Multi-Head Attention**: Дозволяє моделі зосереджуватися на різних частинах вхідного тексту одночасно.
|
||
- **Layer Normalization**: Крок нормалізації для стабілізації та покращення навчання.
|
||
- **Feed Forward Layer**: Відповідає за обробку інформації з шару уваги та прогнозування наступного токена.
|
||
- **Dropout Layers**: Ці шари запобігають перенавчанню, випадковим чином відключаючи одиниці під час навчання.
|
||
4. **Final Output Layer**: Модель виводить **тензор розміром 4x50,257**, де **50,257** представляє розмір словника. Кожен рядок у цьому тензорі відповідає вектору, який модель використовує для прогнозування наступного слова в послідовності.
|
||
5. **Goal**: Мета полягає в тому, щоб взяти ці вбудовування та перетворити їх назад у текст. Конкретно, останній рядок виходу використовується для генерації наступного слова, представленого як "вперед" у цій діаграмі.
|
||
|
||
### Code representation
|
||
```python
|
||
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**
|
||
```python
|
||
# 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 плавно відображає вхідні дані на виходи, дозволяючи невеликі, ненульові значення для негативних вхідних даних.
|
||
- **Математичне Визначення:**
|
||
|
||
<figure><img src="../../images/image (2) (1) (1) (1).png" alt=""><figcaption></figcaption></figure>
|
||
|
||
> [!TIP]
|
||
> Мета використання цієї функції після лінійних шарів всередині шару FeedForward полягає в тому, щоб змінити лінійні дані на нелінійні, щоб дозволити моделі вивчати складні, нелінійні зв'язки.
|
||
|
||
### **FeedForward Нейронна Мережа**
|
||
|
||
_Форми були додані як коментарі для кращого розуміння форм матриць:_
|
||
```python
|
||
# 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`.
|
||
|
||
> [!TIP]
|
||
> Як ви можете бачити, мережа Feed Forward використовує 3 шари. Перший - це лінійний шар, який множить розміри на 4, використовуючи лінійні ваги (параметри для навчання всередині моделі). Потім функція GELU використовується у всіх цих вимірах, щоб застосувати нелінійні варіації для захоплення багатших представлень, і нарешті ще один лінійний шар використовується для повернення до початкового розміру вимірів.
|
||
|
||
### **Механізм Багатоголової Уваги**
|
||
|
||
Це вже було пояснено в попередньому розділі.
|
||
|
||
#### **Мета та Функціональність**
|
||
|
||
- **Багатоголова Самоувага:** Дозволяє моделі зосереджуватися на різних позиціях у вхідній послідовності під час кодування токена.
|
||
- **Ключові Компоненти:**
|
||
- **Запити, Ключі, Значення:** Лінійні проекції вхідних даних, які використовуються для обчислення оцінок уваги.
|
||
- **Голови:** Кілька механізмів уваги, що працюють паралельно (`num_heads`), кожен з зменшеною розмірністю (`head_dim`).
|
||
- **Оцінки Уваги:** Обчислюються як скалярний добуток запитів і ключів, масштабовані та замасковані.
|
||
- **Маскування:** Застосовується каузальна маска, щоб запобігти моделі звертатися до майбутніх токенів (важливо для авторегресивних моделей, таких як GPT).
|
||
- **Ваги Уваги:** Softmax замаскованих і масштабованих оцінок уваги.
|
||
- **Контекстний Вектор:** Вагова сума значень відповідно до ваг уваги.
|
||
- **Вихідна Проекція:** Лінійний шар для об'єднання виходів усіх голів.
|
||
|
||
> [!TIP]
|
||
> Мета цієї мережі - знайти відносини між токенами в одному контексті. Більше того, токени діляться на різні голови, щоб запобігти перенавчанню, хоча фінальні відносини, знайдені для кожної голови, об'єднуються в кінці цієї мережі.
|
||
>
|
||
> Більше того, під час навчання застосовується **каузальна маска**, щоб пізні токени не враховувалися при пошуку специфічних відносин до токена, і також застосовується **dropout** для **запобігання перенавчанню**.
|
||
|
||
### **Нормалізація** Шарів
|
||
```python
|
||
# 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` до нормалізованого виходу.
|
||
|
||
> [!TIP]
|
||
> Мета полягає в тому, щоб забезпечити середнє значення 0 з дисперсією 1 по всіх вимірах одного й того ж токена. Мета цього - **стабілізувати навчання глибоких нейронних мереж** шляхом зменшення внутрішнього зміщення коваріат, що відноситься до зміни розподілу активацій мережі через оновлення параметрів під час навчання.
|
||
|
||
### **Блок Трансформера**
|
||
|
||
_Форми були додані як коментарі для кращого розуміння форм матриць:_
|
||
```python
|
||
# 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)
|
||
|
||
```
|
||
#### **Мета та Функціональність**
|
||
|
||
- **Складові Шарів:** Поєднує багатоголову увагу, мережу прямого проходження, нормалізацію шарів та залишкові з'єднання.
|
||
- **Нормалізація Шарів:** Застосовується перед шарами уваги та прямого проходження для стабільного навчання.
|
||
- **Залишкові З'єднання (Швидкі З'єднання):** Додають вхід шару до його виходу для покращення потоку градієнтів та можливості навчання глибоких мереж.
|
||
- **Випадкове Вимкнення:** Застосовується після шарів уваги та прямого проходження для регуляризації.
|
||
|
||
#### **Покрокова Функціональність**
|
||
|
||
1. **Перший Залишковий Шлях (Само-Увага):**
|
||
- **Вхід (`shortcut`):** Зберегти оригінальний вхід для залишкового з'єднання.
|
||
- **Нормалізація Шару (`norm1`):** Нормалізувати вхід.
|
||
- **Багатоголова Увага (`att`):** Застосувати само-увагу.
|
||
- **Випадкове Вимкнення (`drop_shortcut`):** Застосувати випадкове вимкнення для регуляризації.
|
||
- **Додати Залишок (`x + shortcut`):** Поєднати з оригінальним входом.
|
||
2. **Другий Залишковий Шлях (Пряме Проходження):**
|
||
- **Вхід (`shortcut`):** Зберегти оновлений вхід для наступного залишкового з'єднання.
|
||
- **Нормалізація Шару (`norm2`):** Нормалізувати вхід.
|
||
- **Мережа Прямого Проходження (`ff`):** Застосувати трансформацію прямого проходження.
|
||
- **Випадкове Вимкнення (`drop_shortcut`):** Застосувати випадкове вимкнення.
|
||
- **Додати Залишок (`x + shortcut`):** Поєднати з входом з першого залишкового шляху.
|
||
|
||
> [!TIP]
|
||
> Блок трансформера об'єднує всі мережі разом і застосовує деяку **нормалізацію** та **випадкові вимкнення** для покращення стабільності навчання та результатів.\
|
||
> Зверніть увагу, як випадкові вимкнення виконуються після використання кожної мережі, тоді як нормалізація застосовується перед.
|
||
>
|
||
> Більше того, він також використовує швидкі з'єднання, які полягають у **додаванні виходу мережі до її входу**. Це допомагає запобігти проблемі зникнення градієнта, забезпечуючи, щоб початкові шари вносили "стільки ж", скільки останні.
|
||
|
||
### **GPTModel**
|
||
|
||
_Форми були додані як коментарі для кращого розуміння форм матриць:_
|
||
```python
|
||
# 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`):** Перетворює індекси токенів в ембедінги. Нагадаємо, це ваги, які надаються кожному виміру кожного токена в словнику.
|
||
- **Positional Embeddings (`pos_emb`):** Додає позиційну інформацію до ембедінгів, щоб зафіксувати порядок токенів. Нагадаємо, це ваги, які надаються токену відповідно до його позиції в тексті.
|
||
- **Dropout (`drop_emb`):** Застосовується до ембедінгів для регуляризації.
|
||
- **Transformer Blocks (`trf_blocks`):** Стек з `n_layers` трансформерних блоків для обробки ембедінгів.
|
||
- **Final Normalization (`final_norm`):** Нормалізація шару перед вихідним шаром.
|
||
- **Output Layer (`out_head`):** Проектує фінальні приховані стани на розмір словника для отримання логітів для прогнозування.
|
||
|
||
> [!TIP]
|
||
> Мета цього класу полягає в тому, щоб використовувати всі інші згадані мережі для **прогнозування наступного токена в послідовності**, що є основоположним для завдань, таких як генерація тексту.
|
||
>
|
||
> Зверніть увагу, як він **використовуватиме стільки трансформерних блоків, скільки вказано**, і що кожен трансформерний блок використовує одну мережу з багатоголовим увагою, одну мережу прямого проходження та кілька нормалізацій. Тож, якщо використовується 12 трансформерних блоків, помножте це на 12.
|
||
>
|
||
> Більше того, **шар нормалізації** додається **перед** **виходом**, і в кінці застосовується фінальний лінійний шар, щоб отримати результати з правильними розмірами. Зверніть увагу, що кожен фінальний вектор має розмір використаного словника. Це тому, що він намагається отримати ймовірність для кожного можливого токена в словнику.
|
||
|
||
## Кількість параметрів для навчання
|
||
|
||
Маючи визначену структуру GPT, можна дізнатися кількість параметрів для навчання:
|
||
```python
|
||
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`
|
||
```python
|
||
token_embedding_params = 50257 * 768 = 38,597,376
|
||
```
|
||
- **Шар:** `nn.Embedding(context_length, emb_dim)`
|
||
- **Параметри:** `context_length * emb_dim`
|
||
```python
|
||
position_embedding_params = 1024 * 768 = 786,432
|
||
```
|
||
**Загальна кількість параметрів вбудовування**
|
||
```python
|
||
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`:**
|
||
|
||
```python
|
||
qkv_params = emb_dim * emb_dim = 768 * 768 = 589,824
|
||
```
|
||
|
||
Оскільки є три такі шари:
|
||
|
||
```python
|
||
total_qkv_params = 3 * qkv_params = 3 * 589,824 = 1,769,472
|
||
```
|
||
|
||
- **Вихідна проекція (`out_proj`):**
|
||
|
||
```python
|
||
out_proj_params = (emb_dim * emb_dim) + emb_dim = (768 * 768) + 768 = 589,824 + 768 = 590,592
|
||
```
|
||
|
||
- **Загальна кількість параметрів багатоголової уваги:**
|
||
|
||
```python
|
||
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)`
|
||
- **Розрахунки:**
|
||
|
||
- **Перший лінійний шар:**
|
||
|
||
```python
|
||
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
|
||
```
|
||
|
||
- **Другий лінійний шар:**
|
||
|
||
```python
|
||
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
|
||
```
|
||
|
||
- **Загальна кількість параметрів мережі зворотного зв'язку:**
|
||
|
||
```python
|
||
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` параметрів (масштаб і зсув).
|
||
- **Розрахунки:**
|
||
|
||
```python
|
||
layer_norm_params_per_block = 2 * (2 * emb_dim) = 2 * 768 * 2 = 3,072
|
||
```
|
||
|
||
**d. Загальна кількість параметрів на один блок трансформера**
|
||
```python
|
||
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
|
||
```
|
||
**Загальна кількість параметрів для всіх блоків трансформера**
|
||
```python
|
||
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` (масштаб і зсув)
|
||
```python
|
||
pythonCopy codefinal_layer_norm_params = 2 * 768 = 1,536
|
||
```
|
||
**b. Вихідний проекційний шар (`out_head`)**
|
||
|
||
- **Шар:** `nn.Linear(emb_dim, vocab_size, bias=False)`
|
||
- **Параметри:** `emb_dim * vocab_size`
|
||
```python
|
||
pythonCopy codeoutput_projection_params = 768 * 50257 = 38,597,376
|
||
```
|
||
#### **4. Підсумування всіх параметрів**
|
||
```python
|
||
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](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch04/01_main-chapter-code/ch04.ipynb):
|
||
```python
|
||
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]))
|
||
```
|
||
## Посилання
|
||
|
||
- [https://www.manning.com/books/build-a-large-language-model-from-scratch](https://www.manning.com/books/build-a-large-language-model-from-scratch)
|
||
|
||
|
||
{{#include ../../banners/hacktricks-training.md}}
|