docs: научная и практическая документация для всех ключевых модулей LLM

- Улучшены и дополнены docstring базовых компонентов (decoder, cached_decoder, multi_head_attention, head_attention, feed_forward, token_embeddings, positional_embeddings, gelu, silu, swi_glu, rope, rms_norm)
- На русском языке: объяснены алгоритмы архитектур, приведены формулы и ссылки на статьи
- Для всех моделей (GPT, GPT2, LLaMA) добавлены подробные описания классов, методов forward/generate, форматы входа/выхода
- Примеры использования в каждом ключевом классе
- Описаны научные концепции, архитектурные отличия и причины выбора решений

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

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

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. Базовый случай:
+    Научная суть:
+        - Осуществляет посимвольное предсказание: каждый токен видит только предыдущие (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)  # [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])
+        >>> x = torch.randn(1, 10, 512)
+        >>> out = decoder(x)
+        >>> print(out.shape)  # torch.Size([1, 10, 512])
     """
     def __init__(self, 
         num_heads: int,

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
@@ -25,20 +34,16 @@ class FeedForward(nn.Module):
     - Обеспечивает взаимодействие между различными размерностями эмбеддингов
     - Работает независимо для каждого токена в последовательности
     
-    Примеры использования:
+    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)  # [batch_size, seq_len, emb_size]
+        >>> x = torch.randn(32, 10, 512)
         >>> 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
     """
     def __init__(self, emb_size: int, dropout: float = 0.1, activation: str = "relu"):
         """

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)

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).
+    Одноголовый механизм внимания (scaled dot-product attention) — фундаментальный строительный блок всех современных Transformer.
 
-    Основной алгоритм:
-    1. Линейные преобразования входных данных в Q (query), K (key), V (value)
-    2. Вычисление scores = Q·K^T / sqrt(d_k)
-    3. Применение causal маски (заполнение -inf будущих позиций)
-    4. Softmax для получения весов внимания
-    5. Умножение весов на значения V
+    Научная суть:
+        - Attention учит модель самостоятельно "выбирать" важные связи между словами, независимо от их положения.
+        - Механизм causal mask гарантирует невозможность "заглядывания в будущее" при генерации (авторегрессия).
         
-    Пример использования:
-    >>> 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]
+        Формула:
+            Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) · V
+            (Q — запросы, K — ключи, V — значения; d_k — размерность ключа)
         
-    Параметры:
-        emb_size (int): Размер входного эмбеддинга
-        head_size (int): Размерность выхода головы внимания
-        max_seq_len (int): Максимальная длина последовательности
+        Поддерживает 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__()

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).
         
-    Математическое описание:
-    MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O
-    где head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)
+        Формула внимания для одной головы:
+            Attention(Q, K, V) = softmax(QK^T/sqrt(d_k))·V
+        Мультиголовый:
+            MultiHead(Q, K, V) = Concat([head_i])*W^O
 
-    Примеры использования:
+    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)
+        >>> 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 голов:
@@ -76,12 +77,19 @@ class MultiHeadAttention(nn.Module):
            - Выход: [batch_size, seq_len, emb_size]
         5. Применение dropout
         
-        Аргументы:
-            x (torch.Tensor): Входной тензор формы [batch_size, seq_len, emb_size]
-            mask (torch.Tensor, optional): Маска внимания формы [seq_len, seq_len]
+        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 для ускорения
 
-        Возвращает:
-            torch.Tensor: Выходной тензор формы [batch_size, seq_len, emb_size]
+        Returns:
+            out (Tensor[float]): [batch, seq_len, emb_size] — результат MHA
+            kv_caches (list): списки новых KV-кэшей (если используется)
+
+        Типичный паттерн:
+            Вход: [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 = []

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)

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.
 
-    Args:
-        dim: Размерность измерения для нормализации
-        eps: Малое значение для численной стабильности
+    Научная суть:
+        - Упрощенный вариант LayerNorm без вычисления среднего, только деление на rms.
+        - Лучшая численная стабильность на больших моделях, меньше вычислений.
+        - Применяется в LLaMA, PaLM и др.
         
-    Attributes:
-        _eps: Малое значение для предотвращения деления на ноль
-        _w: Обучаемый параметр масштабирования формы [dim]
+        Формула:
+            RMSNorm(x) = (x / sqrt(mean(x²) + eps)) * w   (w — обучаемый вектор)
+
+    Args:
+        dim (int): размер последнего измерения (обычно emb_size)
+        eps (float): для численной устойчивости
+    
+    Пример:
+        >>> norm = RMSNorm(emb_size)
+        >>> out = norm(x)
     """
     
     def __init__(self, dim: int, eps: float = 1e-6):

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

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 слой (линейная проекция)
     
-    Args:
-        emb_size: Размерность входных эмбеддингов
-        dropout: Вероятность dropout (по умолчанию 0.1)
+    Научная суть:
+        - Сохраняет преимущества GLU (раздельные гейтом и телом) + мощность Swish/SiLU активации.
+        - Дает надежную гладкую активацию, хорошо работает на больших масштабах.
+        - Статья: "GLU Variants Improve Transformer" (Shazeer, 2020).
 
-    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 слой
+        Формула:
+            SwiGLU(x) = SiLU(W_g·x) * (W_u·x)
+        где SiLU(x) = x*sigma(x)
+
+    Args:
+        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):

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): Размерность векторных представлений
+    Научная суть:
+        - Первый шаг для любого NLP-модуля: вместо индекса токена подаём его dense-вектор.
+        - Эти вектора изучаются в процессе обучения и отражают скрытые взаимосвязи между токенами.
+        - Позволяют обрабатывать тексты как матрицу чисел, а не как символы или индексы.
+        - Аналог словарных эмбеддингов в word2vec, но обучаются энд-ту-энд с моделью.
 
-    Форматы данных:
-        - Вход: тензор (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])
+    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__()

llm/src/llm/models/gpt/gpt.py

@@ -1,16 +1,24 @@
 """
-Original GPT (Generative Pre-trained Transformer) модель.
+Классическая GPT (Generative Pre-trained Transformer), OpenAI 2018.
 
-Реализация классической GPT архитектуры из статьи:
-"Improving Language Understanding by Generative Pre-Training"
+Научная суть:
+    - Первая массовая архитектура языка на основе исключительно self-attention механизмов (трансформер-декодер).
+    - Обучается сначала на задаче языкового моделирования (unsupervised), далее дообучается на downstream-задачах (transfer learning).
+    - Обеспечивает длинную память и “глобальный” контекст благодаря attention.
+
+    Ключевые элементы:
+    - 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
 
-Архитектурные особенности:
-- Трансформер-декодер с masked self-attention
-- Layer Normalization применяется после внимания и FFN
-- GELU активационная функция
-- Learned positional embeddings
-- Обучение на задачах языкового моделирования
+    Пример использования:
+        >>> model = GPT({"vocab_size": 50257, ...})
+        >>> logits = model(input_ids)
+        >>> out = model.generate(input_ids, max_length=30)
     """
 
 import torch

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

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