完全分片数据并行

为了在更大的批量大小上加速训练超大模型，我们可以使用完全分片数据并行模型。 这种数据并行范式通过分片优化器状态、梯度和参数，能够容纳更多的数据和更大的模型。 要了解更多信息及其优势，请参阅 Fully Sharded Data Parallel 博客。 我们已集成最新的 PyTorch 完全分片数据并行（FSDP）训练功能。 您只需通过配置启用它即可。

开箱即用的工作原理

在您的机器上只需运行：

accelerate config

and answer the questions asked. This will generate a config file that will be used automatically to properly set the default options when doing

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 包装单元中。使用此策略时，包装操作会针对每个包含 Multi-Head Attention 及其后若干 MLP 层的块进行。其余层（包括共享嵌入层）会被方便地包装在同一个最外层的 FSDP 单元中。因此，请对基于 Transformer 的模型使用此策略。对于 Transformer 模型，您可以通过将 Do you want to use the model's_no_split_modulesto wrap. It will try to use 的回答设为 yes 来使用 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 会在前向传播执行期间显式预取下一个即将进行的 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 时，将 StateDictType 设置为 SHARDED_STATE_DICT。 以下是使用 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),
)

State Dict

accelerator.get_state_dict 将使用 FullStateDictConfig(offload_to_cpu=True, rank0_only=True) 上下文管理器调用底层的 model.state_dict 实现，以仅获取 rank 0 的 state dict，并将其卸载到 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>。在此配置下，优化器状态、梯度和参数将在每个节点内部进行分片，同时每个节点都保留一份完整副本。

需要注意的几个问题

• In case of multiple models, pass the optimizers to the prepare call in the same order as corresponding models else accelerator.save_state() and accelerator.load_state() will result in wrong/unexpected behaviour.
• 此功能与 Transformers 库中 run_translation.py 脚本的 --predict_with_generate 不兼容。

如需更多控制，用户可以利用 FullyShardedDataParallelPlugin。在创建该类的实例后，用户可以将其传递给 Accelerator 类的实例化过程。 有关这些选项的更多信息，请参阅 PyTorch 的 FullyShardedDataParallel 代码。

如果您对 FSDP 和 DeepSpeed 之间的异同感兴趣，请查看此处的概念指南！