From 3a356f5d792f479332e804209c3e436b47b5ac13 Mon Sep 17 00:00:00 2001 From: Sergey Penkovsky Date: Thu, 16 Oct 2025 12:40:46 +0300 Subject: [PATCH] docs(core): improve and expand docstrings for Decoder module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - docs: rewrite and expand docstrings for Decoder class and its methods (__init__, forward) - clarify the block’s architecture, pre-LN logic, flow with residual connections, and attention masking - add mathematical pseudocode, motivation, feature list, usage example, and external references (papers, blog) - no logic or behavior changes This improves readability and makes the codebase easier to understand for transformer/LLM practitioners. --- llm/src/llm/core/decoder.py | 113 ++++++++++++++++++++++-------------- 1 file changed, 70 insertions(+), 43 deletions(-) diff --git a/llm/src/llm/core/decoder.py b/llm/src/llm/core/decoder.py index 3f96215..11de653 100644 --- a/llm/src/llm/core/decoder.py +++ b/llm/src/llm/core/decoder.py @@ -6,37 +6,45 @@ from .multi_head_attention import MultiHeadAttention class Decoder(nn.Module): """ - Базовый автогерессивный блок-декодер трансформера (без кэша KV). + Decoder — базовый transformer decoder block (pre-LN), классический строительный блок современных языковых моделей. - Предназначен для: - - Обработки последовательностей с учетом контекста (самовнимание) - - Постепенного генерирования выходной последовательности - - Учета масок для предотвращения "заглядывания в будущее" + Назначение: + ----------- + - Инкапсулирует архитектуру: norm → multi-head self-attention → residual → norm → feed-forward → residual + - Подходит как для LLM/GPT, так и для любых autoregressive sequence моделей. + - Использует masked self-attention: каждый токен видит только предыдущие (никакого \"заглядывания в будущее\"). + - Стабильность обеспечивается через residual connections и LayerNorm после каждого sub-layer. - Алгоритм работы: - 1. Входной тензор (batch_size, seq_len, emb_size) - 2. Многоголовое внимание с residual connection и LayerNorm - 3. FeedForward сеть с residual connection и LayerNorm - 4. Выходной тензор (batch_size, seq_len, emb_size) + Почему это важно? + ----------------- + - Все современные языковые модели состоят из подобных блоков, соединённых в стек. + - Алгоритм residual+norm позволяет проще обучать очень глубокие сети. + - Разделение на attention+FFN дает и локальные, и глобальные взаимодействия между токенами. - Основные характеристики: - - Поддержка масок внимания - - Residual connections для стабилизации градиентов - - Layer Normalization после каждого sub-layer - - Конфигурируемые параметры внимания + Формула работы (псевдокод): + --------------------------- + y1 = norm1(x) + attn_out = Attention(y1) + x2 = x + attn_out # residual + y2 = norm2(x2) + ffn_out = FFN(y2) + out = x2 + ffn_out # residual + + Архитектурные особенности: + -------------------------- + - Поддержка внимания с маской (causal mask или произвольная attention mask) + - Residual connections для каждого блока (attention, FFN) + - Pre-LN (norm перед каждым подблоком) + - Зависит от переданных блоков self_attention и feed_forward, а не их реализации + + References: + ----------- + - Vaswani et al., \"Attention is All You Need\" (2017): https://arxiv.org/abs/1706.03762 + - Illustrated Transformer: https://jalammar.github.io/illustrated-transformer/ + - Transformer Circuits (дружественное описание): https://transformer-circuits.pub/2021/framework/index.html - Научная суть: - - Осуществляет посимвольное предсказание: каждый токен видит только предыдущие (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) @@ -52,14 +60,27 @@ class Decoder(nn.Module): dropout: float = 0.1, ): """ - Инициализация декодера. + Инициализация стандартного decoder-блока для Transformer. - Параметры: - num_heads: int - количество голов внимания - emb_size: int - размерность эмбеддингов - head_size: int - размерность каждой головы внимания - max_seq_len: int - максимальная длина последовательности - dropout: float (default=0.1) - вероятность dropout + Аргументы: + ---------- + num_heads: int + Количество attention голов (как делить emb_size на heads) + emb_size: int + Размерность эмбеддингов (и входа и выхода) + head_size: int + Размерность одной attention-головы (emb_size = num_heads * head_size) + max_seq_len: int + Максимальная длина последовательности (важно для mask) + dropout: float, default=0.1 + Dropout после внимания и FFN + + Внутри: + ------- + - Создаёт слой MultiHeadAttention (masked/casual) + - Создаёт двухслойный FeedForward (SwiGLU или GELU) + - Применяет 2 слоя LayerNorm для стабилизации градиентов + - Все блоки реализованы как PyTorch-модули """ super().__init__() self._heads = MultiHeadAttention( @@ -75,20 +96,26 @@ class Decoder(nn.Module): def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor: """ - Прямой проход через декодер. + Один прямой проход через Transformer decoder block. - Вход: - x: torch.Tensor - входной тензор [batch_size, seq_len, emb_size] - mask: torch.Tensor (optional) - маска внимания [seq_len, seq_len] + Аргументы: + ---------- + x : torch.Tensor + Входной тензор [batch_size, seq_len, emb_size] + mask : torch.Tensor, optional + Attention/causal mask (по умолчанию None, тогда будет casual mask по длине seq_len) Возвращает: - torch.Tensor - выходной тензор [batch_size, seq_len, emb_size] + ----------- + out : torch.Tensor + Выходной тензор той же формы, что и x - Алгоритм forward: - 1. Применяем MultiHeadAttention к входу - 2. Добавляем residual connection и LayerNorm - 3. Применяем FeedForward сеть - 4. Добавляем residual connection и LayerNorm + Алгоритм: + --------- + - Применяем attention к нормализованному входу (layernorm) + - Добавляем residual-связь (attention + исходный вход) + - Применяем FFN к нормализованному результату (layernorm) + - Добавляем residual-связь (ffn + предыдущий выход) """ # Self-Attention блок attention, _ = self._heads(x, mask, use_cache=False, cache=None)