[源碼分析] Facebook如何訓練超大模型---(1)

[源碼分析] Facebook如何訓練超大模型---(1)

0x00 摘要

我們在前文介紹過,微軟 ZeRO 可以對一個萬億參數模型可以使用 8 路模型並行、64 路管道並行和 8 路數據並行在 4,096 個 NVIDIA A100 GPU 上進行擴展。

而FSDP(Fully Sharded Data Parallel)是Facebook 深度借鑑微軟ZeRO之後提出的PyTorch DDP升級版本,可以認爲是對標微軟 ZeRO,其本質是 parameter sharding。Parameter sharding 就是把模型參數等切分到各個GPU之上。我們會以 Google,微軟和 Facebook 的論文,博客以及代碼來進行學習分析。

本系列其他文章如下:

[源碼解析] PyTorch 分佈式之 ZeroRedundancyOptimizer

[論文翻譯] 分佈式訓練 Parameter sharding 之 ZeRO

[論文翻譯] 分佈式訓練 Parameter Sharding 之 Google Weight Sharding

0x01 簡介

1.1 FAIR & FSDP

大規模訓練人工智能模型並不容易。除了需要大量的計算能力和資源外,訓練非常大的模型背後還有相當大的工程複雜性。Facebook人工智能研究(FAIR)工程部一直致力於構建工具和基礎設施,以使大型人工智能模型的培訓變得更容易。

Fully Sharded Data Parallel(FSDP)是FAIR引入的最新工具。它將AI模型的參數在數據並行worker之間進行切分,並且可以選擇將部分訓練計算卸載到CPU。顧名思義,FSDP是一種數據並行訓練算法。儘管參數被分片到不同的GPU,但每個微批次數據的計算對於每個GPU worker來說仍然是本地的。這種概念上的簡單性使FSDP更易於理解,並且更適用於各種使用場景(與層內並行和流水線並行相比)。與optimizer state+gradient sharding數據並行方法相比,FSDP在訓練過程中通過通信和計算重疊對模型參數進行更均勻的切分,具有更好的性能。

FSDP可以使用更少的GPU更有效地訓練數量級更大的模型。FSDP已在FairScale庫中實現,允許工程師和開發人員使用簡單的API擴展和優化其模型的培訓。在Facebook,FSDP已經被整合並測試,用於訓練一些NLP和Vision模型。

1.2 大規模訓練計算能力需求

大規模模型訓練需要大量的計算資源,比如OpenAI的GPT-3 擁有1750億個參數。其訓練估計需要355年的GPU時間,相當於1000個GPU連續工作4個月以上。

除了需要大量計算和工程資源外,大多數的訓練擴展方法都會帶來額外的通信成本,並且需要工程師仔細評估內存使用和計算效率之間的權衡。例如,典型的數據並行培訓要求在每個GPU上維護模型的冗餘副本,而模型並行培訓爲在worker(GPU)之間移動激活引入了額外的通信成本。

相比之下,FSDP相對而言沒有做任何權衡。它通過在GPU上分割模型參數、梯度和優化器狀態來提高內存效率,並通過分解通信並將其與前向和後向過程重疊來提高計算效率。FSDP產生與標準分佈式數據並行(DDP)培訓相同的結果,並提供易於使用的接口,該接口是PyTorch分佈式數據並行模塊的替代品。Facebook 的早期測試表明,FSDP可以擴展到數萬億個參數。

0x02 FSDP 如何工作

在標準DDP訓練中,每個worker處理一個單獨的批次,並使用all-reduce對worker之間的梯度進行彙總。雖然DDP已經變得非常流行,但它佔用的GPU內存比它實際需要的要多,因爲模型權重和優化器狀態在所有DDP worker中都有一個副本

2.1 全參數分片

減少副本的一種方法是應用全參數分片( full parameter sharding)的過程,其中僅提供局部計算所需的模型參數、梯度和優化器的子集。這種方法的一個實現 ZeRO-3 已經被微軟所普及。

解鎖全參數切分的關鍵是:我們可以把DDP之中的all reduce操作分解爲獨立的 reduce-scatter 和 all-gather 操作

圖來自 :https://engineering.fb.com/wp-content/uploads/2021/07/FSDP-graph-2a.png?w=1024

“All-reduce”是“reduce-scatter”和“all-gather”的組合。聚合梯度的標準 “All-reduce”操作可以分解爲兩個單獨的階段:“reduce-scatter”和“all-gather”。

  • “reduce-scatter”階段,在每個GPU上,會基於rank 索引對 rank 之間相等的塊進行求和。
  • “all-gather”階段,每個GPU上的聚合梯度分片可供所有GPU使用。

通過重新安排reduce scatter和all gather,每個DDP worker只需要存儲一個參數分片和優化器狀態。

2.2 比對

下圖顯示了標準DDP訓練(上半部分)和FSDP訓練(下半部分):

Full Sharded Data Parallel graph

  • 在標準的數據並行訓練方法中,每個GPU上都有一個模型副本,向前和向後傳遞的序列只在自己的數據分片上進行運行。在這些局部計算之後,每個局部過程的參數和優化器與其他GPU共享,以便計算全局權重更新。

  • 在FSDP中:

    • Model shard :每個GPU上僅存在模型的分片
    • All-gather :每個GPU通過all-gather從其他GPU收集所有權重,以在本地計算前向傳播。就是論文思路Pp下劃線部分
    • Forward(local):在本地進行前向操作。前向計算和後向計算都是利用完整模型。
    • All-gather :然後在後向傳播之前再次執行此權重收集。就是論文思路Pp之中的下劃線部分
    • Backward(local):本地進行後向操作。前向計算和後向計算都是利用完整模型,此時每個GPU上也都是全部梯度
    • Reduce-scatter :在向後傳播之後,局部梯度聚合並且通過 reduce-scatter 在各個GPU上分片,每個分片上的梯度是聚合之後本分區對應的那部分,就是論文思路Pg之中的下劃線部分。
    • Update Weight(local):每個GPU更新其局部權重分片。

爲了最大限度地提高內存效率,我們可以在每層向前傳播後丟棄全部權重,爲後續層節省內存。這可以通過將FSDP包裝應用於網絡中的每一層來實現(通過設置reshard_after_forward=True)。

下面是僞代碼實現:

FSDP forward pass:
    for layer_i in layers:
        all-gather full weights for layer_i # 權重
        forward pass for layer_i
        discard full weights for layer_i # 權重

FSDP backward pass:
    for layer_i in layers:
        all-gather full weights for layer_i # 權重
        backward pass for layer_i
        discard full weights for layer_i # 權重
        reduce-scatter gradients for layer_i # 梯度

2.3 梳理

我們結合論文思路再來梳理一下 FSDP。

2.3.1 思路

論文思路如下:

  • Pp: Parameter Partitioning,每個進程只存儲與其分區對應的參數。正向和反向傳播需要其分區外的參數時,會通過broadcast操作從適當的數據並行進程接收這些參數。雖然乍一看,這可能會導致顯著的通信開銷,但我們發現,這種方法只會將基線DP系統的總通信量增加到1.5倍,同時實現與Nd成比例的內存減少。
  • Pos : Optimizer State Partitioning對於一個\(N_d\)並行度的DP來說,我們將優化器狀態分組到\(N_d\)個相等的分區中,這樣第i個數據並行進程只更新與第i個分區對應的優化器狀態。因此,每個數據並行過程只需要存儲和更新總優化器狀態 的$ \frac{1}{N_d}\(,然後只更新\) \frac{1}{N_d}$個參數。在每個訓練步驟結束時,我們會執行一個跨數據並行進程的all-gather操作,以獲得跨所有數據並行進程的完全更新的參數。
  • Pg: Gradient Partitioning由於每個數據並行進程只負責更新其相應的參數分區,因此,每個節點僅僅對自己負責的那部分參數的梯度進行規約。在歸併之後,每個節點只需要自己參數分區對應的梯度,對於其他的梯度不再需要,所以它們的內存可以被釋放。這將梯度的內存佔用從2ψ字節縮減到 \(\frac{2ψ}{N_d}\)。實際上,這是一種 Reduce-Scatter操作,不同參數的梯度被減少到不同的進程之中。

總結一下:因爲模型參數被分區,所以參數梯度(在框架實現中,梯度往往是參數的成員變量)自然就被分區了。分區的參數被設置到優化器之中,所以優化器只會優化本分區的參數,所以優化器狀態自然就是分區之後的。注意,在前向傳播和後向傳播時候,每個GPU都是用全部模型來計算,得到的梯度也是全部的梯度,只是存儲時候只存儲自己分區對應的部分

2.3.2 流程步驟

我們再來展示一下具體流程。假設數據並行度爲 n,則有 n 個GPU,那麼每個GPU之上保存總模型參數的 1/n,同時梯度,優化器狀態就自然被分區了,每個GPU之上還有數據並行。

  • 起始狀態:每個GPU之上是\(P_n, G_n, O_n\)。注意,因爲本GPU上模型是\(P_n\),所以\(O_n\)自然就對應了\(P_n\),就自動分片了。
  • 正向計算時候,每個 \(GPU_n\) 都把自己負責的參數 \(P_n\) 廣播給其他所有的 GPU,前向計算之後,每個 \(GPU_n\) 都得到自己輸入訓練數據 \(data_n\) 的損失 \(loss_n\)
  • 反向計算時候,每個 \(GPU_n\) 也都把自己負責的參數 \(P_n\) 廣播給其他所有的 GPU,最後計算得到對應於數據 \(data_n\)的梯度 \(G_n\)
  • 將梯度 \(G_n\)聚合到對應的\(GPU_n\)上,這時候\(GPU_n\) 上的梯度就是 \(reduce(G_0, ..., G_n)\) 之中自己rank對應的部分。注意,梯度聚合過程則使用了 reduce-scatter,因爲每個gpu只需要更新自己負責的部分\(P_n,G_n,O_n\),所以不需要進行all-gather了。

0x03 How to use FSDP

目前,FAIR 提供四種解決方案來使用FSDP,以適應不同的需求。

3.1 在語言模型中使用FSDP

對於語言模型,可以在通過以下新參數,在 fairseq framework 之中支持 FSDP:

  • –ddp-backend=fully_sharded: 通過FSDP啓用完全切分。
  • –cpu-offload: 將優化器狀態和FP32模型副本卸載到cpu(與–optimizer=cpu_adam結合使用)。
  • –no-reshard-after-forward: 提高大模型訓練速度 (1B+ params) ,類似於 ZeRO stage 2。
  • 其他常見選項 (–fp16, –update-freq, –checkpoint-activations, –offload-activations, etc.) 還是繼續正常工作。

具體請參閱fairseq教程

3.2 在計算機視覺模型之中使用FSDP

對於計算機視覺模型, VISSL 中可以支持FSDP,並在regnet架構上進行了測試。像BatchNorm和ReLU這樣的層已經被無縫地處理並已經測試過其收斂性。可以使用下面選項來啓用 FSDP。

  • config.MODEL.FSDP_CONFIG.AUTO_SETUP_FSDP=True
  • config.MODEL.SYNC_BN_CONFIG.SYNC_BN_TYPE=pytorch
  • config.MODEL.AMP_PARAMS.AMP_TYPE=pytorch

在如下鏈接可以繼續研究 this section

3.3 在PyTorch Lightning使用FSDP

爲了更容易地與更通用的用例集成,PyTorch Lightning已經將FSDP作爲beta功能。[此教程](https://pytorch-lightning.readthedocs.io/en/latest/advanced/advanced_gpu.html#fully-sharded training) 包含一個關於如何將FSDP插件與PyTorch Lightning一起使用的詳細示例。如下所示,添加plugins='fsdp'可以激活它。

model = MyModel()
trainer = Trainer(gpus=4, plugins='fsdp', precision=16)
trainer.fit(model)

trainer.test()
trainer.predict()

3.4 直接從FairScale使用FSDP庫

FSDP的主要開發庫是FairScale.。您可以通過以下示例直接使用FairScale的FSDP,只需更換DDP。

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
...
# sharded_module = DDP(my_module)
sharded_module = FSDP(my_module)
optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
for sample, label in dataload.next_batch:
  out = sharded_module(x=sample, y=3, z=torch.Tensor([1]))
  loss = criterion(out, label)
  loss.backward()
  optim.step()

FairScale中的FSDP庫爲大規模訓練的許多重要方面提供了選項。當你希望使用FSDP的全部功能,你可以自行研究如下方面。

  1. 模型封裝:爲了最大限度地減少短期內的GPU內存需求,用戶需要以嵌套方式封裝模型。這增加了複雜性,但是在移植現有PyTorch模型代碼時非常有用。
  2. 模型初始化:與DDP不同,FSDP不會在GPU工作進程之間自動同步模型權重。這意味着必須小心地進行模型初始化,以便所有GPU worker具有相同的初始權重。
  3. 優化器設置:由於分片和包裝,FSDP只支持某些類型的優化器和優化器設置。特別是,如果模塊被FSDP包裝,並且其參數被展平爲單個張量,則用戶不能對此類模塊中的不同參數組使用不同的超參數。
  4. **混合精度 **:FSDP支持FP16主權重的高級混合精度訓練,以及在梯度上FP16類型的reduce和scatter。但是,模型的某些部分可能只有在使用全精度時才收斂,在這些情況下,需要額外的wrapping,以便有選擇地以全精度運行模型的某些部分。
  5. 狀態檢查點和推斷:當模型規模較大時,保存和加載模型狀態可能會變得很困難。FSDP支持多種方法使該任務成爲可能,但這些方法是有代價的。
  6. 最後,FSDP通常與激活檢查點函數一起使用,如checkpoint_wrapper 。用戶可能需要仔細調整激活檢查點策略,以便在有限GPU內存空間內容納一個大型模型。

0x04 內存管理

我們接下來看看FSDP如何管理內存。

FairScale提供了受ZeRO <https://arxiv.org/pdf/1910.02054.pdf> 啓發的算法:當使用數據並行訓練時,您需要在計算/通信效率方面權衡內存的使用。另一方面,在使用模型並行訓練時,需要爲了內存而權衡計算/通信。

模型訓練的內存使用通常分爲兩類:

  • 模型狀態:優化器狀態、梯度、參數。

  • 剩餘狀態:激活、臨時緩衝區、碎片內存。

爲了減少模型狀態下的冗餘,ZeRO提出了三種不同的算法。這些在FairScale中實現爲優化器狀態分片(Optimizer State Sharding,即OSS)、分片數據並行(Sharded Data Parallel,即SDP)和最終完全分片數據並行(Fully Sharded Data Parallel,即FSDP)。讓我們深入瞭解每一個算法的實際機制,並理解它們爲什麼能夠節省內存。

4.1 Optimizer State Sharding (OSS)

FairScale已經實現了與優化器內存相關的內存優化 OSS。

像Adam這樣的優化器通常需要保持動量、方差。即使可以使用FP16精度的參數和梯度進行訓練,參數和梯度也需要保存爲FP32精度。當每個rank更新完整模型時,這意味着相當大一部分內存被優化器狀態的冗餘表示所佔用。

爲了克服這種冗餘,優化器狀態分片需要將模型優化步驟劃分在不同的rank之間,以便每個rank只負責更新模型的對應分片。這反過來又確保優化器狀態在每個rank上小得多,並且它不包含跨rank的冗餘信息。

4.1.1 訓練流程

訓練流程可以從DDP的執行流程做如下修改:

  1. wrapped optimizer根據參數大小(而不是使用順序)以貪心算法方式來分割優化器狀態。這是爲了確保每個rank具有幾乎相同的優化器內存佔用。

  2. 訓練過程類似於PyTorch的分佈式數據並行(DDP)的過程。在每個rank上完成前向傳播,然後是向後傳播。在後向傳播過程中,使用allreduce同步梯度。

  3. 每個rank只更新它負責的優化器分配狀態參數,然後丟棄其餘的。

  4. 更新後,將執行broadcast或allgather操作,以確保所有rank都收到最新更新的參數值。

當您使用具有附加狀態的優化器(如Adam)時,OSS非常有用。如果您使用的是SGD或任何內存佔用有限的優化器,那麼在使用多個節點時,由於步驟4中的額外通信,您可能會看到速度減慢。在第2步的allreduce過程中,也有一些用於存儲梯度的浪費內存,這些內存隨後被丟棄。

4.1.2 最佳實踐

  • OSS公開了一個broadcast_fp16 flag,您可能應該在多節點作業中使用它。在單節點實驗中通常不需要這樣做。

  • 如果您的模型在大小方面極不平衡(例如,存在一個巨大的張量),那麼這種方法將不會有很大幫助,而張量切分選項,如'fairscale.nn.FullyShardedDataParallel'將更可取。

  • 3.OSS應該是DDP環境中的一個臨時解決方案,其與大多數DDP功能保持兼容。

4.1.3 性能

  • 在單個節點上,OSS應該總是比vanilla PyTorch快,內存節省會因使用的優化器而異

  • 當使用多個節點時,OSS也可以比vanilla PyTorch快或慢,具體取決於所使用的優化器和可選標誌(如上文提到的broadcast_fp16、梯度壓縮、梯度累積)

  • 如果您的實驗可以使用更大的batch size,則採取更大的batch size並減少所涉及的rank數通常是有益的,或者使用梯度累積,因爲這樣可以降低通信成本。

4.2 Optimizer + Gradient State Sharding

雖然OSS解決了優化器中的冗餘問題,但依然存在梯度聚合計算的重複以及存在用於梯度的額外內存。爲了克服冗餘梯度內存,我們可以使用梯度分片或ZeRO-2。這已由FairScale中的分片數據並行(SDP)API實現。

爲了啓用梯度分片,每個 rank 都被分配一組參數,它們負責管理優化器狀態以及梯度聚合。通過將一個模型分片分配給一個給定的rank,我們確保梯度被規約到特定的rank,而這些rank又負責相應的更新。因此這減少了通信和內存使用。

4.2.1 訓練過程

訓練過程如下:

  1. 與之前一樣,包裝的優化器在不同的列組中分割參數。

  2. 該模型現在使用分片數據並行(SDP)包裝器進行包裝,該包裝器允許我們在訓練過程中添加適當的hook並維護狀態。

  3. SDP關注於可訓練的參數,併爲每個參數添加了一個反向hook。

  4. 在反向傳播過程中,梯度將規約到指定rank,rank是在 1 中作爲切分過程的一部分指定的。使用reduce op代替allreduce op,從而減少通信開銷。

  5. 每個rank更新其負責的參數。

  6. 更新後,將執行廣播或allgather,以確保所有rank都收到最新更新的參數值。

OSS和SDPAPI都允許您減少用於梯度和優化器狀態的內存,但是如果網絡緩慢,則可能存在額外的通信成本。當遇到內存不足(OOM)問題時,可以把OSS和SDP作爲第一步嘗試。

4.2.2 最佳實踐

  • 如果使用多個節點,請通過指定reduce_buffer_size 參數確保SDP正在使用reduce buffers。改變它們的大小可能是一個優化目標,最佳配置可能取決於互連狀況。

  • 如果在單個節點上,通常最好不要使用'reduce_buffer_size',因爲它會帶來延遲成本,但不會增加內存。將此值設置爲0表示不使用此功能。

  • 如果您的實驗可以使用更大的batch size,則採取更大的batch size並減少所涉及的rank數通常是有益的,或者使用梯度累積,因爲這樣可以降低通信成本。

4.3 Optimizer + Gradient + Horizontal Model Sharding

爲了進一步優化訓練並實現更大的內存節省,我們需要啓用參數切分。

參數切分類似於梯度和優化器狀態,即,每個數據並行rank負責模型參數的一個分片。FairScale通過完全分片數據並行(FSDP)API實現參數分片,該API深受 ZeRO-3 <https://arxiv.org/pdf/1910.02054.pdf>的啓發。

參數分片有兩個如下關鍵點:

  • Allreduce操作可以分爲reduce和allgather,類似於以前的分片技術(優化器狀態和梯度)。

  • 可以使用FSDP API包裝各個層,該API允許我們在給定實例中將單個層所需的所有參數引入給定GPU,計算前向傳遞,然後丟棄不屬於該rank的參數。

使用FSDP很簡單,只需要在代碼中簡單地替換原來的DDP即可。注意:FSDP目前要求模型是一個nn.Sequential模型。

from torch.utils.data.dataloader import DataLoader
from torchvision.datasets import FakeData
from torchvision.transforms import ToTensor

from fairscale.experimental.nn.offload import OffloadModel

num_inputs = 8
num_outputs = 8
num_hidden =  4
num_layers =  2
batch_size =  8

transform = ToTensor()
dataloader = DataLoader(
    FakeData(
        image_size=(1, num_inputs, num_inputs),
        num_classes=num_outputs,
        transform=transform,
    ),
    batch_size=batch_size,
)

model = torch.nn.Sequential(
    torch.nn.Linear(num_inputs * num_inputs, num_hidden),
    *([torch.nn.Linear(num_hidden, num_hidden) for _ in range(num_layers)]),
    torch.nn.Linear(num_hidden, num_outputs),
)

4.3.1 訓練過程

具體訓練過程如下:

  • 在開始計算特定層之前,allgather模型每個層的正向傳播所需的參數。

  • 計算向前計算。

  • 在特定層開始反向傳遞之前,allgather模型每個層反向傳播所需的參數。

  • 計算向後傳播。

  • 規約梯度,以便在負責相應參數的rank上累積聚合梯度。

  • 讓每個rank使用聚合梯度更新已分配給它的參數。

有了FSDP,在使用API進行檢查點設置和保存優化器狀態時,需要做一些小的更改。鑑於優化器狀態和參數的分片性質,任何旨在保存模型狀態以供訓練或推理的API都需要考慮保存所有worker的權重。FSDP實現所需的管道(required plumbing)以保存所有worker的權重、保存單個worker的權重以及保存所有worker的優化器狀態。

FSDP還支持混合精度訓練,其中計算和通信均以FP16精度進行。如果要減少在FP32中執行的操作(這是DDP的默認行爲),則必須設置 fp32_reduce_scatter=True

爲了進一步節省內存,FSDP支持將當前未使用的參數和梯度卸載到CPU上。這可以通過將“move_params_to_cpu”和“move_grads_to_cpu”設置爲True來啓用。

4.3.2 最佳實踐

  • 對於FSDP,最好使用 model.zero_grad(set_to_none=True) ,因爲它在單步執行後節省了大量內存。

  • torch.cuda.amp.autocast與FSDP完全兼容。您需要將'mixed_precision'arg設置爲True。

  • 如果與激活檢查點相結合,則最好使用 FSDP(checkpoint_wrapper(module))而不是checkpoint_wrapper(FSDP(module)).。後者將導致更多的通信,速度也會變慢。

  • FSDP與使用pointwise優化器的DDP兼容,例如Adam、AdamW、ADADDelta、Adamax、SGD等。當使用non-pointwise優化器(例如Adagrad、Adafactor、LAMB等)時,sharding將導致略有不同的結果。

4.3.3 性能

  • 爲了獲得最佳內存效率,請使用“auto_wrap”將網絡中的每一層用FSDP進行封裝,並將 reshard_after_forward 設置爲True。這樣速度會慢,但是顯存開銷最小。

  • 爲了獲得最佳訓練速度,請將 reshard_after_forward 設置爲False(不需要包裝每一層,但如果設置,則會進一步提高速度)。

支持,FSDP基本原理和如何使用我們已經介紹完畢,下一篇我們介紹其代碼細節,看看究竟如何做到最大程度減少內存使用。

0xFF 參考

Fully Sharded Data Parallel: faster AI training with fewer GPUs

ZeRO & DeepSpeed:可以讓訓練模型擁有超過1000億個參數的優化(微軟)

Fully Sharded Data Parallel: faster AI training with fewer GPUs

https://github.com/microsoft/DeepSpeed

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章