diff --git a/llm/src/llm/core/base_model.py b/llm/src/llm/core/base_model.py index a1f1c45..0560b53 100644 --- a/llm/src/llm/core/base_model.py +++ b/llm/src/llm/core/base_model.py @@ -1,34 +1,43 @@ -# llm/core/base_model.py """ -Базовый абстрактный класс для всех языковых моделей (LLM). +Базовый абстрактный класс для всех больших языковых моделей (LLM). -Реализует общий интерфейс для прямого прохода и генерации текста. -Все конкретные модели должны наследоваться от этого класса. +Научная суть: +Модели типа LLM строятся по модульному принципу — конкретные GPT, LLaMA и др. должны наследоваться от этого класса и реализовывать базовый набор интерфейсов для совместимости с training loop, генерацией, инференсом и т.д. + +Пользовательский уровень: +Базовый интерфейс минимизирует дублирование кода и позволяет быстро добавлять новые архитектуры. + +Использование: + class MyModel(BaseModel): + ... + model = MyModel(config) + logits = model.forward(input_ids) + tokens = model.generate(input_ids) """ - 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): Конфигурация модели с параметрами архитектуры + config (dict): Параметры архитектуры (размерность эмбеддингов, число слоев, heads и т.д.) Attributes: - config (dict): Конфигурационные параметры модели + config (dict): Конфиг модели """ - def __init__(self, config: dict): """ - Инициализация базовой модели. + Инициализация модели. Args: - config: Словарь с параметрами конфигурации модели + config (dict): Настройки архитектуры модели (размеры слоев, типы блоков и т.д.) """ super().__init__() self.config = config @@ -36,27 +45,28 @@ class BaseModel(nn.Module, ABC): @abstractmethod 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] - + input_ids (Tensor[int]): Индексы токенов [batch, seq_len] + attention_mask (Optional[Tensor[bool]]): Маска разрешенных позиций (если требуется) [batch, seq_len] Returns: - Тензор логитов формы [batch_size, seq_len, vocab_size] + logits (Tensor[float]): Логиты словаря [batch, seq_len, vocab_size] """ pass @abstractmethod def generate(self, input_ids: torch.Tensor, max_length: int = 50) -> torch.Tensor: """ - Генерация текста с использованием greedy decoding или sampling. + Генерация текста (авторегрессивно, greedy или sampling). Args: - input_ids: Начальные токены для генерации формы [batch_size, start_len] - max_length: Максимальная длина генерируемой последовательности - + input_ids (Tensor[int]): Начальные токены [batch, start_len] + max_length (int): Максимальная длина последовательности Returns: - Тензор сгенерированных токенов формы [batch_size, generated_len] + output_tokens (Tensor[int]): Сгенерированная последовательность [batch, generated_len] + Пример: + >>> logits = model.forward(input_ids) + >>> generated = model.generate(input_ids, max_length=128) """ pass diff --git a/llm/src/llm/core/cached_decoder.py b/llm/src/llm/core/cached_decoder.py index 59eae7b..dcc5a4e 100644 --- a/llm/src/llm/core/cached_decoder.py +++ b/llm/src/llm/core/cached_decoder.py @@ -8,34 +8,51 @@ from .rope import RoPE class CachedDecoder(nn.Module): """ - Универсальный декодерный блок с dependency injection для поддержки различных архитектур. - - Поддерживает кэширование ключей-значений для ускорения генерации текста. - + Универсальный декодерный блок для современных LLM (GPT, LLaMA, др.), поддерживает кэширование key-value для эффективной генерации. + + Научная идея: + Автопагрессивная авторегрессия в трансформерах требует быстрого доступа к ранее вычисленным self-attention ключам/значениям — этот класс позволяет прозрачно кэшировать такие состояния для быстрой инференс-генерации. + + Алгоритм: + - Input -> LayerNorm -> Многоголовое внимание с кэшем (может быть RoPE) + - Суммируем residual + - LayerNorm -> FeedForward (любой, например SwiGLU) -> Residual + - Возвращается кортеж (output, kvcache) + Args: - feed_forward_layer: Экземпляр слоя прямого распространения (SwiGLU, FeedForward и т.д.) - num_heads: Количество голов механизма внимания - emb_size: Размерность векторных представлений - head_size: Размерность каждой головы внимания - max_seq_len: Максимальная длина последовательности - norm_layer: Класс слоя нормализации (LayerNorm, RMSNorm и т.д.) - dropout: Вероятность dropout - rope: Экземпляр RoPE для позиционного кодирования (опционально) + feed_forward_layer (nn.Module): FeedForward или SwiGLU слой + num_heads (int): Количество голов внимания + emb_size (int): Размерность эмбеддингов + head_size (int): Размерность головы внимания + max_seq_len (int): Максимальная длина + norm_layer (тип nn.Module): Normalization слой (LayerNorm или RMSNorm) + dropout (float): Dropout + rope (RoPE|None): Экземпляр RoPE (для LLaMA) + + Пример (GPT2 style): + >>> decoder = CachedDecoder( + ... feed_forward_layer=FeedForward(...), + ... norm_layer=nn.LayerNorm, + ... num_heads=4, emb_size=256, head_size=64, max_seq_len=128) + >>> out, cache = decoder(x, use_cache=True) """ def __init__( self, - feed_forward_layer: nn.Module, # Обязательный параметр + feed_forward_layer: nn.Module, num_heads: int, emb_size: int, head_size: int, max_seq_len: int, - norm_layer: type = nn.LayerNorm, # Класс + norm_layer: type = nn.LayerNorm, dropout: float = 0.1, rope: RoPE = None, ): """ Инициализация декодера с кэшированием. + Поведение аналогично блоку TransformerDecoderLayer, + но с гибкой возможностью подмены любых подкомпонент (активация, norm, позиции). + Args: feed_forward_layer: Слой feed-forward (должен быть экземпляром, а не классом) num_heads: Количество голов внимания @@ -67,18 +84,19 @@ class CachedDecoder(nn.Module): cache: list = None, ): """ - Прямой проход через декодерный блок. + Прямой проход с поддержкой кэша. Args: - x: Входной тензор формы [batch_size, seq_len, emb_size] - mask: Маска внимания формы [batch_size, seq_len] (опционально) - use_cache: Флаг использования кэширования - cache: Список кэшированных пар (key, value) тензоров - + x (Tensor[float]): [batch, seq_len, emb_size] — скрытые состояния + mask (Optional[Tensor]): маска внимания (или causal mask), shape [seq_len, seq_len] + use_cache (bool): использовать кэширование KV + cache (list): кэш self-attention для быстрого авторегрессива Returns: - Кортеж (output, new_cache) где: - - output: Выходной тензор формы [batch_size, seq_len, emb_size] - - new_cache: Обновленный кэш или None, если use_cache=False + output (Tensor[float]): выходные состояния [batch, seq_len, emb_size] + kv_caches (list): обновленный кэш, если use_cache + Пример: + >>> out, new_cache = decoder(x, use_cache=True, cache=old_cache) + >>> out.shape # [batch, seq_len, emb_size] """ norm1_out = self._norm1(x) # Передаём все cache/use_cache дальше в attention diff --git a/llm/src/llm/core/decoder.py b/llm/src/llm/core/decoder.py index da4f01f..40cb9dd 100644 --- a/llm/src/llm/core/decoder.py +++ b/llm/src/llm/core/decoder.py @@ -5,41 +5,24 @@ from .multi_head_attention import MultiHeadAttention class Decoder(nn.Module): """ - Декодер трансформера - ключевой компонент архитектуры Transformer. - - Предназначен для: - - Обработки последовательностей с учетом контекста (самовнимание) - - Постепенного генерирования выходной последовательности - - Учета масок для предотвращения "заглядывания в будущее" + Базовый автогерессивный блок-декодер трансформера (без кэша KV). - Алгоритм работы: - 1. Входной тензор (batch_size, seq_len, emb_size) - 2. Многоголовое внимание с residual connection и LayerNorm - 3. FeedForward сеть с residual connection и LayerNorm - 4. Выходной тензор (batch_size, seq_len, emb_size) - - Основные характеристики: - - Поддержка масок внимания - - Residual connections для стабилизации градиентов - - Layer Normalization после каждого sub-layer - - Конфигурируемые параметры внимания - - Примеры использования: - - 1. Базовый случай: - >>> decoder = Decoder(num_heads=8, emb_size=512, head_size=64, max_seq_len=1024) - >>> x = torch.randn(1, 10, 512) # [batch, seq_len, emb_size] - >>> output = decoder(x) - >>> print(output.shape) - torch.Size([1, 10, 512]) - - 2. С маской внимания: - >>> mask = torch.tril(torch.ones(10, 10)) # Нижнетреугольная маска - >>> output = decoder(x, mask) - - 3. Инкрементальное декодирование: - >>> for i in range(10): - >>> output = decoder(x[:, :i+1, :], mask[:i+1, :i+1]) + Научная суть: + - Осуществляет посимвольное предсказание: каждый токен видит только предыдущие (masked attention) + - Состоит из self-attention + feedforward + residual + нормализация + - Residual connection и normalization дают стабильность и градиентный “flow” при обучении + - Механизм предложен в Vaswani et al., "Attention is All You Need", 2017 + Args: + num_heads (int): количество attention-голов + emb_size (int): размер эмбеддинга + head_size (int): размер одной attention-головы + max_seq_len (int): максимальная длина последовательности + dropout (float): вероятность dropout + Пример: + >>> decoder = Decoder(num_heads=8, emb_size=512, head_size=64, max_seq_len=1024) + >>> x = torch.randn(1, 10, 512) + >>> out = decoder(x) + >>> print(out.shape) # torch.Size([1, 10, 512]) """ def __init__(self, num_heads: int, diff --git a/llm/src/llm/core/feed_forward.py b/llm/src/llm/core/feed_forward.py index 93fb3d3..3168dea 100644 --- a/llm/src/llm/core/feed_forward.py +++ b/llm/src/llm/core/feed_forward.py @@ -6,12 +6,21 @@ from .gelu import GELU class FeedForward(nn.Module): """ - Слой прямой связи (Feed Forward Network) для архитектуры трансформеров. - + Классический слой прямого распространения (FeedForward, или FFN) для архитектуры Transformer. + Этот слой состоит из двух линейных преобразований с расширением внутренней размерности в 4 раза и механизмом dropout для регуляризации. Между линейными слоями применяется активация ReLU. + Научная суть: + - После внимания каждому токену применяется одинаковая двухслойная нейросеть. + - Дает глубокую нелинейность; позволяет модели не только сопоставлять, но и моделировать сложные связи между токенами. + - Изначально предложен в «Attention is All You Need» (Vaswani et al., 2017). + + Формула: + FFN(x) = Dropout(W2·act(W1·x)) + где act — ReLU, GELU и др., обычно expansion x4. + Алгоритм работы: 1. Входной тензор x (размерность: [batch_size, seq_len, emb_size]) 2. Линейное преобразование: emb_size -> 4*emb_size @@ -24,21 +33,17 @@ class FeedForward(nn.Module): - Добавляет нелинейность в архитектуру трансформера - Обеспечивает взаимодействие между различными размерностями эмбеддингов - Работает независимо для каждого токена в последовательности - - Примеры использования: - >>> # Инициализация слоя - >>> ff = FeedForward(emb_size=512, dropout=0.1) - >>> - >>> # Прямой проход - >>> x = torch.randn(32, 10, 512) # [batch_size, seq_len, emb_size] - >>> output = ff(x) - >>> print(output.shape) # torch.Size([32, 10, 512]) - >>> - >>> # Работа с разными типами данных - >>> x_double = torch.randn(32, 10, 512, dtype=torch.float64) - >>> output_double = ff(x_double) - >>> print(output_double.dtype) # torch.float64 + Args: + emb_size (int): размерность входных эмбеддингов + dropout (float): вероятность(dropout) + activation (str): нелинейная функция (relu, gelu, gelu_exact) + + Пример: + >>> ff = FeedForward(emb_size=512, dropout=0.1) + >>> x = torch.randn(32, 10, 512) + >>> output = ff(x) + >>> print(output.shape) # torch.Size([32, 10, 512]) """ def __init__(self, emb_size: int, dropout: float = 0.1, activation: str = "relu"): """ diff --git a/llm/src/llm/core/gelu.py b/llm/src/llm/core/gelu.py index e311c6e..95772f6 100644 --- a/llm/src/llm/core/gelu.py +++ b/llm/src/llm/core/gelu.py @@ -2,6 +2,21 @@ import torch from torch import nn class GELU(nn.Module): + """ + Гауссовская Эрф-активация (GELU, Gaussian Error Linear Unit). + + Научная суть: + - Одна из самых популярных smooth активаций для трансформеров. + - Дает более гибкие аппроксимации, чем ReLU/SiLU, улучшает flow градиентов для больших LLM. + - Используется в BERT, GPT, GPT2 и почти всех современных NLP-моделях. + Формула: + GELU(x) = 0.5 * x * (1 + tanh(\sqrt{2/π} * (x + 0.044715 x³))) + Подробнее: Hendrycks & Gimpel, "Gaussian Error Linear Units (GELUs)", arXiv:1606.08415 + Пример: + >>> gelu = GELU() + >>> y = gelu(torch.tensor([-1.0, 0.0, 1.0])) + >>> print(y) + """ def __init__(self): super().__init__() self.sqrt_2_over_pi = torch.sqrt(torch.tensor(2.0) / math.pi) diff --git a/llm/src/llm/core/head_attention.py b/llm/src/llm/core/head_attention.py index ff9909c..4b32032 100644 --- a/llm/src/llm/core/head_attention.py +++ b/llm/src/llm/core/head_attention.py @@ -6,30 +6,34 @@ from .rope import RoPE class HeadAttention(nn.Module): """ - Реализация одного головного механизма внимания из архитектуры Transformer. - Выполняет scaled dot-product attention с маскированием будущих позиций (causal attention). - - Основной алгоритм: - 1. Линейные преобразования входных данных в Q (query), K (key), V (value) - 2. Вычисление scores = Q·K^T / sqrt(d_k) - 3. Применение causal маски (заполнение -inf будущих позиций) - 4. Softmax для получения весов внимания - 5. Умножение весов на значения V - - Пример использования: - >>> attention = HeadAttention(emb_size=64, head_size=32, max_seq_len=128) - >>> x = torch.randn(1, 10, 64) # [batch_size, seq_len, emb_size] - >>> output = attention(x) # [1, 10, 32] - - Параметры: - emb_size (int): Размер входного эмбеддинга - head_size (int): Размерность выхода головы внимания - max_seq_len (int): Максимальная длина последовательности - + Одноголовый механизм внимания (scaled dot-product attention) — фундаментальный строительный блок всех современных Transformer. + + Научная суть: + - Attention учит модель самостоятельно "выбирать" важные связи между словами, независимо от их положения. + - Механизм causal mask гарантирует невозможность "заглядывания в будущее" при генерации (авторегрессия). + + Формула: + Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) · V + (Q — запросы, K — ключи, V — значения; d_k — размерность ключа) + + Поддерживает Rotary Position Encoding (RoPE) для относительного позиционного кодирования. + + Args: + emb_size (int): размер входного эмбеддинга + head_size (int): размерность attention-головы + max_seq_len (int): максимальная длина последовательности + rope (RoPE, optional): экземпляр RoPE для позиций + Примечания: - Использует нижнетреугольную маску для предотвращения "заглядывания в будущее" - Автоматически адаптируется к разным версиям PyTorch - Поддерживает batch-обработку входных данных + + Пример использования: + >>> attention = HeadAttention(emb_size=64, head_size=32, max_seq_len=128) + >>> x = torch.randn(1, 10, 64) + >>> output, _ = attention(x) + >>> print(output.shape) # torch.Size([1, 10, 32]) """ def __init__(self, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None): super().__init__() diff --git a/llm/src/llm/core/multi_head_attention.py b/llm/src/llm/core/multi_head_attention.py index b865714..c788359 100644 --- a/llm/src/llm/core/multi_head_attention.py +++ b/llm/src/llm/core/multi_head_attention.py @@ -5,32 +5,32 @@ from .rope import RoPE class MultiHeadAttention(nn.Module): """ - Реализация механизма многоголового внимания (Multi-Head Attention) из архитектуры Transformer. + Мультиголовый (многоголовый) механизм внимания — ключевой компонент любого Transformer. - Основные характеристики: - - Параллельная обработка входных данных несколькими головами внимания - - Поддержка маскирования (causal mask и пользовательские маски) - - Финальная проекция с dropout регуляризацией + Научная суть: + - Модель параллельно агрегирует информацию через несколько подпространств (головы), + чтобы видеть разные связи в последовательности (разный контекст, локально/глобально). + - Каждый attention блок работает независимо, выход конкатенируется. + - Механизм предложен в статье "Attention is All You Need" (Vaswani et al., 2017). + + Формула внимания для одной головы: + Attention(Q, K, V) = softmax(QK^T/sqrt(d_k))·V + Мультиголовый: + MultiHead(Q, K, V) = Concat([head_i])*W^O - Математическое описание: - MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O - где head_i = Attention(QW_i^Q, KW_i^K, VW_i^V) + Args: + num_heads (int): количество attention "голов" + emb_size (int): размерности входа и выхода + head_size (int): размер одной attention-головы (emb_size/num_heads) + max_seq_len (int): максимальная длина последовательности + rope (RoPE, optional): если задан, используется Rotary Positional Encoding + dropout (float): вероятность регуляризации - Примеры использования: - - 1. Базовый пример: - >>> mha = MultiHeadAttention(num_heads=8, emb_size=512, head_size=64, max_seq_len=1024) - >>> x = torch.randn(2, 50, 512) # [batch_size, seq_len, emb_size] - >>> output = mha(x) # [2, 50, 512] - - 2. С использованием маски: - >>> mask = torch.tril(torch.ones(50, 50)) # Causal mask - >>> output = mha(x, mask) - - 3. Интеграция в Transformer: - >>> # В составе Transformer слоя - >>> self.attention = MultiHeadAttention(...) - >>> x = self.attention(x, mask) + Пример использования: + >>> mha = MultiHeadAttention(num_heads=8, emb_size=512, head_size=64, max_seq_len=1024) + >>> x = torch.randn(2, 50, 512) + >>> out, cache = mha(x) + >>> print(out.shape) """ def __init__(self, num_heads: int, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None, dropout: float = 0.1): """ @@ -62,7 +62,8 @@ class MultiHeadAttention(nn.Module): def forward(self, x: torch.Tensor, mask: torch.Tensor = None, use_cache: bool = True, cache: list = None): """ - Прямой проход через слой многоголового внимания. + Прямой проход (forward): + Для каждого токена оценивает "важность" остальных токенов сразу через несколько attention-блоков. Подробное описание преобразований тензоров: 1. Входной тензор [batch_size, seq_len, emb_size] разделяется на N голов: @@ -75,13 +76,20 @@ class MultiHeadAttention(nn.Module): 4. Линейная проекция: - Выход: [batch_size, seq_len, emb_size] 5. Применение dropout + + Args: + x (Tensor[float]): [batch, seq_len, emb_size] — вход + mask (Optional[Tensor[bool]]): маска позиции [seq_len, seq_len] + use_cache (bool): использовать ли key-value кэш (для генерации) + cache (list): предыдущие значения KV для ускорения - Аргументы: - x (torch.Tensor): Входной тензор формы [batch_size, seq_len, emb_size] - mask (torch.Tensor, optional): Маска внимания формы [seq_len, seq_len] + Returns: + out (Tensor[float]): [batch, seq_len, emb_size] — результат MHA + kv_caches (list): списки новых KV-кэшей (если используется) - Возвращает: - torch.Tensor: Выходной тензор формы [batch_size, seq_len, emb_size] + Типичный паттерн: + Вход: [batch, seq, emb] → N голов [batch, seq, head_size] → + → concat [batch, seq, N*head_size] → проекция → dropout Пример преобразований для emb_size=512, num_heads=8: Вход: [4, 100, 512] @@ -90,6 +98,10 @@ class MultiHeadAttention(nn.Module): -> Конкатенация: [4, 100, 512] -> Проекция: [4, 100, 512] -> Dropout: [4, 100, 512] + + Пример: + >>> out, caches = mha(x) + >>> out.shape # [batch, seq_len, emb_size] """ # 1. Вычисляем attention для каждой головы attention_results = [] diff --git a/llm/src/llm/core/positional_embeddings.py b/llm/src/llm/core/positional_embeddings.py index f167c31..fb0e7fb 100644 --- a/llm/src/llm/core/positional_embeddings.py +++ b/llm/src/llm/core/positional_embeddings.py @@ -3,19 +3,19 @@ from torch import nn, Tensor class PositionalEmbeddings(nn.Module): """ - Класс для создания позиционных эмбеддингов через nn.Embedding. - + Обучаемые позиционные эмбеддинги (learnable positional embeddings). + Позиционные эмбеддинги используются в нейросетях для передачи информации о позиции элементов в последовательности (например, в Transformer). - Особенности: - - Создаёт обучаемые позиционные эмбеддинги фиксированной длины - - Поддерживает обработку последовательностей переменной длины - - Автоматически размещает вычисления на том же устройстве, что и параметры - + Научная суть: + - Трансформеры не используют рекуррентность, а значит сами по себе не различают порядок слов. + - Позиционные эмбеддинги добавляются к токеновым, чтобы сеть понимала, в каком месте последовательности находится каждый токен. + - Обычно реализуются как отдельная матрица (nn.Embedding), которая обучается вместе с моделью (это learnable вариант, как в GPT и BERT). + Args: - max_seq_len (int): Максимальная длина последовательности - emb_size (int): Размерность векторного представления позиций + max_seq_len (int): максимальная длина последовательности + emb_size (int): размер вектора позиции Пример использования: >>> pos_encoder = PositionalEmbeddings(max_seq_len=100, emb_size=256) diff --git a/llm/src/llm/core/rms_norm.py b/llm/src/llm/core/rms_norm.py index 305e7f0..47def26 100644 --- a/llm/src/llm/core/rms_norm.py +++ b/llm/src/llm/core/rms_norm.py @@ -24,18 +24,26 @@ from typing import Optional class RMSNorm(nn.Module): """ - Реализация RMS Normalization. - + RMS Normalization (Root Mean Square Layer Normalization). + Нормализует входные данные по последнему измерению используя среднеквадратичное значение вместо среднего, как в стандартном LayerNorm. - + + Научная суть: + - Упрощенный вариант LayerNorm без вычисления среднего, только деление на rms. + - Лучшая численная стабильность на больших моделях, меньше вычислений. + - Применяется в LLaMA, PaLM и др. + + Формула: + RMSNorm(x) = (x / sqrt(mean(x²) + eps)) * w (w — обучаемый вектор) + Args: - dim: Размерность измерения для нормализации - eps: Малое значение для численной стабильности + dim (int): размер последнего измерения (обычно emb_size) + eps (float): для численной устойчивости - Attributes: - _eps: Малое значение для предотвращения деления на ноль - _w: Обучаемый параметр масштабирования формы [dim] + Пример: + >>> norm = RMSNorm(emb_size) + >>> out = norm(x) """ def __init__(self, dim: int, eps: float = 1e-6): diff --git a/llm/src/llm/core/silu.py b/llm/src/llm/core/silu.py index b557280..c23d542 100644 --- a/llm/src/llm/core/silu.py +++ b/llm/src/llm/core/silu.py @@ -2,5 +2,18 @@ import torch from torch import nn class SiLU(nn.Module): - def forward(self, x: torch.Tensor): # [batch_size × seq_len × emb_size] + """ + SiLU (Swish) — современная активационная функция для нейросетей. + + Научная суть: + - Формула: $SiLU(x) = x * \sigm(x)$, где $\sigm(x)$ — сигмоида. + - Более гладкая альтернатива ReLU, улучшает поток градиентов в глубоких сетях. + - Используется во многих «state-of-the-art» архитектурах (SwiGLU, PaLM, LLaMA). + - Также известна как Swish (Ramachandran et al, 2017). + Пример: + >>> act = SiLU() + >>> x = torch.tensor([-1.0, 0.0, 1.0]) + >>> print(act(x)) + """ + def forward(self, x: torch.Tensor): return torch.sigmoid(x) * x \ No newline at end of file diff --git a/llm/src/llm/core/swi_glu.py b/llm/src/llm/core/swi_glu.py index 863fa8f..ace972d 100644 --- a/llm/src/llm/core/swi_glu.py +++ b/llm/src/llm/core/swi_glu.py @@ -24,6 +24,8 @@ from .silu import SiLU class SwiGLU(nn.Module): """ + SwiGLU (Swish-Gated Linear Unit) — современная нелинейность для архитектур LLM (LLaMA, PaLM). + Реализация SwiGLU активационной функции. Состоит из трех линейных слоев и активации SiLU: @@ -32,16 +34,21 @@ class SwiGLU(nn.Module): 3. Element-wise multiplication gate и up 4. Down слой (линейная проекция) + Научная суть: + - Сохраняет преимущества GLU (раздельные гейтом и телом) + мощность Swish/SiLU активации. + - Дает надежную гладкую активацию, хорошо работает на больших масштабах. + - Статья: "GLU Variants Improve Transformer" (Shazeer, 2020). + + Формула: + SwiGLU(x) = SiLU(W_g·x) * (W_u·x) + где SiLU(x) = x*sigma(x) + 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 слой + emb_size (int): размер входов/выходов + dropout (float): после выходной проекции + Пример: + >>> ff = SwiGLU(emb_size=512, dropout=0.1) + >>> y = ff(torch.randn(2,10,512)) """ def __init__(self, emb_size: int, dropout: float = 0.1): diff --git a/llm/src/llm/core/token_embeddings.py b/llm/src/llm/core/token_embeddings.py index e6e2737..5674e6f 100644 --- a/llm/src/llm/core/token_embeddings.py +++ b/llm/src/llm/core/token_embeddings.py @@ -4,29 +4,30 @@ from torch import Tensor class TokenEmbeddings(nn.Module): """ - Модуль PyTorch для преобразования индексов токенов в векторные представления (эмбеддинги). - + Токеновые эмбеддинги — обучаемые векторные представления для каждого токена словаря. + Преобразует целочисленные индексы токенов в обучаемые векторные представления фиксированного размера. Обычно используется как первый слой в нейронных сетях для задач NLP. - Аргументы: - vocab_size (int): Размер словаря (количество уникальных токенов) - emb_size (int): Размерность векторных представлений - - Форматы данных: - - Вход: тензор (batch_size, seq_len) индексов токенов - - Выход: тензор (batch_size, seq_len, emb_size) векторных представлений - - Примеры использования: - >>> embedding_layer = TokenEmbeddings(vocab_size=10000, emb_size=256) - >>> tokens = torch.tensor([[1, 2, 3], [4, 5, 6]]) # batch_size=2, seq_len=3 - >>> embeddings = embedding_layer(tokens) - >>> embeddings.shape - torch.Size([2, 3, 256]) - + Научная суть: + - Первый шаг для любого NLP-модуля: вместо индекса токена подаём его dense-вектор. + - Эти вектора изучаются в процессе обучения и отражают скрытые взаимосвязи между токенами. + - Позволяют обрабатывать тексты как матрицу чисел, а не как символы или индексы. + - Аналог словарных эмбеддингов в word2vec, но обучаются энд-ту-энд с моделью. + + Args: + vocab_size (int): размер словаря (количество уникальных токенов) + emb_size (int): размерность эмбеддинга (длина вектора) + Примечание: - Индексы должны быть в диапазоне [0, vocab_size-1] - Эмбеддинги инициализируются случайно и обучаются в процессе тренировки модели + + Пример: + >>> emb = TokenEmbeddings(vocab_size=10000, emb_size=256) + >>> tokens = torch.tensor([[1, 2, 3]]) + >>> vecs = emb(tokens) + >>> vecs.shape # torch.Size([1, 3, 256]) """ def __init__(self, vocab_size: int, emb_size: int): super().__init__() diff --git a/llm/src/llm/models/gpt/gpt.py b/llm/src/llm/models/gpt/gpt.py index 3aa605e..2121da8 100644 --- a/llm/src/llm/models/gpt/gpt.py +++ b/llm/src/llm/models/gpt/gpt.py @@ -1,17 +1,25 @@ """ -Original GPT (Generative Pre-trained Transformer) модель. +Классическая GPT (Generative Pre-trained Transformer), OpenAI 2018. -Реализация классической GPT архитектуры из статьи: -"Improving Language Understanding by Generative Pre-Training" -https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf +Научная суть: + - Первая массовая архитектура языка на основе исключительно self-attention механизмов (трансформер-декодер). + - Обучается сначала на задаче языкового моделирования (unsupervised), далее дообучается на downstream-задачах (transfer learning). + - Обеспечивает длинную память и “глобальный” контекст благодаря attention. -Архитектурные особенности: -- Трансформер-декодер с masked self-attention -- Layer Normalization применяется после внимания и FFN -- GELU активационная функция -- Learned positional embeddings -- Обучение на задачах языкового моделирования -""" + Ключевые элементы: + - masked self-attention (causal) + - LayerNorm ПОСЛЕ attention и FFN (что отличает от GPT2) + - GELU активация + - Absolute learned positional embeddings + + Подробнее: Radford et al., "Improving Language Understanding by Generative Pre-Training", arXiv:1801.10198 + https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf + + Пример использования: + >>> model = GPT({"vocab_size": 50257, ...}) + >>> logits = model(input_ids) + >>> out = model.generate(input_ids, max_length=30) + """ import torch import torch.nn as nn diff --git a/llm/src/llm/models/gpt/gpt2.py b/llm/src/llm/models/gpt/gpt2.py index 09726b4..54c60a0 100644 --- a/llm/src/llm/models/gpt/gpt2.py +++ b/llm/src/llm/models/gpt/gpt2.py @@ -1,3 +1,23 @@ +""" +GPT-2 — масштабируемый автогерессивный языковой трансформер второго поколения от OpenAI (2019). + +Научная суть: + - В сравнении с классическим GPT, layer normalization теперь применяется ПЕРЕД attention и FFN. + - Позволило сильно увеличить глубину и размер модели (GPT2-модели имеют от 117M до 1.5B параметров). + - Используется GELU активация; эффективное кэширование KV attention для генерации. + +Формула attention-блока: + LN(x) → Attention → рез. связь → LN → FFN → рез. связь + +Подробнее: + Radford et al. "Language Models are Unsupervised Multitask Learners" + https://cdn.openai.com/better-language-models/language-models.pdf + +Пример использования: + >>> model = GPT2({"vocab_size": 50257, ...}) + >>> logits = model(input_ids) + >>> out = model.generate(input_ids, max_length=30) +""" import torch from torch import nn, Tensor import torch.nn.functional as F @@ -8,6 +28,22 @@ from llm.core.cached_decoder import CachedDecoder from llm.core.feed_forward import FeedForward class GPT2(BaseModel): + """ + GPT2 — автогерессивная языковая модель, архитектура Transformer, предложенная OpenAI. + + Научная суть: + - Масштабируемый автогерессивный трансформер для предсказания токенов слева направо. + - Главное отличие от классической GPT: порядок layer normalization ПЕРЕД attention и FFN. + - Используется GELU, efficient KV-cache, несет наследие классической GPT, но делает архитектуру глубже/шире. + + Args: + config (dict): параметры архитектуры (vocab_size, embed_dim, num_heads, num_layers, max_position_embeddings, dropout) + + Пример использования: + >>> model = GPT2({"vocab_size": 50257, ...}) + >>> logits = model(input_ids) + >>> out = model.generate(input_ids, max_length=20) + """ def __init__(self, config): super().__init__(config) @@ -39,6 +75,20 @@ class GPT2(BaseModel): self._linear = nn.Linear(config["embed_dim"], config["vocab_size"]) def forward(self, x: torch.Tensor, use_cache: bool = True, cache: list = None) -> tuple: + """ + Прямой проход GPT2: + - Все слои работают как autoregressive transformer (masked self-attention). + - При use_cache=True возвращает также новый кэш KV attention (ускоряет генерацию). + Args: + x (Tensor): Входные индексы токенов [batch, seq_len] + use_cache (bool): Кэшировать KV attention для ускорения autoregressive генерации + cache (list|None): Список KV-кэшей от предыдущих шагов (или None) + Returns: + logits (Tensor): [batch, seq_len, vocab_size] + cache (list): новый кэш если use_cache=True, иначе None + Пример: + >>> logits, cache = model.forward(x, use_cache=True) + """ # Проверка длины последовательности (только при отсутствии кэша) if cache is None and x.size(1) > self._max_seq_len: raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}") @@ -97,6 +147,24 @@ class GPT2(BaseModel): top_p: float = None, use_cache: bool = True ) -> torch.Tensor: + """ + Генерация текста с использованием autoregressive трансформера (GPT2). + Поддерживаются greedy, sampling, top-k/top-p (nucleus sampling) режимы. + Args: + x (Tensor[int]): начальная последовательность [batch, seq_len] + max_new_tokens (int): сколько токенов сгенерировать + do_sample (bool): использовать стохастическое сэмплирование вместо жадного выбора + temperature (float): коэффициент сглаживания логитов (низкое — более консервативно) + top_k (int|None): ограничить выбор top-k наиболее вероятных токенов + top_p (float|None): ограничить суммарную вероятность (nucleus sampling) + use_cache (bool): ускорять autoregressive инференс + Returns: + output (Tensor[int]): сгенерированный тензор токенов [batch, seq_len + max_new_tokens] + Пример: + >>> prompt = tokenizer.encode('Привет', return_tensors="pt") + >>> output = model.generate(prompt, max_new_tokens=20, do_sample=True) + >>> print(tokenizer.decode(output[0])) + """ cache = None for _ in range(max_new_tokens): diff --git a/llm/src/llm/models/llama/llama.py b/llm/src/llm/models/llama/llama.py index 240d9a3..7e63221 100644 --- a/llm/src/llm/models/llama/llama.py +++ b/llm/src/llm/models/llama/llama.py @@ -12,6 +12,23 @@ from llm.core.cached_decoder import CachedDecoder class Llama(BaseModel): + """ + LLaMA (Large Language Model Meta AI) — высокоэффективная масштабируемая языковая модель, разработанная Meta AI Research. + + Ключевые идеи: + - Rotary Positional Encoding (RoPE) вместо стандартных позиционных эмбеддингов + - RMSNorm (Root Mean Square LayerNorm) вместо LayerNorm + - SwiGLU как нелинейность вместо ReLU/GELU (больше экспрессивности) + - Глубокая оптимизация inference (большая экономия памяти и FLOPs) + Подробнее: https://arxiv.org/abs/2302.13971 + + Args: + config (dict): параметры архитектуры (vocab_size, embed_dim, num_heads, num_layers, max_position_embeddings, dropout) + Пример: + >>> model = Llama({...}) + >>> logits, cache = model(input_ids, use_cache=True) + >>> out = model.generate(input_ids, max_new_tokens=20) + """ def __init__(self,config): super().__init__(config) @@ -44,6 +61,19 @@ class Llama(BaseModel): self._linear = nn.Linear(config["embed_dim"], config["vocab_size"]) def forward(self, x: torch.Tensor, use_cache: bool = True, cache: list = None) -> tuple: + """ + Прямой проход через LLaMA (inference/train): авторегрессионное предсказание токенов. + + Args: + x (Tensor[int]): входные токены [batch, seq_len] + use_cache (bool): использовать ли кэш (ускоряет генерацию) + cache (list|None): ключи и значения attention для autoregressive режима + Returns: + logits (Tensor): [batch, seq_len, vocab_size] + new_cache (list|None): новый кэш attention (если use_cache) + Пример: + >>> logits, cache = model.forward(x, use_cache=True) + """ # Проверка длины последовательности (только при отсутствии кэша) if cache is None and x.size(1) > self._max_seq_len: raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}") @@ -102,6 +132,27 @@ class Llama(BaseModel): top_p: float = None, use_cache: bool = True ) -> torch.Tensor: + """ + Генерация текста c помощью LLaMA (autoregressive Transformer). + Поддерживается: + - greedy и вероятностное сэмплирование (top-k, top-p, temperature) + - кэш attention для ускорения генерации длинных последовательностей + + Args: + x (Tensor[int]): начальная последовательность [batch, seq_len] + max_new_tokens (int): сколько новых токенов сгенерировать + do_sample (bool): использовать стохастику (True) или жадный выбор (False) + temperature (float): масштаб для softmax (важно для sampling) + top_k (int|None): ограничение на количество кандидатов (top-k sampling) + top_p (float|None): nucleus sampling + use_cache (bool): ускоряет autoregressive при длинной генерации + Returns: + output (Tensor[int]): [batch, seq_len + max_new_tokens] + Пример: + >>> prompt = tokenizer.encode('Meta AI', return_tensors="pt") + >>> generated = model.generate(prompt, max_new_tokens=30, do_sample=True) + >>> print(tokenizer.decode(generated[0])) + """ cache = None for _ in range(max_new_tokens):