claudegoodies
Skill

pytorch-lightning

From NousResearch

High-level PyTorch framework with Trainer class, automatic distributed training (DDP/FSDP/DeepSpeed), callbacks system, and minimal boilerplate. Scales from laptop to supercomputer with same code. Use when you want clean training loops with built-in best practices.

Wraps PyTorch training loops in a LightningModule/Trainer API with built-in distributed strategies and callbacks.

Use it when

  • Refactoring manual PyTorch training loops to cut boilerplate
  • Scaling the same training script from single GPU to DDP/FSDP/DeepSpeed multi-node
  • Need built-in checkpointing, early stopping, or LR monitoring via callbacks
  • Want automatic mixed precision, gradient accumulation, or TensorBoard logging without manual setup

Skip it if

  • Need full manual control over the training loop (Lightning imposes its own structure)
  • Project isn't in PyTorch (adds torch, lightning, transformers as dependencies)
  • Only need a single simple training run where plain PyTorch is already sufficient

Facts

Status
Actively maintained
Last commit

Source preview

The instructions Claude Code reads when this skill runs.

# PyTorch Lightning - High-Level Training Framework

## Quick start

PyTorch Lightning organizes PyTorch code to eliminate boilerplate while maintaining flexibility.

**Installation**:
```bash
pip install lightning
```

**Convert PyTorch to Lightning** (3 steps):

```python
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)
        self.log('train_loss', loss)  # Auto-logged to TensorBoard
        return loss

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

# Step 2: Create data
train_loader = DataLoader(train_dataset, batch_size=32)

# Step 3: Train with Trainer (handles everything else!)
trainer = L.Trainer(max_epochs=10, accelerator='gpu', devices=2)
model = LitModel()
trainer.fit(model, train_loader)
```

**That's it!** Trainer handles:
- GPU/TPU/CPU switching
- Distributed training (DDP, FSDP, DeepSpeed)
- Mix
View full source on GitHub →

Other skills