title: “Pytorch Lightning” sidebar_label: “Pytorch Lightning” description: “High-level PyTorch framework with Trainer class, automatic distributed training (DDP/FSDP/DeepSpeed), callbacks system, and minimal boilerplate”

{/* 本页面由 website/scripts/generate-skill-docs.py 根据技能对应的 SKILL.md 文件自动生成。请直接编辑源文件 SKILL.md，而非此页面。 */}

Pytorch Lightning

这是一个高级的 PyTorch 框架，具备 Trainer 类、自动分布式训练功能（DDP/FSDP/DeepSpeed）、回调系统以及极简的代码结构。使用同一份代码即可在从笔记本电脑到超级计算机的各类设备上实现高效训练。当您需要具备内置最佳实践的简洁训练流程时，可选用此框架。

技能元数据

参考：完整的 SKILL.md 文件

:::info 以下是 Hermes 在触发该技能时加载的完整技能定义。当技能处于激活状态时，智能体将看到这些内容作为操作指南。

::

PyTorch Lightning —— 高级训练框架

快速入门

PyTorch Lightning 能够对 PyTorch 代码进行结构化整理，在保持灵活性的同时减少冗余代码。

安装方式：

pip install lightning

将 PyTorch 转换为 Lightning（3个步骤）：

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)

就是这样！ Trainer可处理以下功能：

• GPU/TPU/CPU切换
• 分布式训练（DDP、FSDP、DeepSpeed）
• 混合精度训练（FP16、BF16）
• 梯度累积
• 检查点保存
• 日志记录
• 进度条显示

常见工作流程

工作流程1：从PyTorch转换为Lightning

原始PyTorch代码：

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()

轻量版：

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)

优势：代码行数从40多行缩减至15行，无需进行设备管理，可实现自动分布式处理。

工作流程2：验证与测试

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)

自动功能：

• 默认情况下，每个训练周期都会执行验证
• 将指标记录到 TensorBoard 中
• 根据 val_loss 自动保存最佳模型检查点

# 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)

启动：

# Single command, Lightning handles the rest
python train.py

无需任何更改：

• 自动数据分发
• 梯度同步
• 多节点支持（只需设置 num_nodes=2 即可）

工作流 4：用于监控的回调功能

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)

结果：

• 自动保存最优的3个模型
• 若连续5个训练周期未见改进则提前终止
• 将学习率信息记录至TensorBoard中

工作流程5：学习率调度

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)

何时使用 PyTorch Lightning 及其替代方案

适合使用 PyTorch Lightning 的情况包括：

• 希望拥有结构清晰、条理分明的代码
• 需要可用于生产环境的训练流程
• 在单 GPU、多 GPU 或 TPU 环境之间切换
• 需要内置的回调函数和日志记录功能
• 团队协作（标准化代码结构）

主要优势：

• 结构清晰：将研究代码与工程实现分开
• 操作便捷：仅需一行代码即可启用 DDP、FSDP 和 DeepSpeed
• 回调功能：支持模块化的训练扩展
• 可复现性强：减少样板代码，从而降低错误率
• 经过验证：每月下载量超 100 万次，经实际项目检验

适合选择替代方案的情况：

• Accelerate：对现有代码改动最小，灵活性更高
• Ray Train：支持多节点调度及超参数调优
• 原始 PyTorch：提供最大程度的控制权，适用于学习研究
• Keras：属于 TensorFlow 生态系统

常见问题

问题：损失值不下降

请检查数据与模型设置：

# Add to training_step
def training_step(self, batch, batch_idx):
    if batch_idx == 0:
        print(f"Batch shape: {batch[0].shape}")
        print(f"Labels: {batch[1]}")
    loss = ...
    return loss

问题：内存不足

请减小批量大小或采用梯度累积技术：

trainer = L.Trainer(
    accumulate_grad_batches=4,  # Effective batch = batch_size × 4
    precision='bf16'  # Or 'fp16', reduces memory 50%
)

问题：未运行验证流程

请确保传入了 val_loader 参数：

# WRONG
trainer.fit(model, train_loader)

# CORRECT
trainer.fit(model, train_loader, val_loader)

问题：DDP意外启动了多个进程

Lightning会自动检测GPU。如需手动指定设备，请进行如下设置：

# Test on CPU first
trainer = L.Trainer(accelerator='cpu', devices=1)

# Then GPU
trainer = L.Trainer(accelerator='gpu', devices=1)

高级主题

回调机制：如需了解 EarlyStopping、ModelCheckpoint、自定义回调函数以及回调钩子的相关内容，请参阅 references/callbacks.md。

分布式训练策略：关于 DDP、FSDP、DeepSpeed ZeRO 集成以及多节点部署的详细信息，请查看 references/distributed.md。

超参数调优：若需了解如何与 Optuna、Ray Tune 以及 WandB 的超参数扫描功能进行集成，请参考 references/hyperparameter-tuning.md。

硬件要求

• CPU：可用（适合调试）
• 单 GPU：可用
• 多 GPU：支持 DDP（默认）、FSDP 或 DeepSpeed
• 多节点环境：支持 DDP、FSDP、DeepSpeed
• TPU：支持（需 8 核）
• Apple MPS：支持

精度选项：

• FP32（默认）
• FP16（适用于 V100 及更早版本的 GPU）
• BF16（推荐用于 A100/H100）
• FP8（仅适用于 H100）

相关资源

• 文档：https://lightning.ai/docs/pytorch/stable/
• GitHub 仓库：https://github.com/Lightning-AI/pytorch-lightning ⭐ 29,000+ 次星标
• 版本：2.5.5 及以上
• 示例代码：https://github.com/Lightning-AI/pytorch-lightning/tree/master/examples
• Discord 社群：https://discord.gg/lightning-ai
• 应用场景：Kaggle 冠军团队、研究实验室及生产环境团队