PyTorch Lightning: чистые циклы обучения с поддержкой распределённых вычислений
PyTorch Lightning организует ваш код PyTorch так, чтобы устранить избыточный шаблонный код, при этом сохраняя гибкость; преобразование модели в LightningModule занимает всего три шага.
Что понадобится
- Требуются пакеты lightning, torch и transformers; работает на Linux, macOS и Windows.
Как применить
Используйте PyTorch Lightning, если вам нужен чистый, организованный код, готовые к продакшену циклы обучения, встроенные колбэки и механизмы логирования, стандартизированная структура для командной работы или простая смена режима между одной GPU, несколькими GPU и TPU. Альтернативы: Accelerate — для минимальных изменений с большей гибкостью, Ray Train — для оркестрации нескольких узлов и настройки гиперпараметров, чистый PyTorch — для максимального контроля или обучения, Keras — для экосистемы TensorFlow. Работает на CPU (удобно для отладки), одной GPU, нескольких GPU и нескольких узлах через DDP (по умолчанию), FSDP или DeepSpeed, TPU (8 ядер) и Apple MPS. Доступные варианты точности: FP32 (по умолчанию), FP16 (V100, более старые GPU), BF16 (A100/H100, рекомендуется) и FP8 (H100).
Выбирайте Lightning, если вам нужен чистый, организованный код.
Он также подходит, если вам нужны готовые к продакшену циклы обучения.
Он подходит для переключения между одной GPU, несколькими GPU и TPU.
Он поможет, если вам нужны встроенные колбэки и механизмы логирования.
Он способствует командной работе благодаря стандартизированной структуре.
В таких случаях используйте альтернативы.
Accelerate требует минимальных изменений в существующем коде и обеспечивает большую гибкость.
Ray Train подходит для оркестрации нескольких узлов и настройки гиперпараметров.
Чистый PyTorch предоставляет максимальный контроль и хорошо подходит для обучения.
Keras подходит для экосистемы TensorFlow.
CPU поддерживается и удобен для отладки.
Одна GPU поддерживается.
Для нескольких GPU используйте DDP (по умолчанию), FSDP или DeepSpeed.
Для нескольких узлов используйте DDP, FSDP или DeepSpeed.
Поддерживается TPU с 8 ядрами.
Поддерживается Apple MPS.
Вот доступные варианты точности.
FP32 — это значение по умолчанию.
FP16 подходит для V100 и более старых GPU.
BF16 рекомендуется для A100/H100.
FP8 доступен на H100.
Начните с быстрого старта, установив пакет.
Установка:
pip install lightning
Преобразуйте свой код PyTorch в Lightning за три шага.
import lightning as L
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset
# Step 1: Define LightningModule (organize your PyTorch code)
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) # Auto-logged to TensorBoard
return loss
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=1e-3)
# Step 2: Create data
train_loader = DataLoader(train_dataset, batch_size=32)
# Step 3: Train with Trainer (handles everything else!)
trainer = L.Trainer(max_epochs=10, accelerator='gpu', devices=2)
model = LitModel()
trainer.fit(model, train_loader)
Рабочий процесс 1 охватывает переход от PyTorch к Lightning, начиная с исходного кода PyTorch.
Исходный код 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) # No .to('cuda') needed!
return loss
def configure_optimizers(self):
return torch.optim.Adam(self.parameters())
# Train
trainer = L.Trainer(max_epochs=10, accelerator='gpu')
trainer.fit(LitModel(), train_loader)
Рабочий процесс 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)
# Train with validation
trainer = L.Trainer(max_epochs=10)
trainer.fit(model, train_loader, val_loader)
# Test
trainer.test(model, test_loader)
Рабочий процесс 3 охватывает распределённое обучение с использованием DDP.
# Same code as single GPU!
model = LitModel()
# 8 GPUs with DDP (automatic!)
trainer = L.Trainer(
accelerator='gpu',
devices=8,
strategy='ddp' # Or 'fsdp', 'deepspeed'
)
trainer.fit(model, train_loader)
Запустите его так:
# Single command, Lightning handles the rest
python train.py
Рабочий процесс 4 демонстрирует колбэки для мониторинга.
from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor
# Create callbacks
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')
# Add to Trainer
trainer = L.Trainer(
max_epochs=100,
callbacks=[checkpoint, early_stop, lr_monitor]
)
trainer.fit(model, train_loader, val_loader)
Рабочий процесс 5 охватывает планирование скорости обучения.
class LitModel(L.LightningModule):
# ... (training_step, etc.)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
# Cosine annealing
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=100,
eta_min=1e-5
)
return {
'optimizer': optimizer,
'lr_scheduler': {
'scheduler': scheduler,
'interval': 'epoch', # Update per epoch
'frequency': 1
}
}
# Learning rate auto-logged!
trainer = L.Trainer(max_epochs=100)
trainer.fit(model, train_loader)
Далее следуют советы по устранению неполадок.
Если потери не уменьшаются, проверьте данные и настройки модели.
Проверьте данные и настройки модели:
# Add to 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, # Effective batch = batch_size × 4
precision='bf16' # Or 'fp16', reduces memory 50%
)
Если проверка не запускается, убедитесь, что вы передали val_loader.
Убедитесь, что вы передали val_loader:
# WRONG
trainer.fit(model, train_loader)
# CORRECT
trainer.fit(model, train_loader, val_loader)
Если DDP неожиданно запускает несколько процессов, помните, что Lightning автоматически определяет GPU, поэтому явно укажите устройства.
Lightning автоматически определяет GPU, поэтому явно укажите устройства:
# Test on CPU first
trainer = L.Trainer(accelerator='cpu', devices=1)
# Then GPU
trainer = L.Trainer(accelerator='gpu', devices=1)
Ссылки
Источник: nousresearch/hermes-agent и участники проекта. Лицензия: MIT.
Copyright (c) 2025 Nous Research
Перевод и редакционная адаптация: ClawCases.
>Ресурсы
>Похожие юзкейсы
Арбитр разрешения конфликтов слияния агентов
Разрешает конфликт слияния Git между ветками двух агентов в качестве беспристрастной третьей стороны: нейтральный арбитр, который принимает обе версии изменений и указанные намерения, после чего формирует итоговый результат слияния, не перезаписывая ни одну из сторон.
Автономное управление проектами с помощью подагентов
Координация проектов автономно: главный агент делегирует работу подагентам-менеджерам проектов, которые отслеживают задачи в общем файле STATE.yaml, обновляют его по мере выполнения, параллельно берутся за неблокируемые задачи и сохраняют всё состояние задач в файлах с контролем версий.
Присоединиться к звонку Google Meet, расшифровать его и отправить последующие сообщения
Когда пользователь просит агента присутствовать на встрече Google Meet, этот сценарий подключается к звонку, сохраняет живые субтитры в расшифровку, при явном выборе может говорить в режиме реального времени, а после встречи — подготовить резюме и запрошенные последующие сообщения.