Every PyTorch project eventually grows the same skeleton: a training function that moves batches to the device, zeroes gradients, calls backward, clips gradients, steps the optimizer, logs metrics, validates at intervals, saves checkpoints, and then gets rewritten when someone says the words multi-GPU or mixed precision. Lightning-AI/pytorch-lightning, the Apache-2.0 repository behind the packages you install as lightning and pytorch-lightning, is that skeleton extracted into a framework, tested by an enormous test suite, and scaled to everything from a laptop CPU to multi-node clusters. The project’s README puts the value proposition in one line: if PyTorch is JavaScript, PyTorch Lightning is ReactJS or NextJS. You write the model and the science; the framework owns the engineering around it. As always in this series, what follows is an educational tour of published source code.

The repository is a monorepo serving two packages at once. The src/lightning/pytorch tree holds the high-level Trainer experience; the src/lightning/fabric tree holds Fabric, a thinner expert API that gives you explicit control over devices, precision, and distribution without owning your loop. Both sit on the exact same strategy layer, and a small bridge class inside core/module.py even lets Fabric strategies power the full Trainer. Around these live a loop engine that decomposes training into nestable, restartable pieces, a connector that translates constructor flags into accelerator and device decisions, and a callback system with twenty pluggable behaviors from checkpointing to learning-rate finding. Reading the overview from left to right walks one fit call from the facade, through the loops, down into the hardware layer.

PyTorch Lightning repository architecture overview

Overview of pytorch-lightning: the Trainer facade resolves hardware through the accelerator connector, the loop engine drives your LightningModule, and a Strategy family from single device to DeepSpeed handles the distributed math, with Fabric sharing the same layer.

Reading the overview from left to right:

  • Trainer is the facade at line 89 of a 1797-line file, exposing fit at line 533, validate at 647, test at 765, and predict at 884, all funneling into one run path.
  • LightningModule is the user contract at line 92 of an 1849-line file: training_step at 711, configure_optimizers at 991, and the rank-aware log method at 381.
  • FitLoop is the outer loop at line 59 of 565 lines, nesting epoch loops while tracking progress and checkpointing its own reload state.
  • The training epoch loop at line 51 of 588 lines iterates batches, manages accumulation windows, and routes each one to optimization or evaluation.
  • Automatic optimization is a 329-line driver at line 153 whose Closure at line 92 caches the training-step result so backward and optimizer stepping can be replayed safely.
  • Strategy is the 649-line base contract at line 47; DDPStrategy at line 68 and DeepSpeedStrategy at line 76 of 945 lines are two of its many heirs.
  • The accelerator connector is a 593-line resolver at line 74 that turns constructor flags into a concrete accelerator, device set, and precision plugin.
  • Fabric is the expert API at line 86 of 1260 lines, exposing explicit setup, launch, backward, and step methods over the same strategy layer through a shared 580-line connector.

Why You Need This

The honest answer starts with repetition. A serious PyTorch project does not need three lines to train; it needs hundreds of lines that have nothing to do with the model: device plumbing, gradient accumulation, distributed process groups, precision handling, checkpoint formats, resumption logic, and logging glue. Each of these is easy to write once and painful to maintain across every experiment, and each becomes a bug farm the moment you move from one GPU to eight. Lightning takes that entire category of code and owns it in one audited place, so the diff between your CPU experiment and your cluster run becomes a constructor argument instead of a refactor.

The second reason is scale without rewrites. Because every distributed concern lives behind the Strategy interface, the same LightningModule you debug on a single device runs under DDP, FSDP, DeepSpeed, model parallel, or TPU execution without touching your model code. The framework decides how tensors move, which processes run which code, and how the optimizer states are sharded. When a new parallel scheme appears, the library adds a strategy; your code does not change. That decoupling is why the framework survived the industry’s rapid shifts from DDP to ZeRO-style sharding to fully sharded data parallel approaches over just a few years.

The third reason is the safety net under research code. The sanity check that runs a few validation batches before training starts catches broken dataloaders in seconds instead of after an epoch. The checkpoint system saves model, optimizer, and loop state together, so a crashed run resumes mid-epoch rather than mid-apology. Callbacks like ModelCheckpoint and EarlyStopping watch metrics and act on them, and the Tuner can find a learning rate and a batch size before you waste a week on both. These are unglamorous features, and they are exactly the ones that decide whether a project ships.

PyTorch Lightning detailed architecture

The detailed map: Trainer entry points converge on one run path, the loop engine decomposes into fit, epoch, evaluation, and prediction loops with an automatic optimization driver, the strategy family implements the hardware contract, and Fabric plus the Tuner share the same plumbing.

Follow one call to trainer.fit and the design snaps into focus. The fit method at line 533 validates its arguments, records that the function stage is fitting through the enums in states.py (RunningStage at line 44, TrainerState at 75), and calls _run at line 1008. That setup method is where the accelerator connector’s decisions become real: the strategy is instantiated and set up, loggers and callbacks get attached, and the environment is checked. Then _run_stage at line 1121 dispatches: training goes to the fit loop, evaluation to the shared evaluation loop. Before any real training, _run_sanity_check at line 1138 runs two validation batches by default, a tiny preflight that has saved countless runs from silently broken data.

The loop engine is the most instructive part. FitLoop owns epochs; the _TrainingEpochLoop owns batches and decides, per batch, whether this step belongs to automatic or manual optimization. Under automatic optimization, the Closure at line 92 wraps your training_step call, stores the result in a ClosureResult, and hands a replayable object to the strategy’s optimizer step. That indirection is what lets the framework accumulate gradients across batches, clip them, and step the right optimizer on the right boundary, including when a scheduler needs the closure at step time. Manual optimization flips the contract: override backward at line 1121 and optimizer_step at line 1354 on your LightningModule and the driver steps aside, which is how GAN-style alternating updates are expressed.

The strategy layer is where hardware stops being your problem. The Strategy base defines setup, model and optimizer wiring, backward, optimizer step, and the distributed operations, plus a _ForwardRedirection class at line 602 that keeps nn.Module hooks alive when forward calls must be re-routed. DDPStrategy wraps torch.distributed with launcher support for subprocess and notebook spawns, DeepSpeedStrategy bridges Microsoft’s ZeRO stages and CPU offloading, and FSDP and model parallel strategies cover sharded and layer-split execution. Even checkpoint saving routes through the strategy: trainer.save_checkpoint at line 1450 delegates so that each parallel scheme persists exactly what it needs to restore itself.

Fabric is the same machinery with the scaffolding stripped off. Instead of a Trainer owning the loop, Fabric gives you about a dozen explicit calls: setup the module and optimizers, launch the processes, call backward and step yourself. It shares the strategy layer through the same connector that powers the Trainer, and the _TrainerFabricShim at line 1835 of module.py is the bridge that lets Fabric strategies serve full Trainer runs. The Tuner at line 25 rounds out the toolbox with lr find and batch size scaling. And migration is handled by stub packages: the src/pytorch_lightning directory contains only metadata files that re-export the lightning namespace, so old imports keep working while the monorepo carries both eras.

From Install to Trained Model

The README’s promise is that organized PyTorch needs almost no ceremony:

pip install lightning
pip install torchvision
python main.py

And the smallest meaningful program is a LightningModule plus a Trainer, straight from the repository’s quick start:

import torch, torch.nn as nn, torch.utils.data as data, torchvision as tv
import torch.nn.functional as F
import lightning as L

class LitAutoEncoder(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        return self.encoder(x)

    def training_step(self, batch, batch_idx):
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log("train_loss", loss)
        return loss

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

dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
train, val = data.random_split(dataset, [55000, 5000])

autoencoder = LitAutoEncoder()
trainer = L.Trainer()
trainer.fit(autoencoder, data.DataLoader(train), data.DataLoader(val))

That program runs unchanged on a GPU box with the DDP strategy enabled, accumulates gradients with a flag, resumes from a checkpoint with another, and swaps optimizers’ precision with a third. The science stayed in the class; the engineering moved into the framework.

Honest limits: the abstraction has real weight. Debugging a training bug means learning where your code ends and the loop engine begins, and the hook ordering has enough surface area that newcomers occasionally fight the framework instead of the bug. The rename from pytorch-lightning to lightning left two import namespaces and a long deprecation trail. Expert users who want tight control often graduate to Fabric, which is both a feature and an admission that the Trainer is not for everyone. And the loop engine, restartable as it is, still assumes batch-centric training; genuinely exotic schedules end up in manual optimization mode anyway. None of this dims the core achievement: this repository turned the world’s most repeated training boilerplate into a maintained, tested, hardware-portable foundation, which is why so much modern AI code starts with import lightning.

Watch PyShine on YouTube

Contents