完全分片数据并行 (Fully Sharded Data Parallel)¶
完全分片数据并行 (FSDP) 是一种数据并行方法，它将模型的参数、梯度和优化器状态分布在可用的 GPU（也称为工作进程或 rank）上。与 DistributedDataParallel (DDP) 不同，FSDP 通过在每个 GPU 上复制模型来减少内存使用。这提高了 GPU 的内存效率，并允许你在更少的 GPU 上训练更大的模型。FSDP 集成了 Accelerate 库，这是一个用于轻松管理分布式环境中的训练的库，这意味着你可以在 Trainer 类中使用它。
在开始之前，请确保已安装 Accelerate 并且至少使用 PyTorch 2.1.0 或更新版本。
pip install accelerate
FSDP 配置¶
首先，运行以下命令以创建一个配置文件，用于设置你的训练环境。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 保存中间检查点，因为当 CPU 卸载在 rank 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 支持 TPU 的 FSDP 训练，可以通过修改 accelerate config 生成的 FSDP 配置文件来启用。除了上面指定的分片策略和包装选项，你还可以在文件中添加以下参数。
要启动训练，运行以下命令，它将自动使用你之前使用 accelerate config 创建的配置文件。
accelerate launch my-trainer-script.py
你也可以显式指定 FSDP 配置：
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 模型 博客文章。