李白:你的模型權重很不錯,可惜被我沒收了

語言: CN / TW / HK

影象2022-11-15 11.26.jpeg 撰文|CPFLAME

大噶好,年更樓主今天想推的是,主打分散式訓練的模型庫_李白(LiBai)

https://github.com/Oneflow-Inc/libai

對於目前市面上的模型庫來說,選擇實在是太多了,換了一批又一批,眼睛都挑花了,為什麼要用LiBai? (如果你覺得LiBai萬一某天能用到,或者這篇文章讀下來感覺比較開心,可以去GitHub上點贊,如果能三連就更好了。眾所周知,GitHub點贊其實是個收藏夾功能)。

按照現在的趨勢來說,模型越來越大了,大到一張GPU甚至裝不下完整的模型,必須得上分散式並行技術,但是分散式程式碼在很多框架下都是高度定製化的,對於新手來說根本讀不懂,也不知道應該怎麼使用,導致大家上手非常的困難,讓自己珍貴的髮際線顯得更加珍貴。

針對大模型上述存在的痛點,導致我們必須上分散式(資料並行、模型並行、流水並行)才能跑起來一個大模型。

那麼,LiBai有哪些特點呢? 你坐好,我要發功了。

需要詳細分章介紹的優勢(看上去還不錯,使用者也可以聽得懂,也知道要幹什麼):

  • 簡單易用的分散式程式碼,單機程式碼和分散式程式碼基本一致
  • 可以無縫使用PyTorch、HuggingFace的model權重,並且還可以在LiBai下進行多機多卡的分散式推理
  • 開箱即用,所有的分散式並行配置(Grad Acc,AMP,Checkpointing,ZeRO,Auto Parallel)技術都只需要在config裡面一鍵設定就可以生效,不需要在演算法程式碼model.py中額外新增
  • 支援模型一鍵轉換 ONNX

我擱這兒就要介紹完的優勢(看上去大家也有,很虛的帽子話),為了不讓大家覺得過於虛,在介紹的同時也會插入相關的例子。

  1. 具有高度靈活性和高效率,同時支援動態圖eager模式和靜態圖graph模式,支援一鍵切換,在方便debug和高效性之間反覆橫跳。

  1. 對於分散式並行的支援比較全面,大家可以在裡面盡情地組合各種分散式並行的元件。

  1. LiBai下面有內建的layers直接使用,避免重複造輪子,比如用LiBai下面的Linear層就可以快速地構建一個2D並行(資料並行+模型並行)的MLP。

  1. 採用LazyCall(借鑑自detectron2)的配置系統,基於Python語法構建相比於 argparse 和 yacs-based 方式更靈活,而且每次訓練都會序列化yaml檔案,使用者可以一鍵讀取yaml檔案來複現“上古時期”的實驗結果。

  1. 具有豐富的Projects實現。由於LiBai的分散式並行設計與演算法邏輯進行了解耦,使得在Projects下面的演算法都可以享受到LiBai下面的分散式並行技術,而且隨著分散式並行技術的更新,Projects下面的演算法程式碼不需要任何更新就可以享受到更新後的成果。

  2. 和業界翹楚Megatron比起來,具有不弱於它的吞吐,甚至稍佔優勢,完整的對比實驗在LiBai tutorialhttps://libai.readthedocs.io/en/latest/tutorials/get_started/Benchmark.html )和《大模型訓練難於上青天?效率超群、易用的李白模型庫來了》,這裡給一個GPT2的3D並行資料簡單感受一下。

可能有人會問,怎麼都和Megatron去比,你們各個同行之間有對比資料嗎?主要原因有二:1) 只要呂布不說話反駁,那麼我邢道榮就有不下於呂布的勇武,人均小呂布,這很合理;2) 大家都是國產框架,中國人不卷中國人,咱們薅著一個外國人可勁兒的打。

下面分章詳細說說上述優點。

開源自助

LiBai最基礎的一個功能: 開源自助。也就是除了LiBai訓練出來的模型以外,我們還可以載入PyTorch以及HuggingFace上面的模型進行分散式推理。

由於LiBai的底層是基於OneFlow來實現的,而OneFlow的運算元絕大部分都已經和PyTorch進行了對齊,這能發揮出什麼優勢?

使用前還要看看預備知識。一個完整的模型由兩個部分構成:模型結構,換種說法就是model.py模型權重,再換種說法就是model_best.pth。

假設我們在框架A下面,有modelA.py和model_best_A.pth,我們想在框架B上面跑起來這個框架A下面的模型,應該怎麼做呢?

在框架B下面,用框架B的運算元搭建出一個modelB.py,該modelB的引數名字可以和modelA的不一致,但是前向推理的邏輯運算最好一致,然後去載入model_best_A.pth得到model_A_state_dict(),把model_A_state_dict()裡面的引數格式全部轉換成框架B下面支援的格式,其中可以運用中間格式進行轉換。

舉個例子,比如torch.tensor()->np.numpy()(中間格式)->oneflow.tensor()之前提到了modelB中的引數名字可以和modelA中的不一致,如果不一致,那麼需要把model_A_state_dict()中的key值改一下和modelB的一致。

做完了以後,直接載入我們轉換好的引數modelB.load_state_dict(model_A_state_dict),就可以在框架B下面進行推理。為了保證模型轉換好以後的準確性,可以餵給modelA以及modelB相同的輸入,檢查一下是否能得到相同輸出。

這個預備知識不僅限於LiBai,在任何模型復現或者模型遷移上面都適用。

有了預備知識以後怎麼使用PyTorch或者HuggingFace下面的模型?簡單來說分為以下幾步:

把torch的運算元替換為oneflow: 把torch_model.py下面的torch全部替換為oneflow,得到oneflow_model.py.把oneflow_model.py中的layer儘可能地替換成LiBai中支援的layer,只替換你想要的部分也可以(比如只替換Linear層),LiBai會自動把沒有替換的layer 轉換成分散式並行所需要的格式。

這一步是支援分散式推理的關鍵繼承LiBai提供好的分散式推理的基類basic.py,過載轉換權重的函式,按照PyTorch那樣寫好預處理和後處理,就可以進行分散式推理了。

下面連結裡面有極其詳細的步驟解答,看在作者不僅授人以魚,還授人以漁,還做了程式設計師最討厭的文件活兒,可以順便給LiBai點個star收藏,而且保不齊以後萬一有個什麼復現的任務,這裡面的知識點也用得上,至少可以用別人release出來的預訓練權重來驗證自己復現的model.py是否正確。這叫什麼?這叫call back!

LiBai分散式推理介紹:

http//:github.com/Oneflow-Inc/libai/discussions/386

以MT5應用HuggingFace model為例子,我們在2機4卡下面進行模型並行2X流水並行2的分散式推理,跑起來的程式碼風格如下:

```

test_inference.py

from libai.inference.text_generation import TextGenerationPipeline from libai.utils import distributed as dist

if name == "main": pipeline = TextGenerationPipeline( "projects/MT5/configs/t5_inference.py", data_parallel=1, tensor_parallel=2, pipeline_parallel=2, pipeline_stage_id=[0] * 12 + [1] * 12, pipeline_num_layers=12 * 2, model_path="data_test/t5_inference_model", mode="huggingface", )

text = ["summarize: She is a student, She is tall, She loves study"] dict1 = pipeline(text) if dist.is_main_process(): print(dict1) ```

那麼多機多卡的分散式推理指令碼

在node0上輸入指令:

NODE=2 NODE_RANK=0 ADDR=192.168.0.1 PORT=12345 bash tools/infer.sh test_inference.py 2

在node1上輸入指令:

NODE=2 NODE_RANK=1 ADDR=192.168.0.1 PORT=12345 bash tools/infer.sh test_inference.py 2

細心的朋友已經發現了,LiBai下面可以通過設定pipeline_stage_id, 來讓使用者自己設定每個stage上group的層數是多少,方便在某些極端情況下(比如你的機器0很強,但是機器1很拉胯,或者你的encoder計算量巨大,但是decoder計算量較小)手動實現負載均衡。

大模型訓練

眾所周知,大家都喜歡做點"出格"的事情,比如在上班的時候摸魚,在VScode上面炒股......

那麼LiBai呢?你甚至可以拿它來訓練模型!

在Projects(https://github.com/Oneflow-Inc/libai/tree/main/projects )下支援的模型:

下面來談談模型之外,LiBai有什麼不一樣的地方,換句話說,也就是核心競爭力在哪裡?

分散式配置和演算法邏輯解耦

LiBai進行了模組化的設計,使得分散式的配置和演算法邏輯解耦,這意味著什麼?

這意味著使用者只需要把大部分的注意力專注到演算法邏輯上面,而不用再苦惱怎麼插入各種並行的程式碼了。

簡單來說,下面這些模組都可以在config.py中進行一鍵配置。

```

my_config.py

from libai.config import get_config train = get_config("common/train.py").train optim = get_config("common/optim.py").optim graph = get_config("common/models/graph.py").graph

set dist

train.dist.data_parallel_size = 2 train.dist.tensor_parallel_size = 2 train.dist.pipeline_parallel_size = 2

set model layers for pipeline

train.dist.pipeline_num_layers = 24

set pipeline_stage_id according to your own needs.

if None, LiBai will use its own mode of distribution

train.dist.custom_pipeline_stage_id = [0]14 + [1]10

set auto parallel in LiBai

graph.auto_parallel.enabled = True

enable amp (fp16)

train.amp.enabled = True

enable gradient clipping

optim.params.clip_grad_norm = 1.0 optim.params.clip_grad_norm_type = 2.0

enable grad accumulation for 8 steps

train.num_accumulation_steps = 8

enable activation checkpointing

train.activation_checkpoint.enabled = True

enable zero for leval-2

train.zero_optimization.enabled = True train.zero_optimization.stage = 2 ```

單機和分散式程式碼幾乎一致

下面給一個簡單的2D並行(資料並行+模型並行)的MLP例子, 比如你的Linear層在16384這個維度上面比較大, 需要把它切分在不同的卡上才能裝下, 那麼在LiBai下面只需要如下所示就可以完成了,幾乎跟單機程式碼沒有區別。

``` from libai.layers.linear import Linear from oneflow import nn

write a Simple 2D Parallel MLP

class MLP_2D(nn.Module): def init(self,): super().init() self.linear1 = Linear(in_features=1024, out_features=16384, parallel="col") self.relu = nn.GELU() self.linear2 = Linear(in_features=16384, out_features=1024, parallel="row") self.dropout = nn.Dropout(p=0.5)

def forward(self, x):
    x = self.linear1(x)
    x = self.relu(x)
    x = self.linear2(x)
    x = self.dropout(x)
    return x

```

支援一鍵轉換ONNX

本人對一鍵轉ONNX的執念可謂是相當之深了。同樣以MT5為例子,LiBai支援了一鍵轉換ONNX的功能,點選以下連結就可以體驗:

https://github.com/Oneflow-Inc/libai/tree/main/libai/onnx_export

更詳細的說明和教程會在LiBai中持續釋出。如果這篇文章對你有啟發,請不要吝惜手中的收藏按鈕,歡迎去GitHub上Star、Fork、Watch三連,持續跟進最新進展。

GitHub地址:
https://github.com/Oneflow-Inc/libai

也歡迎你掃碼進技術交流群探討。

引用

1. https://github.com/facebookresearch/detectron2

2. https://github.com/Oneflow-Inc/oneflow

3. https://github.com/Oneflow-Inc/oneflow_convert

4. https://github.com/NVIDIA/Megatron-LM

歡迎下載體驗 OneFlow v0.8.0 最新版本:
https://github.com/Oneflow-Inc/oneflow/