FullyShardedDataParallel

Fully Sharded Data Parallel (FSDP) 是一種結合了資料並行和模型並行優勢的並行化方法，用於分散式訓練。

與 DistributedDataParallel (DDP) 不同，FSDP 更節省記憶體，因為它不會在每個 GPU 上覆制模型。它將模型的引數、梯度和最佳化器狀態分片到各個 GPU 上。每個模型分片處理部分資料，並同步結果以加快訓練速度。

本指南介紹如何使用 FSDP 和 Accelerate（一個用於管理分散式訓練的庫）設定模型訓練。

pip install accelerate

配置選項

始終首先執行 accelerate config 命令，以幫助 Accelerate 設定正確的分散式訓練環境。

accelerate config

以下部分討論了一些重要的 FSDP 配置選項。您可以在 fsdp_config 引數中瞭解更多可用選項。

分片策略

FSDP 提供多種分片策略來分發模型。請參閱下表，以幫助您為您的設定選擇最佳策略。在配置檔案中使用 fsdp_sharding_strategy 引數指定策略。

CPU 解除安裝

當模型引數和梯度不被使用時，將其解除安裝到 CPU，以節省額外的 GPU 記憶體。這對於即使使用 FSDP，模型仍然過大的情況非常有用。

在配置檔案中指定 fsdp_offload_params: true 來啟用解除安裝。

封裝策略

FSDP 透過封裝網路中的每一層來應用。封裝通常以巢狀方式應用，即在每次前向傳播後丟棄完整的權重，以節省下一層的記憶體。

有幾種封裝策略可用，但*自動封裝*策略最簡單，不需要對您的程式碼進行任何更改。指定 fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP 以封裝 Transformer 層，並指定 fsdp_transformer_layer_cls_to_wrap 以確定要封裝的層（例如，BertLayer）。

也提供基於大小的封裝。如果一層超過一定數量的引數，它將被封裝。指定 fsdp_wrap_policy: SIZED_BASED_WRAP 和 min_num_param 來設定要封裝的層的最小引數數量。

檢查點

中間檢查點應儲存為分片狀態字典，因為儲存完整狀態字典（即使有 CPU 解除安裝）也耗時且可能導致 `NCCL Timeout` 錯誤，原因是廣播期間無限期掛起。

在配置檔案中指定 `fsdp_state_dict_type: SHARDED_STATE_DICT` 以儲存分片狀態字典。現在您可以使用 load_state 從分片狀態字典恢復訓練。

accelerator.load_state("directory/containing/checkpoints")

但訓練完成後，您應該儲存完整狀態字典，因為分片狀態字典僅與 FSDP 相容。

if trainer.is_fsdp_enabled:
  trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

trainer.save_model(script_args.output_dir)

TPU

PyTorch XLA 是一個用於在 XLA 裝置上執行 PyTorch 的軟體包，它在 TPU 上啟用了 FSDP。修改配置檔案以包含以下引數。有關 FSDP 可配置的其他 XLA 特定引數，請參閱 xla_fsdp_settings 引數。

xla: True # must be set to True to enable PyTorch/XLA
xla_fsdp_settings: # XLA specific FSDP parameters
xla_fsdp_grad_ckpt: True # enable gradient checkpointing

訓練

執行 accelerate config 後，您的配置檔案應該準備就緒。下面顯示了一個示例配置檔案，它在兩個 GPU 上完全分片了引數、梯度和最佳化器狀態。您的檔案可能因配置設定方式而異。

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_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: true
  fsdp_sharding_strategy: 1
  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 命令，以您在配置檔案中選擇的 FSDP 配置啟動訓練指令碼。

accelerate launch my-training-script.py

也可以直接在命令列中指定一些 FSDP 引數。

accelerate launch --fsdp="full shard" --fsdp_config="path/to/fsdp_config/" my-training-script.py

資源

與其他並行策略相比，FSDP 是一種用於使用更少 GPU 訓練大型模型的強大工具。請參閱以下資源以瞭解有關 FSDP 的更多資訊。

• 請遵循更深入的 Accelerate FSDP 指南。
• 閱讀 Introducing PyTorch Fully Sharded Data Parallel (FSDP) API 部落格文章。
• 閱讀 Scaling PyTorch models on Cloud TPUs with FSDP 部落格文章。