Accelerate 文档

完全分片数据并行

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

完全分片数据并行

为了以更大的批次规模加速训练大型模型,我们可以使用完全分片数据并行模型。这种数据并行范式通过对优化器状态、梯度和参数进行分片,从而能够容纳更多数据和更大的模型。要了解更多相关信息及其好处,请查看完全分片数据并行博客。我们集成了 PyTorch 最新的完全分片数据并行(FSDP)训练功能。您只需通过配置启用它即可。

开箱即用的工作方式

在您的机器上,只需运行

accelerate config

并回答所提出的问题。这将生成一个配置文件,在执行时会自动使用该文件来正确设置默认选项

accelerate launch my_script.py --args_to_my_script

例如,以下是如何启用 FSDP 运行 examples/nlp_example.py(从仓库的根目录)

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_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  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 examples/nlp_example.py

目前,Accelerate 通过 CLI 支持以下配置

fsdp_sharding_strategy: [1] FULL_SHARD (分片优化器状态、梯度和参数),[2] SHARD_GRAD_OP (分片优化器状态和梯度),[3] NO_SHARD (DDP),[4] HYBRID_SHARD (在每个节点内分片优化器状态、梯度和参数,而每个节点都拥有完整的副本),[5] HYBRID_SHARD_ZERO2 (在每个节点内分片优化器状态和梯度,而每个节点都拥有完整的副本)。更多信息,请参阅官方 PyTorch 文档

fsdp_offload_params:决定是否将参数和梯度卸载到 CPU。

fsdp_auto_wrap_policy: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP

fsdp_transformer_layer_cls_to_wrap: 仅适用于 Transformers。当使用 fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP 时,用户可以提供一个逗号分隔的 Transformer 层类名字符串(区分大小写)进行包装,例如 BertLayerGPTJBlockT5BlockBertLayer,BertEmbeddings,BertSelfOutput。这很重要,因为共享权重的子模块(例如嵌入层)不应出现在不同的 FSDP 包装单元中。使用此策略,包装会针对包含多头注意力和几个 MLP 层的每个块进行。包括共享嵌入在内的其余层,会被方便地包装在同一个最外层的 FSDP 单元中。因此,对基于 Transformer 的模型请使用此策略。您可以通过回答“是”来使用模型的 _no_split_modules,以选择“是否要使用模型的 _no_split_modules 进行包装”。它会在可能的情况下尝试使用 model._no_split_modules

fsdp_min_num_params:使用 fsdp_auto_wrap_policy=SIZE_BASED_WRAP 时的最小参数数量。

fsdp_backward_prefetch_policy: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH

fsdp_forward_prefetch:如果为 True,则 FSDP 在前向传播执行期间会显式预取下一个即将到来的 all-gather。只应在静态图模型中使用,因为预取遵循第一次迭代的执行顺序。也就是说,如果子模块的顺序在模型执行期间动态变化,请不要启用此功能。

fsdp_state_dict_type: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT

fsdp_use_orig_params:如果为 True,则允许在初始化期间使用非均匀的 requires_grad,这意味着支持交错的冻结和可训练参数。此设置在参数高效微调等情况下很有用,如这篇文章中所讨论。此选项还允许拥有多个优化器参数组。在使用 FSDP 准备/包装模型之前创建优化器时,此值应为 True

fsdp_cpu_ram_efficient_loading:仅适用于 Transformers 模型。如果为 True,只有第一个进程会加载预训练模型的检查点,而所有其他进程的权重为空。如果在通过 from_pretrained 方法加载预训练 Transformers 模型时遇到错误,则应将其设置为 False。当此设置为 True 时,fsdp_sync_module_states 也必须为 True,否则除了主进程外的所有进程都将具有随机权重,导致训练期间出现意外行为。为使其生效,请确保在调用 Transformers from_pretrained 方法之前已初始化分布式进程组。使用 Trainer API 时,分布式进程组在您创建 TrainingArguments 类的实例时被初始化。

fsdp_sync_module_states:如果为 True,每个单独包装的 FSDP 单元将从 rank 0 广播模块参数。

要进行额外和更精细的控制,您可以通过 FullyShardedDataParallelPlugin 指定其他 FSDP 参数。创建 FullyShardedDataParallelPlugin 对象时,将那些不属于 accelerate 配置的参数传递给它,或者如果您想覆盖它们。FSDP 参数将根据 accelerate 配置文件或启动命令参数进行选择,而您直接通过 FullyShardedDataParallelPlugin 对象传递的其他参数将设置/覆盖这些参数。

以下是一个示例

from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig

fsdp_plugin = FullyShardedDataParallelPlugin(
    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
)

accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

保存和加载

使用 FSDP 模型时,新的推荐检查点方法是在设置 accelerate 配置时使用 SHARDED_STATE_DICT 作为 StateDictType。以下是使用 accelerate 的 save_state 工具进行保存的代码片段。

accelerator.save_state("ckpt")

检查检查点文件夹,可以看到模型和优化器按每个进程分片存储

ls ckpt
# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin

cd ckpt

ls optimizer_0
# __0_0.distcp  __1_0.distcp

ls pytorch_model_0
# __0_0.distcp  __1_0.distcp

要加载它们以恢复训练,请使用 accelerate 的 load_state 工具

accelerator.load_state("ckpt")

当使用 transformers 的 save_pretrained 时,传递 state_dict=accelerator.get_state_dict(model) 来保存模型的状态字典。下面是一个示例

  unwrapped_model.save_pretrained(
      args.output_dir,
      is_main_process=accelerator.is_main_process,
      save_function=accelerator.save,
+     state_dict=accelerator.get_state_dict(model),
)

状态字典

accelerator.get_state_dict 将使用 FullStateDictConfig(offload_to_cpu=True, rank0_only=True) 上下文管理器调用底层的 model.state_dict 实现,以仅为 rank 0 获取状态字典,并且它将被卸载到 CPU。

然后,您可以将 state 传递到 save_pretrained 方法中。您可以使用 StateDictTypeFullStateDictConfig 的几种模式来控制 state_dict 的行为。更多信息,请参阅 PyTorch 文档

如果您选择使用 StateDictType.SHARDED_STATE_DICT,在 Accelerator.save_state 期间,模型的权重将根据模型上的每个子拆分被分割成 n 个文件。要将它们合并回一个单独的字典,以便在训练后加载回模型,您可以使用 merge_weights 工具

from accelerate.utils import merge_fsdp_weights

# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder
merge_fsdp_weights("pytorch_model_fsdp_0", "output_path", safe_serialization=True)

最终的输出将被保存到 model.safetensorspytorch_model.bin(如果传递了 safe_serialization=False)。

这也可以通过 CLI 调用

accelerate merge-weights pytorch_model_fsdp_0/ output_path

FSDP 分片策略与 DeepSpeed ZeRO 阶段的映射关系

  • FULL_SHARD 对应 DeepSpeed ZeRO Stage-3。对优化器状态、梯度和参数进行分片。
  • SHARD_GRAD_OP 对应 DeepSpeed ZeRO Stage-2。对优化器状态和梯度进行分片。
  • NO_SHARD 对应 ZeRO Stage-0。不进行分片,其中每个 GPU 都拥有模型、优化器状态和梯度的完整副本。
  • HYBRID_SHARD 对应 ZeRO++ Stage-3,其中 zero_hpz_partition_size=<num_gpus_per_node>。这里,它将在每个节点内对优化器状态、梯度和参数进行分片,而每个节点都拥有完整的副本。

一些需要注意的事项

  • 在有多个模型的情况下,将优化器按照对应模型的顺序传递给 prepare 调用,否则 accelerator.save_state()accelerator.load_state() 会导致错误/意外行为。
  • 此功能与 Transformers 库中 run_translation.py 脚本的 --predict_with_generate 不兼容。

为了进行更多控制,用户可以利用 FullyShardedDataParallelPlugin。创建此类的实例后,用户可以将其传递给 Accelerator 类的实例化。有关这些选项的更多信息,请参阅 PyTorch FullyShardedDataParallel 代码。

对 FSDP 和 DeepSpeed 的异同感兴趣的读者,请查看这里的概念指南

< > 在 GitHub 上更新