Lightop

简介

Lightop是一个自研的融合算子库，提供了大量的自研算子，用来加速模型的推理和训练。

这些算子通常从torch里面获取，然后进行性能优化，常见的norm类算子比torch native有2-5倍的性能提升。

许可证

Lightop 使用 Apache License 2.0。完整许可证文本见根目录的 LICENSE。

本仓库还包含从第三方项目复制或适配的源码。第三方组件仍受其各自的版权和许可证约束，来源、版本证据、许可证和本仓库中的路径见 THIRD_PARTY_NOTICES.md。

安装

环境依赖： pytorch1.10 \ pytorch1.13 \ pytorch2.0 \ pytorch2.1 \ pytoch2.3

dtk22.04 \ dtk22.10 \ dtk23.04 \ dtk23.10 \ dtk24.04

安装脚本：

python3 setup.py install

对于23.10之前的dtk由于不支持928，因此需要去掉928代码，使用下面的编译命令：

PYTORCH_ROCM_ARCH='gfx906;gfx926' python3 setup.py install

编译安装包：

python3 setup.py install bdist_wheel

编译好之后，安装包在dist目录

算子介绍

核心算子

核心算子主要包括quant_gemm、moe_gemm，具体算子类型和开发负责人如下：

算子	开发负责人	性能表
gemm_w8a8	郑光飞 谢松林 马翔
moe_gemm_w8a8	杨亿男 罗尧	https://r0ddbu55vzx.feishu.cn/sheets/Fr3wsNhrmh777xtl7v4cOodjnH0?sheet=ClHIQQ
moe_gemm_w4a8	王安全 刘振成 桂中化
moe_gemmm_w4a16	刘振成

测试方法：

cd ./test

benchmark_test_kernel.sh

生成对应的算子性能和精度表格

torch算子

主要对应的torch的原生算子，以下时间均为单次前向传播和反向传播的总时间，包括了CPU时间和GPU时间。测试环境：pytorch1.13, dtk23.04, 10.6.10.67服务器z100计算卡。

算子	测试size	优化前	优化后	性能提升	使用说明
droppath	128,3136,96	0.84ms	0.74ms	13.5%	Swin-Transformer
ln_droppath	128, 3136, 96	15.7ms	4.06ms	287%	Swin-Transformer. droppath(ln(input))
layernorm	1, 2048, 3072	0.58ms	0.32ms	81.3%	NLP类
ln_dropout	1, 2048, 3072	0.78ms	0.42ms	85.7%	NLP类。例如transformer,bert。 ln(dropout(input))
ln_add_dropout	1, 2048, 3072	0.97ms	0.47ms	106%	NLP类。例如transformer,bert。 ln(dropout(input)+residual)
ln_add_dropout_bias	1, 2048, 3072	1.16ms	0.53ms	119%	NLP类。例如transformer,bert。 ln(dropout(input+bias)+residual)
rmsnorm	1, 2048, 3072	1.28ms	0.29ms	341%	NLP类。例如llama
rn_dropout	1, 2048, 3072	1.51ms	0.38ms	297%	NLP类。例如llama。 rn(dropout(input))
rn_add_dropout	1, 2048, 3072	1.68ms	0.45ms	273%	NLP类。例如llama。 rn(dropout(input)+residual)
rn_add_dropout_bias	1, 2048, 3072	1.88ms	0.51ms	268%	NLP类。例如llama。rn(dropout(input+bias)+residual)
groupnorm	2,640,8,8	0.037ms	0.017ms	116%	NLP类cv类
celu_gn	2,640,8,8	0.048ms	0.015ms	219%	NLP类cv类gn。例如x-linear。 gn(celu(input))
gn_silu	2,640,8,8	0.080ms	0.015ms	420%	NLP类cv类。例如：SD。 silu(gn(input))
add_dropout	1, 2048, 3072	0.49ms	0.38ms	28.9%	NLP类.dropout(input)+residual
bias_dropout_add	1, 2048, 3072	0.69ms	0.43ms	60.5%	NLP类.dropout(input+bias)+residual
bias_gelu	1, 2048, 3072	0.38ms	0.27ms	40.7%	NLP类,gelu(input+bias)
bias_swiglu	2, 2048, 512	0.50ms	0.28ms	78.5%	NLP类SwiGLU(input+bias)
FusedScaleMask SoftmaxDropout	2,16,128,128	0.26ms	0.34ms	30.7%	NLP类.
bn_act/bn_relu	4,10,128,128	0.56ms	0.33ms	69.7%	CNN，如resnet，都会使用BN。act支持relu,softrelu,abs,power,clipped_relu,leaky_relu,silu_sigmoid,tanh等激活函数。bn_relu是bn_act的特例，减少一个参数，调用更方便。对于bn 2d，推荐NCHW维度中N*H*W>20000，此算子对于小size性能较差。
softmax	512,4096	0.23ms	0.16ms	43.8%	一般分类模型的结尾会有softmax层。推荐size:第一维度大于300，第二维度大于等于1024小于等于4096。
loss_ctc	120,60,240	4.51ms	1.14ms	396%	语音文本等RNN网络的输出层
loss_rnnt	10,30,21,4000	16.52ms	2.93ms	429%	语音文本等RNN网络的输出层

测试方式：

• 进入test文件夹

• 运行任意一个测试脚本，例如：python3 test_ln_add_dropout.py

• 运行python3 all_test.py可以测试全部脚本

• 测试模型时，参考测试脚本里调用算子的方法。

• 若无test文件夹，则运行以下指令验证是否安装成功：

python -c "import lightop; print(lightop.__version__)",若成功打印版本信息则安装成功。