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。修改配置文件以包含以下参数。有关 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 博客文章。