全分片数据并行

完全分片数据并行 (FSDP) 是一种数据并行方法，它将模型的参数、梯度和优化器状态分片到可用 GPU 的数量上（也称为工作进程或rank）。与 分布式数据并行 (DDP) 不同，FSDP 减少了内存使用，因为它不会在每个 GPU 上复制模型。这提高了 GPU 内存效率，并允许您在更少的 GPU 上训练更大的模型。FSDP 与 Accelerate 集成，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 可以降低大规模训练的高成本。如果您有兴趣了解更多信息，以下内容可能会有所帮助

• 按照 Accelerate 的更深入指南了解 FSDP。
• 阅读 介绍 PyTorch 完全分片数据并行 (FSDP) API 博客文章。
• 阅读使用 FSDP 在云 TPU 上扩展 PyTorch 模型 博客文章。