kernl是一个针对Pytorch transformer模型进行GPU加速的一个框架，利用OpenAI的Triton接口，利用Torch2.0的新特性，而不再使用CUDA来实现，加速方法也是非常的简单，只需简单几行即可完成，故而在本地进行了一些尝试，测试结果显示，其对cuda版本以及显卡型号有较高的要求，同时python也需要较高版本的适配，对于transformer加速效果，当len比较短的时候，会取得比较好的效果，当len比较长的时候，不一定比tensorrt会更好，期待起进一步发展吧，持续跟进。
Github：https://github.com/ELS-RD/kernl
当前结论
硬件要求：
- 这个库对GPU的算力要求比较高，代码中强制要求要8.0以上，2080只有7.5不行，所以下面结论都是在3080测试得到的。
速度提升：
- kernl对transformer类模型加速效果还是比较好，比如bert，albert，swin等，部分速度是超过tensorrt的，对resnet也有一定的加速效果，但是相对tensorrt的速度会没有优势，详细可以参考下面的速度对比。
- 在NLP语言模型上，以albert为例子，len越短，kernl的速度相比tensorrt会快一些，当len<128时候kernl的速度是超过tensorrt的，当len>128，tensorrt会更快，详细见下面的速度表格。
- 在CV模型上，kernl没有优势，以resnet50为例子，tensorrt还是比kernl快挺多的。
- 速度的优化主要依赖于实现的算子，目前优化的算子主要是transformer的atten啥的，如果后期提供的优化算子越来越多，效果应该会越来越好，期待未来的发展。
使用场景：
- 使用起来真的挺方便的，一行代码即可，python无脑使用。
- 模型需要在线生成，无法保存到本地，不过在线生成的速度也挺快的。
- 项目还有些许bug：比如batch只能为1，Git提了issue，作者回复确实是个bug。
一、项目地址：
https://github.com/ELS-RD/kernl
二、docker运行：
# build DOCKER_BUILDKIT=1 docker build -t kernl . # run docker run –rm -it –gpus all -v $(pwd):/kernl kernl # 安装torchvision以及torch2.0 pip3 install numpy –pre torch torchvision torchaudio –force-reinstall –index-url https://download.pytorch.org/whl/nightly/cu117
三、环境需求：
- python3.9 （这个是因为openai的triton需要的）
- cuda11.7以上。
- pytorch 2.0以上，需要支持torch.dynamo
- nvidia的算力要求：8.0及以上
四、使用方法：
很简单，参考如下代码即可。
from kernl.model_optimization import optimize_model optimize_model(model_opt) with torch.inference_mode(), torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True): inputs = { “input_ids”: torch.ones((batch,seq_len), device=”cuda”, dtype=torch.long), “attention_mask”: torch.ones((batch, seq_len), device=”cuda”, dtype=torch.long), } outputs = model_opt(**inputs)
五、速度对比:
备注：单位ms， 3080ti机器(10.219.20.239)
| kernl与tensorrt的耗时对比 | ||||||
|---|---|---|---|---|---|---|
| batch*len | 1x256 | 1x128 | 1x64 | 1x32 | swin | resnet50 |
| pytorch（half） | 15.7 | 16 | 15.5 | 16 | 22 | 5.7 |
| kernl（fp16） | 1.95 | 1.6 | 1.02 | 0.84 | 2.95 | 0.96 |
| tensorrt（fp16） | 1.74 | 1.36 | 1.19 | 1.05 | 2.2 | 0.55 |
| kernl提升倍数 | 8.05 | 10.00 | 15.20 | 19.05 | 7.46 | 5.94 |
| tensorrt提升倍数 | 9.02 | 11.76 | 13.03 | 15.24 | 10.00 | 10.36 |
六、输出结果校验：结果基本一致
6.1 kernl：
第一行是pytorch，第二行是kernl，albert模型，fp16加速。
七、依然没解决的问题
- 目前经过kernl优化的模型，batch维度只能为1，不知道为何。。。。尝试了半天没有修复，算了吧，提个issue，看看有没有人帮我。。。
https://github.com/ELS-RD/kernl/issues/286
目前作者已经回复了，确实是个bug，待他们研究一下。
- 无法将优化后的模型保存下来，只能在线优化，git issue给的解释如下：
PyTorch 2.0 有一种方法可以保存在磁盘上并重新加载已编译的 Triton 内核。 使用他们的 API 应该可以在预热期间节省大量时间。 我们刚刚合并了一个最新版本的 PyTorch 2.0，所以我想调用他们的 API 来启动 Triton 内核而不是原来的 Triton 内核是非常可行的。 但是，无法导出和重用 CUDA 图。 主要原因是内核使用它们的参数重新执行。 这些参数包括 gpu 内存地址。 并且您的张量内存地址很可能会在每次启动时发生变化（如果您退出 Python 会话并且 CUDA 池被释放）。 理论上更新图形参数当然是可能的，但这似乎非常困难，所以重新运行预热可能更好（无需重新编译 triton 内核 ofc）。
关注这个issue : https://github.com/ELS-RD/kernl/issues/224
八、遇到的一些已经解决的问题
8.1 git提供的镜像打不起来：
分析原因，github提供的镜像ubuntu版本太高了，本机ubuntu版本比较低，有bug
另外需要注意，kernl需要python3.9， cuda11.7以上
解决方式，把FROM的基础镜像改了，同时，python地方也需要修改，按照如下方式：
1 | # FROM nvidia/cuda:12.0.0-devel-ubuntu22.04 |
8.2 安装timm会遇到问题
在安装timm的时候，需要安装torchvision，默认会安装非最新版，所以需要手动先安装好最新版的torchvision
解决方法：手动安装最新版torchvision后，再安装torchvision
8.3 算力要求8.0以上：
解决方法：换机器
raise RuntimeError(“GPU compute capability 8.0 (Ampere) or higher is required to use Kernl”)
九、一些原理性的东西
kernl库如何实现的加速呢？其实主要有三个东西
- torchdynamo，可以自定义后端，利用优化的内核替换掉原来的。
- cudagraph
- openai triton，可以更方便的写cuda内核代码。
自定义后端在src/kernl/model_optimization.py中定义，代码如下：
def compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]): dynamo_backend_ofi(gm) return cuda_graphs_wrapper(gm, example_inputs, pool=pool)
目前代码支持的自定义后端有如下一些 ：
1 |
|
所以，目前来看，优化还处于初级阶段，提速的本质还是依靠实现的内核，目前支持的都是transformer的内核替换，所以在transformer上的提速较为明显，期待未来的更为完善。