[ PROMPT_NODE_22330 ]
distributed
[ SKILL_DOCUMENTATION ]
# PyTorch Lightning 分布式训练
## 分布式策略
Lightning 通过简单的参数更改支持多种分布式策略。
### 1. DDP (DistributedDataParallel)
**多 GPU 的默认策略**:
python
# 在所有可用 GPU 上自动使用 DDP
trainer = L.Trainer(accelerator='gpu', devices=4, strategy='ddp')
# 或自动检测
trainer = L.Trainer(accelerator='gpu', devices='auto')
**DDP 工作原理**:
- 在每个 GPU 上复制模型
- 每个 GPU 处理不同的批次
- 梯度在 GPU 之间进行全归约 (all-reduced)
- 模型权重同步
**启动**:
bash
# Lightning 自动处理进程生成
python train.py
**DDP 配置**:
python
from lightning.pytorch.strategies import DDPStrategy
strategy = DDPStrategy(
find_unused_parameters=False, # 如果模型有未使用的参数，设为 True
gradient_as_bucket_view=True, # 内存优化
static_graph=False, # 如果图不改变，设为 True
)
trainer = L.Trainer(strategy=strategy)
### 2. FSDP (Fully Sharded Data Parallel)
**用于大模型 (7B+ 参数)**:
python
from lightning.pytorch.strategies import FSDPStrategy
strategy = FSDPStrategy(
sharding_strategy="FULL_SHARD", # 等同于 ZeRO-3
activation_checkpointing=None, # 或指定层类型
cpu_offload=False, # 用于内存的 CPU 卸载
)
trainer = L.Trainer(
accelerator='gpu',
devices=8,
strategy=strategy,
precision='bf16' # 建议与 FSDP 一起使用
)
trainer.fit(model, train_loader)
**FSDP 分片策略**:
python
# FULL_SHARD (内存效率最高，等同于 ZeRO-3)
strategy = FSDPStrategy(sharding_strategy="FULL_SHARD")
# SHARD_GRAD_OP (内存效率较低，等同于 ZeRO-2)
strategy = FSDPStrategy(sharding_strategy="SHARD_GRAD_OP")
# NO_SHARD (无分片，类似于 DDP)
strategy = FSDPStrategy(sharding_strategy="NO_SHARD")
**自动包装策略** (包装 Transformer 块):
python
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} # 对这些块进行检查点处理
)
### 3. DeepSpeed
**用于超大规模模型 (70B+ 参数)**:
python
from lightning.pytorch.strategies import DeepSpeedS
粤公网安备44030002003366号