From c338556cfe0c45743c5c02460caba8a8627370b3 Mon Sep 17 00:00:00 2001 From: Sergey Penkovsky Date: Thu, 16 Oct 2025 12:47:47 +0300 Subject: [PATCH] docs(core): improve and expand docstrings for FeedForward module - docs: rewrite and clarify docstrings for FeedForward class and its methods (__init__, forward) with architectural explanation, pseudocode, motivation, parameter details, usage example, and key references (GELU, SwiGLU, Transformer) - no changes to logic or APIs This makes the feed-forward block more transparent for users and researchers working with transformer models. --- llm/src/llm/core/feed_forward.py | 112 +++++++++++++++++++------------ 1 file changed, 70 insertions(+), 42 deletions(-) diff --git a/llm/src/llm/core/feed_forward.py b/llm/src/llm/core/feed_forward.py index f13018f..f329712 100644 --- a/llm/src/llm/core/feed_forward.py +++ b/llm/src/llm/core/feed_forward.py @@ -6,53 +6,71 @@ from .gelu import GELU class FeedForward(nn.Module): """ - Классический слой прямого распространения (FeedForward, или FFN) для архитектуры Transformer. + FeedForward — классический позиционно-независимый блок для Transformer, применяется к каждому токену отдельно. - Этот слой состоит из двух линейных преобразований с расширением внутренней размерности - в 4 раза и механизмом dropout для регуляризации. Между линейными слоями применяется - активация ReLU. + Назначение и роль: + ------------------ + - Реализует двухслойную (или более сложную) нейронную сеть, которая обрабатывает каждый токен ПОРЯДОЧНО независимо (по последней измерении). + - Дает модели "нелинейную мощность": любой токен может быть переосмыслен вне глобального контекста. + - После слоя внимания (MHA) FFN помогает связать смысл локальных (внутри токена) “скрытых” значений. - Научная суть: - - После внимания каждому токену применяется одинаковая двухслойная нейросеть. - - Дает глубокую нелинейность; позволяет модели не только сопоставлять, но и моделировать сложные связи между токенами. - - Изначально предложен в «Attention is All You Need» (Vaswani et al., 2017). + Архитектурные детали: + --------------------- + - Обычно используется блок: (Linear → Activation → Dropout → Linear → Dropout) + - В современных LLM обычно в 4 раза расширяют скрытый слой (inner_dim = 4 * emb_size). + - Активация часто GELU или SiLU (Swish), иногда SwiGLU, ReGLU, GeGLU (см. PaLM, Llama). - Формула: - FFN(x) = Dropout(W2·act(W1·x)) - где act — ReLU, GELU и др., обычно expansion x4. + Формула (обычная версия): + ------------------------- + FFN(x) = Linear2(Dropout(Activation(Linear1(x)))) + где Linear1: [emb_size → 4*emb_size], Activation: GELU/SiLU, Linear2: [4*emb_size → emb_size] - Алгоритм работы: - 1. Входной тензор x (размерность: [batch_size, seq_len, emb_size]) - 2. Линейное преобразование: emb_size -> 4*emb_size - 3. Активация ReLU - 4. Линейное преобразование: 4*emb_size -> emb_size - 5. Применение dropout - 6. Возврат результата (размерность: [batch_size, seq_len, emb_size]) + Параметры конструктора: + ----------------------- + emb_size: int — размерность входа/выхода токена + inner_dim: int (необязательно) — размер скрытого слоя (по умолчанию 4*emb_size) + activation: str — тип активации ('gelu', 'silu', 'relu', ...), см. варианты ниже + dropout: float — dropout после каждой линейной проекции - Предназначение: - - Добавляет нелинейность в архитектуру трансформера - - Обеспечивает взаимодействие между различными размерностями эмбеддингов - - Работает независимо для каждого токена в последовательности + Пример использования: + --------------------- + >>> ffn = FeedForward(emb_size=256, dropout=0.1, activation='gelu') + >>> x = torch.randn(2, 32, 256) # [batch, seq_len, emb_size] + >>> y = ffn(x) + >>> print(y.shape) # torch.Size([2, 32, 256]) - Args: - emb_size (int): размерность входных эмбеддингов - dropout (float): вероятность(dropout) - activation (str): нелинейная функция (relu, gelu, gelu_exact) + Пояснения: + ---------- + - FeedForward не использует позицию токена — это МLP, применяемый к каждому токену независимо. + - Длина последовательности и размер батча не имеют значения (broadcast/reshape по [-2, -1]). + - Используется во всех декодерах/энкодерах трансформеров. + + Подробнее смотри: + ----------------- + - Vaswani et al., "Attention is All You Need": https://arxiv.org/abs/1706.03762 + - GELU: https://arxiv.org/abs/1606.08415 + - SwiGLU (PaLM, Llama): https://arxiv.org/abs/2002.05202 - Пример: - >>> 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"): """ - Инициализация слоя Feed Forward Network. + Инициализация FeedForward блока для трансформера. - Args: - emb_size: Размерность входных эмбеддингов - dropout: Вероятность dropout для регуляризации (по умолчанию: 0.1) + Аргументы: + ---------- + emb_size: int + Размерность входного и выходного эмбеддинга модели. + dropout: float, по умолчанию 0.1 + Dropout после линии и/или активации (уменьшает переобучение). + activation: str, по умолчанию 'gelu' + Какая нелинейность использовать ('gelu', 'silu', 'relu' и т.д.). + inner_dim: int, опционально + Размер скрытого слоя (по умолчанию 4 * emb_size, как в оригинальном Transformer). + + Внутри: + ------- + - Задает структуру: Linear → Activation → Dropout → Linear → Dropout. """ super().__init__() # Первый линейный слой (расширение размерности) @@ -73,13 +91,23 @@ class FeedForward(nn.Module): def forward(self, x: torch.Tensor): """ - Прямой проход через слой Feed Forward Network. - - Args: - x: Входной тензор размерности [batch_size, seq_len, emb_size] - - Returns: - Тензор той же размерности, что и входной + Прямой проход через FeedForward блок. + + Аргументы: + ---------- + x : torch.Tensor + Входной тензор формы [..., emb_size] (используется на каждом токене отдельно!) + + Возвращает: + ----------- + torch.Tensor — выход такой же формы, как вход (только последняя размерность сохраняется). + + Пример: + ------- + >>> ffn = FeedForward(emb_size=256) + >>> x = torch.randn(8, 16, 256) + >>> y = ffn(x) + >>> y.shape # [8, 16, 256] """ # Сохраняем dtype входных данных input_dtype = x.dtype