Pytorch Lightning

Высокоуровневый фреймворк для PyTorch с классом Trainer, автоматическим распределённым обучением (DDP/FSDP/DeepSpeed), системой колбэков и минимальным шаблонным кодом. Масштабируется от ноутбука до суперкомпьютера без изменения кода. Используйте, когда нужны чистые циклы обучения со встроенными лучшими практиками.

Метаданные навыка


Источник	Опционально — установка: `vibeos skills install official/mlops/pytorch-lightning`
Путь	`optional-skills/mlops/pytorch-lightning`
Версия	`1.0.0`
Автор	Orchestra Research
Лицензия	MIT
Зависимости	`lightning`, `torch`, `transformers`
Платформы	linux, macos, windows
Теги	`PyTorch Lightning`, `Training Framework`, `Distributed Training`, `DDP`, `FSDP`, `DeepSpeed`, `High-Level API`, `Callbacks`, `Best Practices`, `Scalable`

Справочник: полный SKILL.md

к сведению

Ниже приведено полное описание навыка, которое VibeOS загружает при его активации. Это те инструкции, которые видит агент, когда навык активен.

PyTorch Lightning — высокоуровневый фреймворк для обучения

Быстрый старт

PyTorch Lightning организует код PyTorch, устраняя шаблонный код и сохраняя гибкость.

Установка:

pip install lightning

Преобразование PyTorch в Lightning (3 шага):

import lightning as L
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset

# Шаг 1: Определите LightningModule (организуйте код PyTorch)
class LitModel(L.LightningModule):
    def __init__(self, hidden_size=128):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(28 * 28, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 10)
        )

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = nn.functional.cross_entropy(y_hat, y)
        self.log('train_loss', loss)  # Автоматически логируется в TensorBoard
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

# Шаг 2: Создайте данные
train_loader = DataLoader(train_dataset, batch_size=32)

# Шаг 3: Обучайте с Trainer (всё остальное он берёт на себя!)
trainer = L.Trainer(max_epochs=10, accelerator='gpu', devices=2)
model = LitModel()
trainer.fit(model, train_loader)

Вот и всё! Trainer обрабатывает:

Переключение GPU/TPU/CPU
Распределённое обучение (DDP, FSDP, DeepSpeed)
Смешанную точность (FP16, BF16)
Накопление градиентов
Сохранение контрольных точек
Логирование
Индикаторы прогресса

Типовые рабочие процессы

Рабочий процесс 1: От PyTorch к Lightning

Исходный код PyTorch:

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
model.to('cuda')

for epoch in range(max_epochs):
    for batch in train_loader:
        batch = batch.to('cuda')
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()

Версия Lightning:

class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = MyModel()

    def training_step(self, batch, batch_idx):
        loss = self.model(batch)  # .to('cuda') не нужен!
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())

# Обучение
trainer = L.Trainer(max_epochs=10, accelerator='gpu')
trainer.fit(LitModel(), train_loader)

Преимущества: 40+ строк → 15 строк, без управления устройствами, автоматическое распределение

Рабочий процесс 2: Валидация и тестирование

class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = MyModel()

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = nn.functional.cross_entropy(y_hat, y)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        val_loss = nn.functional.cross_entropy(y_hat, y)
        acc = (y_hat.argmax(dim=1) == y).float().mean()
        self.log('val_loss', val_loss)
        self.log('val_acc', acc)

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        test_loss = nn.functional.cross_entropy(y_hat, y)
        self.log('test_loss', test_loss)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

# Обучение с валидацией
trainer = L.Trainer(max_epochs=10)
trainer.fit(model, train_loader, val_loader)

# Тестирование
trainer.test(model, test_loader)

Автоматические возможности:

Валидация запускается каждую эпоху по умолчанию
Метрики логируются в TensorBoard
Сохранение лучшей модели на основе val_loss

Рабочий процесс 3: Распределённое обучение (DDP)

# Тот же код, что и для одного GPU!
model = LitModel()

# 8 GPU с DDP (автоматически!)
trainer = L.Trainer(
    accelerator='gpu',
    devices=8,
    strategy='ddp'  # Или 'fsdp', 'deepspeed'
)

trainer.fit(model, train_loader)

Запуск:

# Одна команда, Lightning делает всё остальное
python train.py

Никаких изменений не требуется:

Автоматическое распределение данных
Синхронизация градиентов
Поддержка нескольких узлов (просто укажите num_nodes=2)

Рабочий процесс 4: Колбэки для мониторинга

from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor

# Создание колбэков
checkpoint = ModelCheckpoint(
    monitor='val_loss',
    mode='min',
    save_top_k=3,
    filename='model-{epoch:02d}-{val_loss:.2f}'
)

early_stop = EarlyStopping(
    monitor='val_loss',
    patience=5,
    mode='min'
)

lr_monitor = LearningRateMonitor(logging_interval='epoch')

# Добавление в Trainer
trainer = L.Trainer(
    max_epochs=100,
    callbacks=[checkpoint, early_stop, lr_monitor]
)

trainer.fit(model, train_loader, val_loader)

Результат:

Автосохранение 3 лучших моделей
Досрочная остановка при отсутствии улучшений в течение 5 эпох
Логирование скорости обучения в TensorBoard

Рабочий процесс 5: Планирование скорости обучения

class LitModel(L.LightningModule):
    # ... (training_step и т.д.)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)

        # Косинусное затухание
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=100,
            eta_min=1e-5
        )

        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'interval': 'epoch',  # Обновление каждую эпоху
                'frequency': 1
            }
        }

# Скорость обучения логируется автоматически!
trainer = L.Trainer(max_epochs=100)
trainer.fit(model, train_loader)

Когда использовать вместо альтернатив

Используйте PyTorch Lightning когда:

Нужен чистый, организованный код
Требуются готовые к продакшену циклы обучения
Планируется переключение между одним GPU, несколькими GPU, TPU
Нужны встроенные колбэки и логирование
Работа в команде (стандартизированная структура)

Ключевые преимущества:

Организованность: Отделяет исследовательский код от инженерного
Автоматизация: DDP, FSDP, DeepSpeed одной строкой
Колбэки: Модульные расширения для обучения
Воспроизводимость: Меньше шаблонного кода = меньше ошибок
Проверенность: 1M+ загрузок в месяц, проверен в бою

Используйте альтернативы:

Accelerate: Минимальные изменения существующего кода, больше гибкости
Ray Train: Оркестрация нескольких узлов, настройка гиперпараметров
Чистый PyTorch: Максимальный контроль, для обучения
Keras: Экосистема TensorFlow

Частые проблемы

Проблема: Потери не уменьшаются

Проверьте данные и настройку модели:

# Добавьте в training_step
def training_step(self, batch, batch_idx):
    if batch_idx == 0:
        print(f"Форма батча: {batch[0].shape}")
        print(f"Метки: {batch[1]}")
    loss = ...
    return loss

Проблема: Не хватает памяти

Уменьшите размер батча или используйте накопление градиентов:

trainer = L.Trainer(
    accumulate_grad_batches=4,  # Эффективный батч = batch_size × 4
    precision='bf16'  # Или 'fp16', уменьшает память на 50%
)

Проблема: Валидация не запускается

Убедитесь, что передаёте val_loader:

# НЕПРАВИЛЬНО
trainer.fit(model, train_loader)

# ПРАВИЛЬНО
trainer.fit(model, train_loader, val_loader)

Проблема: DDP неожиданно запускает несколько процессов

Lightning автоматически определяет GPU. Явно укажите устройства:

# Сначала тест на CPU
trainer = L.Trainer(accelerator='cpu', devices=1)

# Затем GPU
trainer = L.Trainer(accelerator='gpu', devices=1)

Продвинутые темы

Колбэки: См. references/callbacks.md для EarlyStopping, ModelCheckpoint, пользовательских колбэков и хуков колбэков.

Стратегии распределения: См. references/distributed.md для DDP, FSDP, интеграции DeepSpeed ZeRO, настройки нескольких узлов.

Настройка гиперпараметров: См. references/hyperparameter-tuning.md для интеграции с Optuna, Ray Tune и WandB sweeps.

Требования к оборудованию

CPU: Работает (хорошо для отладки)
Один GPU: Работает
Несколько GPU: DDP (по умолчанию), FSDP или DeepSpeed
Несколько узлов: DDP, FSDP, DeepSpeed
TPU: Поддерживается (8 ядер)
Apple MPS: Поддерживается

Варианты точности:

FP32 (по умолчанию)
FP16 (V100, старые GPU)
BF16 (A100/H100, рекомендуется)
FP8 (H100)

Ресурсы

Документация: https://lightning.ai/docs/pytorch/stable/
GitHub: https://github.com/Lightning-AI/pytorch-lightning ⭐ 29,000+
Версия: 2.5.5+
Примеры: https://github.com/Lightning-AI/pytorch-lightning/tree/master/examples
Discord: https://discord.gg/lightning-ai
Используется: победителями Kaggle, исследовательскими лабораториями, продакшен-командами

Метаданные навыка​

Справочник: полный SKILL.md​

PyTorch Lightning — высокоуровневый фреймворк для обучения

Быстрый старт​

Типовые рабочие процессы​

Рабочий процесс 1: От PyTorch к Lightning​

Рабочий процесс 2: Валидация и тестирование​

Рабочий процесс 3: Распределённое обучение (DDP)​

Рабочий процесс 4: Колбэки для мониторинга​

Рабочий процесс 5: Планирование скорости обучения​

Когда использовать вместо альтернатив​

Частые проблемы​

Продвинутые темы​

Требования к оборудованию​

Ресурсы​