Hugging Face 的 LOGO

Accelerate 文档

完全分片数据并行实用工具

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

完全分片数据并行实用工具

enable_fsdp_ram_efficient_loading

accelerate.utils.enable_fsdp_ram_efficient_loading

< >

( )

在环境中为 FSDP 启用 Hugging Face 模型的 RAM 高效加载。

disable_fsdp_ram_efficient_loading

accelerate.utils.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 )

参数

  • checkpoint_dir (str) — 包含 FSDP 检查点的目录（可以是模型或优化器）。
  • output_path (str) — 保存合并后检查点的路径。
  • safe_serialization (bool, *可选*, 默认为 True) — 是否使用 safetensors 保存合并后的权重（推荐）。
  • 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` 时适用。

该插件用于启用完全分片数据并行。

set_auto_wrap_policy

< >

( model )

给定 `model`，根据传入的策略以及是否可以使用 `transformer_cls_to_wrap` 来创建 `auto_wrap_policy`

set_mixed_precision

< >

( mixed_precision buffer_autocast = False override = False )

设置 FSDP 的混合精度策略

set_state_dict_type

< >

( state_dict_type = None )

根据 `StateDictType` 设置状态字典配置。

validate_mixed_precision_policy

< >

( )

验证混合精度策略，抽象出来以避免在不需要时引入导入。

fsdp2_load_full_state_dict

accelerate.utils.fsdp2_load_full_state_dict

< >

( accelerator model: Module full_sd: dict )

参数

  • accelerator (Accelerator) — accelerator 实例
  • model (torch.nn.Module) — 要加载状态字典的模型，预期在元设备上，否则可能发生显存峰值
  • full_sd (dict) — 要加载的完整状态字典，只能在 rank 0 上

将完整的状态字典（可能仅在 rank 0 上）加载到分片模型中。这是通过将参数从 rank 0 广播到所有其他 rank 来完成的。此函数会就地修改模型。

fsdp2_switch_optimizer_parameters

accelerate.utils.fsdp2_switch_optimizer_parameters

< >

( optimizer: Optimizer mapping: dict )

参数

  • optimizer (torch.optim.Optimizer) — 包含原始模型参数的优化器实例
  • mapping (dict) — 从原始参数（由 data_ptr 指定）到分片参数的映射

引发

KeyError

  • KeyError — 如果优化器中的参数无法切换到其分片版本。这种情况永远不应发生，它表示存在一个 bug。如果我们保留原始参数而不是抛出错误，训练在数值上将不正确，并且权重不会得到更新。

将优化器的参数切换为新的参数（通常情况下是分片参数）。此函数会就地修改优化器。

fsdp2_prepare_model

accelerate.utils.fsdp2_prepare_model

< >

( accelerator model: Module ) → torch.nn.Module

参数

  • accelerator (Accelerator) — 加速器实例
  • model (torch.nn.Module) — 要准备的模型

返回

torch.nn.Module

准备好的模型

就地为 FSDP2 准备模型。同时返回模型以避免误用原始模型。

fsdp2_prepare_auto_wrap_policy

< > 在 GitHub 上更新

←Megatron-LM 实用工具
© . This site is unofficial and not affiliated with Hugging Face, Inc.