作者｜BBuf、謝子鵬、馮文
2017 年，Google 提出了 Transformer 架構(gòu)，隨后 BERT 、GPT、T5等預(yù)訓(xùn)練模型不斷涌現(xiàn)，并在各項(xiàng)任務(wù)中都不斷刷新 SOTA 紀(jì)錄。去年，清華提出了 GLM 模型（https://github.com/THUDM/GLM），不同于上述預(yù)訓(xùn)練模型架構(gòu)，它采用了一種自回歸的空白填充方法， 在 NLP 領(lǐng)域三種主要的任務(wù)（自然語言理解、無條件生成、有條件生成）上都取得了不錯的結(jié)果。 ?
很快，清華基于 GLM 架構(gòu)又推出了 GLM-130B（https://keg.cs.tsinghua.edu.cn/glm-130b/zh/posts/glm-130b/），這是一個開源開放的雙語（中文和英文）雙向稠密模型，擁有 1300 億參數(shù)，在語言理解、語言建模、翻譯、Zero-Shot 等方面都更加出色。 ?
(資料圖片)
預(yù)訓(xùn)練模型的背后離不開開源深度學(xué)習(xí)框架的助力。在此之前，GLM 的開源代碼主要是由 PyTorch、DeepSpeed 以及 Apex 來實(shí)現(xiàn)，并且基于 DeepSpeed 提供的數(shù)據(jù)并行和模型并行技術(shù)訓(xùn)練了 GLM-Large（335M），GLM-515M（515M），GLM-10B（10B）等大模型，這在一定程度上降低了 GLM 預(yù)訓(xùn)練模型的使用門檻。
即便如此，對更廣大范圍的普通用戶來說，訓(xùn)練 GLM 這樣的模型依然令人頭禿，同時，預(yù)訓(xùn)練模型的性能優(yōu)化還有更大的提升空間。
為此，我們近期將原始的 GLM 項(xiàng)目移植到了使用 OneFlow 后端進(jìn)行訓(xùn)練的 One-GLM 項(xiàng)目。得益于 OneFlow 和 PyTorch 無縫兼容性，我們快速且平滑地移植了 GLM，并成功跑通了預(yù)訓(xùn)練任務(wù)（訓(xùn)練 GLM-large）。
此外，由于 OneFlow 原生支持 DeepSpeed 和 Apex 的很多功能和優(yōu)化技術(shù)，用戶不再需要這些插件就可訓(xùn)練 GLM 等大模型。更重要的是，針對當(dāng)前 OneFlow 移植的 GLM 模型，在簡單調(diào)優(yōu)后就能在性能以及顯存占用上有大幅提升。
具體是怎么做到的？下文將進(jìn)行揭曉。
One-GLM：https://github.com/Oneflow-Inc/one-glm
OneFlow：https://github.com/Oneflow-Inc/oneflow
首先先展示一下分別使用官方的 GLM 倉庫以及 One-GLM 倉庫訓(xùn)練 GLM-large 網(wǎng)絡(luò)的性能和顯存表現(xiàn)（數(shù)據(jù)并行技術(shù)），硬件環(huán)境為 A100 PCIE 40G，BatchSize 設(shè)置為 8。
可以看到，在 GLM-large 的訓(xùn)練任務(wù)中，相比原始的基于 PyTorch、DeepSpeed、Apex 的 GLM 實(shí)現(xiàn)，OneFlow的性能有 120% - 276% 的加速，并且顯存占用降低了10% -30%（測試結(jié)果均可用 oneflow >=0.9.0 復(fù)現(xiàn)）。
由于 OneFlow 無縫兼容了 PyTorch 的生態(tài)，只需改動幾行代碼，就可以讓用戶輕松遷移 GLM 大模型到 One-GLM： ?
將 import torch 替換為 ?import oneflow as torch
將 import torch.xx 替換為 import oneflow.xx
將 from apex.optimizers import FusedAdam as Adam 替換為 from oneflow.optim import Adam
將 from apex.normalization.fused_layer_norm import FusedLayerNorm as LayerNorm 替換為 from oneflow.nn import LayerNorm
注釋掉 torch.distributed.ReduceOp，torch.distributed.new_group,，torch.distributed.TCPStore，torch.distributed.all_reduce 這些API，它們是 PyTorch DDP 所需要的，但 OneFlow 的數(shù)據(jù)并行是由內(nèi)部的 SBP 和 Global Tensor 機(jī)制實(shí)現(xiàn)，并不需要這些 API。
其它許多模型的遷移更簡單，比如在和 torchvision 對標(biāo)的 flowvision 中，許多模型只需通過在 torchvision 模型文件中加入 import oneflow as torch 即可得到，讓用戶幾乎沒有額外成本。 ?
此外，OneFlow 還提供全局 “mock torch” 功能（https://docs.oneflow.org/master/cookies/oneflow_torch.html），在命令行運(yùn)行 eval $(oneflow-mock-torch) 就可以讓接下來運(yùn)行的所有 Python 腳本里的 import torch 都自動指向 oneflow。 ?
在原始的 GLM 實(shí)現(xiàn)中，loss計算部分使用到了 mpu.vocab_parallel_cross_entropy 這個函數(shù) (https://github.com/THUDM/GLM/blob/main/pretrain_glm.py#L263) 。 ?
通過分析這個函數(shù)，發(fā)現(xiàn)它實(shí)現(xiàn)了 sparse_softmax_cross_entropy 的功能，但在實(shí)現(xiàn)過程中，原始的 GLM 倉庫使用了 PyTorch 的 autograd.Function 模塊，并且使用了大量的小算子來拼接出 sparse_softmax_cross_entropy 整體的功能。而在 OneFlow 的算子庫中，已經(jīng)有 sparse_softmax_cross_entropy 這個算子對應(yīng)的 CUDA 實(shí)現(xiàn)了，也就是 flow.sparse_softmax_cross_entropy 這個 API。
所以，我們將 GLM 對 sparse_softmax_cross_entropy 的 naive 實(shí)現(xiàn)替換為 flow.sparse_softmax_cross_entropy 這個 API，并進(jìn)行了 loss 對齊實(shí)驗(yàn)。
結(jié)果如何？下圖展示了基于 OneFlow 的 Graph 模式訓(xùn)練 GLM-large 模型前 1000 輪的 loss 對齊情況，并分別測試了 FP32 和 AMP 模式：
可以看到，將原始 GLM 的 naive sparse_softmax_cross_entropy 實(shí)現(xiàn)替換為 flow.sparse_softmax_cross_entropy 之后 loss 是完全對齊的，可以保證正確性。
相比原始的 GLM 的單卡性能，這個替換使得 One-GLM 的單卡性能有大幅提升，主要原因是 OneFlow 對 sparse_softmax_cross_entropy 算子做了極致的性能優(yōu)化，并且減少了原始 GLM 中大量的碎算子拼湊帶來的訪存開銷。此外，這樣做也降低了 torch.autograd.Function 本身帶來的一些系統(tǒng)開銷。
除上述優(yōu)化外，GLM 模型本質(zhì)上就是一個編解碼的 Transformer 架構(gòu)，所以我們將之前優(yōu)化 GPT、BERT 的一些 Fuse Pattern 也帶給了 One-GLM 模型。具體包含以下兩個 Fuse Pattern :
fused_bias_add_gelu: 將 bias_add 和 gelu 算子融合在一起。
fused_bias_add_dropout：將 bias_add 和 dropout 算子融合在一起。
這兩個 fuse 都可以顯著改善計算的訪存，并減少 Kernel Launch 帶來的開銷，由于 GLM 模型越大則層數(shù)就會越多，那么這種 Fuse Pattern 帶來的的優(yōu)勢也會不斷放大。
最終，在上述兩方面的優(yōu)化作用下，在 A100 PCIE 40G，batch_size = 8 環(huán)境中的訓(xùn)練 GLM-large 的任務(wù)時，單卡 FP32 模式的性能相比原始的 GLM 取得了 280%（FP32 模式） 和 307%（ AMP 模式） 的訓(xùn)練加速。
當(dāng)模型規(guī)模過于龐大，單個 GPU 設(shè)備無法容納大規(guī)模模型參數(shù)時，便捷好用的分布式訓(xùn)練和推理需求就相繼出現(xiàn)，業(yè)內(nèi)也隨之推出相應(yīng)的工具。
基于 OneFlow 構(gòu)建的 LiBai 模型庫讓分布式上手難度降到最低，用戶不需要關(guān)注模型如何分配在不同的顯卡設(shè)備，只需要修改幾個配置數(shù)據(jù)就可以設(shè)置不同的分布式策略。當(dāng)然，加速性能更是出眾。
LiBai ：https://github.com/Oneflow-Inc/libai
LiBai 相關(guān)介紹：大模型訓(xùn)練之難，難于上青天？預(yù)訓(xùn)練易用、效率超群的「李白」模型庫來了！
GLM：https://github.com/Oneflow-Inc/libai/tree/glm_project/projects/GLM
用 LiBai 搭建的 GLM 可以便捷地實(shí)現(xiàn)model parallel + pipeline parallel推理, 很好地解決單卡放不下大規(guī)模模型的問題。
那么，用戶如何利用大規(guī)模模型訓(xùn)練與推理倉庫 LiBai 來構(gòu)建 GLM 的分布式推理部分？下面用一個小例子解釋一下。
要知道，模型的參數(shù)其實(shí)就是許多 tensor，也就是以矩陣的形式出現(xiàn)，大模型的參數(shù)也就是大矩陣，并行策略就是把大矩陣分為多個小矩陣，并分配到不同的顯卡或不同的設(shè)備上，基礎(chǔ)的 LinearLayer 在LiBai中的實(shí)現(xiàn)代碼如下：
class Linear1D(nn.Module): def __init__(self, in_features, out_features, parallel="data", layer_idx=0, ...): super().__init__() if parallel == "col": weight_sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.split(0)]) elif parallel == "row": weight_sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.split(1)]) elif parallel == "data": weight_sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]) else: raise KeyError(f"{parallel} is not supported! Only support ("data", "row" and "col")") self.weight = flow.nn.Parameter( flow.empty( (out_features, in_features), dtype=flow.float32, placement=dist.get_layer_placement(layer_idx), # for pipeline parallelism placement sbp=weight_sbp, ) ) init_method(self.weight) ... def forward(self, x): ...
在這里，用戶可選擇去如何切分 Linear 層的矩陣，如何切分?jǐn)?shù)據(jù)矩陣，而OneFlow 中的 SBP 控制豎著切、橫著切以及其他拆分矩陣的方案（模型并行、數(shù)據(jù)并行），以及通過設(shè)置 Placement 來控制這個 LinearLayer 是放在第幾張顯卡上（流水并行）。
所以，根據(jù) LiBai 中各種 layer 的設(shè)計原理以及基于 OneFlow 中 tensor 自帶的 SBP 和 Placement 屬性的天然優(yōu)勢，使得用戶搭建的模型能夠很簡單地就實(shí)現(xiàn)數(shù)據(jù)并行、模型并行以及流水并行操作。
這里為用戶展示 LiBai 中 GLM 的單卡和便捷的多卡推理 Demo，模型可在 HuggingFace 上獲取：https://huggingface.co/models?filter=glm
單卡 generate 任務(wù)，我們選擇 glm-10b 模型：
python demo.py
# demo.pyimport oneflow as flowfrom projects.GLM.tokenizer.glm_tokenizer import GLMGPT2Tokenizerfrom libai.utils import distributed as distfrom projects.GLM.configs.glm_inference import cfgfrom projects.GLM.modeling_glm import GLMForConditionalGenerationfrom projects.GLM.utils.glm_loader import GLMLoaderHuggerFacefrom omegaconf import DictConfigtokenizer = GLMGPT2Tokenizer.from_pretrained("/data/home/glm-10b")input_ids = tokenizer.encode( [ "Ng is an adjunct professor at [MASK] (formerly associate professor and Director of its Stanford AI Lab or SAIL ). Also a pioneer in online education, Ng co-founded Coursera and deeplearning.ai." ], return_tensors="of",)inputs = {"input_ids": input_ids, "attention_mask": flow.ones(input_ids.size())}inputs = tokenizer.build_inputs_for_generation(inputs, max_gen_length=512)sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast])placement = dist.get_layer_placement(0)dist.set_device_type("cpu")loader = GLMLoaderHuggerFace(GLMForConditionalGeneration, cfg, "/path/to/glm-10b")model = loader.load()model = model.half().cuda()dist.set_device_type("cuda")outputs = model.generate( inputs=inputs["input_ids"].to_global(sbp=sbp, placement=placement), position_ids=inputs["position_ids"].to_global(sbp=sbp, placement=placement), generation_attention_mask=inputs["generation_attention_mask"].to_global(sbp=sbp, placement=placement).half(), max_length=512)res = tokenizer.decode(outputs[0])print(res)>>> [CLS] Ng is an adjunct professor at [MASK] (formerly associate professor and Director of its Stanford AI Lab or SAIL ). Also a pioneer in online education, Ng co-founded Coursera and deeplearning.ai.<|endoftext|> <|startofpiece|> Stanford University and a co-founder of <|endofpiece|>
4卡 model parallel+pipeline parallel generate 任務(wù)，選擇 glm-10b 模型：
python3 -m oneflow.distributed.launch --nproc_per_node 4 demo.py
# demo.pyimport oneflow as flowfrom projects.GLM.tokenizer.glm_tokenizer import GLMGPT2Tokenizerfrom libai.utils import distributed as distfrom projects.GLM.configs.glm_inference import cfgfrom projects.GLM.modeling_glm import GLMForConditionalGenerationfrom projects.GLM.utils.glm_loader import GLMLoaderHuggerFacefrom omegaconf import DictConfig# 只需簡單配置并行方案parallel_config = DictConfig( dict( data_parallel_size=1, tensor_parallel_size=2, pipeline_parallel_size=2, pipeline_num_layers=2 * 24 ))dist.setup_dist_util(parallel_config)tokenizer = GLMGPT2Tokenizer.from_pretrained("/data/home/glm-10b")input_ids = tokenizer.encode( [ "Ng is an adjunct professor at [MASK] (formerly associate professor and Director of its Stanford AI Lab or SAIL ). Also a pioneer in online education, Ng co-founded Coursera and deeplearning.ai." ], return_tensors="of",)inputs = {"input_ids": input_ids, "attention_mask": flow.ones(input_ids.size())}inputs = tokenizer.build_inputs_for_generation(inputs, max_gen_length=512)sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast])placement = dist.get_layer_placement(0)loader = GLMLoaderHuggerFace(GLMForConditionalGeneration, cfg, "/path/to/glm-10b")model = loader.load()outputs = model.generate( inputs=inputs["input_ids"].to_global(sbp=sbp, placement=placement), position_ids=inputs["position_ids"].to_global(sbp=sbp, placement=placement), generation_attention_mask=inputs["generation_attention_mask"].to_global(sbp=sbp, placement=placement), max_length=512)res = tokenizer.decode(outputs[0])if dist.is_main_process(): print(res)>>> [CLS] Ng is an adjunct professor at [MASK] (formerly associate professor and Director of its Stanford AI Lab or SAIL ). Also a pioneer in online education, Ng co-founded Coursera and deeplearning.ai.<|endoftext|> <|startofpiece|> Stanford University and a co-founder of <|endofpiece|>
使用 One- GLM 訓(xùn)練的模型進(jìn)行推理
LiBai對于OneFlow的模型加載同樣方便，如果你希望使用one-glm訓(xùn)練后的模型進(jìn)行推理，只需簡單的將上述demo中的 GLMLoaderHuggerFace 替換為 GLMLoaderLiBai。 ?
基于 OneFlow 來移植 GLM 大模型非常簡單，相比于原始版本 PyTorch GLM ?訓(xùn)練 GLM-large 模型，OneFlow 能大幅提升性能和節(jié)省顯存。
此外，通過使用 GLM-10B 這個百億級大模型做推理，表明基于 OneFlow 的 LiBai 來做大模型推理可以開箱即用，并實(shí)現(xiàn)更高的推理速度，如果你想配置不同的并行方式來推理大模型，只需要簡單配置文件的幾個參數(shù)即可。
未來，OneFlow團(tuán)隊將探索使用 OneFlow 訓(xùn)練更大的 GLM-130B 千億模型的可行性，相信基于 OneFlow 可以更快地訓(xùn)練 GLM-130B 千億級別模型，加速國產(chǎn)大模型訓(xùn)練和推理任務(wù)。歡迎Star、試用One-GLM：
One-GLM：https://github.com/Oneflow-Inc/one-glm
OneFlow：https://github.com/Oneflow-Inc/oneflow
其他人都在看
35張圖，直觀理解Stable Diffusion
2023年AI十大展望：GPT-4領(lǐng)銜大模型變革
李白：你的模型權(quán)重很不錯，可惜被我沒收了
OpenAI掌門Sam Altman：AI下一個發(fā)展階段
比快更快，開源Stable Diffusion刷新作圖速度
OneEmbedding:單卡訓(xùn)練TB級推薦模型不是夢
“零”代碼改動，靜態(tài)編譯讓太乙Stable Diffusion推理速度翻倍
關(guān)鍵詞：
營業(yè)執(zhí)照公示信息