完全分片数据并行

完全分片数据并行（FSDP）是一种数据并行方法， 它将模型的参数、梯度和优化器状态在可用 GPU（也称为 Worker 或 rank）的数量上进行分片。 与分布式数据并行（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 - 将模型参数、梯度和优化器状态跨 Worker 进行分片；为此选项选择 1
• SHARD_GRAD_OP- 将梯度和优化器状态跨 Worker 进行分片；为此选项选择 2
• NO_SHARD - 不分片任何内容（这等同于 DDP）；为此选项选择 3
• HYBRID_SHARD - 在每个 Worker 中分片模型参数、梯度和优化器状态，其中每个 Worker 也有完整副本；为此选项选择 4
• HYBRID_SHARD_ZERO2 - 在每个 Worker 中分片梯度和优化器状态，其中每个 Worker 也有完整副本；为此选项选择 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 来保存中间检查点， 因为在排名 0 上保存完整状态字典需要很长时间，通常会导致 NCCL Timeout 错误，因为在广播过程中会无限期挂起。 您可以使用 `load_state“ 方法加载分片状态字典以恢复训练。

# 包含检查点的目录
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 支持用于 TPUs 的 FSDP 训练， 可以通过修改由 accelerate config 生成的 FSDP 配置文件来启用。除了上面指定的分片策略和包装选项外， 您还可以将以下参数添加到文件中。

xla: True # 必须设置为 True 以启用 PyTorch/XLA
xla_fsdp_settings: # XLA 特定的 FSDP 参数
xla_fsdp_grad_ckpt: True # 使用梯度检查点

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 模型博文。