完全分片数据并行
Fully Sharded Data Parallel (FSDP) 是一种数据并行方法,它将模型的参数、梯度和优化器状态在可用的GPU数量(也称为工作器或rank)之间进行分片。与DistributedDataParallel (DDP)不同,FSDP减少了内存使用,因为模型在每个GPU上被复制。这提高了GPU的内存效率,并允许你在更少的GPU上训练更大的模型。FSDP与Accelerate集成,这是一个用于轻松管理分布式环境中训练的库,这意味着它可以从Trainer类中使用。
在开始之前,请确保已安装Accelerate并且至少是PyTorch 2.1.0或更新版本。
pip install accelerate
FSDP 配置
首先,运行accelerate config
命令来为您的训练环境创建一个配置文件。Accelerate使用此配置文件根据您在accelerate config
中选择的训练选项自动设置正确的训练环境。
accelerate config
当你运行 accelerate config
时,系统会提示你一系列选项来配置你的训练环境。本节介绍了一些最重要的 FSDP 选项。要了解更多关于其他可用的 FSDP 选项,请查看 fsdp_config 参数。
分片策略
FSDP 提供了多种分片策略供选择:
FULL_SHARD
- 在多个工作节点之间分片模型参数、梯度和优化器状态;选择1
作为此选项SHARD_GRAD_OP
- 跨工作节点分片梯度和优化器状态;选择2
作为此选项NO_SHARD
- 不进行任何分片(这相当于DDP);选择3
作为此选项HYBRID_SHARD
- 在每个工作节点内分片模型参数、梯度和优化器状态,其中每个工作节点还拥有完整的副本;选择4
作为此选项HYBRID_SHARD_ZERO2
- 在每个工作节点内分片梯度和优化器状态,其中每个工作节点还拥有完整的副本;选择5
作为此选项
这是通过fsdp_sharding_strategy
标志启用的。
CPU 卸载
你也可以在不使用参数和梯度时将它们卸载到CPU上,以节省更多的GPU内存,并帮助你适应即使FSDP也可能不足的大型模型。这可以通过在运行accelerate config
时设置fsdp_offload_params: true
来启用。
包装策略
FSDP通过将网络中的每一层进行包装来应用。这种包装通常以嵌套的方式进行,其中在每次前向传递后丢弃完整的权重,以节省内存供下一层使用。自动包装策略是实现这一点的最简单方法,您不需要更改任何代码。您应该选择fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
来包装Transformer层,并使用fsdp_transformer_layer_cls_to_wrap
来指定要包装的层(例如BertLayer
)。
否则,您可以选择基于大小的包装策略,如果某一层的参数超过一定数量,则应用FSDP。这可以通过设置fsdp_wrap_policy: SIZE_BASED_WRAP
和min_num_param
为所需的大小阈值来启用。
检查点
中间检查点应使用fsdp_state_dict_type: SHARDED_STATE_DICT
保存,因为在rank 0上使用CPU卸载保存完整状态字典会花费大量时间,并且由于在广播期间无限挂起,经常会导致NCCL Timeout
错误。您可以使用load_state方法从分片状态字典中恢复训练。
# directory containing checkpoints
accelerator.load_state("ckpt")
然而,当训练结束时,您希望保存完整的状态字典,因为分片状态字典仅与FSDP兼容。
if trainer.is_fsdp_enabled:
trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
trainer.save_model(script_args.output_dir)
TPU
PyTorch XLA 支持 TPU 的 FSDP 训练,可以通过修改由 accelerate config
生成的 FSDP 配置文件来启用。除了上面指定的分片策略和包装选项外,您还可以将以下参数添加到文件中。
xla: True # must be set to True to enable PyTorch/XLA
xla_fsdp_settings: # XLA-specific FSDP parameters
xla_fsdp_grad_ckpt: True # use gradient checkpointing
xla_fsdp_settings
允许您为FSDP配置额外的XLA特定参数。
启动训练
一个示例FSDP配置文件可能如下所示:
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
fsdp_backward_prefetch_policy: BACKWARD_PRE
fsdp_cpu_ram_efficient_loading: true
fsdp_forward_prefetch: false
fsdp_offload_params: true
fsdp_sharding_strategy: 1
fsdp_state_dict_type: SHARDED_STATE_DICT
fsdp_sync_module_states: true
fsdp_transformer_layer_cls_to_wrap: BertLayer
fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
要启动训练,请运行accelerate launch
命令,它将自动使用您之前使用accelerate config
创建的配置文件。
accelerate launch my-trainer-script.py
accelerate launch --fsdp="full shard" --fsdp_config="path/to/fsdp_config/ my-trainer-script.py
下一步
FSDP 可以成为训练真正大型模型的强大工具,前提是你拥有多个 GPU 或 TPU。通过分片模型参数、优化器和梯度状态,甚至在它们不活跃时将它们卸载到 CPU,FSDP 可以降低大规模训练的高成本。如果你有兴趣了解更多,以下内容可能会有所帮助:
- 跟随更深入的加速指南了解FSDP。
- 阅读Introducing PyTorch Fully Sharded Data Parallel (FSDP) API博客文章。
- 阅读Scaling PyTorch models on Cloud TPUs with FSDP博客文章。