Accelerate 文档
完全分片数据并行实用工具
并获得增强的文档体验
开始使用
完全分片数据并行实用工具
enable_fsdp_ram_efficient_loading
在环境中为 FSDP 启用 Hugging Face 模型的 RAM 高效加载。
disable_fsdp_ram_efficient_loading
在环境中为 FSDP 禁用 Hugging Face 模型的 RAM 高效加载。
merge_fsdp_weights
accelerate.utils.merge_fsdp_weights
< 源码 >( checkpoint_dir: str output_path: str safe_serialization: bool = True remove_checkpoint_dir: bool = False )
将分片的 FSDP 模型检查点的权重合并到一个组合检查点中。如果模型使用了 `SHARDED_STATE_DICT`，则应使用此功能。如果 `safe_serialization` 为 `True`，权重将保存到 `{output_path}/model.safetensors`，否则保存到 `pytorch_model.bin`。
注意：这是一个 CPU 密集型过程。
FullyShardedDataParallelPlugin
class accelerate.FullyShardedDataParallelPlugin
< 源码 >( fsdp_version: int = None sharding_strategy: typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy')] = None reshard_after_forward: typing.Union[str, ForwardRef('torch.distributed.fsdp.ShardingStrategy'), bool] = None backward_prefetch: typing.Union[str, ForwardRef('torch.distributed.fsdp.BackwardPrefetch'), NoneType] = None mixed_precision_policy: typing.Union[dict, ForwardRef('torch.distributed.fsdp.MixedPrecision'), ForwardRef('torch.distributed.fsdp.MixedPrecisionPolicy'), NoneType] = None auto_wrap_policy: typing.Union[typing.Callable, typing.Literal['transformer_based_wrap', 'size_based_wrap', 'no_wrap'], NoneType] = None cpu_offload: typing.Union[bool, ForwardRef('torch.distributed.fsdp.CPUOffload'), ForwardRef('torch.distributed.fsdp.CPUOffloadPolicy')] = None ignored_modules: typing.Union[collections.abc.Iterable[torch.nn.modules.module.Module], str, NoneType] = None state_dict_type: typing.Union[str, ForwardRef('torch.distributed.fsdp.StateDictType')] = None state_dict_config: typing.Union[ForwardRef('torch.distributed.fsdp.FullStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedStateDictConfig'), NoneType] = None optim_state_dict_config: typing.Union[ForwardRef('torch.distributed.fsdp.FullOptimStateDictConfig'), ForwardRef('torch.distributed.fsdp.ShardedOptimStateDictConfig'), NoneType] = None limit_all_gathers: bool = True use_orig_params: typing.Optional[bool] = None param_init_fn: typing.Optional[typing.Callable[[torch.nn.modules.module.Module], NoneType]] = None sync_module_states: typing.Optional[bool] = None forward_prefetch: bool = None activation_checkpointing: bool = None cpu_ram_efficient_loading: bool = None transformer_cls_names_to_wrap: typing.Optional[list[str]] = None min_num_params: typing.Optional[int] = None )
参数
- fsdp_version (
int, 默认为1) — 要使用的 FSDP 版本。默认为 1。如果设置为 2，启动器期望配置已转换为 FSDP2 格式。 - sharding_strategy (
Union[str, torch.distributed.fsdp.ShardingStrategy], 默认为'FULL_SHARD') — 要使用的分片策略。应为 `str` 或 `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy` 的实例。已被 `reshard_after_forward` 弃用。 - reshard_after_forward (
Union[str, torch.distributed.fsdp.ShardingStrategy, bool], 对于fsdp_version=1默认为'FULL_SHARD'，对于fsdp_version=2默认为True) — 要使用的分片策略。如果 `fsdp_version` 设置为 2，应为布尔值，否则为 `str` 或 `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy` 的实例。 - backward_prefetch (
Union[str, torch.distributed.fsdp.BackwardPrefetch], 默认为'NO_PREFETCH') — 要使用的反向预取策略。应为 `str` 或 `torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch` 的实例。 - mixed_precision_policy (
Optional[Union[dict, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]], 默认为None) — 用于启用 FullyShardedDataParallel 混合精度训练的配置。如果传入 `dict`，它应包含 `param_dtype`、`reduce_dtype` 和 `buffer_dtype` 键。如果 `fsdp_version` 设置为 2，可以是一个 `torch.distributed.fsdp.MixedPrecisionPolicy` 实例。 - auto_wrap_policy (
Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), 默认为NO_WRAP) — 用于递归地用 FSDP 包装层的可调用对象或字符串策略。如果为字符串，则必须是 `transformer_based_wrap`、`size_based_wrap` 或 `no_wrap` 之一。有关其具体形式，请参阅 `torch.distributed.fsdp.wrap.size_based_wrap_policy`。 - cpu_offload (
Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy], 默认为False) — 是否将参数卸载到 CPU。应为布尔值，或者如果 `fsdp_version` 设置为 2，则为 `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload` 或 `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy` 的实例。 - ignored_modules (
Optional[Union[Iterable[torch.nn.Module], str]], 默认为None) — 在使用 FSDP 包装时要忽略的模块列表。当传入字符串时，将使用正则表达式的 `fullmatch` 按名称匹配模块。 - state_dict_type (
Union[str, torch.distributed.fsdp.StateDictType], 默认为'FULL_STATE_DICT') — 要使用的状态字典类型。如果为字符串，则必须是 `full_state_dict`、`local_state_dict` 或 `sharded_state_dict` 之一。 - state_dict_config (
Optional[Union[torch.distributed.fsdp.FullStateDictConfig, torch.distributed.fsdp.ShardedStateDictConfig], 默认为None) — 要使用的状态字典配置。如果未传入，则根据 `state_dict_type` 确定。 - optim_state_dict_config (
Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig], 默认为None) — 要使用的优化器状态字典配置。如果未传入，则根据 `state_dict_type` 确定。 - limit_all_gathers (
bool, 默认为True) — FSDP 是否应显式同步 CPU 线程以防止过多的正在进行的 all-gather 操作。此布尔值仅影响调度 all-gather 的分片策略。启用此选项有助于减少 CUDA malloc 重试次数。 - use_orig_params (
bool, 默认为False) — 是否为优化器使用原始参数。 - param_init_fn (
Optional[Callable[[torch.nn.Module], None], 默认为None) — 一个 `Callable[torch.nn.Module] -> None`，用于指定当前在元设备上的模块应如何初始化到实际设备上。仅在 `sync_module_states` 为 `True` 时适用。默认是一个在模块上调用 `to_empty` 的 `lambda` 函数。 - sync_module_states (
bool, 默认为False) — 每个独立包装的 FSDP 单元是否应从 rank 0 广播模块参数，以确保它们在初始化后在所有 rank 上保持一致。默认为 `False`，除非 `cpu_ram_efficient_loading` 为 `True`，届时将被强制启用。 - forward_prefetch (
bool, 默认为False) — 是否让 FSDP 在前向传播过程中显式预取下一个即将到来的 all-gather 操作。仅适用于静态图。 - activation_checkpointing (
bool, 默认为False) — 一种通过清除某些层的激活并在反向传播期间重新计算它们来减少内存使用的技术。实际上，这是以额外的计算时间换取减少的内存使用。 - cpu_ram_efficient_loading (
bool, 默认为None) — 如果为 True，只有第一个进程加载预训练模型的检查点，而所有其他进程的权重为空。仅适用于 Transformers。使用此选项时，`sync_module_states` 需要为 `True`。 - transformer_cls_names_to_wrap (
Optional[List[str]], 默认为None) — 要包装的 transformer 层类名列表。仅在 `auto_wrap_policy` 为 `transformer_based_wrap` 时适用。 - min_num_params (
Optional[int], 默认为None) — 一个模块必须拥有的最小参数数量才能被包装。仅在 `auto_wrap_policy` 为 `size_based_wrap` 时适用。
该插件用于启用完全分片数据并行。
给定 `model`，根据传入的策略以及是否可以使用 `transformer_cls_to_wrap` 来创建 `auto_wrap_policy`
设置 FSDP 的混合精度策略
根据 `StateDictType` 设置状态字典配置。
验证混合精度策略，抽象出来以避免在不需要时引入导入。
fsdp2_load_full_state_dict
accelerate.utils.fsdp2_load_full_state_dict
< 源码 >( accelerator model: Module full_sd: dict )
将完整的状态字典（可能仅在 rank 0 上）加载到分片模型中。这是通过将参数从 rank 0 广播到所有其他 rank 来完成的。此函数会就地修改模型。
fsdp2_switch_optimizer_parameters
accelerate.utils.fsdp2_switch_optimizer_parameters
< 源码 >( optimizer: Optimizer mapping: dict )
将优化器的参数切换为新的参数（通常情况下是分片参数）。此函数会就地修改优化器。
fsdp2_prepare_model
accelerate.utils.fsdp2_prepare_model
< source >( accelerator model: Module ) → torch.nn.Module
就地为 FSDP2 准备模型。同时返回模型以避免误用原始模型。