PyTorch Lightning is a high-level wrapper around PyTorch that removes boilerplate from training loops while keeping you in control of the actual model logic. Vanilla PyTorch gives you maximum flexibility and transparency — every part of training is code you wrote. The decision between them is not about which is better in absolute terms: it is about which fits your workflow at the current stage of your project. Understanding what Lightning does and does not abstract helps you make that call without guessing.
What Lightning Actually Removes
Lightning takes ownership of the training loop scaffolding: moving data and models to the right device, calling optimizer.zero_grad(), loss.backward(), and optimizer.step() in the right order, running validation after each epoch, logging metrics, saving checkpoints, and handling multi-GPU distribution. Your job is to define training_step(), validation_step(), configure_optimizers(), and the model itself. Lightning runs everything else.
Side-by-Side: The Same Model
# ── Vanilla PyTorch ──────────────────────────────────────────
import torch, torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(nn.Flatten(), nn.Linear(784,256), nn.ReLU(), nn.Linear(256,10))
def forward(self, x): return self.net(x)
model = MLP().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
train_loader = DataLoader(datasets.MNIST('./data', train=True, download=True,
transform=transforms.ToTensor()), batch_size=256, shuffle=True)
for epoch in range(10):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1} done")# ── PyTorch Lightning ────────────────────────────────────────
import torch, torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
class LitMLP(pl.LightningModule):
def __init__(self):
super().__init__()
self.net = nn.Sequential(nn.Flatten(), nn.Linear(784,256), nn.ReLU(), nn.Linear(256,10))
self.criterion = nn.CrossEntropyLoss()
def forward(self, x): return self.net(x)
def training_step(self, batch, batch_idx):
x, y = batch
loss = self.criterion(self(x), y)
self.log('train_loss', loss)
return loss
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=1e-3)
train_loader = DataLoader(datasets.MNIST('./data', train=True, download=True,
transform=transforms.ToTensor()), batch_size=256, shuffle=True)
model = LitMLP()
trainer = pl.Trainer(max_epochs=10, accelerator='auto')
trainer.fit(model, train_loader)Both do the same thing. The Lightning version is shorter and gains automatic device placement (accelerator='auto' picks GPU/MPS/CPU), built-in logging, and easy multi-GPU support — but it requires you to learn Lightning’s callback and hook system to do anything non-standard.
When Vanilla PyTorch Wins
Vanilla PyTorch is better when you are implementing something non-standard. Custom training loops that mix multiple models (GANs, meta-learning inner/outer loops, RLHF reward model training) are easier to write and debug when you own every line of the loop. Research that modifies gradient flow in unusual ways — gradient surgery, gradient checkpointing at specific layers, mixed-precision with custom scaling — is clearer in vanilla PyTorch where there is no abstraction layer to reason through. Debugging is also easier: a training loop you wrote is a training loop you can step through line by line. Lightning’s training_step runs inside the trainer’s orchestration and errors can surface with confusing stack traces.
When Lightning Wins
Lightning is the right choice when your training loop is standard and you want the infrastructure without writing it. Multi-GPU training with DDP requires significant boilerplate in vanilla PyTorch; in Lightning it is trainer = pl.Trainer(devices=4, strategy='ddp'). Logging to TensorBoard, Weights and Biases, or MLflow is built in — call self.log('metric', value) and Lightning routes it to whichever logger you configured. Gradient clipping, learning rate scheduling, early stopping, and checkpoint saving are all Trainer arguments rather than manual code. For teams where multiple people run similar experiments, Lightning enforces a consistent structure that makes reading others’ code easier.
Figure 1 — What Lightning abstracts vs what you control
Multi-GPU: The Clearest Lightning Win
# Vanilla PyTorch DDP — significant setup required
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])
sampler = DistributedSampler(dataset)
# ... many more lines
# Lightning DDP — three characters different from single GPU
trainer = pl.Trainer(max_epochs=10, devices=4, strategy='ddp', accelerator='gpu')Escape Hatches: When You Need to Go Custom
Lightning is not a black box — it exposes hooks at every stage of the training loop. If you need custom behaviour, override the relevant hook rather than abandoning Lightning entirely:
class CustomModel(pl.LightningModule):
def training_step(self, batch, batch_idx):
# Standard step — you write this
x, y = batch
loss = self.criterion(self(x), y)
return loss
def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
# Custom gradient manipulation before step
for group in optimizer.param_groups:
for p in group['params']:
if p.grad is not None:
p.grad.data.clamp_(-1.0, 1.0) # gradient clipping
super().optimizer_step(epoch, batch_idx, optimizer, optimizer_closure)
def on_train_epoch_end(self):
# Custom logic at epoch end
avg_loss = self.trainer.callback_metrics.get('train_loss')
print(f"Epoch {self.current_epoch}: {avg_loss:.4f}")The Practical Decision
Start with vanilla PyTorch for new research ideas — the loop transparency helps you understand what is happening and spot bugs faster. Switch to Lightning when the experiment is validated and you need to scale: multi-GPU runs, systematic logging, reproducible checkpoints for sharing with collaborators. Many practitioners write an initial prototype in vanilla PyTorch, then port to Lightning for the scaling runs. The port is straightforward — move your forward() into a LightningModule, split the training loop into training_step() and configure_optimizers(), and the Trainer handles the rest. The two approaches are not rivals; they are different tools for different phases of the same workflow.
If you are joining a team that already uses Lightning, learn Lightning — the shared structure is a genuine productivity benefit. If you are a solo researcher exploring novel ideas, vanilla PyTorch keeps every decision visible. For standard supervised learning tasks (classification, regression, fine-tuning) where the loop is unremarkable, Lightning reduces the lines you have to write and maintain without taking anything meaningful away.
Figure 1 — What Lightning abstracts vs what you control
Lightning Callbacks and Plugins
Lightning’s callback system is where much of its power lives. Callbacks are objects that hook into named points in the training lifecycle — on_train_epoch_start, on_validation_end, on_before_optimizer_step, and so on. Built-in callbacks handle the most common needs:
from pytorch_lightning.callbacks import (
ModelCheckpoint, EarlyStopping, LearningRateMonitor, RichProgressBar
)
trainer = pl.Trainer(
max_epochs=100,
accelerator='auto',
callbacks=[
ModelCheckpoint(monitor='val_loss', save_top_k=3, mode='min'),
EarlyStopping(monitor='val_loss', patience=10, mode='min'),
LearningRateMonitor(logging_interval='epoch'),
RichProgressBar(),
],
logger=pl.loggers.WandbLogger(project='my-project'), # or TensorBoardLogger
)Each of these would be 20–50 lines of manual code in vanilla PyTorch. The callback system also makes it easy to add custom behaviour without modifying the model — write a Callback subclass, override the hooks you need, pass it to the Trainer.
Installation and Version Notes
pip install pytorch-lightning
# or the newer rebranded package
pip install lightningThe package was rebranded from pytorch-lightning to lightning in version 2.0. Both work; lightning is the current name. Import as import lightning as L or import pytorch_lightning as pl — both are supported. Lightning 2.0 introduced some breaking changes from 1.x (Trainer argument names, callback signatures); check the migration guide if upgrading an existing project from 1.x.
LightningDataModule
Lightning’s DataModule encapsulates all data-related logic — downloading, preprocessing, splitting, and creating DataLoaders — into a reusable object that can be shared across experiments:
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
class MNISTDataModule(LightningDataModule):
def __init__(self, batch_size=256):
super().__init__()
self.batch_size = batch_size
def setup(self, stage=None):
full = datasets.MNIST('./data', train=True, download=True,
transform=transforms.ToTensor())
self.train_set, self.val_set = random_split(full, [55000, 5000])
self.test_set = datasets.MNIST('./data', train=False,
transform=transforms.ToTensor())
def train_dataloader(self):
return DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True)
def val_dataloader(self):
return DataLoader(self.val_set, batch_size=self.batch_size)
def test_dataloader(self):
return DataLoader(self.test_set, batch_size=self.batch_size)
dm = MNISTDataModule(batch_size=256)
trainer.fit(model, dm)The DataModule pattern makes experiments fully reproducible — the same data splits, transforms, and DataLoader settings are always applied, and you can swap datasets by swapping DataModules without touching the model code.
Profiling and Debugging in Lightning
Lightning includes a built-in profiler that measures time spent in each part of the training loop:
from pytorch_lightning.profilers import PyTorchProfiler
trainer = pl.Trainer(
max_epochs=5,
profiler=PyTorchProfiler(filename='profiler-output')
)
trainer.fit(model, datamodule)For debugging, set fast_dev_run=True to run a single batch through train, val, and test in one pass — a quick sanity check before committing to a full run:
trainer = pl.Trainer(fast_dev_run=True) # 1 batch train + val + test
trainer.fit(model, datamodule)
# Or run N batches
trainer = pl.Trainer(fast_dev_run=5) # 5 batches eachHydra Integration for Config Management
For larger projects, Lightning pairs well with Hydra for configuration management. Hydra lets you define experiment configs in YAML files and override them from the command line, making hyperparameter sweeps and experiment tracking clean:
python train.py trainer.max_epochs=50 model.lr=0.001 data.batch_size=512The Lightning CLI (from lightning.pytorch.cli import LightningCLI) provides this integration out of the box with minimal setup — your model and datamodule are automatically configurable from YAML without writing argparse code.
Mixed Precision and Gradient Clipping
Both features are single Trainer arguments in Lightning versus multiple lines in vanilla PyTorch:
# Lightning: one-line AMP + gradient clipping
trainer = pl.Trainer(
precision='16-mixed', # automatic mixed precision
gradient_clip_val=1.0, # gradient clipping by value
gradient_clip_algorithm='norm' # or 'value'
)
# Vanilla PyTorch equivalent:
from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler()
# ... then in each step:
with autocast():
loss = model(x)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()The Lightning version is not hiding complexity — it is the same operations under the hood. The difference is that Lightning has already debugged the ordering (unscale before clip, step after unscale) so you do not have to. Getting the AMP + clipping + DDP interaction right in vanilla PyTorch is a known source of subtle bugs that Lightning avoids by enforcing the correct sequence internally.
Common Lightning Gotchas
A few things that trip up newcomers. Logging inside training_step: call self.log('metric', value, on_step=True, on_epoch=True) rather than print() — the Trainer routes logged values to whatever logger you configured and handles averaging over the epoch. Returning loss from training_step: you must return the loss tensor so Lightning can call loss.backward(); returning None silently skips the backward pass. validation_step does not update gradients: Lightning wraps it in torch.no_grad() automatically, so you do not need to add it. Epoch-level metrics: use self.log(..., on_epoch=True) rather than manually accumulating; Lightning averages the logged values across batches for you.
Start with vanilla PyTorch for new research ideas — loop transparency helps you spot bugs faster. Switch to Lightning when the experiment is validated and you need to scale: multi-GPU runs, systematic logging, reproducible checkpoints. Many practitioners write an initial prototype in vanilla PyTorch, then port to Lightning for scaling runs. The port is straightforward — move your forward() into a LightningModule, split the training loop into training_step() and configure_optimizers(), and the Trainer handles the rest. The two are not rivals; they are different tools for different phases of the same workflow. The real risk to avoid is choosing Lightning for a project where you need to understand every step of training, and choosing vanilla PyTorch for a project where you end up rewriting checkpoint management and multi-GPU setup from scratch — both of those are avoidable costs once you know which phase you are in.
If you are on the fence, a useful heuristic: would you rather read a bug in your training loop code or a bug in Lightning’s trainer? If the former — stay in vanilla PyTorch. If the latter feels less likely because Lightning is well-tested and your own loop code is not — use Lightning.