Hugging Face 的 logo

加速文档

完全分片数据并行

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 层类名称字符串（区分大小写）来包装，例如，BertLayer、GPTJBlock、T5Block、BertLayer,BertEmbeddings,BertSelfOutput。这很重要，因为共享权重的子模块（例如，嵌入层）不应最终位于不同的 FSDP 包装单元中。使用此策略，每个包含多头注意力后跟几个 MLP 层的块都会进行包装。包括共享嵌入在内的其余层方便地包装在同一个最外层的 FSDP 单元中。因此，将其用于基于 Transformer 的模型。您可以通过对您是否要使用模型的_no_split_modules进行包装？回答是来使用 Transformer 模型的model._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 在前向传递中执行时显式预取下一个即将到来的全收集。仅应用于静态图模型，因为预取遵循第一次迭代的执行顺序。即，如果子模块的顺序在模型执行期间动态更改，则不要启用此功能。

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 单元将从排名 0 广播模块参数。

为了进行额外的和更细致的控制，您可以通过FullyShardedDataParallelPlugin指定其他 FSDP 参数。在创建FullyShardedDataParallelPlugin对象时，将不属于加速配置的参数传递给它，或者如果要覆盖它们。FSDP 参数将根据加速配置文件或启动命令参数以及您将直接通过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 模型时，推荐的检查点方法是使用SHARDED_STATE_DICT作为StateDictType来设置加速配置。以下是使用加速的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

要加载它们以恢复训练，请使用加速的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实现以仅获取排名 0 的状态字典，并且它将被卸载到 CPU。

然后，您可以将state传递到save_pretrained方法中。StateDictType和FullStateDictConfig有几种模式，您可以使用它们来控制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.safetensors或pytorch_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 上更新

←DDP 通信 Hook Megatron-LM→