完全分片資料並行

為了以更大的批次規模加速訓練大型模型，我們可以使用完全分片資料並行模型。這種資料並行正規化透過對最佳化器狀態、梯度和引數進行分片，從而能夠容納更多資料和更大的模型。要了解更多相關資訊及其好處，請檢視完全分片資料並行部落格。我們集成了 PyTorch 最新的完全分片資料並行（FSDP）訓練功能。您只需透過配置啟用它即可。

開箱即用的工作方式

在您的機器上，只需執行

accelerate config

並回答所提出的問題。這將生成一個配置檔案，在執行時會自動使用該檔案來正確設定預設選項

accelerate launch my_script.py --args_to_my_script

例如，以下是如何啟用 FSDP 執行 examples/nlp_example.py（從倉庫的根目錄）

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

accelerate launch examples/nlp_example.py

目前，Accelerate 透過 CLI 支援以下配置

fsdp_sharding_strategy: [1] FULL_SHARD (分片最佳化器狀態、梯度和引數)，[2] SHARD_GRAD_OP (分片最佳化器狀態和梯度)，[3] NO_SHARD (DDP)，[4] HYBRID_SHARD (在每個節點內分片最佳化器狀態、梯度和引數，而每個節點都擁有完整的副本)，[5] HYBRID_SHARD_ZERO2 (在每個節點內分片最佳化器狀態和梯度，而每個節點都擁有完整的副本)。更多資訊，請參閱官方 PyTorch 文件。

fsdp_offload_params：決定是否將引數和梯度解除安裝到 CPU。

fsdp_auto_wrap_policy: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP

fsdp_transformer_layer_cls_to_wrap: 僅適用於 Transformers。當使用 fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP 時，使用者可以提供一個逗號分隔的 Transformer 層類名字串（區分大小寫）進行包裝，例如 BertLayer、GPTJBlock、T5Block、BertLayer,BertEmbeddings,BertSelfOutput。這很重要，因為共享權重的子模組（例如嵌入層）不應出現在不同的 FSDP 包裝單元中。使用此策略，包裝會針對包含多頭注意力和幾個 MLP 層的每個塊進行。包括共享嵌入在內的其餘層，會被方便地包裝在同一個最外層的 FSDP 單元中。因此，對基於 Transformer 的模型請使用此策略。您可以透過回答“是”來使用模型的 _no_split_modules，以選擇“是否要使用模型的 _no_split_modules 進行包裝”。它會在可能的情況下嘗試使用 model._no_split_modules。

fsdp_min_num_params：使用 fsdp_auto_wrap_policy=SIZE_BASED_WRAP 時的最小引數數量。

fsdp_backward_prefetch_policy: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH

fsdp_forward_prefetch：如果為 True，則 FSDP 在前向傳播執行期間會顯式預取下一個即將到來的 all-gather。只應在靜態圖模型中使用，因為預取遵循第一次迭代的執行順序。也就是說，如果子模組的順序在模型執行期間動態變化，請不要啟用此功能。

fsdp_state_dict_type: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT

fsdp_use_orig_params：如果為 True，則允許在初始化期間使用非均勻的 requires_grad，這意味著支援交錯的凍結和可訓練引數。此設定在引數高效微調等情況下很有用，如這篇文章中所討論。此選項還允許擁有多個最佳化器引數組。在使用 FSDP 準備/包裝模型之前建立最佳化器時，此值應為 True。

fsdp_cpu_ram_efficient_loading：僅適用於 Transformers 模型。如果為 True，只有第一個程序會載入預訓練模型的檢查點，而所有其他程序的權重為空。如果在透過 from_pretrained 方法載入預訓練 Transformers 模型時遇到錯誤，則應將其設定為 False。當此設定為 True 時，fsdp_sync_module_states 也必須為 True，否則除了主程序外的所有程序都將具有隨機權重，導致訓練期間出現意外行為。為使其生效，請確保在呼叫 Transformers from_pretrained 方法之前已初始化分散式程序組。使用 Trainer API 時，分散式程序組在您建立 TrainingArguments 類的例項時被初始化。

fsdp_sync_module_states：如果為 True，每個單獨包裝的 FSDP 單元將從 rank 0 廣播模組引數。

要進行額外和更精細的控制，您可以透過 FullyShardedDataParallelPlugin 指定其他 FSDP 引數。建立 FullyShardedDataParallelPlugin 物件時，將那些不屬於 accelerate 配置的引數傳遞給它，或者如果您想覆蓋它們。FSDP 引數將根據 accelerate 配置檔案或啟動命令引數進行選擇，而您直接透過 FullyShardedDataParallelPlugin 物件傳遞的其他引數將設定/覆蓋這些引數。

以下是一個示例

from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig

fsdp_plugin = FullyShardedDataParallelPlugin(
    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
)

accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

儲存和載入

使用 FSDP 模型時，新的推薦檢查點方法是在設定 accelerate 配置時使用 SHARDED_STATE_DICT 作為 StateDictType。以下是使用 accelerate 的 save_state 工具進行儲存的程式碼片段。

accelerator.save_state("ckpt")

檢查檢查點資料夾，可以看到模型和最佳化器按每個程序分片儲存

ls ckpt
# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin

cd ckpt

ls optimizer_0
# __0_0.distcp  __1_0.distcp

ls pytorch_model_0
# __0_0.distcp  __1_0.distcp

要載入它們以恢復訓練，請使用 accelerate 的 load_state 工具

accelerator.load_state("ckpt")

當使用 transformers 的 save_pretrained 時，傳遞 state_dict=accelerator.get_state_dict(model) 來儲存模型的狀態字典。下面是一個示例

  unwrapped_model.save_pretrained(
      args.output_dir,
      is_main_process=accelerator.is_main_process,
      save_function=accelerator.save,
+     state_dict=accelerator.get_state_dict(model),
)

狀態字典

accelerator.get_state_dict 將使用 FullStateDictConfig(offload_to_cpu=True, rank0_only=True) 上下文管理器呼叫底層的 model.state_dict 實現，以僅為 rank 0 獲取狀態字典，並且它將被解除安裝到 CPU。

然後，您可以將 state 傳遞到 save_pretrained 方法中。您可以使用 StateDictType 和 FullStateDictConfig 的幾種模式來控制 state_dict 的行為。更多資訊，請參閱 PyTorch 文件。

如果您選擇使用 StateDictType.SHARDED_STATE_DICT，在 Accelerator.save_state 期間，模型的權重將根據模型上的每個子拆分被分割成 n 個檔案。要將它們合併回一個單獨的字典，以便在訓練後加載回模型，您可以使用 merge_weights 工具

from accelerate.utils import merge_fsdp_weights

# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder
merge_fsdp_weights("pytorch_model_fsdp_0", "output_path", safe_serialization=True)

最終的輸出將被儲存到 model.safetensors 或 pytorch_model.bin（如果傳遞了 safe_serialization=False）。

這也可以透過 CLI 呼叫

accelerate merge-weights pytorch_model_fsdp_0/ output_path

FSDP 分片策略與 DeepSpeed ZeRO 階段的對映關係

• FULL_SHARD 對應 DeepSpeed ZeRO Stage-3。對最佳化器狀態、梯度和引數進行分片。
• SHARD_GRAD_OP 對應 DeepSpeed ZeRO Stage-2。對最佳化器狀態和梯度進行分片。
• NO_SHARD 對應 ZeRO Stage-0。不進行分片，其中每個 GPU 都擁有模型、最佳化器狀態和梯度的完整副本。
• HYBRID_SHARD 對應 ZeRO++ Stage-3，其中 zero_hpz_partition_size=<num_gpus_per_node>。這裡，它將在每個節點內對最佳化器狀態、梯度和引數進行分片，而每個節點都擁有完整的副本。

一些需要注意的事項

• 在有多個模型的情況下，將最佳化器按照對應模型的順序傳遞給 prepare 呼叫，否則 accelerator.save_state() 和 accelerator.load_state() 會導致錯誤/意外行為。
• 此功能與 Transformers 庫中 run_translation.py 指令碼的 --predict_with_generate 不相容。

為了進行更多控制，使用者可以利用 FullyShardedDataParallelPlugin。建立此類的例項後，使用者可以將其傳遞給 Accelerator 類的例項化。有關這些選項的更多資訊，請參閱 PyTorch FullyShardedDataParallel 程式碼。