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.