PyTorch Lightning 完全指南:一套代码从单卡到 DDP/FSDP/DeepSpeed 多机分布式训练
PyTorch Lightning 完全指南:一套代码从单卡到 DDP/FSDP/DeepSpeed 多机分布式训练
PyTorch Lightning 是一个建立在 PyTorch 之上的高层训练框架,核心价值在于通过 LightningModule 与 Trainer 的分工,把「研究代码」与「工程代码」解耦,从而消除样板代码,让同一套代码无需改动即可在 CPU、单卡 GPU、多卡 GPU、多机乃至 TPU 上运行,并内置 DDP / FSDP / DeepSpeed 分布式策略、混合精度、梯度累积、断点续训、日志与进度条等最佳实践。本文以 08-distributed-training/pytorch-lightning/SKILL.md 为主干,结合该技能目录下的 callbacks.md、distributed.md、hyperparameter-tuning.md 三份深度参考文档,完整覆盖从「PyTorch 代码改造」到「回调系统」「分布式策略」「超参数调优」的实战全流程。读完本文,你将能独立把一个普通 PyTorch 训练脚本改造成 Lightning 工程,并熟练配置分布式训练、模型保存、早停与超参搜索。
快速上手:三步把 PyTorch 代码改造成 Lightning
PyTorch Lightning 的组织思路非常直接:保留 PyTorch 的模型与数据代码,只改变组织方式。安装只需一条命令:
pip install lightning
安装后,将现有 PyTorch 代码转换为 Lightning 只需三步:
import lightning as L
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset
# Step 1: 定义 LightningModule(组织你的 PyTorch 代码)
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: 创建数据
train_loader = DataLoader(train_dataset, batch_size=32)
# Step 3: 用 Trainer 训练(其余一切交给它!)
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)
- 梯度累积
- Checkpointing(断点保存)
- 日志记录
- 进度条
从这套 API 可以看出,LightningModule 只负责「定义模型、单步计算、配置优化器」,而所有涉及多进程、多卡、调度、日志的复杂逻辑全部下沉到 Trainer。这正是它「从笔记本到超级计算机,同一套代码」承诺的实现基础。
核心机制:LightningModule 的 Hook 方法与 Trainer 的分工
要理解 Lightning,关键是理解它的「约定式 Hook」设计。training_step、validation_step、test_step、configure_optimizers 等只是 LightningModule 暴露的钩子方法,Trainer 会在正确的时机自动调用它们,并替开发者处理设备迁移(.to('cuda') 不再需要手写)、梯度归零、反向传播与参数更新。
以最小实现为例,模型只需声明四个方法即可完成训练闭环:
| Hook 方法 | 职责 | 调用时机 |
|---|---|---|
training_step(batch, batch_idx) |
前向计算并返回 loss | 每个训练 batch |
validation_step(batch, batch_idx) |
验证集评估 | 每个验证 batch(默认每个 epoch 末) |
test_step(batch, batch_idx) |
测试集评估 | 调用 trainer.test() 时 |
configure_optimizers() |
返回优化器 / 学习率调度器 | 训练开始前 |
self.log() 是 Lightning 统一的指标记录入口,写入的指标会自动进入 TensorBoard 等 logger,并可作为 ModelCheckpoint、EarlyStopping 等回调的监控对象。这套设计的好处是:研究逻辑(怎么算 loss、怎么评估)与工程逻辑(何时保存、何时早停、怎么同步梯度)完全分离,团队协作时每个人只需关注自己负责的那一层。
五个常用工作流
工作流 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)
对比可见,手动设备管理(model.to('cuda')、batch.to('cuda'))和手写训练循环(zero_grad / backward / step)全部消失。原文档记录的经验值是:40 多行代码缩减到 15 行,且自动获得分布式能力。这段改造逻辑对应本文开头「三步转换法」的实际落地场景:把循环体内的单步逻辑挪进 training_step,把优化器配置挪进 configure_optimizers。
工作流 2:验证与测试
Lightning 把训练、验证、测试三阶段统一定义在同一个模块里:
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自动保存最佳模型 checkpoint
这里要注意 trainer.fit 的第三个参数是验证 DataLoader;只有在传入 val_loader 的前提下,验证阶段和后续的模型监控(如早停、最佳模型保存)才会生效。
工作流 3:分布式训练(DDP)
分布式是 Lightning 的核心卖点之一——多卡训练的代码与单卡完全相同:
# 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)
启动方式依然是单条命令,Lightning 负责其余一切:
# Single command, Lightning handles the rest
python train.py
完全无需改动的地方包括:
- 数据自动分发(Lightning 自动为 DataLoader 包装
DistributedSampler) - 梯度同步(跨 GPU all-reduce)
- 多节点支持(只需设置
num_nodes=2)
关于 strategy 各选项(ddp / fsdp / deepspeed)的选型与底层细节,见下文「分布式训练策略」一节,更完整的说明见 references/distributed.md。
工作流 4:用回调做监控
回调(Callback)是 Lightning 可插拔的扩展机制,用来在不改动模型代码的前提下添加监控、保存、早停等功能:
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
回调系统的完整深度讲解(内置回调的参数详解、自定义回调、全部 Hook 列表)见下文「回调系统深度解析」,对应仓库文档为 references/callbacks.md。
工作流 5:学习率调度
Lightning 支持在 configure_optimizers 里同时返回优化器和学习率调度器:
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)
lr_scheduler 字典里的 interval 可选 'epoch' 或 'step'(按 step 更新时适合 warmup 类调度器),frequency 表示每 N 个 interval 更新一次。配合 LearningRateMonitor,学习率曲线会被自动记录,无需手动 log。
回调系统深度解析
回调捕获的是非核心逻辑(checkpoint、早停、日志等),在不修改 LightningModule 的前提下为训练注入功能。以下是本技能目录 references/callbacks.md 中记录的完整回调能力。
内置回调 1:ModelCheckpoint
按验证指标保存最佳模型:
from lightning.pytorch.callbacks import ModelCheckpoint
# Save top 3 models based on validation loss
checkpoint = ModelCheckpoint(
dirpath='checkpoints/',
filename='model-{epoch:02d}-{val_loss:.2f}',
monitor='val_loss',
mode='min',
save_top_k=3,
save_last=True, # Also save last epoch
verbose=True
)
trainer = L.Trainer(callbacks=[checkpoint])
trainer.fit(model, train_loader, val_loader)
全部常用配置项:
checkpoint = ModelCheckpoint(
monitor='val_acc', # Metric to monitor
mode='max', # 'max' for accuracy, 'min' for loss
save_top_k=5, # Keep best 5 models
save_last=True, # Save last epoch separately
every_n_epochs=1, # Save every N epochs
save_on_train_epoch_end=False, # Save on validation end instead
filename='best-{epoch}-{val_acc:.3f}', # Naming pattern
auto_insert_metric_name=False # Don't auto-add metric to filename
)
加载 checkpoint:
# Load best model
best_model_path = checkpoint.best_model_path
model = LitModel.load_from_checkpoint(best_model_path)
# Resume training
trainer = L.Trainer(callbacks=[checkpoint])
trainer.fit(model, train_loader, val_loader, ckpt_path='checkpoints/last.ckpt')
需要注意:filename 中的 {epoch}、{val_loss:.2f} 等占位符会按当前 epoch 与监控指标值动态填充;mode='min' 表示指标越小越好(loss),mode='max' 表示越大越好(accuracy)。
内置回调 2:EarlyStopping
指标不再提升时提前停止训练,节省算力:
from lightning.pytorch.callbacks import EarlyStopping
early_stop = EarlyStopping(
monitor='val_loss',
patience=5, # Wait 5 epochs
mode='min',
min_delta=0.001, # Minimum change to qualify as improvement
verbose=True,
strict=True, # Crash if monitored metric not found
check_on_train_epoch_end=False # Check on validation end
)
trainer = L.Trainer(callbacks=[early_stop])
trainer.fit(model, train_loader, val_loader)
# Stops automatically if no improvement for 5 epochs
进阶用法:
early_stop = EarlyStopping(
monitor='val_loss',
patience=10,
min_delta=0.0,
verbose=True,
mode='min',
stopping_threshold=0.1, # Stop if val_loss < 0.1
divergence_threshold=5.0, # Stop if val_loss > 5.0
check_finite=True # Stop on NaN/Inf
)
stopping_threshold 与 divergence_threshold 分别用于「达到目标立即停止」与「发散立即停止」两种场景,check_finite=True 则能在指标出现 NaN/Inf 时中断训练。
内置回调 3:LearningRateMonitor
自动记录学习率:
from lightning.pytorch.callbacks import LearningRateMonitor
lr_monitor = LearningRateMonitor(
logging_interval='epoch', # Or 'step'
log_momentum=True # Also log momentum
)
trainer = L.Trainer(callbacks=[lr_monitor])
# Learning rate automatically logged to TensorBoard/WandB
内置回调 4:TQDMProgressBar
自定义进度条刷新行为:
from lightning.pytorch.callbacks import TQDMProgressBar
progress_bar = TQDMProgressBar(
refresh_rate=10, # Update every 10 batches
process_position=0
)
trainer = L.Trainer(callbacks=[progress_bar])
内置回调 5:GradientAccumulationScheduler
训练过程中动态调整梯度累积步数:
from lightning.pytorch.callbacks import GradientAccumulationScheduler
# Accumulate more gradients as training progresses
accumulator = GradientAccumulationScheduler(
scheduling={
0: 8, # Epochs 0-4: accumulate 8 batches
5: 4, # Epochs 5-9: accumulate 4 batches
10: 2 # Epochs 10+: accumulate 2 batches
}
)
trainer = L.Trainer(callbacks=[accumulator])
与 Trainer(accumulate_grad_batches=4) 的固定配置相比,这个回调允许按 epoch 阶段动态变化累积数——典型场景是训练早期用大累积模拟大 batch,后期减小以加快收敛。
内置回调 6:StochasticWeightAveraging(SWA)
对权重做平均以获得更好的泛化:
from lightning.pytorch.callbacks import StochasticWeightAveraging
swa = StochasticWeightAveraging(
swa_lrs=1e-2, # SWA learning rate
swa_epoch_start=0.8, # Start at 80% of training
annealing_epochs=10, # Annealing period
annealing_strategy='cos' # 'cos' or 'linear'
)
trainer = L.Trainer(callbacks=[swa])
自定义回调
基础自定义回调——只需继承 Callback 并覆写对应 Hook:
from lightning.pytorch.callbacks import Callback
class PrintingCallback(Callback):
def on_train_start(self, trainer, pl_module):
print("Training is starting!")
def on_train_end(self, trainer, pl_module):
print("Training is done!")
def on_epoch_end(self, trainer, pl_module):
print(f"Epoch {trainer.current_epoch} ended")
# Use it
trainer = L.Trainer(callbacks=[PrintingCallback()])
进阶:带状态的指标记录回调:
class MetricsCallback(Callback):
"""Logs custom metrics every N batches."""
def __init__(self, log_every_n_batches=100):
self.log_every_n_batches = log_every_n_batches
self.metrics = []
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
if batch_idx % self.log_every_n_batches == 0:
# Compute custom metric
metric = self.compute_metric(outputs)
self.metrics.append(metric)
# Log to Lightning
pl_module.log('custom_metric', metric)
def compute_metric(self, outputs):
# Your custom logic
return outputs['loss'].item()
def state_dict(self):
"""Save callback state in checkpoint."""
return {'metrics': self.metrics}
def load_state_dict(self, state_dict):
"""Restore callback state from checkpoint."""
self.metrics = state_dict['metrics']
实现 state_dict / load_state_dict 后,回调的自定义状态会随 checkpoint 一起保存与恢复,保证断点续训后回调状态不丢失。
梯度监控回调——用于检测梯度爆炸:
class GradientMonitorCallback(Callback):
"""Monitor gradient norms."""
def on_after_backward(self, trainer, pl_module):
# Compute gradient norm
total_norm = 0.0
for p in pl_module.parameters():
if p.grad is not None:
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
# Log
pl_module.log('grad_norm', total_norm)
# Warn if exploding
if total_norm > 100:
print(f"Warning: Large gradient norm: {total_norm:.2f}")
模型检查回调——通过注册 forward hook 观测中间层激活:
class ModelInspectionCallback(Callback):
"""Inspect model activations during training."""
def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
if batch_idx == 0: # First batch of epoch
# Register hooks
self.activations = {}
def get_activation(name):
def hook(model, input, output):
self.activations[name] = output.detach()
return hook
# Attach to specific layers
pl_module.model.layer1.register_forward_hook(get_activation('layer1'))
pl_module.model.layer2.register_forward_hook(get_activation('layer2'))
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
if batch_idx == 0:
# Log activation statistics
for name, activation in self.activations.items():
mean = activation.mean().item()
std = activation.std().item()
pl_module.log(f'{name}_mean', mean)
pl_module.log(f'{name}_std', std)
全部回调 Hook 一览
回调的生命周期覆盖 fit / validate / test / predict 的完整过程,常用 Hook 如下:
class MyCallback(Callback):
# Setup/Teardown
def setup(self, trainer, pl_module, stage):
"""Called at beginning of fit/test/predict."""
pass
def teardown(self, trainer, pl_module, stage):
"""Called at end of fit/test/predict."""
pass
# Training
def on_train_start(self, trainer, pl_module):
pass
def on_train_epoch_start(self, trainer, pl_module):
pass
def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
pass
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
pass
def on_train_epoch_end(self, trainer, pl_module):
pass
def on_train_end(self, trainer, pl_module):
pass
# Validation
def on_validation_start(self, trainer, pl_module):
pass
def on_validation_epoch_start(self, trainer, pl_module):
pass
def on_validation_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
pass
def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
pass
def on_validation_epoch_end(self, trainer, pl_module):
pass
def on_validation_end(self, trainer, pl_module):
pass
# Test (same structure as validation)
def on_test_start(self, trainer, pl_module):
pass
# ... (test_epoch_start, test_batch_start, etc.)
# Predict
def on_predict_start(self, trainer, pl_module):
pass
# ... (predict_epoch_start, predict_batch_start, etc.)
# Backward
def on_before_backward(self, trainer, pl_module, loss):
pass
def on_after_backward(self, trainer, pl_module):
pass
# Optimizer
def on_before_optimizer_step(self, trainer, pl_module, optimizer):
pass
# Checkpointing
def on_save_checkpoint(self, trainer, pl_module, checkpoint):
"""Add data to checkpoint."""
pass
def on_load_checkpoint(self, trainer, pl_module, checkpoint):
"""Restore data from checkpoint."""
pass
组合多个回调与最佳实践
多回调组合——按需叠加,互不影响:
from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor
# Create all callbacks
checkpoint = ModelCheckpoint(monitor='val_loss', mode='min', save_top_k=3)
early_stop = EarlyStopping(monitor='val_loss', patience=5)
lr_monitor = LearningRateMonitor(logging_interval='epoch')
custom_callback = MyCustomCallback()
# Add all to Trainer
trainer = L.Trainer(
callbacks=[checkpoint, early_stop, lr_monitor, custom_callback]
)
trainer.fit(model, train_loader, val_loader)
执行顺序:回调按加入列表的顺序依次执行。
最佳实践要点(来自参考文档):
- 保持回调独立:不要在回调里假设其他回调一定存在。例如通过
trainer.checkpoint_callback.best_model_path访问保存路径是脆弱的写法(Bad),更稳妥的方式是遍历trainer.callbacks用isinstance判断后再访问(Good)。 - 用 state_dict 持久化:有状态回调务必实现
state_dict/load_state_dict,否则断点续训会丢失回调内部状态。 - 处理分布式环境:多卡训练时,用
trainer.is_global_zero判断是否只在主进程执行打印等操作,避免日志重复输出 8 份。
分布式训练策略
Lightning 支持通过一个参数切换多种分布式策略。以下内容对应 references/distributed.md 的完整说明。
策略 1:DDP(DistributedDataParallel)
多卡默认策略:
# Automatic DDP on all available GPUs
trainer = L.Trainer(accelerator='gpu', devices=4, strategy='ddp')
# Or auto-detect
trainer = L.Trainer(accelerator='gpu', devices='auto')
DDP 的工作原理:
- 每个 GPU 上复制一份模型
- 每个 GPU 处理不同的 batch
- 梯度跨 GPU 做 all-reduce 归约
- 模型权重保持同步
启动方式——Lightning 自动负责进程的 spawn:
# Lightning handles spawning processes automatically
python train.py
DDP 高级配置:
from lightning.pytorch.strategies import DDPStrategy
strategy = DDPStrategy(
find_unused_parameters=False, # Set True if model has unused params
gradient_as_bucket_view=True, # Memory optimization
static_graph=False, # Set True if graph doesn't change
)
trainer = L.Trainer(strategy=strategy)
其中 gradient_as_bucket_view=True 通过让梯度视图复用 bucket 内存来降低显存拷贝;static_graph=True 适用于计算图结构固定不变的模型(典型如 Transformer),可以提升通信效率。
策略 2:FSDP(Fully Sharded Data Parallel)
面向大模型(7B 以上参数):
from lightning.pytorch.strategies import FSDPStrategy
strategy = FSDPStrategy(
sharding_strategy="FULL_SHARD", # ZeRO-3 equivalent
activation_checkpointing=None, # Or specify layer types
cpu_offload=False, # CPU offload for memory
)
trainer = L.Trainer(
accelerator='gpu',
devices=8,
strategy=strategy,
precision='bf16' # Recommended with FSDP
)
trainer.fit(model, train_loader)
三种分片策略对比:
# FULL_SHARD (most memory efficient, equivalent to ZeRO-3)
strategy = FSDPStrategy(sharding_strategy="FULL_SHARD")
# SHARD_GRAD_OP (less memory efficient, equivalent to ZeRO-2)
strategy = FSDPStrategy(sharding_strategy="SHARD_GRAD_OP")
# NO_SHARD (no sharding, like DDP)
strategy = FSDPStrategy(sharding_strategy="NO_SHARD")
自动包装策略(把 Transformer Block 作为 FSDP 分片单元):
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.gpt2.modeling_gpt2 import GPT2Block
import functools
auto_wrap_policy = functools.partial(
transformer_auto_wrap_policy,
transformer_layer_cls={GPT2Block}
)
strategy = FSDPStrategy(
auto_wrap_policy=auto_wrap_policy,
activation_checkpointing_policy={GPT2Block} # Checkpoint these blocks
)
对 Transformer 类模型,把每个 block 作为 FSDP 分片单位、并对 block 做激活检查点(activation checkpointing),是训练大模型的标准内存优化组合。
策略 3:DeepSpeed
面向超大模型(70B 以上参数):
from lightning.pytorch.strategies import DeepSpeedStrategy
# DeepSpeed ZeRO-3 with CPU offload
strategy = DeepSpeedStrategy(
stage=3, # ZeRO-3
offload_optimizer=True, # CPU offload optimizer
offload_parameters=True, # CPU offload parameters
cpu_checkpointing=True, # Checkpoint to CPU
)
trainer = L.Trainer(
accelerator='gpu',
devices=8,
strategy=strategy,
precision='bf16'
)
trainer.fit(model, train_loader)
DeepSpeed 配置文件(deepspeed_config.json):
{
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"gradient_accumulation_steps": "auto",
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"reduce_bucket_size": 5e8,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e6
},
"bf16": {
"enabled": true
}
}
使用配置文件:
strategy = DeepSpeedStrategy(config='deepspeed_config.json')
trainer = L.Trainer(strategy=strategy)
注意配置中的 "auto" 值:DeepSpeed 会从 Trainer 和 DataLoader 自动推断 batch 大小与梯度累积步数,避免手填出错。
策略 4:DDP Spawn
Windows 与 Jupyter 环境兼容的 DDP 变体:
# Use when DDP doesn't work (e.g., Windows, Jupyter)
trainer = L.Trainer(
accelerator='gpu',
devices=2,
strategy='ddp_spawn' # Spawns new processes
)
注意:由于进程 spawn 有额外开销,ddp_spawn 比标准 DDP 慢,仅在不支持 fork 启动方式的环境中才需要。
多节点训练
节点 0(master):
export MASTER_ADDR=192.168.1.100
export MASTER_PORT=12355
export WORLD_SIZE=16 # 2 nodes × 8 GPUs
export NODE_RANK=0
python train.py
节点 1(worker):
export MASTER_ADDR=192.168.1.100
export MASTER_PORT=12355
export WORLD_SIZE=16
export NODE_RANK=1
python train.py
训练脚本——只需声明 num_nodes:
trainer = L.Trainer(
accelerator='gpu',
devices=8, # GPUs per node
num_nodes=2, # Total nodes
strategy='ddp'
)
trainer.fit(model, train_loader)
SLURM 与 Kubernetes 集成
SLURM 作业脚本:
#!/bin/bash
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8
#SBATCH --time=24:00:00
# Lightning auto-detects SLURM environment
srun python train.py
训练脚本无需改动,Lightning 自动读取 SLURM 环境变量:
# Lightning automatically reads SLURM environment variables
trainer = L.Trainer(
accelerator='gpu',
devices=8,
num_nodes=4, # From SBATCH --nodes
strategy='ddp'
)
Kubernetes 环境:Lightning 同样自动检测 K8s 环境变量:
import os
# Lightning auto-detects Kubernetes
trainer = L.Trainer(
accelerator='gpu',
devices=int(os.getenv('WORLD_SIZE', 1)),
strategy='ddp'
)
混合精度训练
BF16(A100/H100,推荐):
trainer = L.Trainer(
precision='bf16', # Or 'bf16-mixed'
accelerator='gpu'
)
BF16 的优势:不需要梯度缩放器(gradient scaler)、动态范围与 FP32 相同、可带来约 2 倍加速与 50% 显存节省。
FP16(V100 及更早 GPU):
trainer = L.Trainer(
precision='16-mixed', # Or just '16'
accelerator='gpu'
)
FP16 的自动梯度缩放由 Lightning 处理。
FP8(H100):
# Requires transformer_engine
# pip install transformer-engine[pytorch]
trainer = L.Trainer(
precision='transformer-engine',
accelerator='gpu'
)
FP8 需要额外安装 transformer-engine,在 H100 上可取得比 BF16 更快的训练速度。
梯度累积
用累积模拟更大的 batch size:
trainer = L.Trainer(
accumulate_grad_batches=4, # Accumulate 4 batches
precision='bf16'
)
# Effective batch = batch_size × accumulate_grad_batches × num_gpus
# Example: 32 × 4 × 8 = 1024
动态累积(训练早期累积更多):
# Accumulate more early in training
trainer = L.Trainer(
accumulate_grad_batches={
0: 8, # Epochs 0-4: accumulate 8
5: 4, # Epochs 5-9: accumulate 4
10: 2 # Epochs 10+: accumulate 2
}
)
分布式下的 Checkpoint
保存——默认只有 rank 0 进程保存,避免多进程同时写盘冲突:
from lightning.pytorch.callbacks import ModelCheckpoint
# Only rank 0 saves by default
checkpoint = ModelCheckpoint(
dirpath='checkpoints/',
filename='model-{epoch:02d}',
save_top_k=3
)
trainer = L.Trainer(callbacks=[checkpoint], strategy='ddp')
trainer.fit(model, train_loader)
手动保存:
class MyModel(L.LightningModule):
def training_step(self, batch, batch_idx):
# Training...
loss = ...
# Save every 1000 steps (only rank 0)
if batch_idx % 1000 == 0 and self.trainer.is_global_zero:
self.trainer.save_checkpoint(f'checkpoint_step_{batch_idx}.ckpt')
return loss
加载:
# Resume training
trainer = L.Trainer(strategy='ddp')
trainer.fit(model, train_loader, ckpt_path='checkpoints/last.ckpt')
# Load for inference
model = MyModel.load_from_checkpoint('checkpoints/best.ckpt')
model.eval()
策略对比一览
| Strategy | Memory Efficiency | Speed | Use Case |
|---|---|---|---|
| DDP | Low | Fast | Small models (<7B), single node |
| FSDP | High | Medium | Large models (7-70B) |
| DeepSpeed ZeRO-2 | Medium | Fast | Medium models (1-13B) |
| DeepSpeed ZeRO-3 | Very High | Slower | Massive models (70B+) |
| DDP Spawn | Low | Slow | Windows, debugging |
分布式最佳实践
1. 按模型规模选策略:
# Model size guide
if model_params < 1e9: # <1B
strategy = 'ddp'
elif model_params < 7e9: # 1-7B
strategy = 'ddp' or DeepSpeedStrategy(stage=2)
elif model_params < 70e9: # 7-70B
strategy = FSDPStrategy(sharding_strategy="FULL_SHARD")
else: # 70B+
strategy = DeepSpeedStrategy(stage=3, offload_optimizer=True)
trainer = L.Trainer(strategy=strategy)
2. 避免同步问题——分布式下每个进程都会执行 training_step,日志操作要限制在主进程:
class MyModel(L.LightningModule):
def training_step(self, batch, batch_idx):
# WRONG: This runs on all GPUs independently
if batch_idx % 100 == 0:
self.log_something() # Logged 8 times on 8 GPUs!
# CORRECT: Use is_global_zero
if batch_idx % 100 == 0 and self.trainer.is_global_zero:
self.log_something() # Logged once
loss = ...
return loss
3. 高效数据加载——Lightning 在 DDP 下会自动包装 DistributedSampler,无需手写:
from torch.utils.data import DataLoader, DistributedSampler
# Lightning handles DistributedSampler automatically
train_loader = DataLoader(
dataset,
batch_size=32,
num_workers=4, # 4 workers per GPU
pin_memory=True,
persistent_workers=True
)
# Lightning automatically wraps with DistributedSampler in DDP
trainer.fit(model, train_loader)
4. 降低通信开销:
from lightning.pytorch.strategies import DDPStrategy
strategy = DDPStrategy(
gradient_as_bucket_view=True, # Reduce memory copies
static_graph=True, # If model graph doesn't change (faster)
)
trainer = L.Trainer(strategy=strategy)
分布式常见问题
NCCL 超时:表现为训练卡住并报 NCCL timeout。
方案一:增大超时时间
export NCCL_TIMEOUT=3600 # 1 hour
python train.py
方案二:检查网络连通性
# Test inter-node communication
nvidia-smi nvlink -s
# Verify all nodes can ping each other
ping <node-2-ip>
FSDP 下 OOM:开启 CPU offload
strategy = FSDPStrategy(
sharding_strategy="FULL_SHARD",
cpu_offload=True # Offload to CPU
)
DDP 下结果不一致:原因是各 GPU 随机种子不同,需统一设置种子
class MyModel(L.LightningModule):
def __init__(self):
super().__init__()
L.seed_everything(42, workers=True) # Same seed everywhere
DeepSpeed 配置报错:改用 Lightning 自动生成配置
strategy = DeepSpeedStrategy(
stage=3,
# Don't specify config file, Lightning generates automatically
)
超参数调优
Lightning 与主流调参框架(Ray Tune、Optuna、WandB Sweeps、Hyperopt)均有开箱即用的集成,同时自带 LR Finder 与 Batch Size Finder。以下内容对应 references/hyperparameter-tuning.md。
集成 1:Ray Tune
安装:
pip install ray[tune]
pip install lightning
基础示例:
import lightning as L
from ray import tune
from ray.tune.integration.pytorch_lightning import TuneReportCallback
class LitModel(L.LightningModule):
def __init__(self, lr, batch_size):
super().__init__()
self.lr = lr
self.batch_size = batch_size
self.model = nn.Sequential(nn.Linear(10, 128), nn.ReLU(), nn.Linear(128, 1))
def training_step(self, batch, batch_idx):
loss = self.model(batch).mean()
self.log('train_loss', loss)
return loss
def validation_step(self, batch, batch_idx):
val_loss = self.model(batch).mean()
self.log('val_loss', val_loss)
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=self.lr)
def train_fn(config):
"""Training function for Ray Tune."""
model = LitModel(lr=config["lr"], batch_size=config["batch_size"])
# Add callback to report metrics to Tune
trainer = L.Trainer(
max_epochs=10,
callbacks=[TuneReportCallback({"loss": "val_loss"}, on="validation_end")]
)
trainer.fit(model, train_loader, val_loader)
# Define search space
config = {
"lr": tune.loguniform(1e-5, 1e-1),
"batch_size": tune.choice([16, 32, 64, 128])
}
# Run hyperparameter search
analysis = tune.run(
train_fn,
config=config,
num_samples=20, # 20 trials
resources_per_trial={"gpu": 1}
)
# Best hyperparameters
best_config = analysis.get_best_config(metric="loss", mode="min")
print(f"Best config: {best_config}")
核心衔接点是 TuneReportCallback:它把 Lightning 的 val_loss 指标实时上报给 Tune,Tune 据此进行搜索与剪枝。
进阶:Population-Based Training(PBT):
from ray.tune.schedulers import PopulationBasedTraining
# PBT scheduler
scheduler = PopulationBasedTraining(
time_attr='training_iteration',
metric='val_loss',
mode='min',
perturbation_interval=5, # Perturb every 5 epochs
hyperparam_mutations={
"lr": tune.loguniform(1e-5, 1e-1),
"batch_size": [16, 32, 64, 128]
}
)
analysis = tune.run(
train_fn,
config=config,
num_samples=8, # Population size
scheduler=scheduler,
resources_per_trial={"gpu": 1}
)
集成 2:Optuna
安装:
pip install optuna
pip install optuna-integration
示例:
import optuna
from optuna.integration import PyTorchLightningPruningCallback
def objective(trial):
# Suggest hyperparameters
lr = trial.suggest_loguniform('lr', 1e-5, 1e-1)
batch_size = trial.suggest_categorical('batch_size', [16, 32, 64, 128])
n_layers = trial.suggest_int('n_layers', 1, 3)
hidden_size = trial.suggest_int('hidden_size', 64, 512, step=64)
# Create model
model = LitModel(lr=lr, n_layers=n_layers, hidden_size=hidden_size)
# Pruning callback (early stopping for bad trials)
pruning_callback = PyTorchLightningPruningCallback(trial, monitor="val_loss")
trainer = L.Trainer(
max_epochs=20,
callbacks=[pruning_callback],
enable_progress_bar=False,
logger=False
)
trainer.fit(model, train_loader, val_loader)
return trainer.callback_metrics["val_loss"].item()
# Create study
study = optuna.create_study(
direction='minimize',
pruner=optuna.pruners.MedianPruner() # Prune bad trials early
)
# Optimize
study.optimize(objective, n_trials=50, timeout=3600)
# Best params
print(f"Best trial: {study.best_trial.params}")
print(f"Best value: {study.best_value}")
# Visualization
optuna.visualization.plot_optimization_history(study).show()
optuna.visualization.plot_param_importances(study).show()
PyTorchLightningPruningCallback 把早停决策交给 Optuna 的 pruner,配合 MedianPruner 可提前终止表现差的 trial,显著节约算力。
Optuna 分布式调参(共享数据库):
import optuna
# Shared database for distributed optimization
storage = optuna.storages.RDBStorage(
url='postgresql://user:pass@localhost/optuna'
)
study = optuna.create_study(
study_name='distributed_study',
storage=storage,
load_if_exists=True,
direction='minimize'
)
# Run on multiple machines
study.optimize(objective, n_trials=50)
集成 3:Weights & Biases(WandB)Sweeps
安装:
pip install wandb
sweep 配置(sweep.yaml):
program: train.py
method: bayes
metric:
name: val_loss
goal: minimize
parameters:
lr:
distribution: log_uniform_values
min: 0.00001
max: 0.1
batch_size:
values: [16, 32, 64, 128]
optimizer:
values: ['adam', 'sgd', 'adamw']
dropout:
distribution: uniform
min: 0.0
max: 0.5
训练脚本(train.py):
import wandb
import lightning as L
from lightning.pytorch.loggers import WandbLogger
def train():
# Initialize wandb
wandb.init()
config = wandb.config
# Create model with sweep params
model = LitModel(
lr=config.lr,
batch_size=config.batch_size,
optimizer=config.optimizer,
dropout=config.dropout
)
# WandB logger
wandb_logger = WandbLogger(project='hyperparameter-sweep')
trainer = L.Trainer(
max_epochs=20,
logger=wandb_logger
)
trainer.fit(model, train_loader, val_loader)
if __name__ == '__main__':
train()
启动 sweep:
# Initialize sweep
wandb sweep sweep.yaml
# Output: wandb: Created sweep with ID: abc123
# Run agent (can run on multiple machines)
wandb agent your-entity/your-project/abc123
集成 4:Hyperopt
安装:
pip install hyperopt
示例:
from hyperopt import hp, fmin, tpe, Trials
def objective(params):
model = LitModel(
lr=params['lr'],
batch_size=int(params['batch_size']),
hidden_size=int(params['hidden_size'])
)
trainer = L.Trainer(
max_epochs=10,
enable_progress_bar=False,
logger=False
)
trainer.fit(model, train_loader, val_loader)
# Return loss (minimize)
return trainer.callback_metrics["val_loss"].item()
# Define search space
space = {
'lr': hp.loguniform('lr', np.log(1e-5), np.log(1e-1)),
'batch_size': hp.quniform('batch_size', 16, 128, 16),
'hidden_size': hp.quniform('hidden_size', 64, 512, 64)
}
# Optimize
trials = Trials()
best = fmin(
fn=objective,
space=space,
algo=tpe.suggest, # Tree-structured Parzen Estimator
max_evals=50,
trials=trials
)
print(f"Best hyperparameters: {best}")
注意 Hyperopt 的 quniform 返回浮点数,构造模型时需 int() 转换。
内置调参:自动学习率搜索与 Batch Size 搜索
Auto Learning Rate Finder:
class LitModel(L.LightningModule):
def __init__(self, lr=1e-3):
super().__init__()
self.lr = lr
self.model = nn.Linear(10, 1)
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=self.lr)
def training_step(self, batch, batch_idx):
loss = self.model(batch).mean()
return loss
# Find optimal learning rate
model = LitModel()
trainer = L.Trainer(auto_lr_find=True)
# This runs LR finder before training
trainer.tune(model, train_loader)
# Or manually
from lightning.pytorch.tuner import Tuner
tuner = Tuner(trainer)
lr_finder = tuner.lr_find(model, train_loader)
# Plot results
fig = lr_finder.plot(suggest=True)
fig.show()
# Get suggested LR
suggested_lr = lr_finder.suggestion()
print(f"Suggested LR: {suggested_lr}")
# Update model
model.lr = suggested_lr
# Train with optimal LR
trainer.fit(model, train_loader)
Auto Batch Size Finder:
class LitModel(L.LightningModule):
def __init__(self, batch_size=32):
super().__init__()
self.batch_size = batch_size
self.model = nn.Linear(10, 1)
def train_dataloader(self):
return DataLoader(dataset, batch_size=self.batch_size)
model = LitModel()
trainer = L.Trainer(auto_scale_batch_size='binsearch')
# Find optimal batch size
trainer.tune(model)
print(f"Optimal batch size: {model.batch_size}")
# Train with optimal batch size
trainer.fit(model, train_loader)
auto_scale_batch_size='binsearch' 使用二分搜索在显存允许范围内寻找最大 batch size;trainer.tune() 会先执行 finder,再把结果写回模型/训练器供后续 fit 使用。
进阶调参策略
多保真度优化——Successive Halving(ASHA):
from ray.tune.schedulers import ASHAScheduler
# ASHA: Asynchronous Successive Halving Algorithm
scheduler = ASHAScheduler(
max_t=100, # Max epochs
grace_period=10, # Min epochs before stopping
reduction_factor=2 # Halve resources each round
)
analysis = tune.run(
train_fn,
config=config,
num_samples=64,
scheduler=scheduler,
resources_per_trial={"gpu": 1}
)
其工作方式:启动 64 个 trial → 10 epoch 后淘汰表现最差的 50%(剩 32)→ 20 epoch 后再淘汰一半(剩 16)→ 依次类推,最后把剩余 4 个 trial 跑到 100 epoch。用少量资源快速淘汰差 trial,是调参效率的核心手段。
贝叶斯优化:
from ray.tune.search.bayesopt import BayesOptSearch
search = BayesOptSearch(
metric="val_loss",
mode="min"
)
analysis = tune.run(
train_fn,
config=config,
num_samples=50,
search_alg=search,
resources_per_trial={"gpu": 1}
)
网格搜索(穷举,trial 数 = 各维度取值数的乘积):
from ray import tune
# Exhaustive grid search
config = {
"lr": tune.grid_search([1e-5, 1e-4, 1e-3, 1e-2]),
"batch_size": tune.grid_search([16, 32, 64, 128]),
"optimizer": tune.grid_search(['adam', 'sgd', 'adamw'])
}
# Total trials: 4 × 4 × 3 = 48
analysis = tune.run(train_fn, config=config)
随机搜索:
config = {
"lr": tune.loguniform(1e-5, 1e-1),
"batch_size": tune.choice([16, 32, 64, 128]),
"dropout": tune.uniform(0.0, 0.5),
"hidden_size": tune.randint(64, 512)
}
# Random sampling
analysis = tune.run(
train_fn,
config=config,
num_samples=100 # 100 random samples
)
调参最佳实践
1. 由粗到细分阶段搜索:
# Phase 1: Coarse search (fast)
coarse_config = {
"lr": tune.loguniform(1e-5, 1e-1),
"batch_size": tune.choice([32, 64])
}
coarse_analysis = tune.run(train_fn, config=coarse_config, num_samples=10, max_epochs=5)
# Phase 2: Fine-tune around best (slow)
best_lr = coarse_analysis.best_config["lr"]
fine_config = {
"lr": tune.uniform(best_lr * 0.5, best_lr * 2),
"batch_size": tune.choice([16, 32, 64, 128])
}
fine_analysis = tune.run(train_fn, config=fine_config, num_samples=20, max_epochs=20)
2. 配合 Checkpoint 做可恢复调参:
def train_fn(config, checkpoint_dir=None):
model = LitModel(lr=config["lr"])
trainer = L.Trainer(
max_epochs=100,
callbacks=[
TuneReportCheckpointCallback(
metrics={"loss": "val_loss"},
filename="checkpoint",
on="validation_end"
)
]
)
# Resume from checkpoint if exists
ckpt_path = None
if checkpoint_dir:
ckpt_path = os.path.join(checkpoint_dir, "checkpoint")
trainer.fit(model, train_loader, val_loader, ckpt_path=ckpt_path)
3. 监控资源占用:
import GPUtil
def train_fn(config):
# Before training
GPUs = GPUtil.getGPUs()
print(f"GPU memory before: {GPUs[0].memoryUsed} MB")
# Train
model = LitModel(lr=config["lr"], batch_size=config["batch_size"])
trainer.fit(model, train_loader)
# After training
GPUs = GPUtil.getGPUs()
print(f"GPU memory after: {GPUs[0].memoryUsed} MB")
调参常见问题
Trial OOM:降低并发 trial 数或 batch size
analysis = tune.run(
train_fn,
config=config,
resources_per_trial={"gpu": 0.5}, # 2 trials per GPU
max_concurrent_trials=2 # Limit concurrent trials
)
调参太慢:使用早期停止调度器
from ray.tune.schedulers import ASHAScheduler
scheduler = ASHAScheduler(
max_t=100,
grace_period=5, # Stop bad trials after 5 epochs
reduction_factor=3
)
无法复现最优 trial:在训练函数里固定种子
def train_fn(config):
L.seed_everything(42, workers=True)
# Rest of training...
何时使用 Lightning:与替代方案的取舍
适合选择 PyTorch Lightning 的场景:
- 想要干净、组织良好的代码结构
- 需要生产级训练循环
- 需要在单卡 GPU、多卡 GPU、TPU 之间切换
- 需要内置的回调与日志能力
- 团队协作(标准化代码结构,方便代码评审与复用)
核心优势总结:
- 组织清晰:研究代码与工程代码分离,研究者只需实现 Hook 方法
- 自动化:DDP、FSDP、DeepSpeed 一行切换
- 回调机制:模块化的训练扩展点
- 可复现:样板代码更少,意味着出 bug 的面更小
- 久经考验:SKILL 文档的资源节记录其参考版本为 2.5.5+,并提到被 Kaggle 竞赛、研究实验室与生产团队广泛使用
适合选择其他方案的场景:
- HuggingFace Accelerate:对现有代码改动最小、更灵活(本仓库对应的技能见 accelerate/SKILL.md,其思路是给任意 PyTorch 脚本加 4 行代码即可获得分布式能力)
- Ray Train:多节点编排与超参数调优一体化的场景
- 原生 PyTorch:追求最大控制力、学习底层原理的场景
- Keras:TensorFlow 生态下的场景
常见问题排查(来自 SKILL.md)
问题:loss 不下降
先检查数据和模型设置,在 training_step 里打印 batch 信息:
# 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
问题:显存不足(OOM)
减小 batch size 或开启梯度累积与混合精度:
trainer = L.Trainer(
accumulate_grad_batches=4, # Effective batch = batch_size × 4
precision='bf16' # Or 'fp16', reduces memory 50%
)
问题:验证没有运行
确保把 val_loader 传给了 trainer.fit:
# 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)
硬件与精度支持
硬件支持矩阵(来自 SKILL.md):
- CPU:支持(适合调试)
- 单卡 GPU:支持
- 多卡 GPU:DDP(默认)、FSDP 或 DeepSpeed
- 多节点:DDP、FSDP、DeepSpeed
- TPU:支持(8 核)
- Apple MPS:支持
精度选项:
- FP32(默认)
- FP16(V100 及更早 GPU)
- BF16(A100/H100,推荐)
- FP8(H100)
在 AI Research Skills 仓库中的定位与延伸阅读
本技能位于仓库 08-distributed-training/pytorch-lightning/ 目录,属于「08-distributed-training」分布式训练分类(该分类共 6 个技能,还包括 Megatron-Core、DeepSpeed、FSDP2、Accelerate、Ray Train)。它是 AI 研究全生命周期中「训练工程」环节的高层封装层——上接模型架构定义,下接分布式底层(DDP/FSDP/DeepSpeed),并与本仓库的 accelerate(轻量接入)形成互补。
三份参考文档是本文内容的深度延伸:
- references/callbacks.md:回调 API 全解,覆盖全部内置回调与自定义回调 Hook
- references/distributed.md:分布式策略全解,覆盖 DDP/FSDP/DeepSpeed 配置、多节点、混合精度与策略选型
- references/hyperparameter-tuning.md:超参调优全解,覆盖 Ray Tune/Optuna/WandB/Hyperopt 集成与内置 Finder
这三个文件与本 SKILL.md 共同构成了一个完整的 Lightning 实战知识包,配合仓库中其他技能(如 10-optimization 下的量化与显存优化、13-mlops 下的实验跟踪),即可支撑从单机调试到大模型多机分布式训练再到系统化调参的完整研发链路。