Accelerate 文档
FSDP 对比 DeepSpeed
并获得增强的文档体验
开始使用
FSDP 对比 DeepSpeed
Accelerate 通过集成 PyTorch FSDP 和 Microsoft DeepSpeed 这两个功能强大的分布式训练工具,提供了训练框架的灵活性。本教程旨在比较两者的异同,以便用户能够在这两个框架之间无缝切换。
要在框架之间切换,我们建议使用 accelerate launch
命令,通过 --config_file
传入正确的配置文件,或者直接为 FSDP 和 DeepSpeed 传入相应的参数。
DeepSpeed 的 Accelerate 配置示例请参见此处,FSDP 的配置示例请参见此处,或者在示例库中的“启动配置”部分查找。
本教程仅适用于单节点、多 GPU 场景。
配置功能
模型张量被分成不同的 GPU 以便扩展模型大小;这在 FSDP 中称为“分片”(*sharding*),在 DeepSpeed 中称为“分区”(*partitioning*)。FSDP 分片和 DeepSpeed ZeRO(分区)阶段分别通过 --fsdp_sharding_strategy
和 --zero_stage
进行配置。特别是,FSDP FULL_SHARD
对应于 DeepSpeed ZeRO 阶段 3
;有关 FSDP 分片和 DeepSpeed ZeRO 设置之间的全面映射,请参阅此页面。下表总结并分组了类似设置:
组 | 框架 | 配置 | 示例 | 限制(如果有) |
---|---|---|---|---|
分片 / 分区 | FSDP DeepSpeed | --fsdp_sharding_strategy --zero_stage | 1 (FULL_SHARD )3 | |
卸载 | FSDP DeepSpeed | --fsdp_offload_params --offload_param_device --offload_optimizer_device | true cpu cpu | 全部或不卸载 |
模型加载 | FSDP DeepSpeed | --fsdp_cpu_ram_efficient_loading --zero3_init_flag | true true | 仅限 ZeRO 3 |
高效检查点 | FSDP DeepSpeed | --fsdp_state_dict_type --zero3_save_16bit_model | SHARDED_STATE_DICT true | 仅限 ZeRO 3 |
权重预取 | FSDP DeepSpeed | --fsdp_forward_prefetch --fsdp_backward_prefetch 无 | true BACKWARD_PRE | |
模型 | FSDP DeepSpeed | --fsdp_auto_wrap_policy --fsdp_transformer_layer_cls_to_wrap 无 | TRANSFORMER_BASED_WRAP <层类> | 通常不需要 对用户透明。 |
参数召唤 | FSDP DeepSpeed | --fsdp_use_orig_params 无 | true | torch.compile 所必需对用户透明 |
参数同步 | FSDP DeepSpeed | --fsdp_sync_module_states 无 | true | |
训练 | FSDP DeepSpeed | 无--gradient_accumulation_steps --gradient_clipping | 自动 自动 | 对用户透明 |
有关上述内容的详细说明,请参阅 Accelerate
启动文档。
要访问其他 DeepSpeed 配置,例如混合精度设置,您需要传入 --deepspeed_config_file
,请参阅文档。
DeepSpeed 也可以通过 DeepSpeedPlugin 配置,例如,DeepSpeedPlugin.zero_stage
等同于 --zero_stage
,DeepSpeedPlugin.hf_ds_config
可用于传递 --deepspeed_config_file
。
FSDP 也可以通过 FullyShardedDataParallelPlugin 配置,例如,FullyShardedDataParallelPlugin.sharding_strategy
等同于 --fsdp_sharding_strategy
。
检查点
请注意,FSDP 可以通过 --fsdp_state_dict_type
配置来保存完整或分片检查点。
对于 DeepSpeed Zero3,可以传入 --zero3_save_16bit_model true
,它方便地将模型整合到一个 rank 中并保存;这相当于 FSDP 中的 fsdp_state_dict_type: FULL_STATE_DICT
。
对于大型模型,将模型整合到一个 rank 中可能会非常慢。
为了加快检查点保存速度,对于 FSDP,请使用 fsdp_state_dict_type: SHARDED_STATE_DICT
;对于 DeepSpeed Zero3,请使用 zero_to_fp32.py
脚本来后转换分片检查点。
卸载
FSDP 只允许*全部或不卸载*(即,要么卸载所有参数、梯度和优化器,要么将它们全部保留在 GPU 中),但 DeepSpeed 可以不同地卸载参数和优化器。此外,DeepSpeed 还支持卸载到 NVME。
预取
FSDP 允许两种预取配置:--fsdp_forward_prefetch
和 --fsdp_backward_prefetch
,以牺牲额外内存为代价提高通信/计算重叠,参见FSDP 文档。对于 DeepSpeed,预取会在需要时开启,并且它会根据某些超参数(如 stage3_param_persistence_threshold
、stage3_max_reuse_distance
等)开启,这些超参数可以为 Zero3 配置;如果您未在 deepspeed 配置文件中明确设置这些超参数,accelerate
可能会自动设置它们。
对于 FSDP,如果内存允许,请将 fsdp_backward_prefetch: BACKWARD_PRE
设置为提高吞吐量。
模型加载
FSDP 需要显式设置 --fsdp_cpu_ram_efficient_loading true
来激活高效模型加载,而 transformers
在使用 DeepSpeed Zero3 时会自动激活类似功能。
对于 FSDP,当设置 --fsdp_cpu_ram_efficient_loading true
时,accelerate
会自动将 sync_module_states
设置为 true。对于 RAM 高效加载,权重将只加载到单个 rank 中,因此需要 sync_module_states
将权重广播到其他 rank。
模型
FSDP 需要显式的 --fsdp_auto_wrap_policy
标志,以便算法决定如何调度 all-gather 和 reduce-scatter 操作。但对于 DeepSpeed,这对用户是透明的。
对于 FSDP,只需设置 fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
。使用最新的 transformers
版本,我们尽力找出适合 HF transformers 模型的 fsdp_transformer_layer_cls_to_wrap
。但是,如果遇到相关错误,请指定此参数。
参数召唤
如果使用 torch.compile
,FSDP 需要显式设置 --fsdp_use_orig_params
标志,请参阅 pytorch 文档。对于 DeepSpeed,这对用户是透明的。
对于 FSDP,当使用 torch.compile
时,请设置 fsdp_use_orig_params: True
。
训练
DeepSpeed 需要显式的 --gradient_accumulation_steps
和 --gradient_clipping
标志。对于 FSDP,这对用户是透明的。
使用 DeepSpeed 时,设置 gradient_accumulation_steps: "auto"
和 gradient_clipping: "auto"
将自动选择在 Accelerator 或 TrainingArguments
(如果使用 transformers
)中设置的值。
关于数据精度处理的差异
为了讨论 FSDP 和 DeepSpeed 中数据精度如何处理,首先概述这些框架中模型参数的处理方式是有益的。在模型/优化器参数分发到 GPU 之前,参数准备涉及先将它们“展平”为一维 torch.Tensor
。FSDP / DeepSpeed 的实现在这些“展平”参数的 dtype
存储方面有所不同,这会影响 torch.Optimizer
如何分配其 dtype
。下表概述了这两个框架的过程;“本地”列表示在每个 GPU 级别上发生的过程,因此向上转换导致的任何内存开销都应被使用的 GPU 数量摊销。
作为经验法则,为了自动混合精度的稳定训练,所有可训练参数必须是 torch.float32
。
过程 | 本地 | 框架 | 详情 |
---|---|---|---|
加载,即 AutoModel.from_pretrained(..., torch_dtype=torch_dtype) | |||
准备,即创建“展平参数” | ✅ | FSDP DeepSpeed | 以 torch_dtype 创建。忽略 torch_dtype ,以 float32 创建。 |
优化器初始化 | ✅ | FSDP DeepSpeed | 以 torch_dtype 创建参数以 float32 创建参数 |
训练步骤,即前向、反向、规约 | FSDP DeepSpeed | 遵循 MixedPrecision 遵循 deepspeed_config_file 混合精度设置。 | |
优化器(预步骤) | ✅ | FSDP DeepSpeed | 向上转换(如果有)到 torch_dtype 向上转换为 float32 |
优化器(实际步骤) | ✅ | FSDP DeepSpeed | 发生在 torch_dtype 发生在 float32 。 |
因此,当使用少量 GPU 的 DeepSpeed 时,请注意由于准备过程中的向上转换而可能导致的显着内存开销。
使用 FSDP 时,在没有混合精度的情况下,可以将 torch.Optimizer
以低精度 torch_dtype
运行,这在使用少量 GPU 时可能会有所帮助。
在混合精度下,FSDP 和 DeepSpeed 将在模型准备步骤中向上转换(参见上表)。但请注意,FSDP 将以向上转换后的精度保存检查点;如果指定了 --zero3_save_16bit_model
,DeepSpeed 仍可能保存低精度检查点。
为阐明上表,请考虑下面的具体示例;为简洁起见,优化器预步骤和实际步骤合并。FSDP 可以以以下两种模式运行,而 DeepSpeed 只能以一种模式运行。
框架 | 模型加载(torch_dtype ) | 混合精度 | 准备(本地) | 训练 | 优化器(本地) |
---|---|---|---|---|---|
FSDP | bf16 | 默认(无) | bf16 | bf16 | bf16 |
FSDP | bf16 | bf16 | fp32 | bf16 | fp32 |
DeepSpeed | bf16 | bf16 | fp32 | bf16 | fp32 |