pytorch-lightning-distributed

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

Lightning 版本：

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 保存最佳模型检查点

工作流 3：分布式训练（DDP）

# 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 的情况：

• 希望代码整洁、有条理
• 需要可用于生产环境的训练循环
• 需要在单 GPU、多 GPU 和 TPU 之间切换
• 希望使用内置回调和日志记录
• 团队协作（采用标准化结构）

主要优势：

• 有条理：将研究代码与工程代码分离
• 自动化：只需 1 行即可使用 DDP、FSDP、DeepSpeed
• 回调：模块化的训练扩展
• 可复现：样板代码更少，意味着 Bug 更少
• 经过验证：每月下载量超过 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。请显式设置 devices：

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

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

高级主题

回调：请参阅 references/callbacks.md，了解 EarlyStopping、ModelCheckpoint、自定义回调和回调钩子。

分布式策略：请参阅 references/distributed.md，了解 DDP、FSDP、DeepSpeed ZeRO 集成和多节点设置。

超参数调优：请参阅 references/hyperparameter-tuning.md，了解与 Optuna、Ray Tune 和 WandB sweeps 的集成。

硬件要求

• 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 获奖者、研究实验室、生产团队

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 明确输入。 根据数据或用户确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 进行冒烟测试：使用小型模型或子集、少量步骤，并确认损失下降且检查点成功保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果发生这种情况	请这样做
CUDA 显存不足	减小批次大小、启用梯度累积/检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据是否包含无效值，并启用梯度裁剪。
这些说明中的函数、标志或端点在已安装版本中缺失	查阅已安装版本自己的文档（help()、--help、官方文档），调整做法并告知用户。切勿编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所采用的假设。

诚信规则

• 切勿伪造结果、参数、标识符、引用或统计数据。如果无法运行或验证某项内容，请如实说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前应估算计算成本。
• 此处的特定版本信息可能已过时：请对照已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• pytorch-lightning：使用 lightning 包（PyTorch Lightning）组织 PyTorch 训练代码：LightningModule、LightningDataModule、Trainer、回调等……
• ray-train：在集群间进行分布式训练编排。
• huggingface-accelerate：使用 HuggingFace Accelerate（Accelerator 类、accelerate config、accelerate launch）封装现有 PyTorch 训练脚本，让同一份代码……