FullyShardedDataParallel

完全分片数据并行 (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 超时错误。

在配置文件中指定 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。修改配置文件以包含以下参数。有关您可以为 FSDP 配置的其他 XLA 特定参数，请参阅 xla_fsdp_settings 参数。

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”博文。