参数

• checkpoint_dir (str) — 包含 FSDP 检查点（可以是模型或优化器）的目录。
• output_path (str) — 保存合并检查点的路径。
• safe_serialization (bool, 可选, 默认为 True) — 是否使用 safetensors 保存合并后的权重（推荐）。
• remove_checkpoint_dir (bool, 可选, 默认为 False) — 是否在合并后删除检查点目录。

参数

• 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，则应为 bool 值，否则应为 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]], defaults to 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"]]), defaults to 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], defaults to False) — 是否将参数卸载到 CPU。应为 bool 或 torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload 或 torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffloadPolicy 的实例，如果 fsdp_version 设置为 2。
• ignored_modules (Optional[Iterable[torch.nn.Module]], defaults to None) — 使用 FSDP 包装时要忽略的模块列表。
• state_dict_type (Union[str, torch.distributed.fsdp.StateDictType], defaults to '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], defaults to None) — 要使用的状态字典配置。如果未传入，则根据 state_dict_type 确定。
• optim_state_dict_config (Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig], defaults to None) — 要使用的优化器状态字典配置。如果未传入，则根据 state_dict_type 确定。
• limit_all_gathers (bool, defaults to True) — 是否让 FSDP 显式同步 CPU 线程，以防止过多的正在进行中的 all-gather 操作。此布尔值仅影响调度 all-gather 操作的分片策略。启用此功能可以帮助减少 CUDA malloc 重试的次数。
• use_orig_params (bool, defaults to False) — 是否对优化器使用原始参数。
• param_init_fn (Optional[Callable[[torch.nn.Module], None], defaults to None) — 一个 Callable[torch.nn.Module] -> None，用于指定当前在 meta 设备上的模块应如何在实际设备上初始化。仅当 sync_module_states 为 True 时适用。默认情况下是一个 lambda，它在模块上调用 to_empty。
• sync_module_states (bool, defaults to False) — 每个单独包装的 FSDP 单元是否应从 rank 0 广播模块参数，以确保初始化后它们在所有 rank 中都相同。默认为 False，除非 cpu_ram_efficient_loading 为 True，届时将强制启用。
• forward_prefetch (bool, defaults to False) — 是否让 FSDP 在前向传递中显式预取下一个即将到来的 all-gather 操作。仅用于静态图。
• activation_checkpointing (bool, defaults to False) — 一种通过清除某些层的激活并在反向传递期间重新计算它们来减少内存使用量的技术。实际上，这用额外的计算时间换取了减少的内存使用量。
• cpu_ram_efficient_loading (bool, defaults to None) — 如果为 True，则只有第一个进程加载预训练模型检查点，而所有其他进程都具有空权重。仅适用于 Transformers。使用此选项时，sync_module_states 需要为 True。
• transformer_cls_names_to_wrap (Optional[List[str]], defaults to None) — 要包装的 transformer 层类名称列表。仅当 auto_wrap_policy 为 transformer_based_wrap 时适用。
• min_num_params (Optional[int], defaults to None) — 模块必须具有的要包装的最小参数数量。仅当 auto_wrap_policy 为 size_based_wrap 时适用。

参数

• accelerator (Accelerator) — 加速器实例
• model (torch.nn.Module) — 要将状态字典加载到的模型
• full_sd (dict) — 要加载的完整状态字典，只能在 rank 0 上

参数

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

• KeyError — 如果优化器中的参数无法切换到其分片版本。这不应该发生，并表示存在错误。如果我们保留原始参数而不是引发异常，则训练在数值上将不正确，并且权重将不会得到更新。

参数

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

准备好的模型