Pytorch Lightning

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

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


Источник	Опциональный — установка командой `hermes 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

инфо

Ниже приведено полное определение навыка, которое Hermes загружает при его срабатывании. Именно это агент видит как инструкции, когда навык активен.

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

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

Проблема: потери (loss) не падают

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

# Добавьте в training_step
def training_step(self, batch, batch_idx):
    if batch_idx == 0:
        print(f"Batch shape: {batch[0].shape}")
        print(f"Labels: {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. Задайте devices явно:

# Сначала проверьте на 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.

Требования к железу

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, исследовательские лаборатории, продакшен-команды