使用Colossal-AI分散式訓練BERT模型
前言
最近幾周在研究分散式訓練中的模型並行技術。為了直觀感受和加深記憶，閱讀相關論文的同時，動手用開源的大模型訓練框架Colossal-AI逐步改寫出了一個數據並行+模型並行的BERT來幫助理解。在這裡想介紹一下藉助Colossal-AI提供的零冗餘優化器、張量並行、流水線並行等技術一點點縮小BERT模型記憶體佔用的過程。
文章內容:
- 大規模模型對分散式訓練帶來了什麼挑戰?什麼是Colossal-AI？
- 用Colossal-AI提供的分散式技術訓練BERT模型
- 資料並行
- 零冗餘優化器
- 張量並行
- 流水線並行
- 實驗
- 總結
背景知識
大規模模型訓練
NLP領域中新的預訓練模型不斷對各種語言任務的效果做出了突破。這些預訓練模型依賴於深而寬的網路結構來“記憶”某些語言表徵，往往把模型的Layer增多增寬（引數量也隨之變多）能進一步提升模型表現，因此近年來NLP模型的模型引數也越來越多。比如BERT論文中提到的BERT-base有1億引數和BERT-large有3億引數；最近的GPT-3和PaLM的large更是高達1750億引數和5400億引數。 大規模模型為分散式訓練帶來了新的挑戰：過去單卡就能放下的小規模模型僅用資料並行就能達到優秀的效能和可擴充套件性；而如今單個模型甚至運算元的引數量就超過了單卡記憶體，需要用更復雜的並行技術來將引數分佈到各個節點多張卡上，使分散式訓練能支援更大規模的模型訓練。
Colossal-AI
Colossal-AI是一個專注於大規模模型訓練的深度學習系統，Colossal-AI基於PyTorch開發，旨在支援完整的高效能分散式訓練生態。Colossal-AI已在GitHub上開源，且多次登頂GitHub trending榜單，感興趣的同學可以訪問我們的GitHub主頁（http://github.com/hpcaitech/ColossalAI）。在Colossal-AI中，我們支援了不同的分散式加速方式，包括張量並行、流水線並行、零冗餘資料並行、異構計算等。
簡單的來說，就像我們藉助PyTorch、Tensorflow等訓練框架提供的方法來寫單機模型訓練一樣，Colossal-AI旨在幫助使用者像寫單機訓練一樣去寫大規模模型的分散式訓練，並提供了PyTorch-like的介面和使用方式，使使用者儘量無痛遷移目前的單機模型。
分散式加速方式
關於上面提到的各種優化技術，Colossal-AI官方文件有原論文地址和詳細的探討，這裡不贅述，想了解細節的朋友移步我們的文件Paradigms of Parallelism | Colossal-AI。
Colossal-AI:從資料並行的BERT模型到模型並行的BERT模型
下面我們用一個PyTorch開發的BERT模型由易到難逐步實驗各種優化技術的效果。 原始模型我們直接採用huggingface-BERT的BERTForMaskedLM，該模型用Masked Language Model單任務預訓練BERT。Colossal-AI使用v0.1.0。參考Colossal-AI官方的初始化和engine文件，我們定義配置檔案config.py，並用colossalai.launch和colossalai.initialize介面啟動一個engine來執行BERTForMaskedLM訓練。 下面展示一些 內聯程式碼片。
1 import colossalai
2
3 colossalai.launch(config='./config.py', ...)
4
5 # define model, optimizer, dataloader, criterion just like using PyTorch
6 ...
7
8 engine, trainloader, testloader, _ = colossalai.initialize(
9 model=model,
10 optimizer=optimizer,
11 criterion=criterion,
12 train_dataloader=train_dataloader,
13 test_dataloader=test_dataloader,
14 )
15
16 for data, label in trainloader:
17 data, label = data.cuda(), label.cuda()
18 engine.zero_grad()
19 output = engine(data)
20 loss = engine.criterion(output, label)
21 engine.backward(loss)
22 engine.step()
資料並行
首先實驗資料並行。資料並行對於模型程式碼是非侵入式的，只需要正常啟動colossal-AI engine即可，Colossal-AI預設會自動配置資料並行。由於GPU的數量等於資料並行大小 x 張量並行大小(default=1) x 流水線並行大小(default=1)，Colossal-AI會根據張量並行和流水線並行配置自動配置資料並行。不傳入自己實現的梯度處理器時，預設會使用PyTorch自帶的DistributedDataParallel來做資料並行。所以我們的訓練不額外新增任何config，在啟動並傳入了GPU Number=8後，已經自動啟動了8卡資料並行。
零冗餘優化器
其次實驗零冗餘優化器ZeRO。ZeRO也不需要修改任何原生模型程式碼，只需要我們配置config和在colossalai.zero.init_ctx.ZeroInitContext的上下文內建立模型，Colossal-AI會自動檢測ZeRO相關配置然後管理模型訓練過程中需要減少冗餘的引數。
第一步是修改config檔案，增加如下ZeRO配置表明我們要啟動ZeRO，引數解釋見註釋。ZeRO更全面的設定參見ZeRO文件。
1 from colossalai.zero.shard_utils import TensorShardStrategy
2
3 zero = dict(
4 model_config=dict( #模型引數
5 offload_config=dict(device="cpu"), #在不參與計算時將模型引數解除安裝到CPU上，進一步減少視訊記憶體開銷
6 shard_strategy=TensorShardStrategy() #指定使用的切片策略，這裡我們使用Colossalai預設策略、將每個張量均勻地分片到所有rank上
7 ),
8 optimizer_config=dict( #優化器引數
9 cpu_offload=True, #將優化器狀態從 GPU 解除安裝到 CPU，以節省 GPU 的記憶體使用
10 initial_scale=2**5, #自動混合精度訓練的初始scale
11 )
12 )
第二步是在Colossalai ZeRO上下文內建立模型，讓Colossalai能管理原生的PyTorch模型。
1 from colossalai.zero.init_ctx import ZeroInitContext
2
3 zero_ctx = ZeroInitContext(
4 target_device=torch.cuda.current_device(),
5 #gpc.config is all things defined in config.py as a dict
6 shard_strategy=gpc.config.zero.model_config.shard_strategy,
7 shard_param=True,
8 )
9
10 ...
11
12 with zero_ctx:
13 model = build_model()
14
15 ...
16 # colossalai.initialize(...)
17 # start your train
這樣我們就能啟動Colossalai的零冗餘優化器來減少原生模型的視訊記憶體使用，也是非常便捷。
張量並行
然後實驗張量並行。張量並行比起前兩個優化技術稍微複雜一些，需要在config中配置和改動模型程式碼來使用。 第一步仍然是修改config檔案，增加如下設定表明我們希望在傳入的8卡上啟動一個數據並行大小為4(自動配置)、張量並行大小為2的1D張量並行。
1 parallel = dict(
2 tensor=dict(size=2, mode='1d')
3 )
第二步我們需要更改模型程式碼。
-
首先是將原生BertForMaskedLM中的
torch.nn.Linear/LayerNorm/Embedding/Dropout以及損失函式torch.nn.CrossEntropyLoss替換成colossalai.nn.Linear/LayerNorm/Embedding/Dropout/CrossEntropyLoss。Colossalai提供的這些運算元會根據當前配置的張量並行模式自動切分模型引數：如張量並行大小為2的1D模式下，一個Linear層6x8的weight張量會根據所在位置被切分為2個3x8或者6x4的張量，其輸出結果張量也會相應的切分，並在forward和backward的過程中處理合並邏輯，從而完成自動張量並行。同樣，只要在配置檔案中修改合適的size和mode，方法會自動根據配置檔案變更張量並行方式，非常方便。 -
因為Colossalai要求
colossalai.nn運算元定義的順序和forward呼叫的順序一致來保證多個運算元間輸入輸出切分的維度匹配，我們需要嚴格將所有colossalai.nn運算元定義順序按照實際使用順序排序。 -
對於注意力運算元中的qkv Linear運算元做特殊處理。目前預設狀態下，Colossal 1D Linear運算元會嘗試將qkv三個連續定義的
colossalai.nn.Linear運算元劃分成：q-按列切(col)|k-按行切(row)|v-col，預期三個Linear運算元將會是q(k(v(input)))的計算路徑，這樣的計算路徑下col-row-col的切分能保證運算形狀正確。但這樣不符合我們對qkv的實際使用期望。我們可以直接呼叫colossalai.nn.Linear1D_col手動指定1D按列切分來處理，這樣可以在1D mode正確計算注意力。不過為了讓模型程式碼能保持根據config適配1D、2D、2.5D和3D的能力，我們可以選擇將qkv直接合併為一個size*3的colossalai.nn.Linear運算元，然後在forward計算中再重新chunk成3份來表示qkv，避免k在1D下被誤切分為row。
1 def __init__(self, ...):
2 ...
3 self.query_key_value = colossalai.nn.Linear(hidden_size,
4 num_attention_heads * attention_head_size * 3)
5 ...
6
7 def forward(self, ...):
8 ...
9 qkv = self.query_key_value(hidden_states)
10 q, k, v = torch.chunk(qkv, 3, dim=-1)
11 ...
-
對於模型最後一層Head使用的Linear層，呼叫特殊運算元
colossalai.nn.Classifier替換。Classifier對於Bert的MLM任務輸出會做vocab切分並行，也就是將num_class維度切分。得到的輸出再交給同樣對vocab做切分並行的colossalai.nn.CrossEntropyLoss計算loss。到這裡，大部分使用Colossal做張量並行的工作已經完成，已經可以執行切分邏輯最簡單的1D張量並行。 -
接下來我們對attention mask的維度做微調，來確保1D、2D、2.5D和3D張量並行時候注意力計算都能正確。對於多維(>1D)的張量並行，為了保證模型切分後計算的正確性，需要將attention mask先用
colossalai.nn.partition_batch從batch維度切分，並從[partition_batch_size, seq_length]reshape為[partition_batch_size, 1, 1, seq_length]，這樣計算注意力時可以廣播成[partition_batch_size, num_heads, from_seq_length, to_seq_length]，以便於多維張量平行計算。最後的轉化dtype是因為我們這裡的初始化賦值有可能改變attention mask的dtype，需要轉化為原本的dtype。
1 extended_attention_mask = attention_mask.view(batch_size, -1)
2 extended_attention_mask = col_nn.partition_batch(extended_attention_mask)
3 extended_attention_mask = extended_attention_mask.unsqueeze(1).unsqueeze(2)
4 extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
5 extended_attention_mask = extended_attention_mask.to(dtype=embedding_output.dtype)
這樣我們就改寫了一個1D、2D、2.5D和3D支援張量並行的Colossal Bert模型，並可以通過config檔案靈活更改張量並行大小和方式。通過張量並行，我們能將運算元和模型分佈到多張GPU上，減少視訊記憶體開銷。
流水線並行
最後實驗流水線並行。流水線並行需要配置config、修改模型、使用Colossalai封裝的高階API trainer代替engine訓練。新版本Colossalai流水線並行的介面做了簡化，請參考Pipeline Parallel | Colossal-AI。 配置config，啟動一個2階段的流水線並行：
1 parallel = dict(
2 pipeline=2,
3 )
流水線並行對原生模型的主要修改是要切分每個流水線階段應該執行的邏輯。這裡的BertForMaskedLM修改可以通過讓Embedding只被first執行，輸出token分類結果只被last執行，然後將中間的Encoder Layers切分給各流水線階段完成：
1 class PipelineBertForMaskedLM(nn.Module):
2 def __init__(..., first, last):
3 ...
4 self.first = first
5 self.last = last
6
7 if self.first:
8 ...
9 # The logic will be executed only in the first stage
10
11 ...
12 # The logic will be executed in every pipeline stage
13
14 if self.last:
15 ...
16 # The logic will be executed only in the last stage
17 ...
18
19 def forward(
20 ...
21 ):
22 ...
23 if self.first:
24 ...
25 # The logic will be executed only in the first stage
26
27 ...
28 # The logic will be executed in every pipeline stage
29
30 if self.last:
31 ...
32 # The logic will be executed only in the last stage
33
34 ...
使用trainer替代engine來啟動Colossalai訓練。trainer是Colossalai提供的高階API，可以簡化訓練過程，並且為流水線並行的自動排程提供支援：
1 from colossalai.logging import get_dist_logger
2 from colossalai.trainer import Trainer, hooks
3
4 # build components and initialize with colossalai.initialize
5 ...
6
7 # create a logger so that trainer can log on the console
8 logger = get_dist_logger()
9
10 # create a trainer object
11 trainer = Trainer(
12 engine=engine,
13 logger=logger
14 )
15
16 # define the hooks to attach to the trainer
17 hook_list = [
18 hooks.LossHook(),
19 hooks.LRSchedulerHook(lr_scheduler=lr_scheduler, by_epoch=True),
20 hooks.AccuracyHook(accuracy_func=Accuracy()),
21 hooks.LogMetricByEpochHook(logger),
22 ]
23
24 # start training
25 trainer.fit(
26 train_dataloader=train_dataloader,
27 epochs=NUM_EPOCHS,
28 test_dataloader=test_dataloader,
29 test_interval=1,
30 hooks=hook_list,
31 display_progress=True
32 )
實驗
實驗環境為搭載8張GPU的小型伺服器，每張顯示卡視訊記憶體為16GB。 因為Colossal-AI可以通過配置檔案靈活修改並行方式，我們可以通過構建不同的config來測試視訊記憶體佔用。config的取值如下：
1 資料並行大小DP size取值{1, 2, 4, 8}
2 張量並行大小TP mode和size取值{1, 1d: {2, 4, 8}, 2d: {4, 8}, 2.5d|depth=2: {8}, 3d: {8}}
3 流水線並行大小PP size取值{1, 2, 4, 8}
4 DP size * TP size * PP size = 8
5 零冗餘優化器根據不使用，使用切片且解除安裝到CPU上，取值{F, T}
測試BERT-large在不同Colossal-AI並行配置下的單張GPU記憶體佔用峰值： 
實驗結果表明通過使用Colossal-AI提供的各類並行技術，能有效控制模型訓練中的GPU記憶體使用，從而支援更大規模的規模的模型訓練。
總結
- 大規模模型訓練對分散式訓練框架提出了更高的要求。
- Colossal-AI基於PyTorch提供了強大的資料並行、零冗餘優化器、張量並行、流水線並行等訓練技術，幫助使用者儘可能簡單地從單機訓練模型遷移至分散式的訓練模型。
- Colossal-AI能通過配置檔案靈活地配置不同的混合並行策略，有效支援更大規模的模型訓練。
專案團隊
潞晨技術團隊的核心成員均來自美國加州大學伯克利分校，斯坦福大學，清華大學，北京大學，新加坡國立大學，新加坡南洋理工大學等國內外知名高校；擁有Google Brain、IBM、Intel、 Microsoft、NVIDIA等知名廠商工作經歷。公司成立即獲得創新工場、真格基金等多家頂尖VC機構種子輪投資。
目前，潞晨科技還在廣納英才，招聘全職/實習AI分散式系統、架構、編譯器、網路、CUDA、SaaS、k8s等核心系統研發人員，開源社群運營、銷售人員。
潞晨科技提供有競爭力的薪資回報，特別優秀的，還可以申請遠端工作。也歡迎各位向潞晨科技引薦優秀人才，如果您推薦優秀人才成功簽約潞晨科技，我們將為您提供數千元至數萬元的推薦費。
工作地點：中國北京，新加坡，美國。(可相互轉崗) 簡歷投遞郵箱：[email protected]
傳送門
BERT專案地址： http://github.com/hpcaitech/ColossalAI-Examples
Colossal-AI專案地址： http://github.com/hpcaitech/ColossalAI
Colossal-AI文件地址： http://www.colossalai.org/
- 開源方案復現ChatGPT流程！1.62GB視訊記憶體即可體驗，單機訓練提速7.73倍
- 硬體預算最高直降46倍！低成本上手AIGC和千億大模型，一行程式碼自動並行，Colossal-AI再升級
- 潞晨科技完成600萬美元種子及天使輪融資，藍馳領投天使輪
- 無縫支援Hugging Face社群，Colossal-AI低成本輕鬆加速大模型
- 推理加速效能超越英偉達FasterTransformer 50%，開源方案打通大模型落地關鍵路徑
- 在個人電腦用單塊GPU帶動180億引數GPT！熱門開源專案再添新特性
- 使用Colossal-AI分散式訓練BERT模型
- 使用Colossal-AI復現Pathways Language Model
- 霸榜GitHub熱門第一多日後，Colossal-AI正式版釋出
- 助力藥物研發，低成本加速AlphaFold訓練從11天到67小時，11倍推理加速——開源解決方案FastFold