完全分片数据并行(FSDP)是一种数据并行方法,
它将模型的参数、梯度和优化器状态在可用 GPU(也称为 Worker 或 rank)的数量上进行分片。
与分布式数据并行(DDP)不同,
FSDP 减少了内存使用量,因为模型在每个 GPU 上都被复制了一次。这就提高了 GPU 内存效率,
使您能够用较少的 GPU 训练更大的模型。FSDP 已经集成到 Accelerate 中,
这是一个用于在分布式环境中轻松管理训练的库,这意味着可以从 [Trainer
] 类中调用这个库。
在开始之前,请确保已安装 Accelerate,并且至少使用 PyTorch 2.1.0 或更高版本。
pip install accelerate
首先,运行 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 上,以节省更多 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
错误,因为在广播过程中会无限期挂起。
您可以使用 [~accelerate.Accelerator.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)
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 模型博文。