From d947b7beb3ff5a2dda3a953ccb8b96fdb81fc877 Mon Sep 17 00:00:00 2001 From: Sergey Penkovsky Date: Fri, 17 Oct 2025 16:23:43 +0300 Subject: [PATCH] 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 | 64 ++++++++++++++- llm/src/llm/training/scheduler.py | 54 ++++++++++++- llm/src/llm/training/trainer.py | 116 ++++++++++++++++++++++++--- llm/tests/training/test_optimizer.py | 35 ++++++++ llm/tests/training/test_scheduler.py | 62 ++++++++++++++ llm/tests/training/test_trainer.py | 62 ++++++++++++++ 6 files changed, 380 insertions(+), 13 deletions(-) create mode 100644 llm/tests/training/test_optimizer.py create mode 100644 llm/tests/training/test_scheduler.py create mode 100644 llm/tests/training/test_trainer.py diff --git a/llm/src/llm/training/optimizer.py b/llm/src/llm/training/optimizer.py index 0ee6359..2ea312a 100644 --- a/llm/src/llm/training/optimizer.py +++ b/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) diff --git a/llm/src/llm/training/scheduler.py b/llm/src/llm/training/scheduler.py index 9105a12..cf7a165 100644 --- a/llm/src/llm/training/scheduler.py +++ b/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) diff --git a/llm/src/llm/training/trainer.py b/llm/src/llm/training/trainer.py index 5ed16e9..77f90e3 100644 --- a/llm/src/llm/training/trainer.py +++ b/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 @@ -113,4 +211,4 @@ class Trainer: total_loss += loss.item() avg_loss = total_loss / len(self.val_loader) - print(f"Validation loss: {avg_loss:.4f}") + print(f"Validation loss: {avg_loss:.4f}") \ No newline at end of file diff --git a/llm/tests/training/test_optimizer.py b/llm/tests/training/test_optimizer.py new file mode 100644 index 0000000..9d696c2 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/llm/tests/training/test_scheduler.py b/llm/tests/training/test_scheduler.py new file mode 100644 index 0000000..c010805 --- /dev/null +++ b/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 diff --git a/llm/tests/training/test_trainer.py b/llm/tests/training/test_trainer.py new file mode 100644 index 0000000..3c51a23 --- /dev/null +++ b/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