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 を切り替えたい
• 組み込みのコールバックとロギングを使いたい
• チームで共同作業する（標準化された構成が必要）

主な利点：

• 整理しやすい：研究コードとエンジニアリングを分離します
• 自動化：DDP、FSDP、DeepSpeed を 1 行で利用できます
• コールバック：トレーニングをモジュール単位で拡張できます
• 再現性：ボイラープレートが減るため、バグも減ります
• 実績：月間ダウンロード数 100 万以上で、十分に実運用されています

代わりに以下を使う場合：

• Accelerate：既存コードへの変更を最小限に抑え、柔軟性を高めたい場合
• Ray Train：マルチノードのオーケストレーションやハイパーパラメーター調整を行いたい場合
• Raw 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)

高度なトピック

コールバック：EarlyStopping、ModelCheckpoint、カスタムコールバック、コールバックフックについては、references/callbacks.md を参照してください。

分散ストラテジー：DDP、FSDP、DeepSpeed ZeRO の統合、マルチノードの設定については、references/distributed.md を参照してください。

ハイパーパラメーター調整：Optuna、Ray Tune、WandB sweeps との統合については、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 の優勝者、研究機関、本番環境のチーム

エージェントの運用手順

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：既存の PyTorch 学習スクリプトを HuggingFace Accelerate（Accelerator クラス、accelerate config、accelerate launch）でラップし、同じ c…