Lightning Training Framework Reference

Overview

The Lightning training framework organizes PyTorch code into a high-level structure that eliminates boilerplate while maintaining flexibility. It provides production-ready training loops, built-in callbacks, and automatic logging. Use it when you want clean, organized code, need to switch between single GPU, multi-GPU, or TPU, or require standardized structure for team collaboration.

Prerequisites

Install Lightning with the following command:

pip install lightning

You also need PyTorch installed. For GPU training, ensure you have a CUDA-compatible GPU and drivers. For TPU, you need a TPU runtime. For Apple MPS, use macOS with an M-series chip. Lightning version 2.5.5 or later is required.

Core Concepts

LightningModule

A LightningModule subclass organizes your PyTorch code into methods such as __init__, training_step, and configure_optimizers. This eliminates manual device management and boilerplate loops.

Trainer

The Trainer object handles training loops, distributed strategies, mixed precision, checkpointing, and logging. You configure it with parameters like max_epochs, accelerator, devices, and strategy.

Procedures

Convert PyTorch to Lightning

Define a LightningModule subclass with methods like __init__, training_step, and configure_optimizers.
Create data using DataLoader with a Dataset.
Train with Trainer: instantiate L.Trainer with desired settings, then call trainer.fit(model, train_loader).

For example, a raw PyTorch loop like this:

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()

Becomes this Lightning version:

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)
        return loss

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

# Step 2: Create data
dataset = torch.randn(100, 28 * 28)  # dummy data
train_loader = DataLoader(dataset, batch_size=32)

# Step 3: Train
trainer = L.Trainer(max_epochs=10, accelerator='gpu')
trainer.fit(LitModel(), train_loader)

The Lightning version is about 15 lines compared to 40+ lines of raw PyTorch.

Add Validation and Testing

In your LightningModule, define validation_step(self, batch, batch_idx) and test_step(self, batch, batch_idx). Log metrics using self.log('metric_name', value). Pass val_loader and test_loader to trainer.fit and trainer.test respectively.

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)

Validation runs every epoch by default. Metrics are automatically logged to TensorBoard.

Distributed Training with DDP

Define your model as usual with a LightningModule. Create a Trainer with accelerator='gpu', devices=8, and strategy='ddp' (or 'fsdp', 'deepspeed'). Run the script normally; Lightning handles multi-process launch.

# 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)

Use a single command to launch:

# Single command, Lightning handles the rest
python train.py

For multi-node training, set num_nodes=2. For debugging, test on CPU first with accelerator='cpu' and devices=1.

Use Callbacks for Monitoring

Import callbacks from lightning.pytorch.callbacks: ModelCheckpoint, EarlyStopping, LearningRateMonitor. Instantiate them with desired parameters and pass the list to the Trainer via the callbacks parameter.

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)

The best model checkpointing is based on val_loss by default.

Learning Rate Scheduling

In configure_optimizers, create an optimizer and a scheduler. Return a dictionary with the optimizer and lr_scheduler configuration.

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)

The learning rate is automatically logged by the LearningRateMonitor callback.

Parameters

Parameter	Meaning	Required
`accelerator`	Hardware accelerator type (`'gpu'`, `'cpu'`, `'tpu'`, `'mps'`)	No (default: auto-detect)
`devices`	Number of devices to use (e.g., 1, 2, 8) or `'auto'`	No (default: 1)
`strategy`	Distributed strategy (`'ddp'`, `'fsdp'`, `'deepspeed'`, etc.)	No (default: `'ddp'` for multi-GPU)
`max_epochs`	Maximum number of training epochs	No (default: 1000)
`accumulate_grad_batches`	Number of batches to accumulate gradients over (effective batch size = batch_size x this value)	No
`precision`	Precision mode: `'fp32'`, `'fp16'`, `'bf16'`, `'fp8'`	No (default: `'fp32'`)
`callbacks`	List of callback instances	No
`num_nodes`	Number of nodes for multi-node training	No (default: 1)

Mixed Precision and Gradient Accumulation

Use accumulate_grad_batches to simulate larger batch sizes. Use precision to reduce memory usage. For example, setting accumulate_grad_batches=4 with precision='bf16' reduces memory by 50%:

trainer = L.Trainer(
    accumulate_grad_batches=4,  # Effective batch = batch_size x 4
    precision='bf16'  # Or 'fp16', reduces memory 50%
)

FP16 works on V100 and older GPUs. BF16 is recommended for A100/H100. FP8 works on H100. TPU supports 8 cores. Apple MPS is also supported.

Failure Modes and Debugging

Loss not decreasing: Check data shapes and labels by printing in training_step at batch_idx==0:

# 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

Out of memory: Reduce batch size, use gradient accumulation (accumulate_grad_batches), and lower precision (precision='bf16' or 'fp16').

Validation not running: Ensure val_loader is passed to trainer.fit. The correct call is:

# WRONG
trainer.fit(model, train_loader)

# CORRECT
trainer.fit(model, train_loader, val_loader)

DDP spawns multiple processes unexpectedly: Test on CPU first:

# Test on CPU first
trainer = L.Trainer(accelerator='cpu', devices=1)

# Then GPU
trainer = L.Trainer(accelerator='gpu', devices=1)

Hyperparameter Tuning

For hyperparameter tuning, integrate with Optuna, Ray Tune, or WandB sweeps. See the hyperparameter tuning reference for details. For advanced callbacks and distributed strategies, see the respective reference files.

PyTorch Lightning: Clean Training Loops with Built-in Distributed Support