参数

• 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, str, torch.distributed.fsdp.MixedPrecision, torch.distributed.fsdp.MixedPrecisionPolicy]]，默认为 None) — 用于启用 FullyShardedDataParallel 混合精度训练的配置。如果传入 dict，它应包含以下键：param_dtype、reduce_dtype 和 buffer_dtype；如果 fsdp_version 设置为 2，则可以是 torch.distributed.fsdp.MixedPrecisionPolicy 的实例。如果传入 str，它应为以下值之一：fp8、fp16、bf16、fp32，用于设置 param_dtype、reduce_dtype 和 buffer_dtype。
• auto_wrap_policy (Optional(Union[Callable, Literal["transformer_based_wrap", "size_based_wrap", "no_wrap"]]), defaults to NO_WRAP) -- A callable or string specifying a policy to recursively wrap layers with FSDP. If a string, it must be one of transformer_based_wrap, size_based_wrap, or no_wrap. See torch.distributed.fsdp.wrap.size_based_wrap_policy` for a direction on what it should look like.
• cpu_offload (Union[bool, torch.distributed.fsdp.CPUOffload, torch.distributed.fsdp.CPUOffloadPolicy]，默认为 False) — 是否将参数卸载到 CPU。应为 bool 类型，或者当 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 按名称匹配模块。如果 fsdp_version 设置为 2，这些模块将被转换为参数并使用。
• state_dict_type (Union[str, torch.distributed.fsdp.StateDictType]，默认为 'FULL_STATE_DICT') — 要使用的 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 配置。如果未传入，则根据 state_dict_type 确定。
• optim_state_dict_config (Optional[Union[torch.distributed.fsdp.FullOptimStateDictConfig, torch.distributed.fsdp.ShardedOptimStateDictConfig]，默认为 None) — 要使用的 Optim state dict 配置。如果未传入，则根据 state_dict_type 确定。
• limit_all_gathers (bool，默认为 True) — 是否让 FSDP 显式同步 CPU 线程，以防止过多的 in-flight all-gathers。此 bool 仅影响调度 all-gathers 的分片策略。启用此选项有助于减少 CUDA malloc 重试次数。
• use_orig_params (bool，默认为 False) — 是否使用原始参数作为优化器的参数。
• param_init_fn (Optional[Callable[[torch.nn.Module], None]，默认为 None) — 一个 Callable[torch.nn.Module] -> None，用于指定当前位于 meta 设备上的模块应如何初始化到实际设备上。仅当 sync_module_states 为 True 时适用。默认是一个 lambda，它会对模块调用 to_empty。
• sync_module_states (bool，默认为 False) — 是否应让每个单独包装的 FSDP 单元从 rank 0 广播模块参数，以确保初始化后所有 rank 上的参数一致。除非 cpu_ram_efficient_loading 为 True，否则默认为 False；若为 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 时适用。

参数

• accelerator (Accelerator) — 加速器实例
• model (torch.nn.Module) — 要加载状态字典的模型，预期位于 meta 设备上，否则可能会发生 VRAM 峰值
• full_sd (dict) — 要加载的完整状态字典，只能位于 rank 0
• cpu_offload (bool，默认为 False) — 如果为 True，则在分发后将分片参数移至 CPU。启用 FSDP CPU offloading 时必需。

参数

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

参数

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