FSDP1 與 FSDP2 對比

本指南解釋了 `FSDP1` 和 `FSDP2` 之間的主要區別，並幫助您以最小的改動將現有程式碼遷移到 `FSDP2`。

FSDP2 比 FSDP1 好在哪裡？

首先，我們需要了解 `FSDP1` 和 `FSDP2` 的內部工作原理，以便理解它們之間的差異。這也有助於我們理解 `FSDP1` 的侷限性以及 `FSDP2` 如何解決這些問題。

我們將討論一個場景：我們有一個包含 3 個 `Linear` 層的 `Layer`，並使用 `FSDP` 將其包裝以在 2 個 GPU 上進行分片。

FSDP1

首先，我們必須瞭解最初的 `FSDP1` 及其帶來的侷限性。它將每個 `FSDP` 模組表示為單個 `FlatParameter`，這是一個包含所有模組引數的一維張量，然後在不同的程序（ranks）之間進行分片。也就是說，如果用 `FSDP1` 包裝 `Layer`，會得到如下結果：

您可能會注意到一個問題。整個 `Layer` 被展平成一個 `FlatParameter`，然後在各個程序之間分片。但如果它是一個單獨的 `FlatParameter` 物件，我們如何儲存元資料呢？這是其侷限性之一。如果不使用一些不優雅的技巧，就無法正確儲存每個引數的元資料，例如 `dtype`、`requires_grad` 等。

FSDP2

這就是為什麼引入了 `FSDP2`。它不使用 `FlatParameter`，而是使用 `DTensor`，即“分散式張量”（Distributed Tensor）。每個 `DTensor` 基本上代表一個已在不同程序間分片的普通 `torch.Tensor`。它包含關於原始 `torch.Tensor` 的元資料以及它是如何分片的，放置型別是什麼等等。這就是為什麼它被稱為“按引數分片”（per-parameter sharding）。下圖顯示了差異：

原始 `Layer` 的每個引數都在第 0 維上進行分片，並在 2 個 GPU 之間分割。現在，每個 `Linear` 層都是一個單獨的 `DTensor`，按引數儲存元資料變得可能且直接。

在上圖中，為了適應螢幕顯示，張量是在第 1 維上分片的，實際上，如上所述，它們是在第 0 維上分片的。

FSDP2 提供了什麼？

`FSDP2` 是 PyTorch 完全分片資料並行訓練 API 的新改進版本。它的主要優勢是使用 `DTensor` 來表示分片引數。與 `FSDP1` 相比，它提供了：

更簡單的內部實現，其中每個 `Parameter` 都是一個單獨的 `DTensor`
由於上述原因，可以輕鬆實現部分引數凍結，這使得像 `LORA` 這樣的方法可以直接使用
透過 `DTensor`，`FSDP2` 支援在同一模型中混合使用 `fp8` 和其他引數型別
使用 `SHARDED_STATE_DICT` 和 `torch.distributed.checkpoint`，可以實現更快、更簡單的檢查點，無需跨程序進行額外通訊，這樣，每個程序只儲存自己的分片和相應的元資料
對於載入，它使用分片模型的 `state_dict` 直接載入分片引數
支援非同步檢查點，其中引數首先被複制到 CPU 記憶體，之後主執行緒繼續訓練，而另一個執行緒將引數儲存到磁碟
記憶體效率和確定性的記憶體使用，`FSDP2` 不再使用 `recordStream`，而是使用流到流的同步（更多技術細節請參閱此論壇帖子和此 issue）
未來，計劃透過 `torch.compile` 最佳化通訊模式，進一步提高效能和記憶體效率

API 差異

我們已經討論了內部差異，現在讓我們討論一下您作為使用者需要了解的差異。

以下是透過 `accelerate` CLI 使用 `FSDP2` 時配置選項的主要變化：

舊版 (`FSDP1`)	新版 (`FSDP2`)	變化說明
`--fsdp_sharding_strategy`	`--fsdp_reshard_after_forward`	取代 `--fsdp_sharding_strategy`，改為 `true`（之前是 `FULL_SHARD`）或 `false`（之前是 `SHARD_GRAD_OP`）
`--fsdp_backward_prefetch`	**已移除**	`FSDP2` 預設使用之前的 `BACKWARD_PRE` 選項，因為只有這樣才能實現通訊和計算的重疊
`--fsdp_forward_prefetch`	**尚未實現**	如何實現此功能正在積極討論中，目前在 `FSDP2` 中尚不支援
`--fsdp_sync_module_states`	**已移除**	使用 `FSDP2` 後，此引數變得多餘
`--fsdp_cpu_ram_efficient_loading`	`--fsdp_cpu_ram_efficient_loading`	如果為 `true`，`FSDP2` 同樣只在 rank 0 上載入模型，然後引數同步到其他 rank，這與 `FSDP1` 的行為相同，但是不再需要設定 `--fsdp_sync_module_states`
`--fsdp_state_dict_type`	`--fsdp_state_dict_type`	`LOCAL_STATE_DICT` 已過時，在 `FSDP2` 中 `SHARDED_STATE_DICT` 是預設選項，這導致沒有額外的通訊，並且每個 rank 儲存自己的分片；另一個可能的選項是 `FULL_STATE_DICT`，它會導致額外的通訊和記憶體使用激增，但會從 rank 0 儲存完整的模型。
`--fsdp_use_orig_params`	**已移除**	`FSDP2` 在後臺使用 `DTensor` 類，這意味著它預設總是使用原始引數
新增	`--fsdp_version`	預設選項為 `1`，以免破壞現有程式碼，設定為 `2` 以使用 `FSDP2`

對於所有其他未更改的選項，請參閱 `FSDP` 文件。

如何切換到 FSDP2

如果使用 Python 程式碼：

在建立外掛時，只需設定 `fsdp_version=2`，並根據上表替換選項即可。

from accelerate import FullyShardedDataParallelPlugin, Accelerator

fsdp_plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2
    # other options...
)
accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

如果使用 YAML 配置：

使用我們的轉換工具

accelerate to-fsdp2 --config_file config.yaml --output_file new_config.yaml

這將自動將所有 FSDP1 設定轉換為其 FSDP2 的等效設定。使用 `--overwrite` 來更新現有檔案，而不是建立一個新檔案。

< > 在 GitHub 上更新

Accelerate