深度学习框架的流行（如PyTorch，Tensorflow等）极大方便了我们研发设计各种各样的AI模型，而在实际落地的环节中，孵化于实验室里的模型代码往往在生产环境上面临着性能、准确度、资源等各种各样的问题。随着AI与现实世界的业务结合愈加密切，高性能且易使用的深度学习算子愈发受到AI研究者和开发者们的青睐。
幻方AI依托萤火二号集群，自研了一系列AI基础设施：3FS，hfreduce，ffrecord等。如之前系列文章中介绍的，他们极大地提升了模型训练中数据读取的性能，充分利用GPU的算力。而对于计算环节，幻方AI同样对Pytorch框架进行了深度优化，结合萤火二号的集群特点，对一些常用的AI算子重新研发，提升效率，进一步提升了模型整体的训练效率。
本期文章将为大家介绍，幻方AI的深度学习算子，hfai.nn，供大家使用。
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 dist
ip = os.environ["MASTER_ADDR"]
port = os.environ["MASTER_PORT"]
hosts = int(os.environ["WORLD_SIZE"]) # number of nodes
rank = int(os.environ["RANK"]) # node id
gpus = 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_hfai
to_hfai是方便一键使用hfai算子的工具，会对你的model进行一次深度遍历，将其中能做转化Module替换成hfai的。转化后的算子在training以及inference的时候都会有可观的性能提升。
to_torch
to_torch是将使用hfai算子的model转化回使用torch算子的工具，以便在非 hfai 下使用训练出来的模型。设计思路同to_hfai。
使用案例：
from hfai.nn import to_hfai,to_torch
# get model from hfai && to_hfai
model = hfai.models.bert_large_MLM(vocab_size=vocab_size, max_length=max_length).cuda()
model = to_hfai(model)性能测试
那么，幻方AI优化的算子具体能达到多少的加速效果呢？这里，我们把性能测试数据给出：
LSTM
参数配置：
input_size: 985
hidden_size: 985
bias: True
num_layers: 1
batch_first: True
dropout: 0
bidirectional: False
batch_size: 4
seq_len: 1000连续测试10轮，统计整体每轮的耗时，测试结果如下：
-
torch.nn.LSTM
- forward: 31.78 ms/epoch
- backard: 93.60 ms/epoch
- all: 125.38 ms/epoch
-
hfai.nn.LSTM
- forward: 4.82 ms/epoch
- backard: 9.25 ms/epoch
- all: 14.07 ms/epoch
forward 提升6倍，backward 提升了10倍，整体训练提升了8倍！
Attention
参数配置：
batch_size: 200
seq_len: 256
embed_dim: 1024
attn_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: 5000
seq_len: 64
embed_dim: 512
attn_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 = 200
seq_len = 128
hidden_size = 1024连续测试10轮，统计forward，backward，还有inference的耗时，测试结果如下：
- 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基础设施结合，让模型训练高效迅速，计算资源充分利用，加速更多场景的探索与研究。算子优化一直持续中，非常欢迎更多优秀同行加入我们一同共建。