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

@@ -6,53 +6,71 @@ from .gelu import GELU
 
 class FeedForward(nn.Module):
     """
-    Классический слой прямого распространения (FeedForward, или FFN) для архитектуры Transformer.
+    FeedForward — классический позиционно-независимый блок для Transformer, применяется к каждому токену отдельно.
 
-    Этот слой состоит из двух линейных преобразований с расширением внутренней размерности
+    Назначение и роль:
-    в 4 раза и механизмом dropout для регуляризации. Между линейными слоями применяется
+    ------------------
-    активация ReLU.
+    - Реализует двухслойную (или более сложную) нейронную сеть, которая обрабатывает каждый токен ПОРЯДОЧНО независимо (по последней измерении).
+    - Дает модели "нелинейную мощность": любой токен может быть переосмыслен вне глобального контекста.
+    - После слоя внимания (MHA) FFN помогает связать смысл локальных (внутри токена) “скрытых” значений.
 
-    Научная суть:
+    Архитектурные детали:
-        - После внимания каждому токену применяется одинаковая двухслойная нейросеть.
+    ---------------------
-        - Дает глубокую нелинейность; позволяет модели не только сопоставлять, но и моделировать сложные связи между токенами.
+    - Обычно используется блок: (Linear → Activation → Dropout → Linear → Dropout)
-        - Изначально предложен в «Attention is All You Need» (Vaswani et al., 2017).
+    - В современных 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
+    emb_size: int — размерность входа/выхода токена
-    3. Активация ReLU
+    inner_dim: int (необязательно) — размер скрытого слоя (по умолчанию 4*emb_size)
-    4. Линейное преобразование: 4*emb_size -> emb_size
+    activation: str — тип активации ('gelu', 'silu', 'relu', ...), см. варианты ниже
-    5. Применение dropout
+    dropout: float — dropout после каждой линейной проекции
-    6. Возврат результата (размерность: [batch_size, seq_len, emb_size])
 
-    Предназначение:
+    Пример использования:
-    - Добавляет нелинейность в архитектуру трансформера
+    ---------------------
-    - Обеспечивает взаимодействие между различными размерностями эмбеддингов
+        >>> 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)
+    - FeedForward не использует позицию токена — это МLP, применяемый к каждому токену независимо.
-        activation (str): нелинейная функция (relu, gelu, gelu_exact)
+    - Длина последовательности и размер батча не имеют значения (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.
+        Прямой проход через FeedForward блок.
     
-        Args:
+        Аргументы:
-            x: Входной тензор размерности [batch_size, seq_len, emb_size]
+        ----------
+        x : torch.Tensor
+            Входной тензор формы [..., emb_size] (используется на каждом токене отдельно!)
     
-        Returns:
+        Возвращает:
-            Тензор той же размерности, что и входной
+        -----------
+        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