Labsco
firecrawl logo

pytorch-lightning

✓ Official★ 11

by firecrawl · part of firecrawl/ai-research-skills

High-level PyTorch framework with Trainer class, automatic distributed training (DDP/FSDP/DeepSpeed), callbacks system, and minimal boilerplate. Scales from…

🔥🔥🔥✓ VerifiedFreeQuick setup
#automation#calendar#scraping#git#monitoring#project-mgmt
🧩 One of 7 skills in the firecrawl/ai-research-skills package — works on its own, and pairs well with its siblings.

High-level PyTorch framework with Trainer class, automatic distributed training (DDP/FSDP/DeepSpeed), callbacks system, and minimal boilerplate. Scales from…

Inspect the full instructions your agent will receiveExpand

This is the exact playbook injected into your agent when the skill activates — shown here so you can audit it before installing. You don't need to read it to use the skill.

by firecrawl

High-level PyTorch framework with Trainer class, automatic distributed training (DDP/FSDP/DeepSpeed), callbacks system, and minimal boilerplate. Scales from… npx skills add https://github.com/firecrawl/ai-research-skills --skill pytorch-lightning Download ZIPGitHub11

PyTorch Lightning - High-Level Training Framework

Common workflows

Workflow 1: From PyTorch to Lightning

Original PyTorch code:

Copy & paste — that's it
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()

Lightning version:

Copy & paste — that's it
class LitModel(L.LightningModule):
 def __init__(self):
 super().__init__()
 self.model = MyModel()

 def training_step(self, batch, batch_idx):
 loss = self.model(batch) # No .to('cuda') needed!
 return loss

 def configure_optimizers(self):
 return torch.optim.Adam(self.parameters())

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

Benefits: 40+ lines → 15 lines, no device management, automatic distributed

Workflow 2: Validation and testing

Copy & paste — that's it
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)

Automatic features:

  • Validation runs every epoch by default

  • Metrics logged to TensorBoard

  • Best model checkpointing based on val_loss

Workflow 3: Distributed training (DDP)

Copy & paste — that's it
# 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)

Launch:

Copy & paste — that's it
# Single command, Lightning handles the rest
python train.py

No changes needed:

  • Automatic data distribution

  • Gradient synchronization

  • Multi-node support (just set num_nodes=2)

Workflow 4: Callbacks for monitoring

Copy & paste — that's it
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)

Result:

  • Auto-saves best 3 models

  • Stops early if no improvement for 5 epochs

  • Logs learning rate to TensorBoard

Workflow 5: Learning rate scheduling

Copy & paste — that's it
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)

When to use vs alternatives

Use PyTorch Lightning when:

  • Want clean, organized code

  • Need production-ready training loops

  • Switching between single GPU, multi-GPU, TPU

  • Want built-in callbacks and logging

  • Team collaboration (standardized structure)

Key advantages:

  • Organized: Separates research code from engineering

  • Automatic: DDP, FSDP, DeepSpeed with 1 line

  • Callbacks: Modular training extensions

  • Reproducible: Less boilerplate = fewer bugs

  • Tested: 1M+ downloads/month, battle-tested

Use alternatives instead:

  • Accelerate: Minimal changes to existing code, more flexibility

  • Ray Train: Multi-node orchestration, hyperparameter tuning

  • Raw PyTorch: Maximum control, learning purposes

  • Keras: TensorFlow ecosystem

Advanced topics

Callbacks: See references/callbacks.md for EarlyStopping, ModelCheckpoint, custom callbacks, and callback hooks.

Distributed strategies: See references/distributed.md for DDP, FSDP, DeepSpeed ZeRO integration, multi-node setup.

Hyperparameter tuning: See references/hyperparameter-tuning.md for integration with Optuna, Ray Tune, and WandB sweeps.

Resources