lightning

**PyTorch Lightning** is a **lightweight wrapper around PyTorch that eliminates boilerplate code while preserving full flexibility** — organizing the messy training loop (optimizer.zero_grad(), loss.backward(), optimizer.step(), logging, checkpointing, multi-GPU, mixed precision) into a clean, standardized LightningModule structure where you define only what matters (training_step, configure_optimizers) and Lightning handles everything else, enabling research code to scale from a laptop to a 100-GPU cluster with a single flag change. **What Is PyTorch Lightning?** - **Definition**: An open-source framework (pip install lightning) that restructures PyTorch code into a standardized LightningModule class — separating research logic (model architecture, loss, training step) from engineering boilerplate (device management, distributed training, logging, checkpointing). - **The Problem**: Raw PyTorch training loops are 200+ lines of repetitive code — move tensors to GPU, zero gradients, compute loss, backpropagate, step optimizer, log metrics, save checkpoints, handle multi-GPU. Every researcher rewrites this identically, introducing bugs each time. - **The Philosophy**: "Reorganize, don't abstract." Lightning doesn't hide PyTorch — it organizes it. You still write pure PyTorch inside training_step(). Lightning handles the engineering around it. **What You Write vs What Lightning Handles** | You Write | Lightning Handles | |-----------|------------------| | `training_step(batch, batch_idx)` | Training loop, batching, epochs | | `validation_step(batch, batch_idx)` | Validation loop, metric aggregation | | `configure_optimizers()` | Optimizer stepping, LR scheduling | | Model architecture (`__init__`) | Device placement (CPU/GPU/TPU) | | | Multi-GPU/multi-node distribution | | | Mixed precision (16-bit training) | | | Gradient accumulation/clipping | | | Checkpointing (best + last) | | | Logging (TensorBoard, WandB) | | | Early stopping | | | Profiling | **Code Comparison** ```python # Raw PyTorch: ~50 lines of boilerplate per training loop for epoch in range(num_epochs): model.train() for batch in train_loader: x, y = batch[0].to(device), batch[1].to(device) optimizer.zero_grad() output = model(x) loss = criterion(output, y) loss.backward() optimizer.step() # PyTorch Lightning: Define only what matters class LitModel(L.LightningModule): def training_step(self, batch, batch_idx): x, y = batch output = self.model(x) loss = self.criterion(output, y) self.log("train_loss", loss) return loss def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=1e-3) trainer = L.Trainer(max_epochs=10, accelerator="gpu", devices=4) trainer.fit(model, train_dataloader) ``` **Scaling With One Flag** | Task | Lightning Flag | |------|---------------| | Single GPU | `Trainer(accelerator="gpu", devices=1)` | | Multi-GPU (4 GPUs) | `Trainer(accelerator="gpu", devices=4)` | | Multi-Node (8 nodes × 8 GPUs) | `Trainer(num_nodes=8, devices=8)` | | Mixed Precision (16-bit) | `Trainer(precision=16)` | | Gradient Accumulation | `Trainer(accumulate_grad_batches=4)` | | TPU | `Trainer(accelerator="tpu", devices=8)` | **PyTorch Lightning is the standard way to write scalable, organized PyTorch code** — eliminating hundreds of lines of boilerplate while preserving full PyTorch flexibility, enabling researchers to focus on model innovation rather than engineering plumbing, and scaling seamlessly from a single GPU to multi-node clusters with zero code changes.

Go deeper with CFSGPT

Get AI-powered deep-dives, save terms, and run advanced simulations — free account.

Create Free Account