撰文｜CPFLAME
大噶好，年更樓主今天想推的是，主打分布式訓練的模型庫_李白（LiBai）。?
https://github.com/Oneflow-Inc/libaihttps://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和高效性之間反復橫跳。
2. 對于分布式并行的支持比較全面，大家可以在里面盡情地組合各種分布式并行的組件。
3. LiBai下面有內置的layers直接使用，避免重復造輪子，比如用LiBai下面的Linear層就可以快速地構建一個2D并行（數據并行+模型并行）的MLP。?
4. 采用LazyCall（借鑒自detectron2）的配置系統，基于Python語法構建相比于 argparse 和 yacs-based 方式更靈活，而且每次訓練都會序列化yaml文件，用戶可以一鍵讀取yaml文件來復現“上古時期”的實驗結果。
5. 具有豐富的Projects實現。由于LiBai的分布式并行設計與算法邏輯進行了解耦，使得在Projects下面的算法都可以享受到LiBai下面的分布式并行技術，而且隨著分布式并行技術的更新，Projects下面的算法代碼不需要任何更新就可以享受到更新后的成果。
6. 和業界翹楚Megatron比起來，具有不弱于它的吞吐，甚至稍占優勢，完整的對比實驗在LiBai tutorial（https://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.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)
那么多機多卡的分布式推理腳本
在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.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
支持一鍵轉換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
其他人都在看
機器學習編譯器的前世今生
TPU演進十年：Google的十大經驗教訓
更快的YOLOv5問世，附送全面中文解析教程
中文Stable Diffusion開源；PyTorch優化技巧
開源吞噬AI界？從Stable Diffusion的爆火說起
OneEmbedding:單卡訓練TB級推薦模型不是夢
大模型訓練難？效率超群、易用的“李白”模型庫來了
歡迎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
關鍵詞：
營業執照公示信息