深度学习框架的流行（如 PyTorch，Tensorflow 等）极大方便了我们研发设计各种各样的 AI 模型，而在实际落地的环节中，孵化于实验室里的模型代码往往在生产环境上面临着性能、准确度、资源等各种各样的问题。随着 AI 与现实世界的业务结合愈加密切，高性能且易使用的深度学习算子愈发受到 AI 研究者和开发者们的青睐。
幻方 AI 依托「萤火二号」超算集群，自研了一系列 AI 基础设施：3FS，hfreduce，ffrecord等。如之前系列文章中介绍的，他们极大地提升了模型训练中数据读取的性能，充分利用 GPU 算力。而对于计算环节，幻方 AI 同样对 Pytorch 框架进行了深度优化，结合「萤火二号」的集群特点，对一些常用的 AI 算子重新研发，提升效率，进一步提升了模型整体的训练效率。
本期文章将为大家介绍，幻方 AI 的深度学习算子 hfai.nn。点击文末链接，即可申请免费使用幻方的「萤火」AI 训练平台。
hfai.nn.parallel在之前的文章《幻方萤火 | 模型并行训练工具 hfreduce》中介绍了幻方 AI 自研的并行训练工具，设计的全新的通讯模式替代 NCCL，在「萤火」超算集群上取得非常好的加速效果。这里，我们对接口按照 PyTorch 框架风格进一步封装，降低大家的使用门槛。
hfai DDP我们提供了 hfai.nn.parallel.DistributedDataParallel 接口，其在 hfreduce 的基础上进行封装，使用方法和 PyTorch 的 DDP 基本相同。在「萤火二号」集群的环境下，hfai DDP 多机并行的性能比 PyTorch DDP 更好。
使用案例：
from hfai.nn.parallel import DistributedDataParallel
# init distip = os.environ["MASTER_ADDR"]port = os.environ["MASTER_PORT"]hosts = int(os.environ["WORLD_SIZE"]) # number of nodesrank = int(os.environ["RANK"]) # node idgpus = torch.cuda.device_count() # gpus per node
dist.init_process_group( backend="nccl", init_method=f"tcp://{ip}:{port}", world_size=hosts * gpus, rank=rank * gpus + local_rank)torch.cuda.set_device(local_rank)
model = MyModel().cuda()model = DistributedDataParallel(model, device_ids=[local_rank])
hfai.nn.to_hfai / to_torch
有了上述分布式训练管理工具，接下来就是具体的优化算子了。为了让模型计算的更快，我们优化 CUDA kernel，融合多运算，优化模型结构，使得深度学习算子计算更加的高效。
具体优化的算子包括：
-
torch.nn.lstm - hfai.nn.lstm
-
torch.nn.MultiheadAttention- hfai.nn.MultiheadAttention
-
torch.nn.LayerNorm- hfai.nn.LayerNorm
使用方法非常简单，幻方 AI 提供了工具可以一行代码进行算子的切换，即 to_hfai 和 to_torch。
to_hfaito_hfai 是方便一键使用 hfai 算子的工具，会对你的 model 进行一次深度遍历，将其中能做转化 Module 替换成 hfai 的。转化后的算子在 training 以及 inference 的时候都会有可观的性能提升。
to_torchto_torch 是将使用 hfai 算子的 model 转化回使用 torch 算子的工具，以便在非 hf_env 下使用训练出来的模型。设计思路同 to_hfai。
使用案例：
from hfai.nn import to_hfai,to_torch # get model from hfai && to_hfaimodel = hfai.models.bert_large_MLM(vocab_size=vocab_size, max_length=max_length).cuda()model = to_hfai(model)性能测试
那么，幻方 AI 优化的算子具体能达到多少的加速效果呢？这里，我们把性能测试数据给出：
LSTM参数配置：
input_size: 985hidden_size: 985bias: Truenum_layers: 1batch_first: Truedropout: 0bidirectional: Falsebatch_size: 4seq_len: 1000
连续测试10轮，统计forward，backward，还有整体的耗时，测试结果如下：
-
torch.nn.lstm:
-
forward: 31.78 ms/epoch
-
backward: 93.60 ms/epoch
-
all: 125.38 ms/epoch
-
-
hfai.nn.lstm:
-
forward: 4.82 ms/epoch
-
backward: 9.25 ms/epoch
-
all: 14.07 ms/epoch
-
forward提升了6倍，backward提升了10倍，整体训练提升了8倍！
Attention参数配置：
batch_size: 200seq_len: 256embed_dim: 1024attn_heads: 16
连续测试10轮，统计forward，backward，还有 inference 的耗时，测试结果如下：
-
torch.nn.MultiheadAttention
-
forward: 17.581 ms/epoch
-
backard: 21.986 ms/epoch
-
all: 39.567 ms/epoch
-
inference: 17.372 ms
-
-
hfai.nn.MultiheadAttention
-
forward: 12.165 ms/epoch
-
backard: 16.515 ms/epoch
-
all: 28.680 ms/epoch
-
inference: 12.112 ms
-
forward 提升 44%，backward 提升了 33%，整体训练提升了 37%，inference 也提升 43%！
同时，当 batch_size 比较大，seq_len 比较小，以及输入 qkv 一样的情况下，hfai.nn.MultiheadAttention 同样具备优异性能。
参数配置：
batch_size: 5000seq_len: 64embed_dim: 512attn_heads: 8
连续测试10轮，统计forward，backward，还有 inference 的耗时，测试结果如下：
-
torch.nn.MultiheadAttention
-
forward: 22.713 ms/epoch
-
backard: 41.831 ms/epoch
-
all: 64.543 ms/epoch
-
inference: 22.925 ms
-
-
hfai.nn.MultiheadAttention
-
forward: 19.599 ms/epoch
-
backard: 24.379 ms/epoch
-
all: 43.978 ms/epoch
-
inference: 19.037 ms
-
forward 提升 15%，backward 提升了 71%，整体训练提升了 46%，inference 也提升 20%！
LayerNorm参数配置：
batch_size = 200seq_len = 128hidden_size = 1024
连续测试10轮，统计forward，backward，及总体耗时，测试结果如下：
-
torch.nn.LayerNorm
-
forward: 0.514 ms/epoch
-
backard: 0.217 ms/epoch
-
all: 2.467 ms/epoch
-
-
hfai.nn.LayerNorm
-
forward: 0.217 ms/epoch
-
backard: 0.121 ms/epoch
-
all: 1.306 ms/epoch
-
forward 提升 137%，backward 提升了 79%，整体训练提升了 88%！
总结通过优化 CUDA kernel 和模型结构，融合多运算操作，幻方 AI 优化的深度学习算子在实际测试中表现出惊人的加速效果，优化的算子与幻方 AI 基础设施结合，让模型训练高效迅速，计算资源充分利用，加速更多场景的探索与研究。算子优化一直持续中，非常欢迎更多优秀同行加入我们一同共建。
幻方 AI 团队， 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力 AI 实验落地与价值创造，降低超大算力需求的使用门槛，欢迎各方数据科学家与开发者们一同共建。
戳下方链接，申请使用萤火超算
High-flyer AIHPC 萤火二号使用申请
High-Flyer AI
万张A100随时待命
点击下方链接
幻方AI BLOG 更 多 干 货 奉 上
幻方 | 技术博客
