作者｜BBuf、谢子鹏、冯文

2017 年，Google 提出了 Transformer 架构，随后 BERT 、GPT、T5等预练习模型不断涌现，并在各项使命中都不断刷新 SOTA 纪录。上一年，清华提出了 GLM 模型（github.com/THUDM/GLM ），不同于上述预练习模型架构，它采用了一种自回归的空白填充办法，在 NLP 领域三种首要的使命（自然语言了解、无条件生成、有条件生成）上都取得了不错的成果。

很快，清华根据 GLM 架构又推出了 GLM-130B（keg.cs.tsinghua.edu.cn/glm-130b/zh… ），这是一个开源开放的双语（中文和英文）双向稠密模型，具有 1300 亿参数，在语言了解、语言建模、翻译、Zero-Shot 等方面都更加超卓。

预练习模型的背后离不开开源深度学习结构的助力。在此之前，GLM 的开源代码首要是由 PyTorch、DeepSpeed 以及 Apex 来完成，而且根据 DeepSpeed 供给的数据并行和模型并行技能练习了 GLM-Large（335M），GLM-515M（515M），GLM-10B（10B）等大模型，这在一定程度上降低了 GLM 预练习模型的运用门槛。

即便如此，对更广大范围的普通用户来说，练习 GLM 这样的模型仍然令人头秃，一起，预练习模型的功能优化还有更大的提高空间。

为此，咱们近期将原始的 GLM 项目移植到了运用 OneFlow 后端进行练习的 One-GLM 项目。得益于 OneFlow 和 PyTorch 无缝兼容性，咱们快速且平滑地移植了 GLM，并成功跑通了预练习使命（练习 GLM-large）。

此外，因为 OneFlow 原生支持 DeepSpeed 和 Apex 的许多功能和优化技能，用户不再需求这些插件就可练习 GLM 等大模型。更重要的是，针对当前 OneFlow 移植的 GLM 模型，在简略调优后就能在功能以及显存占用上有大幅提高。

具体是怎么做到的？下文将进行揭晓。

One-GLM：github.com/Oneflow-Inc…
OneFlow：github.com/Oneflow-Inc…

1、GLM-large 练习功能和显存的体现

首要先展现一下别离运用官方的 GLM 库房以及 One-GLM 库房练习 GLM-large 网络的功能和显存体现（数据并行技能），硬件环境为 A100 PCIE 40G，BatchSize 设置为 8。

能够看到，在 GLM-large 的练习使命中，比较原始的根据 PyTorch、DeepSpeed、Apex 的 GLM 完成，OneFlow的功能有 120% – 276% 的加快，而且显存占用降低了10% -30% （测验成果均可用 oneflow >=0.9.0 复现）。

2、GLM 搬迁，只需修改几行代码

因为 OneFlow 无缝兼容了 PyTorch 的生态，只需改动几行代码，就能够让用户轻松搬迁 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 的数据并行是由内部的 SBP 和 Global Tensor 机制完成，并不需求这些 API。

其它许多模型的搬迁更简略，比如在和 torchvision 对标的 flowvision 中，许多模型只需经过在 torchvision 模型文件中参加 import oneflow as torch 即可得到，让用户几乎没有额外成本。

此外，OneFlow 还供给全局 “mock torch” 功能（docs.oneflow.org/master/cook… ），在命令行运转 eval $(oneflow-mock-torch) 就能够让接下来运转的一切 Python 脚本里的 import torch 都自动指向 oneflow。

3、两大调优手段

loss 核算部分的优化

在原始的 GLM 完成中，loss核算部分运用到了 mpu.vocab_parallel_cross_entropy 这个函数 (github.com/THUDM/GLM/b…) 。

经过分析这个函数，发现它完成了
sparse_softmax_cross_entropy 的功能，但在完成过程中，原始的 GLM 库房运用了 PyTorch 的 autograd.Function 模块，而且运用了大量的小算子来拼接出sparse_softmax_cross_entropy 整体的功能。而在 OneFlow 的算子库中，已经有sparse_softmax_cross_entropy 这个算子对应的 CUDA 完成了，也便是flow.sparse_softmax_cross_entropy 这个 API。

所以，咱们将 GLM 对
sparse_softmax_cross_entropy 的 naive 完成替换为flow.sparse_softmax_cross_entropy 这个 API，并进行了 loss 对齐试验。

成果怎么？下图展现了根据 OneFlow 的 Graph 方式练习 GLM-large 模型前 1000 轮的 loss 对齐状况，并别离测验了 FP32 和 AMP 方式：

能够看到，将原始 GLM 的 naive
sparse_softmax_cross_entropy 完成替换为flow.sparse_softmax_cross_entropy 之后 loss 是完全对齐的，能够确保正确性。

比较原始的 GLM 的单卡功能，这个替换使得 One-GLM 的单卡功能有大幅提高，首要原因是 OneFlow 对 sparse_softmax_cross_entropy 算子做了极致的功能优化，而且削减了原始 GLM 中大量的碎算子凑集带来的访存开销。此外，这样做也降低了 torch.autograd.Function 自身带来的一些系统开销。

CUDA Kernel Fuse

除上述优化外，GLM 模型本质上便是一个编解码的 Transformer 架构，所以咱们将之前优化 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 模型越大则层数就会越多，那么这种 Fuse Pattern 带来的的优势也会不断放大。

最终，在上述两方面的优化作用下，在 A100 PCIE 40G，batch_size = 8 环境中的练习 GLM-large 的使命时，单卡 FP32 方式的功能比较原始的 GLM 取得了 280%（FP32 方式）和 307%（ AMP 方式）的练习加快。

4、LiBai 也能轻松搞定 GLM 推理

当模型规划过于巨大，单个 GPU 设备无法包容大规划模型参数时，快捷好用的分布式练习和推理需求就相继呈现，业界也随之推出相应的东西。

根据 OneFlow 构建的 LiBai 模型库让分布式上手难度降到最低，用户不需求重视模型怎么分配在不同的显卡设备，只需求修改几个配置数据就能够设置不同的分布式战略。当然，加快功能更是拔尖。

LiBai ：github.com/Oneflow-Inc…
LiBai 相关介绍：大模型练习之难，难于上青天？预练习易用、功率出众的「李白」模型库来了！
GLM：github.com/Oneflow-Inc…

用 LiBai 建立的 GLM 能够快捷地完成model parallel + pipeline parallel推理, 很好地解决单卡放不下大规划模型的问题。

那么，用户怎么利用大规划模型练习与推理库房 LiBai 来构建 GLM 的分布式推理部分？下面用一个小例子解释一下。

分布式推理具有天然优势

要知道，模型的参数其实便是许多 tensor，也便是以矩阵的方式呈现，大模型的参数也便是大矩阵，并行战略便是把大矩阵分为多个小矩阵，并分配到不同的显卡或不同的设备上，基础的 LinearLayer 在LiBai中的完成代码如下：

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 层的矩阵，怎么切分数据矩阵，而OneFlow 中的 SBP 控制竖着切、横着切以及其他拆分矩阵的计划（模型并行、数据并行），以及经过设置 Placement 来控制这个 LinearLayer 是放在第几张显卡上（流水并行）。

所以，根据 LiBai 中各种 layer 的规划原理以及根据 OneFlow 中 tensor 自带的 SBP 和 Placement 属性的天然优势，使得用户建立的模型能够很简略地就完成数据并行、模型并行以及流水并行操作。

GLM 推理的 Demo 演示

这里为用户展现 LiBai 中 GLM 的单卡和快捷的多卡推理 Demo，模型可在 HuggingFace 上获取：huggingface.co/models?filt…

单卡 generate 使命，咱们挑选 glm-10b 模型：

python demo.py

# demo.py
import oneflow as flow
from projects.GLM.tokenizer.glm_tokenizer import GLMGPT2Tokenizer
from libai.utils import distributed as dist
from projects.GLM.configs.glm_inference import cfg
from projects.GLM.modeling_glm import GLMForConditionalGeneration
from projects.GLM.utils.glm_loader import GLMLoaderHuggerFace
from omegaconf import DictConfig
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)
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 使命，挑选 glm-10b 模型：

python3 -m oneflow.distributed.launch --nproc_per_node 4 demo.py

# demo.py
import oneflow as flow
from projects.GLM.tokenizer.glm_tokenizer import GLMGPT2Tokenizer
from libai.utils import distributed as dist
from projects.GLM.configs.glm_inference import cfg
from projects.GLM.modeling_glm import GLMForConditionalGeneration
from projects.GLM.utils.glm_loader import GLMLoaderHuggerFace
from 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 练习的模型进行推理

LiBai关于OneFlow的模型加载相同便利，如果你希望运用one-glm练习后的模型进行推理，只需简略的将上述demo中的 GLMLoaderHuggerFace 替换为 GLMLoaderLiBai。

5、结语

根据 OneFlow 来移植 GLM 大模型非常简略，比较于原始版别 PyTorch GLM 练习 GLM-large 模型，OneFlow 能大幅提高功能和节约显存。

此外，经过运用 GLM-10B 这个百亿级大模型做推理，表明根据 OneFlow 的 LiBai 来做大模型推理能够开箱即用，并完成更高的推理速度，如果你想配置不同的并行方式来推理大模型，只需求简略配置文件的几个参数即可。

未来，OneFlow团队将探究运用 OneFlow 练习更大的 GLM-130B 千亿模型的可行性，相信根据 OneFlow 能够更快地练习 GLM-130B 千亿等级模型，加快国产大模型练习和推理使命。
欢迎Star、试用One-GLM：

One-GLM：github.com/Oneflow-Inc…
OneFlow：github.com/Oneflow-Inc…

欢迎 Star、试用 OneFlow 最新版别：
github.com/Oneflow-Inc…

GLM国产大模型训练加速：性能最高提升3倍，显存节省1/3，低成本上手

1、GLM-large 练习功能和显存的体现

2、GLM 搬迁，只需修改几行代码

3、两大调优手段

4、LiBai 也能轻松搞定 GLM 推理

5、结语

作者信息

GLM国产大模型训练加速：性能最高提升3倍，显存节省1/3，低成本上手

1、GLM-large 练习功能和显存的体现

2、GLM 搬迁，只需修改几行代码

3、两大调优手段

4、LiBai 也能轻松搞定 GLM 推理

5、结语

相关文章

手把手教你做一个天猫精灵（一）

Presto自定义函数@SqlNullable血泪史

梳理拯救烂怂代码？我是这么做的

与Go的初次见面 ｜ 青训营笔记

作者信息

与Go的初次见面｜青训营笔记