Pytorch模型的
高性能部署
一直是大家讨论的问题,有两点比较重要:
高度优化的算子不用多说,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实现一样的功能,比如:
#include at::Tensor a = at::ones({2 , 2 }, at::kInt); at::Tensor b = at::randn({2 , 2 });auto c = a + b.to(at::kInt);
转化为Pytorch就是:
import torch a = torch.ones((2 , 2 ), dtype=torch.int32) b = torch.randn((2 , 2 )) c = a + b.to(torch.int32)
而torchscript则用于
trace或者script我们的模型到C++环境中
部署,速度方面变化不大,主要是通过torchscript导出的模型可以在C++环境中加载并运行,不需要依赖python了,可以减少一些python的over head:
#include // One-stop header. #include #include int main (int argc, const char * argv[]) { if (argc != 2 ) { std ::cerr <"usage: example-app \n"; return -1 ; } torch::jit::script::Module module ; try { // Deserialize the ScriptModule from a file using torch::jit::load(). module = torch::jit::load(argv[1 ]); } catch (const c10::Error& e) { std ::cerr <"error loading the model\n"; return -1 ; } std ::cout <"ok\n"; }
关于torchscript的解读有不少,这里不赘述了,感兴趣的可以参阅:
https://zhuanlan.zhihu.com/p/486914187
https://zhuanlan.zhihu.com/p/489090393
https://zhuanlan.zhihu.com/p/363319763
https://zhuanlan.zhihu.com/p/652193676
https://zhuanlan.zhihu.com/p/410507557
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 def opt_foo (x, y) : a = torch.sin(x) b = torch.cos(y) return a + bfor _ in range(100 ): opt_foo(torch.randn(10 ).cuda(), torch.randn(10 ).cuda())
定义好一个函数后,加上
@torch.compile
装饰器,执行几次即可得到优化后的模型,默认使用的优化器是TorchInductor,借助
depyf
[5]
,我们可以看到优化好后生成的triton代码(GPU端):
import tritonimport triton.language as tlfrom torch._inductor.ir import ReductionHintfrom torch._inductor.ir import TileHintfrom torch._inductor.triton_heuristics import AutotuneHint, pointwisefrom torch._inductor.utils import instance_descriptorfrom torch._inductor import triton_helpers@pointwise( size_hints=[16 ], filename=__file__, triton_meta={'signature' : {0 : '*fp32' , 1 : '*fp32' , 2 : '*fp32' , 3 : 'i32' }, 'device' : 0 , 'device_type' : 'cuda' , 'constants' : {}, 'configs' : [instance_descriptor(divisible_by_16=(0 , 1 , 2 ), equal_to_1=(), ids_of_folded_args=(), divisible_by_8=())]}, inductor_meta={'autotune_hints' : set(), 'kernel_name' : 'triton_poi_fused_add_cos_sin_0' , 'mutated_arg_names' : []}, min_elem_per_thread=0 )@triton.jit def triton_ (in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr) : xnumel = 10 xoffset = tl.program_id(0 ) * XBLOCK xindex = xoffset + tl.arange(0 , XBLOCK)[:] xmask = xindex x0 = xindex tmp0 = tl.load(in_ptr0 + (x0), xmask) tmp2 = tl.load(in_ptr1 + (x0), xmask) tmp1 = tl.sin(tmp0) tmp3 = tl.cos(tmp2) tmp4 = tmp1 + tmp3 tl.store(out_ptr0 + (x0), tmp4, xmask)
这个triton代码可以直接调用,但是依赖python环境,如果想要切换到C++端,则修改下config:
import torch._inductor.config as config config.cpp_wrapper = True
后重新执行几次,可以得到生成的cpp调用代码:
#include #include #include #include #include #define reinterpret_tensor torch::inductor::_reinterpret_tensor #define alloc_from_pool torch::inductor::_alloc_from_pool #include [[maybe_unused]] static int64_t align (int64_t nbytes) { return (nbytes + 64 - 1 ) & -64 ; }#include #include #include #define CUDA_DRIVER_CHECK(EXPR) \ do { \ CUresult code = EXPR; \ const char *msg; \ cuGetErrorString(code, &msg); \ if (code != CUDA_SUCCESS) { \ throw std::runtime_error( \ std::string("CUDA driver error: " ) + \ std::string(msg)); \ } \ } while (0); namespace {struct Grid { Grid(uint32_t x, uint32_t y, uint32_t z) : grid_x(x), grid_y(y), grid_z(z) {} uint32_t grid_x; uint32_t grid_y; uint32_t grid_z; bool is_non_zero () { return grid_x > 0 && grid_y > 0 && grid_z > 0 ; } }; } // anonymous namespace static inline CUfunction loadKernel ( std ::string