PyTorch Lightning: clean training loops with distributed support
PyTorch Lightning organizes your PyTorch code to eliminate boilerplate while maintaining flexibility, so converting a model to a LightningModule takes just three steps.
Requirements
- Requires the lightning, torch, and transformers packages, and runs on Linux, macOS, and Windows.
How to use
Use PyTorch Lightning when you want clean, organized code, production-ready training loops, built-in callbacks and logging, standardized structure for team collaboration, or easy switching between single GPU, multi-GPU, and TPU. Alternatives: Accelerate for minimal changes with more flexibility, Ray Train for multi-node orchestration and hyperparameter tuning, raw PyTorch for maximum control or learning, and Keras for the TensorFlow ecosystem. It works on CPU (good for debugging), single GPU, multi-GPU and multi-node via DDP (default), FSDP, or DeepSpeed, TPU (8 cores), and Apple MPS. Precision options are FP32 (default), FP16 (V100, older GPUs), BF16 (A100/H100, recommended), and FP8 (H100).
Choose Lightning when you want clean, organized code.
It also fits when you need production-ready training loops.
It suits switching between single GPU, multi-GPU, and TPU.
It helps if you want built-in callbacks and logging.
It supports team collaboration through a standardized structure.
Use alternatives instead in these cases.
Accelerate requires minimal changes to existing code and gives more flexibility.
Ray Train fits multi-node orchestration and hyperparameter tuning.
Raw PyTorch offers maximum control and is good for learning purposes.
Keras fits the TensorFlow ecosystem.
CPU works and is good for debugging.
Single GPU works.
For multi-GPU, use DDP (default), FSDP, or DeepSpeed.
For multi-node, use DDP, FSDP, or DeepSpeed.
TPU is supported with 8 cores.
Apple MPS is supported.
These are the precision options.
FP32 is the default.
FP16 suits V100 and older GPUs.
BF16 is recommended for A100/H100.
FP8 is available on H100.
Start with the quick start by installing the package.
Installation:
pip install lightning
Convert your PyTorch code to Lightning in three steps.
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)
Workflow 1 covers moving from PyTorch to Lightning, starting with the original PyTorch code.
Original PyTorch code:
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()
Here is the Lightning version.
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)
Workflow 2 covers validation and testing.
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)
Workflow 3 covers distributed training with 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)
Launch it like this:
# Single command, Lightning handles the rest
python train.py
Workflow 4 shows callbacks for monitoring.
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)
Workflow 5 covers learning rate scheduling.
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)
Troubleshooting tips follow.
If the issue is loss not decreasing, check your data and model setup.
Check data and model setup:
# 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
If the issue is running out of memory, reduce batch size or use gradient accumulation.
Reduce batch size or use gradient accumulation:
trainer = L.Trainer(
accumulate_grad_batches=4, # Effective batch = batch_size × 4
precision='bf16' # Or 'fp16', reduces memory 50%
)
If the issue is validation not running, make sure you pass val_loader.
Ensure you pass val_loader:
# WRONG
trainer.fit(model, train_loader)
# CORRECT
trainer.fit(model, train_loader, val_loader)
If DDP spawns multiple processes unexpectedly, remember that Lightning auto-detects GPUs, so set devices explicitly.
Lightning auto-detects GPUs, so explicitly set devices:
# Test on CPU first
trainer = L.Trainer(accelerator='cpu', devices=1)
# Then GPU
trainer = L.Trainer(accelerator='gpu', devices=1)
Links
Source: nousresearch/hermes-agent and its contributors. License: MIT.
Copyright (c) 2025 Nous Research
Translation and editorial adaptation by ClawCases.
>Resources
>Related Use Cases
Agent Merge-Conflict Arbiter
Resolve a git merge conflict between two agents' branches as an impartial third party: a neutral arbiter that takes both diffs and their stated intents and produces a merged result without overwriting either side.
Autonomous Project Management with Subagents
Coordinate projects autonomously: a main agent delegates work to PM subagents that track tasks in a shared STATE.yaml, update it as they progress, pick up unblocked work in parallel, and keep all task state in version-controlled files.
Weekly Review and Planning
Run a bounded weekly review and planning reset that audits commitments, stalled projects, and waiting items, then produces a next-week plan fitted to real calendar capacity.