撰文｜CPFLAME
大噶好，年更樓主今天想推的是，主打分布式訓(xùn)練的模型庫_李白（LiBai）。?
https://github.com/Oneflow-Inc/libaihttps://github.com/Oneflow-Inc/libai
【資料圖】
對于目前市面上的模型庫來說，選擇實在是太多了，換了一批又一批，眼睛都挑花了，為什么要用LiBai？（如果你覺得LiBai萬一某天能用到，或者這篇文章讀下來感覺比較開心，可以去GitHub上點贊，如果能三連就更好了。眾所周知，GitHub點贊其實是個收藏夾功能）。
按照現(xiàn)在的趨勢來說，模型越來越大了，大到一張GPU甚至裝不下完整的模型，必須得上分布式并行技術(shù)，但是分布式代碼在很多框架下都是高度定制化的，對于新手來說根本讀不懂，也不知道應(yīng)該怎么使用，導(dǎo)致大家上手非常的困難，讓自己珍貴的發(fā)際線顯得更加珍貴。
針對大模型上述存在的痛點，導(dǎo)致我們必須上分布式（數(shù)據(jù)并行、模型并行、流水并行）才能跑起來一個大模型。
那么，LiBai有哪些特點呢？你坐好，我要發(fā)功了。
需要詳細(xì)分章介紹的優(yōu)勢（看上去還不錯，用戶也可以聽得懂，也知道要干什么）：
簡單易用的分布式代碼，單機(jī)代碼和分布式代碼基本一致
可以無縫使用PyTorch、HuggingFace的model權(quán)重，并且還可以在LiBai下進(jìn)行多機(jī)多卡的分布式推理
開箱即用，所有的分布式并行配置(Grad Acc，AMP，Checkpointing，ZeRO，Auto Parallel）技術(shù)都只需要在config里面一鍵設(shè)置就可以生效，不需要在算法代碼model.py中額外添加
支持模型一鍵轉(zhuǎn)換 ONNX
我擱這兒就要介紹完的優(yōu)勢（看上去大家也有，很虛的帽子話)，為了不讓大家覺得過于虛，在介紹的同時也會插入相關(guān)的例子。?
1. 具有高度靈活性和高效率，同時支持動態(tài)圖eager模式和靜態(tài)圖graph模式，支持一鍵切換，在方便debug和高效性之間反復(fù)橫跳。
2. 對于分布式并行的支持比較全面，大家可以在里面盡情地組合各種分布式并行的組件。
3. LiBai下面有內(nèi)置的layers直接使用，避免重復(fù)造輪子，比如用LiBai下面的Linear層就可以快速地構(gòu)建一個2D并行（數(shù)據(jù)并行+模型并行）的MLP。?
4. 采用LazyCall（借鑒自detectron2）的配置系統(tǒng)，基于Python語法構(gòu)建相比于 argparse 和 yacs-based 方式更靈活，而且每次訓(xùn)練都會序列化yaml文件，用戶可以一鍵讀取yaml文件來復(fù)現(xiàn)“上古時期”的實驗結(jié)果。
5. 具有豐富的Projects實現(xiàn)。由于LiBai的分布式并行設(shè)計與算法邏輯進(jìn)行了解耦，使得在Projects下面的算法都可以享受到LiBai下面的分布式并行技術(shù)，而且隨著分布式并行技術(shù)的更新，Projects下面的算法代碼不需要任何更新就可以享受到更新后的成果。
6. 和業(yè)界翹楚Megatron比起來，具有不弱于它的吞吐，甚至稍占優(yōu)勢，完整的對比實驗在LiBai tutorial（https://libai.readthedocs.io/en/latest/tutorials/get_started/Benchmark.html）和《大模型訓(xùn)練難于上青天？效率超群、易用的李白模型庫來了》，這里給一個GPT2的3D并行數(shù)據(jù)簡單感受一下。
可能有人會問，怎么都和Megatron去比，你們各個同行之間有對比數(shù)據(jù)嗎？主要原因有二：1) 只要呂布不說話反駁，那么我邢道榮就有不下于呂布的勇武，人均小呂布，這很合理；2) 大家都是國產(chǎn)框架，中國人不卷中國人，咱們薅著一個外國人可勁兒的打。
下面分章詳細(xì)說說上述優(yōu)點。?開源自助
LiBai最基礎(chǔ)的一個功能:?開源自助。也就是除了LiBai訓(xùn)練出來的模型以外，我們還可以加載PyTorch以及HuggingFace上面的模型進(jìn)行分布式推理。?
由于LiBai的底層是基于OneFlow來實現(xiàn)的，而OneFlow的算子絕大部分都已經(jīng)和PyTorch進(jìn)行了對齊，這能發(fā)揮出什么優(yōu)勢?
使用前還要看看預(yù)備知識。一個完整的模型由兩個部分構(gòu)成：模型結(jié)構(gòu)，換種說法就是model.py模型權(quán)重，再換種說法就是model_best.pth。
假設(shè)我們在框架A下面，有modelA.py和model_best_A.pth，我們想在框架B上面跑起來這個框架A下面的模型，應(yīng)該怎么做呢?
在框架B下面，用框架B的算子搭建出一個modelB.py，該modelB的參數(shù)名字可以和modelA的不一致，但是前向推理的邏輯運(yùn)算最好一致，然后去加載model_best_A.pth得到model_A_state_dict()，把model_A_state_dict()里面的參數(shù)格式全部轉(zhuǎn)換成框架B下面支持的格式，其中可以運(yùn)用中間格式進(jìn)行轉(zhuǎn)換。
舉個例子，比如torch.tensor()->np.numpy()(中間格式)->oneflow.tensor()之前提到了modelB中的參數(shù)名字可以和modelA中的不一致，如果不一致，那么需要把model_A_state_dict()中的key值改一下和modelB的一致。
做完了以后，直接加載我們轉(zhuǎn)換好的參數(shù)modelB.load_state_dict(model_A_state_dict)，就可以在框架B下面進(jìn)行推理。為了保證模型轉(zhuǎn)換好以后的準(zhǔn)確性，可以喂給modelA以及modelB相同的輸入，檢查一下是否能得到相同輸出。
這個預(yù)備知識不僅限于LiBai，在任何模型復(fù)現(xiàn)或者模型遷移上面都適用。
有了預(yù)備知識以后怎么使用PyTorch或者HuggingFace下面的模型？簡單來說分為以下幾步：
把torch的算子替換為oneflow: 把torch_model.py下面的torch全部替換為oneflow，得到oneflow_model.py.把oneflow_model.py中的layer盡可能地替換成LiBai中支持的layer，只替換你想要的部分也可以（比如只替換Linear層），LiBai會自動把沒有替換的layer 轉(zhuǎn)換成分布式并行所需要的格式。
這一步是支持分布式推理的關(guān)鍵繼承LiBai提供好的分布式推理的基類basic.py，重載轉(zhuǎn)換權(quán)重的函數(shù)，按照PyTorch那樣寫好預(yù)處理和后處理，就可以進(jìn)行分布式推理了。
下面鏈接里面有極其詳細(xì)的步驟解答，看在作者不僅授人以魚，還授人以漁，還做了程序員最討厭的文檔活兒，可以順便給LiBai點個star收藏，而且保不齊以后萬一有個什么復(fù)現(xiàn)的任務(wù)，這里面的知識點也用得上，至少可以用別人release出來的預(yù)訓(xùn)練權(quán)重來驗證自己復(fù)現(xiàn)的model.py是否正確。這叫什么？這叫call back!
LiBai分布式推理介紹：
http//:github.com/Oneflow-Inc/libai/discussions/386
以MT5應(yīng)用HuggingFace model為例子，我們在2機(jī)4卡下面進(jìn)行模型并行2X流水并行2的分布式推理，跑起來的代碼風(fēng)格如下：
# test_inference.pyfrom libai.inference.text_generation import TextGenerationPipelinefrom libai.utils import distributed as distif __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)
那么多機(jī)多卡的分布式推理腳本
在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
細(xì)心的朋友已經(jīng)發(fā)現(xiàn)了，LiBai下面可以通過設(shè)置pipeline_stage_id， 來讓用戶自己設(shè)置每個stage上group的層數(shù)是多少，方便在某些極端情況下（比如你的機(jī)器0很強(qiáng)，但是機(jī)器1很拉胯，或者你的encoder計算量巨大，但是decoder計算量較小）手動實現(xiàn)負(fù)載均衡。大模型訓(xùn)練眾所周知，大家都喜歡做點"出格"的事情，比如在上班的時候摸魚，在VScode上面炒股......那么LiBai呢？你甚至可以拿它來訓(xùn)練模型！?
在Projects（https://github.com/Oneflow-Inc/libai/tree/main/projects）下支持的模型：
下面來談?wù)勀Ｐ椭猓琇iBai有什么不一樣的地方，換句話說，也就是核心競爭力在哪里？
分布式配置和算法邏輯解耦
LiBai進(jìn)行了模塊化的設(shè)計，使得分布式的配置和算法邏輯解耦，這意味著什么??
這意味著用戶只需要把大部分的注意力專注到算法邏輯上面，而不用再苦惱怎么插入各種并行的代碼了。
簡單來說，下面這些模塊都可以在config.py中進(jìn)行一鍵配置。
# my_config.pyfrom libai.config import get_configtrain = get_config("common/train.py").trainoptim = get_config("common/optim.py").optimgraph = get_config("common/models/graph.py").graph# set disttrain.dist.data_parallel_size = 2train.dist.tensor_parallel_size = 2train.dist.pipeline_parallel_size = 2# set model layers for pipelinetrain.dist.pipeline_num_layers = 24# set pipeline_stage_id according to your own needs.# if `None`, LiBai will use its own mode of distributiontrain.dist.custom_pipeline_stage_id = [0]*14 + [1]*10# set auto parallel in LiBaigraph.auto_parallel.enabled = True# enable amp (fp16)train.amp.enabled = True # enable gradient clippingoptim.params.clip_grad_norm = 1.0optim.params.clip_grad_norm_type = 2.0# enable grad accumulation for 8 stepstrain.num_accumulation_steps = 8# enable activation checkpointingtrain.activation_checkpoint.enabled = True# enable zero for leval-2train.zero_optimization.enabled = Truetrain.zero_optimization.stage = 2
from libai.layers.linear import Linearfrom oneflow import nn # write a Simple 2D Parallel MLPclass 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
支持一鍵轉(zhuǎn)換ONNX
本人對一鍵轉(zhuǎn)ONNX的執(zhí)念可謂是相當(dāng)之深了。同樣以MT5為例子，LiBai支持了一鍵轉(zhuǎn)換ONNX的功能，點擊以下鏈接就可以體驗：
https://github.com/Oneflow-Inc/libai/tree/main/libai/onnx_export
更詳細(xì)的說明和教程會在LiBai中持續(xù)發(fā)布。如果這篇文章對你有啟發(fā)，請不要吝惜手中的收藏按鈕，歡迎去GitHub上Star、Fork、Watch三連，持續(xù)跟進(jìn)最新進(jìn)展。?
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
其他人都在看
機(jī)器學(xué)習(xí)編譯器的前世今生
TPU演進(jìn)十年：Google的十大經(jīng)驗教訓(xùn)
更快的YOLOv5問世，附送全面中文解析教程
中文Stable Diffusion開源；PyTorch優(yōu)化技巧
開源吞噬AI界？從Stable Diffusion的爆火說起
OneEmbedding:單卡訓(xùn)練TB級推薦模型不是夢
大模型訓(xùn)練難？效率超群、易用的“李白”模型庫來了
歡迎Star、試用OneFlow最新版本：GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient. - GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.https://github.com/Oneflow-Inc/oneflow
關(guān)鍵詞：
營業(yè)執(zhí)照公示信息