update and expand scientific docstrings for optimizer, scheduler, trainer

- Expanded module-level and function/class docstrings in optimizer.py, scheduler.py, and trainer.py
- Described mathematical foundations, theoretical motivations, and provided detailed usage examples for students
- All docstrings in Russian, clear scientific style

test(training): add comprehensive tests for optimizer, scheduler, and trainer modules

- Added new test files for get_optimizer, get_linear_schedule_with_warmup, and Trainer
- Tests cover parameter handling, edge cases, and expected learning dynamics (lr schedules and loss behavior)
- Trainer now logs average epoch losses to self.loss_history for testability and analysis

refactor(training/trainer): log epoch loss to loss_history for downstream analysis and tests

BREAKING CHANGE: Trainer.loss_history is a new attribute consolidating average losses per epoch, enabling robust learning dynamics assertions in tests

llm/src/llm/training/optimizer.py

@@ -1,9 +1,71 @@
+"""
+Модуль оптимизации для обучения нейронных сетей.
+
+В данном модуле реализована функция выбора и инициализации оптимизаторов, наиболее популярных при обучении глубоких нейросетей:
+- AdamW
+- Adam
+- SGD
+
+Теоретическое обоснование:
+--------------------------
+Задача оптимизации в обучении нейросети заключается в минимизации функции потерь (Loss) по параметрам модели W. Современные методы базируются на стохастическом градиентном спуске (SGD), а также на его адаптивных модификациях (Adam, AdamW).
+
+**SGD** (Stochastic Gradient Descent) — стохастический градиентный спуск:
+  W_{t+1} = W_t - \eta \nabla_W L(W_t)
+  Здесь \eta — шаг обучения, \nabla_W — градиент по параметрам. SGD позволяет случайно выбирать подмножество обучающих данных для каждой итерации, что ускоряет процесс и уменьшает избыточную корреляцию между примерами.
+
+**Adam** (Adaptive Moment Estimation) — адаптивный алгоритм, который использует скользящую среднюю не только градиентов, но и их квадратов:
+  m_t = \beta_1 m_{t-1} + (1-\beta_1) \nabla_W L(W_t)
+  v_t = \beta_2 v_{t-1} + (1-\beta_2) (\nabla_W L(W_t))^2
+  W_{t+1} = W_t - \eta m_t/(\sqrt{v_t}+\epsilon)
+  Где \beta_1, \beta_2 — коэффициенты экспоненциального сглаживания.
+
+**AdamW** — модификация Adam, в которой weight decay (имплицитная L2-регуляризация) вводится корректно, отдельно от шага градиента, что улучшает обобщающую способность моделей:
+  W_{t+1} = W_t - \eta [ m_t/(\sqrt{v_t}+\epsilon) + \lambda W_t ]
+  Где \lambda — коэффициент weight decay.
+
+Детальное описание: https://arxiv.org/abs/1711.05101
+
+Пример использования:
+---------------------
+>>> optimizer = get_optimizer(model, lr=3e-4, weight_decay=0.01, optimizer_type="adamw")
+>>> for batch in dataloader:
+...     loss = model(batch)
+...     loss.backward()
+...     optimizer.step()
+...     optimizer.zero_grad()
+
+"""
 import torch.optim as optim
 
 
 def get_optimizer(model, lr=3e-4, weight_decay=0.01, optimizer_type="adamw"):
     """
-    Возвращает оптимизатор для обучения модели.
+    Фабричная функция для создания оптимизатора PyTorch по выбранному типу.
+    
+    Параметры
+    ---------
+    model : torch.nn.Module
+        Модель, параметры которой требуется оптимизировать.
+    lr : float, по умолчанию 3e-4
+        Шаг обучения (learning rate).
+    weight_decay : float, по умолчанию 0.01
+        Коэффициент weight decay (L2-регуляризации).
+    optimizer_type : str, по умолчанию 'adamw'
+        Тип оптимизатора: 'adamw', 'adam' или 'sgd'.
+    
+    Возвращаемое значение
+    ---------------------
+    torch.optim.Optimizer
+        Объект-оптимизатор, готовый к использованию.
+
+    Исключения
+    ----------
+    ValueError: Если передан неизвестный тип оптимизатора.
+
+    Пример использования:
+    ---------------------
+    >>> optimizer = get_optimizer(model, lr=1e-3, optimizer_type='sgd')
     """
     if optimizer_type.lower() == "adamw":
         return optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)

llm/src/llm/training/scheduler.py

@@ -1,18 +1,66 @@
-from torch.optim.lr_scheduler import LambdaLR
+"""
+Модуль для управления динамикой шага обучения (learning rate scheduling) при обучении нейронных сетей.
 
+Теоретическое обоснование:
+--------------------------
+Плавная динамика шага обучения существенно влияет на сходимость и итоговое качество моделей. Введение этапа "разогрева" (warmup) — техники, при которой шаг обучения начинается с нуля и постепенно увеличивается до целевого значения, снижает вероятность неустойчивых градиентов на старте обучения. Подобная стратегия показала свою эффективность для крупных нейронных сетей, особенно в трансформерах (Vaswani et al, 2017, https://arxiv.org/abs/1706.03762).
+
+Линейный scheduler с warmup задаёт динамику learning rate по формуле:
+  - если current_step < num_warmup_steps:
+        lr = lr_init * (current_step / num_warmup_steps)
+  - иначе:
+        lr = lr_init * max(0, (num_training_steps - current_step) / (num_training_steps - num_warmup_steps))
+
+Пример использования:
+---------------------
+>>> optimizer = get_optimizer(model, lr=3e-4)
+>>> scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=1000, num_training_steps=10000)
+>>> for step in range(num_training_steps):
+...     optimizer.step()
+...     scheduler.step()
+"""
+
+from torch.optim.lr_scheduler import LambdaLR
 
 def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
     """
-    Линейный планировщик обучения с warmup.
+    Создаёт линейный планировщик изменения шага обучения (learning rate) с этапом warmup для оптимизатора PyTorch.
+
+    Аргументы
+    ---------
+    optimizer : torch.optim.Optimizer
+        Оптимизатор, для которого применяется scheduler.
+    num_warmup_steps : int
+        Количество шагов разогрева (warmup) — начиная с нулевого шага и плавного увеличения lr до номинального значения.
+    num_training_steps : int
+        Общее количество шагов (эпох/итераций) обучения модели.
+
+    Возвращаемое значение
+    ---------------------
+    torch.optim.lr_scheduler.LambdaLR
+        Планировщик lr, который следует вызывать после каждого optimizer.step() во время обучения.
+
+    Теоретическая справка
+    ---------------------
+    Такой scheduler позволяет повысить стабильность и устойчивость обучения крупных моделей (особенно трансформеров), предотвращая резкие скачки градиентов в начале.
+    
+    Пример:
+    -------
+    >>> optimizer = get_optimizer(model, lr=3e-4)
+    >>> scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=1000, num_training_steps=10000)
+    >>> for step in range(num_training_steps):
+    ...     optimizer.step()
+    ...     scheduler.step()
     """
 
     def lr_lambda(current_step):
+        # Линейный рост lr на этапе разогрева
         if current_step < num_warmup_steps:
             return float(current_step) / float(max(1, num_warmup_steps))
+        # Линейное затухание lr после разогрева
         return max(
             0.0,
             float(num_training_steps - current_step)
             / float(max(1, num_training_steps - num_warmup_steps)),
         )
-
     return LambdaLR(optimizer, lr_lambda)

llm/src/llm/training/trainer.py

@@ -1,3 +1,22 @@
+"""
+Модуль для организации процесса обучения больших языковых моделей (LLM).
+
+Научное и техническое обоснование
+----------------------------------
+Эффективное обучение современных трансформеров (GPT, LLaMA, Mistral и др.) опирается на принципы языкового моделирования (Language Modeling):
+- Предсказание вероятности следующего токена на основе предыдущих.
+- Использование функции потерь кросс-энтропии (cross-entropy) с маскированием паддингов.
+- Циклы обратного распространения ошибки (backpropagation), оптимизационные алгоритмы (например, AdamW), управление шагом обучения (scheduler с warmup), обрезка градиентов (grad clipping).
+
+Реализация объединяет лучшие практики обучения LLM, универсальный API к моделям, датасетам, оптимизаторам и lr-схемам.
+
+Подробнее: Vaswani et al. "Attention is All You Need" (2017), Radford et al. "Language Models are Unsupervised Multitask Learners" (2019)
+
+Пример использования
+--------------------
+>>> trainer = Trainer(model, train_dataset, val_dataset, lr=3e-4, batch_size=8, num_epochs=3, warmup_steps=100)
+>>> trainer.train()
+"""
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
@@ -8,7 +27,33 @@ from llm.training.scheduler import get_linear_schedule_with_warmup
 
 class Trainer:
     """
-    Универсальный класс обучения LLM (GPT, LLaMA, Mistral и т.д.)
+    Универсальный и расширяемый класс для обучения больших языковых моделей (Large Language Models, LLM).
+
+    Поддерживаются архитектуры семейства GPT, LLaMA, Mistral и другие автогрессивные модели.
+    Объединяет:
+      - Тренировку по задаче языкового моделирования (Causal LM)
+      - Cross-entropy loss с автоматическим сдвигом логитов/меток
+      - Поддержку Grad Clipping, Scheduler, Validation
+      - Унифицированный даталоадер, автоматический выбор устройства (CPU/GPU)
+
+    Атрибуты
+    --------
+    model : torch.nn.Module
+        Модель для обучения языковому моделированию
+    train_loader : torch.utils.data.DataLoader
+        Даталоадер обучающего набора
+    val_loader : torch.utils.data.DataLoader или None
+        Даталоадер валидационного набора (если задан)
+    optimizer : torch.optim.Optimizer
+        Оптимизатор параметров модели
+    scheduler : torch.optim.lr_scheduler.LambdaLR
+        Планировщик learning rate (инициализируется в train)
+    device : torch.device
+        Устройство (CPU или CUDA), куда помещается модель
+    num_epochs : int
+        Количество эпох обучения
+    warmup_steps : int
+        Число шагов warmup для scheduler
     """
 
     def __init__(
@@ -21,6 +66,26 @@ class Trainer:
         num_epochs=3,
         warmup_steps=100,
     ):
+        """
+        Инициализация обучающего класса Trainer.
+
+        Аргументы
+        ---------
+        model : torch.nn.Module
+            Модель для обучения (например, GPT, LLaMA, Mistral).
+        train_dataset : torch.utils.data.Dataset
+            Обучающий датасет с полями input_ids и labels.
+        val_dataset : torch.utils.data.Dataset, optional
+            Валидационный датасет для контроля качества обучения.
+        lr : float, default=3e-4
+            Начальный шаг обучения.
+        batch_size : int, default=8
+            Размер обучающего мини-батча.
+        num_epochs : int, default=3
+            Количество эпох обучения.
+        warmup_steps : int, default=100
+            Количество шагов разогрева (warmup) learning rate.
+        """
         self.model = model
         self.train_loader = DataLoader(
             train_dataset, batch_size=batch_size, shuffle=True
@@ -37,26 +102,52 @@ class Trainer:
 
     def compute_lm_loss(self, logits, labels):
         """
-        Вычисляет loss для языкового моделирования.
+        Вычисляет функцию потерь (loss) для задачи автогрессивного языкового моделирования.
-        Сдвигает логиты и метки для предсказания следующего токена.
+
+        Производит сдвиг логитов и меток: предсказания делаются для следующего токена.
+        Используется кросс-энтропия (CrossEntropyLoss), что соответствует максимизации логарифма правдоподобия:
+            L = -log P(w_{t+1} | w_1,...,w_t)
+
+        Аргументы
+        ---------
+        logits : torch.Tensor
+            Логиты модели: (batch_size, seq_len, vocab_size)
+        labels : torch.Tensor
+            Правильные метки: (batch_size, seq_len)
+        Возвращаемое значение
+        ---------------------
+        loss : torch.Tensor
+            Средний loss по batch.
         """
-        # Сдвигаем логиты и метки для языкового моделирования
+        # Сдвигаем логиты и метки для языкового моделирования (автогрессия)
         shift_logits = logits[..., :-1, :].contiguous()
         shift_labels = labels[..., 1:].contiguous()
 
-        # Вычисляем cross-entropy loss
+        # CrossEntropyLoss (игнорируем паддинги: ignore_index=-100)
         loss = F.cross_entropy(
             shift_logits.view(-1, shift_logits.size(-1)),
             shift_labels.view(-1),
-            ignore_index=-100,  # Игнорируем padding tokens
+            ignore_index=-100,  # Padding токены не участвуют в loss
         )
         return loss
 
     def train(self):
+        """
+        Запускает процесс обучения модели по заданному числу эпох.
+
+        В процессе:
+        - Применяет optimizer, scheduler с warmup и decay, grad clipping (обрезка градиентов)
+        - Вызывает функцию потерь для языкового моделирования
+        - Показывает динамику процесса (tqdm)
+        - После каждой эпохи возможно проведение валидации
+
+        Параметры задаются на этапе инициализации Trainer.
+        """
         total_steps = len(self.train_loader) * self.num_epochs
         self.scheduler = get_linear_schedule_with_warmup(
             self.optimizer, self.warmup_steps, total_steps
         )
+        self.loss_history = []  # добавлено: лог средних потерь
 
         for epoch in range(self.num_epochs):
             self.model.train()
@@ -71,14 +162,14 @@ class Trainer:
                 input_ids = batch["input_ids"].to(self.device)
                 labels = batch["labels"].to(self.device)
 
-                # Универсально обрабатываем выход (tuple/logits)
+                # Универсально обрабатываем выходы модели: tuple или просто tensor (logits)
                 outputs = self.model(input_ids)
                 if isinstance(outputs, tuple):
                     logits = outputs[0]
                 else:
                     logits = outputs
 
-                # Trainer вычисляет loss
+                # Вычисляем loss автогрессивной LM-задачи
                 loss = self.compute_lm_loss(logits, labels)
                 loss.backward()
 
@@ -90,12 +181,19 @@ class Trainer:
                 progress_bar.set_postfix(loss=loss.item())
 
             avg_loss = total_loss / len(self.train_loader)
+            self.loss_history.append(avg_loss)  # добавлено: запоминаем loss
             print(f"Epoch {epoch+1} finished — avg loss: {avg_loss:.4f}")
 
             if self.val_loader:
                 self.evaluate()
 
     def evaluate(self):
+        """
+        Оценивает модель на валидационном датасете (если задан).
+
+        В режиме eval() модели отключается dropout и все стохастические элементы.
+        Возвращает среднее значение функции потерь (loss) по всему validation set.
+        """
         self.model.eval()
         total_loss = 0
 

llm/tests/training/test_optimizer.py

@@ -0,0 +1,35 @@
+import pytest
+import torch.nn as nn
+from llm.training.optimizer import get_optimizer
+
+class DummyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 1)
+
+def test_get_optimizer_adamw():
+    model = DummyModel()
+    optimizer = get_optimizer(model, lr=1e-3, weight_decay=0.02, optimizer_type="adamw")
+    assert optimizer.__class__.__name__ == 'AdamW'
+    assert optimizer.defaults['lr'] == 1e-3
+    assert optimizer.defaults['weight_decay'] == 0.02
+
+def test_get_optimizer_adam():
+    model = DummyModel()
+    optimizer = get_optimizer(model, lr=1e-4, weight_decay=0.01, optimizer_type="adam")
+    assert optimizer.__class__.__name__ == 'Adam'
+    assert optimizer.defaults['lr'] == 1e-4
+    assert optimizer.defaults['weight_decay'] == 0.01
+
+def test_get_optimizer_sgd():
+    model = DummyModel()
+    optimizer = get_optimizer(model, lr=0.1, optimizer_type="sgd")
+    assert optimizer.__class__.__name__ == 'SGD'
+    assert optimizer.defaults['lr'] == 0.1
+    # SGD: weight_decay по умолчанию 0 для этого вызова
+    assert optimizer.defaults['momentum'] == 0.9
+
+def test_get_optimizer_invalid():
+    model = DummyModel()
+    with pytest.raises(ValueError):
+        get_optimizer(model, optimizer_type="nonexistent")

llm/tests/training/test_scheduler.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+from llm.training.scheduler import get_linear_schedule_with_warmup
+from llm.training.optimizer import get_optimizer
+
+class DummyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(2, 2)
+
+def test_scheduler_warmup_and_decay():
+    model = DummyModel()
+    base_lr = 0.1
+    warmup_steps = 5
+    total_steps = 20
+    optimizer = get_optimizer(model, lr=base_lr, optimizer_type="sgd")
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
+
+    lrs = [optimizer.param_groups[0]['lr']]  # lr до первого .step()
+    for _ in range(total_steps):
+        optimizer.step()
+        scheduler.step()
+        lrs.append(optimizer.param_groups[0]['lr'])
+
+    # Проверяем warmup: lr должен расти линейно в первых warmup_steps (начиная с шага 1)
+    for i in range(warmup_steps + 1):
+        expected = base_lr * min(i, warmup_steps) / max(1, warmup_steps)
+        assert abs(lrs[i] - expected) < 1e-6, f"Warmup step {i}: lr={lrs[i]}, expected={expected}"
+    # Проверяем decay: после warmup lr затухает
+    for i in range(warmup_steps + 1, total_steps + 1):
+        expected = base_lr * max(0.0, (total_steps - (i - 0)) / max(1, total_steps - warmup_steps))
+        assert abs(lrs[i] - expected) < 1e-6, f"Decay step {i}: lr={lrs[i]}, expected={expected}"
+    assert lrs[-1] == 0.0
+
+def test_scheduler_no_warmup():
+    model = DummyModel()
+    base_lr = 0.1
+    warmup_steps = 0
+    total_steps = 10
+    optimizer = get_optimizer(model, lr=base_lr, optimizer_type="adam")
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
+    lrs = [optimizer.param_groups[0]['lr']]
+    for _ in range(total_steps):
+        optimizer.step()
+        scheduler.step()
+        lrs.append(optimizer.param_groups[0]['lr'])
+
+    for i in range(total_steps + 1):
+        expected = base_lr * max(0.0, (total_steps - i) / max(1, total_steps - warmup_steps))
+        assert abs(lrs[i] - expected) < 1e-6, f"Step {i}: lr={lrs[i]}, expected={expected}"
+    assert lrs[-1] == 0.0
+
+def test_scheduler_full_decay_to_zero():
+    model = DummyModel()
+    optimizer = get_optimizer(model, lr=1.0, optimizer_type="adamw")
+    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=2, num_training_steps=2)
+    scheduler.step()
+    scheduler.step()
+    for param_group in optimizer.param_groups:
+        assert param_group['lr'] == 0.0

llm/tests/training/test_trainer.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+from llm.training.trainer import Trainer
+
+# Синтетический небольшой датасет для автогрессивной LM задачи
+class ToyLMDataset(Dataset):
+    def __init__(self, num_samples=16, seq_len=8, vocab_size=16):
+        self.data = torch.randint(1, vocab_size, (num_samples, seq_len))
+    def __len__(self):
+        return len(self.data)
+    def __getitem__(self, idx):
+        # labels == input_ids (identity task)
+        return {"input_ids": self.data[idx], "labels": self.data[idx]}
+
+# Простая dummy-модель — 1 слой linear over vocab
+class TinyModel(nn.Module):
+    def __init__(self, vocab_size=16, seq_len=8):
+        super().__init__()
+        self.linear = nn.Linear(seq_len, vocab_size)
+    def forward(self, x):
+        # logits: (batch, seq_len, vocab_size)
+        # Для простоты делаем транспонирование
+        return self.linear(x.float()).unsqueeze(1).expand(-1, x.shape[1], -1)
+
+def test_train_runs_without_errors():
+    train_data = ToyLMDataset(num_samples=16, seq_len=8, vocab_size=16)
+    model = TinyModel(vocab_size=16, seq_len=8)
+    trainer = Trainer(model, train_data, lr=1e-3, batch_size=4, num_epochs=1, warmup_steps=2)
+    trainer.train()
+
+def test_trainer_evaluate_runs():
+    train_data = ToyLMDataset(num_samples=8)
+    val_data = ToyLMDataset(num_samples=8)
+    model = TinyModel()
+    trainer = Trainer(model, train_data, val_data, lr=1e-3, batch_size=4, num_epochs=1, warmup_steps=2)
+    trainer.train()
+    trainer.evaluate()
+
+def test_trainer_tuple_output():
+    # Модель, возвращающая кортеж (logits, extra)
+    class TupleModel(nn.Module):
+        def __init__(self, vocab_size=16, seq_len=8):
+            super().__init__()
+            self.linear = nn.Linear(seq_len, vocab_size)
+        def forward(self, x):
+            logits = self.linear(x.float()).unsqueeze(1).expand(-1, x.shape[1], -1)
+            extra = torch.zeros(1)
+            return logits, extra
+
+    train_data = ToyLMDataset(num_samples=8)
+    model = TupleModel()
+    trainer = Trainer(model, train_data, lr=1e-3, batch_size=2, num_epochs=1, warmup_steps=1)
+    trainer.train()
+
+def test_trainer_loss_decreases():
+    train_data = ToyLMDataset(num_samples=32, seq_len=8, vocab_size=8)
+    model = TinyModel(vocab_size=8, seq_len=8)
+    trainer = Trainer(model, train_data, lr=0.05, batch_size=8, num_epochs=2, warmup_steps=1)
+    trainer.train()
+    avg_losses = trainer.loss_history
+    assert avg_losses[-1] <= avg_losses[0] or abs(avg_losses[-1] - avg_losses[0]) < 1e-3