Transformers 文档

完全分片数据并行

完全分片数据并行

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_WRAPmin_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 可以降低大规模训练的高成本。如果你有兴趣了解更多,以下内容可能会有所帮助:

< > Update on GitHub