PyTorch Lightning

Overview

PyTorch Lightning is an open-source library that organises PyTorch code. It does not replace PyTorch. Your layers, tensors and losses are still plain PyTorch. What Lightning takes over is the engineering around them: the training loop, moving data to the right device, mixed precision, checkpointing, logging, early stopping and running on many GPUs or machines.

The core idea is a split between research code and engineering code. You describe what one training step does and which optimizer to use. The Lightning Trainer decides how to run that step thousands of times, on which hardware, with which precision, and what to save along the way. Changing from one GPU to eight usually means changing an argument, not rewriting the loop.

Think of a restaurant kitchen. The head chef writes the recipe and decides how each dish should taste. The kitchen staff handle the ovens, timing, stock and cleaning. Lightning plays the role of the kitchen staff, so the researcher can spend their time on the recipe rather than on oven temperatures.

What it is

Lightning is built from a small number of pieces.

A LightningModule is still a torch.nn.Module, so anything that works on a PyTorch model works on it. The Trainer is the part that knows about hardware. The callbacks and loggers let you add behaviour without editing the model class.

Where it came from

William Falcon started PyTorch Lightning while doing research, and the first public release came in May 2019 under the Apache 2.0 licence. The same year, the NeurIPS reproducibility challenge adopted it as a recommended way to structure PyTorch submissions, because a fixed layout makes other people's code easier to read and rerun.

Falcon founded a company, first called Grid.ai, to build products around the library. In June 2022 it rebranded as Lightning AI. On 15 March 2023 the company released PyTorch Lightning 2.0, which promised a stable API, together with Lightning Fabric, a new library for people who want to keep their own training loop but still scale it easily. The project now ships as the lightning package, which contains both the Trainer-based API and Fabric, while the older pytorch_lightning import name remains available.

How it works

Step 1: define the module

You subclass LightningModule. In __init__ you create layers. In training_step you take a batch, compute the loss and return it. In configure_optimizers you return the optimizer and, optionally, a learning-rate scheduler. Validation and test steps follow the same pattern.

import lightning as L
import torch
from torch import nn

class Classifier(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 2))

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = nn.functional.cross_entropy(self.net(x), y)
        self.log("train_loss", loss)
        return loss

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

Step 2: hand it to the Trainer

You create a Trainer with settings such as max_epochs, accelerator, devices and precision, then call trainer.fit(model, datamodule) or pass data loaders directly.

Step 3: what the Trainer does for you

Inside fit, the Trainer moves the model and each batch to the right device, calls zero_grad, runs your step, calls backward, steps the optimizer and scheduler, runs validation at set intervals, writes logs and saves checkpoints. With precision="bf16-mixed" it runs mixed precision. With devices=8 and strategy="ddp" it launches distributed data-parallel training and inserts the right samplers.

Step 4: hooks and callbacks

At dozens of points, such as the start of an epoch or the end of a validation run, the Trainer calls hooks. You can override hooks in the module or write callbacks that live outside it. This is how early stopping, custom checkpointing and learning-rate monitoring are added without touching the core logic.

Fabric for custom loops

Some work does not fit the standard loop: reinforcement learning, unusual multi-model training, or research on the loop itself. Fabric gives you a few calls, such as fabric.setup() and fabric.backward(), that handle devices, precision and distribution while you keep writing the loop yourself.

Key concepts and terms

A worked example

This scenario is illustrative. A research group trains a text classifier on 500,000 labelled reviews. Their first version is a 300-line PyTorch script with a hand-written loop, manual device moves and print statements.

1. They move the model and loss into a LightningModule with training_step, validation_step and configure_optimizers. The class shrinks to about 80 lines.
2. They move tokenisation, splitting and loaders into a DataModule so another project can reuse it.
3. They run Trainer(fast_dev_run=True) on a laptop. One training and one validation batch run end to end, which catches a shape error in seconds.
4. On a single GPU they train with precision="bf16-mixed" and a ModelCheckpoint callback that keeps the best three checkpoints by validation accuracy.
5. To finish faster, they move to a node with eight GPUs and set devices=8, strategy="ddp". The model code does not change. With a per-device batch of 32, the effective batch becomes 256, so they raise the learning rate and add warm-up.
6. An EarlyStopping callback stops training when validation loss has not improved for three evaluations.

The time saved is less about speed and more about fewer bugs and easier handover to colleagues.

Strengths and limitations

Strengths:

• Removes repeated boilerplate for loops, devices and checkpoints.
• A shared structure makes projects easier to read, review and reproduce.
• Scaling to many GPUs, mixed precision and sharded strategies is mostly configuration.
• Callbacks and loggers integrate with common tools.
• The model stays plain PyTorch, so nothing is locked away.

Limitations:

• Another layer to learn, with its own hooks and ordering rules.
• Unusual training loops can fight the structure; Fabric or plain PyTorch may suit them better.
• When something fails inside the Trainer, debugging requires knowing what it does behind the scenes.
• It does not make individual operations faster; speed still comes from PyTorch itself.

How it compares with alternatives

Practical guidance and best practices

1. Keep model logic in the LightningModule and all data handling in a DataModule.
2. Use fast_dev_run=True before every long run to catch errors early.
3. Log metrics with self.log and set sync_dist=True for metrics that should be averaged across devices.
4. Use ModelCheckpoint with a clear monitored metric and keep more than one checkpoint.
5. Call self.save_hyperparameters() so checkpoints record the settings that produced them.
6. Seed runs with L.seed_everything and record library versions.
7. When scaling up, recalculate the effective batch size and adjust the learning rate.
8. Choose Fabric rather than bending the Trainer when your loop is genuinely unusual.

Common mistakes

• Calling .cuda() or .to(device) by hand inside the module, which breaks device handling.
• Calling backward() or optimizer.step() yourself while automatic optimisation is on.
• Logging metrics per step on every device without syncing, giving misleading numbers.
• Forgetting that the effective batch size grows with the number of devices.
• Putting data downloads in setup rather than prepare_data, so every process downloads at once.

Frequently asked questions

Is Lightning a separate framework from PyTorch?

No. It is a library on top of PyTorch. A LightningModule is a normal PyTorch module, and you can load its weights into plain PyTorch code without Lightning installed if you save the state_dict.

Does Lightning make training faster?

Not directly. Each operation runs the same PyTorch code. It makes speed features such as mixed precision, multi-GPU training and compiled models easier to switch on, which often leads to faster training in practice.

What is the difference between the Trainer and Fabric?

The Trainer owns the whole loop and gives you many features automatically. Fabric only handles devices, precision and distribution, and you write the loop. Use the Trainer for standard supervised training and Fabric for custom loops.

Can Lightning train very large models?

Yes, through strategies such as FSDP and DeepSpeed, which shard model state across GPUs. You choose the strategy in the Trainer, though large models still need careful memory planning.

Which package should I install?

Current documentation uses the unified lightning package, imported as import lightning as L. Older projects import pytorch_lightning; both point to the same Trainer API.

Does it work with experiment trackers?

Yes. Built-in loggers cover TensorBoard, MLflow, Weights & Biases and others, and several can be used at once.

Glossary

• Automatic optimisation: the default mode in which the Trainer calls backward and the optimizer for you.
• Callback: a reusable object that adds behaviour at fixed points in training.
• DDP: distributed data-parallel training with one model copy per device.
• Effective batch size: per-device batch times number of devices times gradient accumulation steps.
• Gradient accumulation: summing gradients over several batches before one optimizer step.
• Hook: a named method called by the Trainer at a set moment.
• Logger: a connector that sends metrics to a tracking tool.
• Strategy: the plug-in that decides how computation is split across devices.

Summary

PyTorch Lightning, started by William Falcon in 2019 and now maintained by Lightning AI, organises PyTorch code into a LightningModule and hands the engineering to a Trainer. It removes repeated loop code, makes mixed precision and multi-GPU training a configuration choice, and gives teams a common structure that is easier to share and reproduce. It adds an abstraction to learn and suits standard training best; Fabric or plain PyTorch remain the better fit for unusual loops.

Related

It belongs to Frameworks in my stack, under Deep learning. Others in the same group: