借着triton inference server聊一下各种batching方法

本文链接：https://blog.csdn.net/sinat_37574187/article/details/140480476

借着triton inference server聊一下各种batching方法

原创 Oldpan oldpan博客 2024年06月03日 09:17 北京

在实际的模型部署场景中，我们一般会先优化模型的性能，这也是最直接提升模型服务性能的方式。但如果从更全局方面考虑的话，除了模型的性能，整体的调度和pipeline优化对服务的性能影响也是很大。

比如LLM中提的很多的`Continuous batching`[1]，对整体LLM推理的性能影响就很大，这个不光光是提升kernel性能能够解决的问题。

这里总结下各种batching策略，以及各种batch策略对整体性能的影响，可能不够全面，也希望能够抛砖引玉，一起交流。

单batch

单batch就是不组batch，也就是一个图片或者一个sentence传过来，直接送入模型进行推理。

对于普通CV模型来说，我们的输入tensor大小可以是[1,3,512,512]，以NCHW维度举例子，这里的N是1，即batch=1。对于LLM来说，可能是一个input_ids，维度是[1,1]，比如：

input_ids
tensor([[   0,  376, 1366,  338,  263, 3017,  775, 6160]], device='cuda:0')
input_ids.shape
torch.Size([1, 8])

这种情况比较简单，在搭建服务的时候不需要额外处理什么，单batch比较适合dymamic shape的场景，模型尺寸是否是dynamic也就是NCHW中的HW是否是可变的。如果你的模型的推理场景需要dynamic shape，那么一般无法组batch（不过可以用ragged batch，后续介绍），只能设置为batch=1，一般我们batch的时候HW都需要固定，比如[8,3,256,256]。

对于dynamic shape，不管是优化还是转模型都需要额外注意，比如tensorrt中转换dynamic shape需要设置动态范围：

polygraphy run model_sim.onnx --trt --onnxrt  --fp16 \
        --trt-min-shapes input:[1,64,64,3]   \
        --trt-opt-shapes input:[1,1024,1024,3] \
        --trt-max-shapes input:[1,1920,1920,3]

dynamic shape的好处就是不需要padding了，避免了额外计算，适合那种请求shape变化特别剧烈的场景，不管是传输图片还是推理，少了无效的padding像素，自然变快了不少。