Pytorch模型的高性能部署一直是大家讨论的问题,有两点比较重要:
- 高度优化的算子
- 可以高效率运行计算图的架构和runtime
高度优化的算子不用多说,TensorRT为什么那么快,因为engine在构建的时候,在每个平台(A10、A100、T4等)上搜索到了最优最快的kernel(实现了一些op)。高效率运行计算图也是很关键的一点,TensorRT构建好engine后,需要libnvinfer.so来驱动,其中实现了什么,在使用过程中很容易猜到:
- 序列化和反序列化,也就是所谓的生成engine,读取engine
- 推理engine、多stream运行计算图,管理engine所需要的一些环境,比如显存和中间变量等
为了达到极致的性能,TensorRT的整个运行时都是在C++环境中,虽然提供了Python-API,但实际调用执行的操作都是在C++中,Python只提供包了一层的作用,算子和执行整个计算图的地方都是C++。
c++ api vs python api
python有快速开发以及验证的优点,但是相比C++来说速度较慢而且比较费内存,一般高性能场景都是使用C++去部署,尽量避免使用python环境。
TORCH 1.x时期的C++部署
torch1.x的实际场景中,一般是搭配使用libtorch + torchscript,这俩在很多生产环境中已经验证过了。
libtorch可以使用C++ API去完成和python中使用pytorch-op实现一样的功能,比如:
转化为Pytorch就是:
而torchscript则用于trace或者script我们的模型到C++环境中部署,速度方面变化不大,主要是通过torchscript导出的模型可以在C++环境中加载并运行,不需要依赖python了,可以减少一些python的over head:
TORCH 2.x的C++部署
torch2.0出来的时候,最主要的就是torch.compile的新API,可以直接优化模型。
torch.compile核心是dynamo,dynamo相比torch.jit.trace和torch.jit.script,是一个功能更强大的trace工具[1],trace模型从而优化模型。dynamo出现后,我也很好奇torchscript是否会被废弃?
torchscript
目前看来torchscript还是会继续存在,只是freeze了,功能还会维护,bug还会修,但不会有新功能了。
基于torch.jit.trace的模型导出路径成为过去式了,那么新的基于pt2.0的C++导出方案是啥?
torch官方前一周发了一篇新blog,正式提到了cpp wrapper,核心就是torch.export[2] + cpp wrapper[3]:
- PyTorch 2.1 Contains New Performance Features for AI Developers[4]
使用cpp wrapper去invoke the generated kernels and external kernels in TorchInductor,可以减少python的overhead,实际测试中,模型速度越快,python overhead占比越大,提升也就越大:
cpp wrapper benchmark
我们都知道torch2.0可以基于triton生成高性能的kernel,例如:
定义好一个函数后,加上@torch.compile
装饰器,执行几次即可得到优化后的模型,默认使用的优化器是TorchInductor,借助depyf[5],我们可以看到优化好后生成的triton代码(GPU端):
这个triton代码可以直接调用,但是依赖python环境,如果想要切换到C++端,则修改下config:
后重新执行几次,可以得到生成的cpp调用代码:
其中调用的cubin就是上述生成triton代码编译出来的/tmp/torchinductor_oldpan/rg/xxx.cubin
,这样的话就可以直接拿这个cpp代码去no-python环境跑起来了。
不过实际中我们更多用的是整个模型,例如resnet50,并且带有权重参数,当然这种也是支持的。torch官方也提供了aot工具可以导出整个模型为so:
通过这个aot_compile
可以直接导出带有模型入口的so,因为是aot,需要提前指定输入的一些维度信息,对于支持dynamic来说是必要的。
导出的so可以通过以下C++方式读取:
核心就是AOTIModelContainerRunnerCpu
。
torch针对inductor设计了AOTIModelContainerRunnerCpu类去加载和运行生成的so,so中会包含一些计算图的执行步骤。
具体例子在pytorch/test/cpp/aot_inductor中,有两个例子,一个是
torch::CustomClassHolder
包一层AOTIModelContainerRunnerCuda
跑,也就是和torchscript的结合,另一个是单独的AOTIModelContainerRunnerCuda
去跑,搭配API直接C++调用。
对于一些常见的op,比如全连接self.fc = torch.nn.Linear(64, 10)
,可以直接调用外部算子不需要triton去codegen,上述例子中直接调用的是torch.ops.aten.addmm
,更多细节可以看pytorch/torch/_inductor/select_algorithm.py
。
整体来说,这种导出方式也比较符合常识,常见的op可以直接调用已经高度优化的版本,未见过的一些算子可以使用triton去生成,fuse等图融合操作可以通过fx pass去做,导出c++也可以通过aot的方式导出,还有一些提升性能的runtime细节设计,整体潜力还是蛮大的。
还有很多细节没有来得及看,比如模型中的某些可以并行的op是如何多stream运行的,dynamic的情况是怎么处理的,中间变量如何存放的,显存是如何管理的,都需要花时间去看看。
参考
- https://pytorch.org/tutorials/prototype/inductor_cpp_wrapper_tutorial.html
- https://www.youtube.com/watch?v=eN5fqBNrjOo&list=PL_lsbAsL_o2BivkGLiDfHY9VqWlaNoZ2O&index=33
- https://pytorch.org/blog/new-features-for-ai/
- https://github.com/pytorch/pytorch/pull/111124
- https://github.com/pytorch/pytorch/pull/88167
- https://discuss.pytorch.org/t/pytorch-2-and-the-c-interface/168034/4
- https://github.com/pytorch/TensorRT/discussions/1743
- https://github.com/pytorch/TensorRT/discussions/1557
- https://github.com/pytorch/TensorRT/issues/1404
- https://github.com/pytorch/TensorRT/discussions/1372
- https://discuss.pytorch.org/t/pytorch-2-and-the-c-interface/168034/2
- https://discuss.pytorch.org/t/torch-compiles-deployment-story-to-non-python-host-processes/180943
- https://dev-discuss.pytorch.org/t/using-nsight-systems-to-profile-gpu-workload/59/10
参考资料
[1] dynamo相比torch.jit.trace和torch.jit.script,是一个功能更强大的trace工具: https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html?highlight=torch%20compile
[2] torch.export: https://pytorch.org/docs/main/export.html
[3] cpp wrapper: https://pytorch.org/tutorials/prototype/inductor_cpp_wrapper_tutorial.html
[4] PyTorch 2.1 Contains New Performance Features for AI Developers: https://pytorch.org/blog/new-features-for-ai/
[5] depyf: https://github.com/thuml/depyf