FullyShardedDataParallel

Fully Sharded Data Parallel (FSDP) 是一种结合了数据并行和模型并行优点的分布式训练并行方法。

与DistributedDataParallel (DDP)不同，FSDP通过不在每个GPU上复制模型来节省更多内存。它将模型参数、梯度和优化器状态分片到各个GPU上。每个模型分片处理一部分数据，然后同步结果以加速训练。

本指南涵盖了如何设置使用FSDP和Accelerate（一个用于管理分布式训练的库）来训练模型。

pip install accelerate

配置选项

始终首先运行accelerate config命令，以帮助Accelerate设置正确的分布式训练环境。

accelerate config

下面部分将讨论一些更重要的FSDP配置选项。请参阅fsdp_config参数以了解其他可用选项。

分片策略

FSDP提供了几种分片策略来分发模型。请参阅下表以帮助您为您的设置选择最佳策略。使用配置文件中的fsdp_sharding_strategy参数指定策略。

CPU 卸载

在模型参数和梯度不使用时将其卸载到CPU，以节省额外的GPU内存。这对于模型即使使用FSDP仍然过大的情况很有用。

在配置文件中指定fsdp_offload_params: true以启用卸载。

包裹策略

FSDP通过包裹网络中的每个层来应用。包裹通常以嵌套的方式应用，其中在每个前向传播后丢弃完整的权重，以便为下一层节省内存。

有几种可用的包裹策略，但自动包裹策略最简单，不需要对您的代码进行任何更改。指定fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP来包裹Transformer层，并指定fsdp_transformer_layer_cls_to_wrap来确定要包裹的层（例如，BertLayer）。

也提供基于大小的包裹。如果一个层超过特定数量的参数，它就会被包裹。指定fsdp_wrap_policy: SIZED_BASED_WRAP和min_num_param来设置一个层被包裹的最小参数数量。

检查点

中间检查点应保存为分片状态字典，因为保存完整的状态字典（即使有CPU卸载）也非常耗时，并且可能会因广播过程中的无限挂起而导致NCCL Timeout错误。

在配置文件中指定fsdp_state_dict_type: SHARDED_STATE_DICT以保存分片状态字典。现在您可以使用load_state从分片状态字典恢复训练。

accelerator.load_state("directory/containing/checkpoints")

但是，一旦训练完成，您应该保存完整的状态字典，因为分片状态字典仅与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（一个用于在XLA设备上运行PyTorch的包）支持在TPU上使用FSDP。修改配置文件以包含以下参数。请参阅xla_fsdp_settings参数，了解您可以为FSDP配置的其他XLA特定参数。

xla: True # must be set to True to enable PyTorch/XLA
xla_fsdp_settings: # XLA specific FSDP parameters
xla_fsdp_grad_ckpt: True # enable gradient checkpointing

训练

运行accelerate config后，您的配置文件应已准备就绪。下面展示了一个在两台GPU上完全分片参数、梯度和优化器状态的示例配置文件。根据您设置配置的方式，您的文件可能会有所不同。

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命令来启动使用您在配置文件中选择的FSDP配置的训练脚本。

accelerate launch my-training-script.py

也可以直接在命令行中指定一些FSDP参数。

accelerate launch --fsdp="full shard" --fsdp_config="path/to/fsdp_config/" my-training-script.py

资源

与其他并行策略相比，FSDP是一个强大的工具，可用于使用更少的GPU训练大型模型。请参阅下面的资源，以了解更多关于FSDP的信息。

• 请按照更深入的Accelerate指南，了解FSDP。
• 阅读Introducing PyTorch Fully Sharded Data Parallel (FSDP) API博文。
• 阅读Scaling PyTorch models on Cloud TPUs with FSDP博文。