From 73ee3e16ec85149807ba6f47ae0a8c382b26b86c Mon Sep 17 00:00:00 2001 From: Sergey Penkovsky Date: Mon, 6 Oct 2025 20:34:02 +0300 Subject: [PATCH] docs: update and enhance documentation for all core components and models - Added detailed documentation for GPT, GPT2 and LLaMA models - Enhanced docstrings in base_model.py, rope.py, rms_norm.py, swi_glu.py - Updated README with architectural differences and usage examples - Added scientific references and mathematical foundations - Improved type hints and parameter descriptions --- README.md | 20 ++- llm/README.md | 251 +++++++++++++++++++++++++++++ llm/src/llm/core/base_model.py | 54 ++++++- llm/src/llm/core/cached_decoder.py | 48 ++++-- llm/src/llm/core/rms_norm.py | 66 +++++++- llm/src/llm/core/rope.py | 86 ++++++++-- llm/src/llm/core/swi_glu.py | 85 +++++++++- llm/src/llm/models/gpt/gpt.py | 46 ++++-- 8 files changed, 605 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index 6e150fa..7a63cc3 100644 --- a/README.md +++ b/README.md @@ -24,16 +24,26 @@ llm-arch-research/ │ └── src/llm/ │ ├── core/ # базовые компоненты │ │ ├── base_model.py +│ │ ├── cached_decoder.py # Декодер с кэшированием │ │ ├── decoder.py │ │ ├── multi_head_attention.py │ │ ├── head_attention.py │ │ ├── feed_forward.py │ │ ├── token_embeddings.py -│ │ └── positional_embeddings.py -│ ├── models/gpt/ # GPT и GPT-2 реализация -│ │ ├── gpt.py -│ │ ├── gpt2.py -│ │ └── __init__.py +│ │ ├── positional_embeddings.py +│ │ ├── rope.py # Rotary Positional Embeddings +│ │ ├── rms_norm.py # RMS Normalization +│ │ ├── swi_glu.py # SwiGLU активация +│ │ ├── silu.py # SiLU активация +│ │ └── gelu.py # GELU активация +│ ├── models/ # Реализации моделей +│ │ ├── gpt/ # GPT и GPT-2 архитектуры +│ │ │ ├── gpt.py +│ │ │ ├── gpt2.py +│ │ │ └── __init__.py +│ │ └── llama/ # LLaMA архитектура +│ │ ├── llama.py +│ │ └── __init__.py │ ├── training/ # утилиты обучения │ │ ├── dataset.py │ │ ├── trainer.py diff --git a/llm/README.md b/llm/README.md index e69de29..ab35c48 100644 --- a/llm/README.md +++ b/llm/README.md @@ -0,0 +1,251 @@ +# LLM Framework - Фреймворк для языковых моделей + +Модульная библиотека для создания, обучения и использования больших языковых моделей (LLM) с поддержкой различных архитектур (GPT, LLaMA и др.). + +## 🏗️ Архитектура + +Библиотека построена по модульному принципу с четким разделением ответственности: + +``` +llm/ +├── core/ # Базовые компоненты +│ ├── base_model.py # Абстрактный базовый класс моделей +│ ├── cached_decoder.py # Универсальный декодер с кэшированием +│ ├── decoder.py # Базовый декодер +│ ├── multi_head_attention.py # Многоголовое внимание +│ ├── head_attention.py # Одно-головое внимание +│ ├── feed_forward.py # Стандартный FFN слой +│ ├── token_embeddings.py # Векторные представления токенов +│ ├── positional_embeddings.py # Абсолютные позиционные эмбеддинги +│ ├── rope.py # Rotary Positional Embeddings (RoPE) +│ ├── rms_norm.py # RMS Normalization +│ ├── swi_glu.py # SwiGLU активация +│ ├── silu.py # SiLU активация +│ └── gelu.py # GELU активация +├── models/ # Конкретные реализации моделей +│ ├── gpt/ # GPT архитектуры +│ │ ├── gpt.py # Базовая GPT +│ │ ├── gpt2.py # GPT-2 реализация +│ │ └── __init__.py +│ └── llama/ # LLaMA архитектура +│ ├── llama.py # LLaMA реализация +│ └── __init__.py +├── tokenizers/ # Токенизаторы +│ ├── base_tokenizer.py # Базовый интерфейс +│ └── bpe_tokenizer.py # BPE токенизатор +└── training/ # Утилиты обучения + ├── dataset.py # Датасеты + ├── trainer.py # Тренировочный цикл + ├── optimizer.py # Оптимизаторы + └── scheduler.py # Планировщики обучения +``` + +## 🧩 Ключевые компоненты + +### BaseModel (`core/base_model.py`) +**Абстрактный базовый класс** для всех языковых моделей с единым интерфейсом. + +```python +class BaseModel(nn.Module, ABC): + @abstractmethod + def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """Прямой проход модели.""" + + @abstractmethod + def generate(self, input_ids: torch.Tensor, max_length: int = 50) -> torch.Tensor: + """Генерация текста.""" +``` + +### CachedDecoder (`core/cached_decoder.py`) +**Универсальный декодер** с поддержкой dependency injection и кэширования KV-памяти. + +```python +CachedDecoder( + feed_forward_layer=FeedForward(...), # или SwiGLU + norm_layer=nn.LayerNorm, # или RMSNorm + rope=RoPE(...), # опционально + # ... другие параметры +) +``` + +### RoPE (`core/rope.py`) +**Rotary Positional Embeddings** - ротационные позиционные эмбеддинги. + +**Математическая основа:** +``` +θ_i = base^(-2i/d) +q'_m = q_m * cos(mθ_i) + rotate(q_m) * sin(mθ_i) +``` + +### RMSNorm (`core/rms_norm.py`) +**Root Mean Square Normalization** - упрощенная нормализация без среднего. + +**Формула:** +``` +RMSNorm(x) = (x / RMS(x)) * w +где RMS(x) = sqrt(mean(x²) + eps) +``` + +### SwiGLU (`core/swi_glu.py`) +**Swish-Gated Linear Unit** - современная активация с gating mechanism. + +**Формула:** +``` +SwiGLU(x) = Swish(xW_g + b_g) ⊙ (xW_u + b_u) * W_d + b_d +``` + +## 🚀 Примеры использования + +### Создание классической GPT модели +```python +from llm.models.gpt import GPT + +config = { + "vocab_size": 50257, + "embed_dim": 768, + "num_heads": 12, + "num_layers": 12, + "max_position_embeddings": 1024, + "dropout": 0.1 +} + +model = GPT(config) +``` + +### Создание GPT2 модели +```python +from llm.models.gpt import GPT2 + +config = { + "vocab_size": 50257, + "embed_dim": 768, + "num_heads": 12, + "num_layers": 12, + "max_position_embeddings": 1024, + "dropout": 0.1 +} + +model = GPT2(config) +``` + +### Создание LLaMA модели +```python +from llm.models.llama import Llama +from llm.core.swi_glu import SwiGLU +from llm.core.rms_norm import RMSNorm + +config = { + "vocab_size": 32000, + "embed_dim": 4096, + "num_heads": 32, + "num_layers": 32, + "max_position_embeddings": 2048, + "dropout": 0.1 +} + +model = Llama(config) +``` + +### Генерация текста +```python +# Прямой проход +output = model(input_ids, attention_mask) + +# Генерация текста +generated = model.generate(input_ids, max_length=100) +``` + +## 📊 Входные и выходные данные + +### Входные данные: +- `input_ids`: `Tensor[int64]` формы `[batch_size, seq_len]` - индексы токенов +- `attention_mask`: `Tensor[bool]` формы `[batch_size, seq_len]` - маска внимания +- `cache`: `List[Tuple[Tensor, Tensor]]` - кэш ключей-значений для генерации + +### Выходные данные: +- `logits`: `Tensor[float32]` формы `[batch_size, seq_len, vocab_size]` - вероятности токенов +- `cache`: `List[Tuple[Tensor, Tensor]]` - обновленный кэш (при использовании) + +## 🏆 Поддерживаемые архитектуры + +### GPT (Original) Особенности +- ✅ Многоголовое внимание +- ✅ Layer Normalization (после внимания и FFN) +- ✅ GELU активация +- ✅ Learned positional embeddings +- ✅ Базовая архитектура трансформер-декодера + +### GPT-2 Особенности +- ✅ Улучшенная версия оригинальной GPT +- ✅ Layer Normalization (перед вниманием и FFN) +- ✅ GELU активация +- ✅ Learned positional embeddings +- ✅ Кэширование для эффективной генерации +- ✅ Оптимизированные веса инициализации + +### LLaMA Особенности +- ✅ Rotary Positional Embeddings (RoPE) +- ✅ RMS Normalization вместо LayerNorm +- ✅ SwiGLU активация вместо GELU +- ✅ Оптимизированная структура декодера +- ✅ Эффективное кэширование KV-памяти + +## 🧪 Тестирование + +Запуск всех тестов: +```bash +cd llm +python -m pytest tests/ -v +``` + +**Статус тестов:** ✅ 101 тест пройден + +## 📚 Научные концепции + +### Трансформерная архитектура +Основана на механизме **внимания**, позволяющем модели взвешивать важность разных частей входной последовательности. + +**Формула внимания:** +``` +Attention(Q, K, V) = softmax(Q·Kᵀ/√d_k)·V +``` + +### RoPE (Rotary Positional Embeddings) +Инновационный метод кодирования позиционной информации через **вращение векторов** в комплексном пространстве. + +**Преимущества:** +- Относительное позиционное кодирование +- Лучшая экстраполяция на длинные последовательности +- Сохранение нормы векторов + +### RMSNorm vs LayerNorm +**RMSNorm** устраняет вычитание среднего, что делает его более стабильным и эффективным при обучении больших моделей. + +### SwiGLU vs GELU +**SwiGLU** с gating mechanism показывает лучшую производительность благодаря способности выборочно передавать информацию. + +## 🔧 Настройка и расширение + +Библиотека разработана с учетом **расширяемости**. Для добавления новой архитектуры: + +1. **Наследоваться** от `BaseModel` +2. **Реализовать** обязательные методы `forward()` и `generate()` +3. **Использовать** модульные компоненты из `core/` +4. **Добавить** конфигурацию модели + +### Пример расширения: +```python +class NewModel(BaseModel): + def __init__(self, config): + super().__init__(config) + # Использование готовых компонентов + self.decoder = CachedDecoder(...) + + def forward(self, input_ids, attention_mask=None): + # Реализация прямого прохода + pass +``` + +## 📄 Лицензия + +Проект распространяется под MIT License. diff --git a/llm/src/llm/core/base_model.py b/llm/src/llm/core/base_model.py index da362c4..a1f1c45 100644 --- a/llm/src/llm/core/base_model.py +++ b/llm/src/llm/core/base_model.py @@ -1,20 +1,62 @@ # llm/core/base_model.py +""" +Базовый абстрактный класс для всех языковых моделей (LLM). + +Реализует общий интерфейс для прямого прохода и генерации текста. +Все конкретные модели должны наследоваться от этого класса. +""" + import torch.nn as nn from abc import ABC, abstractmethod +from typing import Optional, Tuple +import torch + class BaseModel(nn.Module, ABC): - """Базовый класс для всех LLM.""" + """ + Абстрактный базовый класс для больших языковых моделей. + + Args: + config (dict): Конфигурация модели с параметрами архитектуры + + Attributes: + config (dict): Конфигурационные параметры модели + """ - def __init__(self, config): + def __init__(self, config: dict): + """ + Инициализация базовой модели. + + Args: + config: Словарь с параметрами конфигурации модели + """ super().__init__() self.config = config @abstractmethod - def forward(self, input_ids, attention_mask=None): - """Прямой проход модели.""" + def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Прямой проход модели. + + Args: + input_ids: Тензор индексов токенов формы [batch_size, seq_len] + attention_mask: Опциональная маска внимания формы [batch_size, seq_len] + + Returns: + Тензор логитов формы [batch_size, seq_len, vocab_size] + """ pass @abstractmethod - def generate(self, input_ids, max_length=50): - """Генерация текста (greedy или sampling).""" + def generate(self, input_ids: torch.Tensor, max_length: int = 50) -> torch.Tensor: + """ + Генерация текста с использованием greedy decoding или sampling. + + Args: + input_ids: Начальные токены для генерации формы [batch_size, start_len] + max_length: Максимальная длина генерируемой последовательности + + Returns: + Тензор сгенерированных токенов формы [batch_size, generated_len] + """ pass diff --git a/llm/src/llm/core/cached_decoder.py b/llm/src/llm/core/cached_decoder.py index 8658683..59eae7b 100644 --- a/llm/src/llm/core/cached_decoder.py +++ b/llm/src/llm/core/cached_decoder.py @@ -8,8 +8,19 @@ from .rope import RoPE class CachedDecoder(nn.Module): """ - Универсальный декодер с поддержкой кэша для autoregressive использования (GPT, LLAMA и пр). - - Поддерживает использование past_key_values для быстрого генеративного инференса. + Универсальный декодерный блок с dependency injection для поддержки различных архитектур. + + Поддерживает кэширование ключей-значений для ускорения генерации текста. + + Args: + feed_forward_layer: Экземпляр слоя прямого распространения (SwiGLU, FeedForward и т.д.) + num_heads: Количество голов механизма внимания + emb_size: Размерность векторных представлений + head_size: Размерность каждой головы внимания + max_seq_len: Максимальная длина последовательности + norm_layer: Класс слоя нормализации (LayerNorm, RMSNorm и т.д.) + dropout: Вероятность dropout + rope: Экземпляр RoPE для позиционного кодирования (опционально) """ def __init__( self, @@ -18,11 +29,23 @@ class CachedDecoder(nn.Module): emb_size: int, head_size: int, max_seq_len: int, - dropout: float = 0.1, norm_layer: type = nn.LayerNorm, # Класс + dropout: float = 0.1, rope: RoPE = None, - activation: str = "gelu", ): + """ + Инициализация декодера с кэшированием. + + Args: + feed_forward_layer: Слой feed-forward (должен быть экземпляром, а не классом) + num_heads: Количество голов внимания + emb_size: Размерность эмбеддингов + head_size: Размерность каждой головы + max_seq_len: Максимальная длина последовательности + norm_layer: Класс нормализации (по умолчанию LayerNorm) + dropout: Вероятность dropout + rope: Rotary Positional Embeddings (опционально) + """ super().__init__() self._heads = MultiHeadAttention( num_heads=num_heads, @@ -44,11 +67,18 @@ class CachedDecoder(nn.Module): cache: list = None, ): """ - x: [batch, seq_len, emb_size] - mask: (optional) - use_cache: использовать ли кэширование KV-слоев (инкрементальный генератив, GPT-style) - cache: список кэшей для голов (или None) - Возвращает: (output, new_cache) если use_cache=True, иначе (output, None) + Прямой проход через декодерный блок. + + Args: + x: Входной тензор формы [batch_size, seq_len, emb_size] + mask: Маска внимания формы [batch_size, seq_len] (опционально) + use_cache: Флаг использования кэширования + cache: Список кэшированных пар (key, value) тензоров + + Returns: + Кортеж (output, new_cache) где: + - output: Выходной тензор формы [batch_size, seq_len, emb_size] + - new_cache: Обновленный кэш или None, если use_cache=False """ norm1_out = self._norm1(x) # Передаём все cache/use_cache дальше в attention diff --git a/llm/src/llm/core/rms_norm.py b/llm/src/llm/core/rms_norm.py index 37fafe8..305e7f0 100644 --- a/llm/src/llm/core/rms_norm.py +++ b/llm/src/llm/core/rms_norm.py @@ -1,13 +1,75 @@ +""" +RMSNorm (Root Mean Square Normalization) - нормализация по среднеквадратичному значению. + +Упрощенная версия LayerNorm без вычисления среднего значения. Широко используется +в современных архитектурах типа LLaMA благодаря лучшей стабильности и производительности. + +Научная статья: "Root Mean Square Layer Normalization" +https://arxiv.org/abs/1910.07467 + +Формула: +RMSNorm(x) = (x / RMS(x)) * w +где RMS(x) = sqrt(mean(x²) + eps) + +Преимущества: +- Меньше вычислений (нет вычитания среднего) +- Лучшая стабильность при обучении +- Сохранение масштаба сигнала +""" + import torch from torch import nn +from typing import Optional + class RMSNorm(nn.Module): + """ + Реализация RMS Normalization. + + Нормализует входные данные по последнему измерению используя среднеквадратичное + значение вместо среднего, как в стандартном LayerNorm. + + Args: + dim: Размерность измерения для нормализации + eps: Малое значение для численной стабильности + + Attributes: + _eps: Малое значение для предотвращения деления на ноль + _w: Обучаемый параметр масштабирования формы [dim] + """ + def __init__(self, dim: int, eps: float = 1e-6): + """ + Инициализация RMSNorm слоя. + + Args: + dim: Размерность нормализуемого измерения + eps: Малое значение для численной стабильности (по умолчанию 1e-6) + """ super().__init__() self._eps = eps self._w = nn.Parameter(torch.ones(dim)) - def forward(self, x: torch.Tensor): # [batch_size × seq_len × emb_size] + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Прямой проход через RMSNorm слой. + + Args: + x: Входной тензор формы [..., dim] + + Returns: + Нормализованный тензор той же формы, что и входной + + Формула: + output = w * (x / sqrt(mean(x²) + eps)) + """ + # Вычисление RMS (Root Mean Square) по последнему измерению rms = (x.pow(2).mean(-1, keepdim=True) + self._eps) ** 0.5 + + # Нормализация и масштабирование norm_x = x / rms - return self._w * norm_x \ No newline at end of file + return self._w * norm_x + + def extra_repr(self) -> str: + """Строковое представление для отладки.""" + return f'dim={self._w.shape[0]}, eps={self._eps}' \ No newline at end of file diff --git a/llm/src/llm/core/rope.py b/llm/src/llm/core/rope.py index 966b981..d07b348 100644 --- a/llm/src/llm/core/rope.py +++ b/llm/src/llm/core/rope.py @@ -1,43 +1,103 @@ +""" +Rotary Positional Embeddings (RoPE) - ротационные позиционные эмбеддинги. + +Реализация ротационного позиционного кодирования, которое кодирует позиционную +информацию через вращение векторов запросов и ключей в комплексном пространстве. + +Научная статья: "RoFormer: Enhanced Transformer with Rotary Position Embedding" +https://arxiv.org/abs/2104.09864 + +Математическая основа: +Для позиции m и измерения i: +θ_i = base^(-2i/d) +q'_m = q_m * cos(mθ_i) + rotate(q_m) * sin(mθ_i) + +Преимущества: +- Относительное позиционное кодирование +- Лучшая экстраполяция на длинные последовательности +- Сохранение нормы векторов +""" + import torch from torch import nn +from typing import Optional + class RoPE(nn.Module): + """ + Rotary Positional Embeddings (RoPE) для механизма внимания. + + Кодирует позиционную информацию через вращение векторов запросов и ключей + в многомерном пространстве с использованием синусов и косинусов. + + Args: + head_size: Размерность головы внимания (должен быть четным) + max_seq_len: Максимальная длина последовательности + base: Базовое значение для вычисления частот (по умолчанию 10000) + + Attributes: + cos_matrix: Буферизованная матрица косинусов формы [max_seq_len, head_size//2] + sin_matrix: Буферизованная матрица синусов формы [max_seq_len, head_size//2] + """ + def __init__(self, head_size: int, max_seq_len: int, base: int = 10_000): + """ + Инициализация RoPE эмбеддингов. + + Args: + head_size: Размерность головы внимания (должен быть четным) + max_seq_len: Максимальная поддерживаемая длина последовательности + base: Базовое значение для вычисления частот (типично 10000) + + Raises: + AssertionError: Если head_size не четный + """ super().__init__() assert head_size % 2 == 0, "head_size должен быть четным" - - # Обратные частоты + + # Вычисление частот: θ_i = base^(-2i/d) для i ∈ [0, d/2-1] freqs = 1.0 / (base ** (2 * torch.arange(head_size // 2).float() / head_size)) - # Позиции + # Позиции от 0 до max_seq_len-1 positions = torch.arange(max_seq_len).float() - # Матрица частот (внешнее произведение) - #freq_matrix = torch.outer(positions, freqs) + # Внешнее произведение: m * θ_i для всех позиций и частот freq_matrix = positions.unsqueeze(1) * freqs.unsqueeze(0) - # Матрицы косинусов и синусов + # Предвычисление матриц косинусов и синусов self.register_buffer('cos_matrix', torch.cos(freq_matrix)) self.register_buffer('sin_matrix', torch.sin(freq_matrix)) - - def forward(self, x: torch.Tensor): # Получает на вход тензор x (тип float) размером [batch_size × seq_len × head_size] + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Применение ротационного позиционного кодирования к входному тензору. + + Args: + x: Входной тензор формы [batch_size, seq_len, head_size] + + Returns: + Тензор с примененным RoPE формы [batch_size, seq_len, head_size] + + Алгоритм: + 1. Разделение векторов на четные и нечетные компоненты + 2. Применение вращения через синусы и косинусы + 3. Объединение компонент обратно + """ seq_len = x.size(1) + # Берем нужную часть матриц и приводим к типу x cos = self.cos_matrix[:seq_len].to(x.dtype) # [seq_len, head_size//2] sin = self.sin_matrix[:seq_len].to(x.dtype) # [seq_len, head_size//2] - - # Разделяем на четные и нечетные + # Разделяем на четные и нечетные компоненты x_even = x[:, :, 0::2] # [batch_size, seq_len, head_size//2] x_odd = x[:, :, 1::2] # [batch_size, seq_len, head_size//2] - # Применяем поворот + # Применяем поворот: q' = q * cos(mθ) + rotate(q) * sin(mθ) x_rotated_even = x_even * cos - x_odd * sin x_rotated_odd = x_even * sin + x_odd * cos - - # Объединяем обратно + # Объединяем обратно в исходную размерность x_rotated = torch.stack([x_rotated_even, x_rotated_odd], dim=-1) x_rotated = x_rotated.flatten(-2) # [batch_size, seq_len, head_size] diff --git a/llm/src/llm/core/swi_glu.py b/llm/src/llm/core/swi_glu.py index 0c3ade7..863fa8f 100644 --- a/llm/src/llm/core/swi_glu.py +++ b/llm/src/llm/core/swi_glu.py @@ -1,21 +1,94 @@ +""" +SwiGLU (Swish-Gated Linear Unit) - активационная функция с gating mechanism. + +Комбинация Swish активации и Gating Linear Unit. Широко используется в современных +моделях типа LLaMA и PaLM благодаря улучшенной производительности. + +Научная статья: "GLU Variants Improve Transformer" +https://arxiv.org/abs/2002.05202 + +Формула: +SwiGLU(x) = Swish(xW_g + b_g) ⊙ (xW_u + b_u) * W_d + b_d + +Преимущества: +- Лучшая производительность чем у ReLU/GELU +- Gating mechanism позволяет модели лучше выбирать информацию +- Хорошо масштабируется для больших моделей +""" + import torch from torch import nn +from typing import Optional from .silu import SiLU -class SwiGLU(nn.Module): - def __init__(self, emb_size: int, dropout: float = 0.1): - super().__init__() +class SwiGLU(nn.Module): + """ + Реализация SwiGLU активационной функции. + + Состоит из трех линейных слоев и активации SiLU: + 1. Gate слой + SiLU активация + 2. Up слой (линейное преобразование) + 3. Element-wise multiplication gate и up + 4. Down слой (линейная проекция) + + Args: + emb_size: Размерность входных эмбеддингов + dropout: Вероятность dropout (по умолчанию 0.1) + + Attributes: + _gate: Линейный слой для gate ветви [emb_size -> 4*emb_size] + _up: Линейный слой для up ветви [emb_size -> 4*emb_size] + _down: Линейный слой проекции [4*emb_size -> emb_size] + _activation: Функция активации SiLU + _dropout: Dropout слой + """ + + def __init__(self, emb_size: int, dropout: float = 0.1): + """ + Инициализация SwiGLU слоя. + + Args: + emb_size: Размерность входных/выходных эмбеддингов + dropout: Вероятность dropout (по умолчанию 0.1) + """ + super().__init__() self._gate = nn.Linear(emb_size, 4 * emb_size) self._up = nn.Linear(emb_size, 4 * emb_size) self._down = nn.Linear(4 * emb_size, emb_size) self._activation = SiLU() self._dropout = nn.Dropout(dropout) - def forward(self, x: torch.Tensor): # [batch_size × seq_len × emb_size]. + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Прямой проход через SwiGLU слой. + + Args: + x: Входной тензор формы [batch_size, seq_len, emb_size] + + Returns: + Выходной тензор формы [batch_size, seq_len, emb_size] + + Алгоритм: + 1. gate = SiLU(linear_gate(x)) + 2. up = linear_up(x) + 3. output = linear_down(gate ⊙ up) + 4. apply dropout + """ + # Gate ветвь: линейное преобразование + активация gate_out = self._gate(x) # [batch, seq, 4*emb] activation_out = self._activation(gate_out) # [batch, seq, 4*emb] + + # Up ветвь: линейное преобразование up_out = self._up(x) # [batch, seq, 4*emb] - out = up_out * activation_out # поэлементное! + + # Element-wise multiplication (gating mechanism) + out = up_out * activation_out # поэлементное умножение! + + # Final projection and dropout out = self._down(out) # [batch, seq, emb] - return self._dropout(out) \ No newline at end of file + return self._dropout(out) + + def extra_repr(self) -> str: + """Строковое представление для отладки.""" + return f'emb_size={self._gate.in_features}, dropout={self._dropout.p}' \ No newline at end of file diff --git a/llm/src/llm/models/gpt/gpt.py b/llm/src/llm/models/gpt/gpt.py index c184a27..3aa605e 100644 --- a/llm/src/llm/models/gpt/gpt.py +++ b/llm/src/llm/models/gpt/gpt.py @@ -1,24 +1,50 @@ -# llm/models/gpt/gpt2.py +""" +Original GPT (Generative Pre-trained Transformer) модель. + +Реализация классической GPT архитектуры из статьи: +"Improving Language Understanding by Generative Pre-Training" +https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf + +Архитектурные особенности: +- Трансформер-декодер с masked self-attention +- Layer Normalization применяется после внимания и FFN +- GELU активационная функция +- Learned positional embeddings +- Обучение на задачах языкового моделирования +""" + import torch import torch.nn as nn import torch.nn.functional as F +from typing import Optional, Dict from llm.core.base_model import BaseModel from llm.core.decoder import Decoder from llm.core.token_embeddings import TokenEmbeddings from llm.core.positional_embeddings import PositionalEmbeddings + class GPT(BaseModel): - """GPT-like трансформер для генерации текста + """ + Original GPT (Generative Pre-trained Transformer) модель. + + Первая версия трансформерной архитектуры от OpenAI, предназначенная + для генеративного предобучения на текстовых данных. Args: - vocab_size: Размер словаря - max_seq_len: Макс. длина последовательности - emb_size: Размерность эмбеддингов - num_heads: Количество голов внимания - head_size: Размерность голов внимания - num_layers: Количество слоёв декодера - dropout: Вероятность dropout (default=0.1) - device: Устройство (default='cpu') + config: Словарь конфигурации с параметрами: + - vocab_size: Размер словаря токенов + - embed_dim: Размерность векторных представлений + - num_heads: Количество голов внимания + - num_layers: Количество декодерных слоев + - max_position_embeddings: Максимальная длина последовательности + - dropout: Вероятность dropout + + Attributes: + _token_embeddings: Слой векторных представлений токенов + _position_embeddings: Слой позиционных эмбеддингов + _decoders: Список декодерных слоев + _norm: Финальный слой нормализации + _linear: Выходной линейный слой """ def __init__(self, config): super().__init__(config)