PyTorch Lightning

高級 PyTorch 框架，提供 Trainer 類、自動分佈式訓練（DDP/FSDP/DeepSpeed）、回調系統以及極少的樣板代碼。使用相同的代碼即可從筆記本電腦擴展至超級計算機。當你希望獲得內置最佳實踐的簡潔訓練循環時，請使用此框架。

技能元數據


來源	可選 — 使用 `hermes skills install official/mlops/pytorch-lightning` 安裝
路徑	`optional-skills/mlops/pytorch-lightning`
版本	`1.0.0`
作者	Orchestra Research
許可證	MIT
依賴項	`lightning`, `torch`, `transformers`
標籤	`PyTorch Lightning`, `Training Framework`, `Distributed Training`, `DDP`, `FSDP`, `DeepSpeed`, `High-Level API`, `Callbacks`, `Best Practices`, `Scalable`

參考：完整 SKILL.md

信息

以下是 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()

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)

自動功能：

默認每個 epoch 運行驗證
指標記錄到 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 個 epoch 內沒有改進則提前停止
將學習率記錄到 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)

何時使用 vs 替代方案

在以下情況使用 PyTorch Lightning：

希望代碼整潔、有條理
需要生產就緒的訓練循環
在單 GPU、多 GPU、TPU 之間切換
希望擁有內置的回調和日誌記錄功能
團隊協作（標準化結構）

主要優勢：

有組織：將研究代碼與工程代碼分離
自動化：僅需 1 行代碼即可啟用 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)

高級主題

回調：請參閱 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 獲獎者、研究實驗室、生產團隊

技能元數據​

參考：完整 SKILL.md​

PyTorch Lightning - 高級訓練框架

快速開始​

常見工作流​

工作流 1：從 PyTorch 到 Lightning​

工作流 2：驗證與測試​

工作流 3：分佈式訓練（DDP）​

工作流 4：用於監控的回調​

工作流 5：學習率調度​

何時使用 vs 替代方案​

常見問題​

高級主題​

硬件要求​

資源​