专栏名称: 极市平台

极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台，为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯，行业动态，在线分享信息，线下活动等。网站: http://cvmart.net/

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

极市平台 · 公众号 · · 2024-06-03 18:37

正文

↑ 点击蓝字关注极市平台

作者丨oldpan

来源丨oldpan博客

编辑丨极市平台

极市导读

总结各种batching策略，以及各种batch策略对整体性能的影响。 >> 加入极市CV技术交流群，走在计算机视觉的最前沿

在实际的模型部署场景中，我们一般会先优化模型的性能，这也是最直接提升模型服务性能的方式。但如果从更全局方面考虑的话，除了模型的性能，整体的调度和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像素，自然变快了不少。

Static Batching

上述batch=1的情况虽然简单而且灵活，不过因为每次只能处理一张图，对GPU资源的利用率还是不如大batch，增大batch一般可以提升模型的FLOPS，同时也可以更多地利用tensor core，实际场景中表现一般就是这样：

# centernet res50  
[1,3,1024,1024]  # trt inference 需要10ms  
[8,3,1024,1024]  # trt inference 需要50ms

很显然增大batch可以提升throughput，所以某些场景，对于时延（latency）要求没有那么高的时候，可以尝试大batch来提升吞吐。

静态batch比较简单，对于图像来说，我们可以在模型推理前，输入将图像cat到一起，比如 [4,3,1024,1024] ，送给模型推理，这时的batch=4。

对于非图像场景也一样，将输入tensor合并后就相当利用了batch。不过比较尴尬的是，LLM场景因为decode阶段batch中每个case最终结束时输出的长度不一样，所以会有有些case吐完字儿了，有些还在继续吐，早吐完的需要等还没吐完的。这种情况不管是速度还是GPU利用率都是比较低的：

LLM-static 来自https://www.anyscale.com/blog/continuous-batching-llm-inference

不过当然有解决方案：

https://www.anyscale.com/blog/continuous-batching-llm-inference

Dynamic Batching

除了在客户端直接送给模型大batch数据，也可以在triton服务端组batch，

Dynamic batching, i.e., server-side batching of incoming queries, significantly improves both latency and performance.

动态组batch在图像推理场景很常用，也可以配合static batch使用。

比如你的模型支持最大batch是16，然后客户端可以每次发送batch=x过来，服务端在收到这些请求后，可以根据规则去组batch，组成更大的batch送给模型：

当模型实际拿到的tensor batch越大，模型的性能就越强，找个例子压测看下：

$ perf_analyzer -m inception_graphdef --percentile=95 --concurrency-range 1:8  
...  
Inferences/Second vs. Client p95 Batch Latency  
Concurrency: 1, throughput: 66.8 infer/sec, latency 19785 usec  
Concurrency: 2, throughput: 80.8 infer/sec, latency 30732 usec  
Concurrency: 3, throughput: 118 infer/sec, latency 32968 usec  
Concurrency: 4, throughput: 165.2 infer/sec, latency 32974 usec  
Concurrency: 5, throughput: 194.4 infer/sec, latency 33035 usec  
Concurrency: 6, throughput: 217.6 infer/sec, latency 34258 usec  
Concurrency: 7, throughput: 249.8 infer/sec, latency 34522 usec  
Concurrency: 8, throughput: 272 infer/sec, latency 35988 usec

随着请求并行度的增加，triton服务端模型组的batch越来越大，服务的QPS也随之增加，不过时延也会有所增加，这个时候就要trade off了。

Continuing Batching

Continuing Batching，也可以叫做inflight batching或者Iteration batching。不同于static batching，当一个输入的生成结束了，就将新的输入插进来，所以batch size是动态的。下图第三个输入先生成完，新的输入S5就插入进来了，一直到输出最长的S2的输出结束的时候，batch size由4变成了7，任意case跑完就能返回：

当然实际情况更为复杂，这里不进行详细讨论。另外Continuing Batching也有人称为dynamic batching，和上述不是一回事儿哈

之前static batching的时候也提过，对于LLM场景，假如batch中某个case已经提前结束了，但是其余batch还在decode，那么这个case只能等其他case而无法直接返回，导致某些请求时延变长了。

另外，假如遇到特别长的case，为了等这个case跑完，其他case都需要等它，其他请求进来也得等他，等来等去待请求队列就长了，服务也就崩了。

所以Continuing Batching这个调度策略很重要，相比模型本身的kernel性能，调度对整体性能的影响是很大的，夸张点，拿23倍的吞吐量轻而易举：

How continuous batching enables 23x throughput in LLM inference while reducing p50 latency ^[2]

很多优秀的LLM推理框架都用到了这个技术，几乎是必用的：TensorRT-LLM、vLLM、Imdeploy等。

Ragged Batching

上文我们提到 Triton有动态批处理功能 ，它可以将多个相同模型执行的请求合并以提供更大的吞吐量。

默认情况下，只有在每个请求中的输入具有 相同shape时，才能进行动态批处理 。如果我们想同时用上dynamic batching和dynamic shape，一般则需要客户端将请求中的输入tensor填充到相同的形状。

举个例子，比如输入 [1,3,768,932] 和 [1,3,1024,768] ，合并为batch的时候需要将两个tensor都padding到同样尺寸，比如 [1,3,1024,1024] ，然后再cat起来为 [2,3,1024,1024] 才可以传给服务端模型；又或者客户端分别发了两个padding后的 [1,3,1024,1024] 过来，服务端自动组batch为 [2,3,1024,1024] 送给模型。

咱们之前也说过，padding会带来不必要的传输和计算量，所以我们可以稍微修改下调度逻辑从而实现dynamic shape + dynamic batching，也就是所谓的ragged batching。

在triton中，ragged batching是一种避免显式padding的功能，它允许用户指定哪些输入不需要进行形状检查。用户可以通过在模型配置中设置allow_ragged_batch字段来指定这样的输入（不规则输入）：

input [  
  {  
    name: "input0"  
    data_type: TYPE_FP32  
    dims: [ 16 ]  
    allow_ragged_batch: true  
  }  
]

举个triton官方的例子。

如果我们有一个接受变长输入tensor INPUT 的模型，INPUT 的形状为 [ -1, -1 ]。第一个维度是 batch 维度，第二个维度是变长内容。当客户端发送三个形状为 [ 1, 3 ]、[ 1, 4 ]、[ 1, 5 ] 的请求时，为了利用动态 batching，最直接的实现方法是期望 INPUT 的形状为 [ -1, -1 ] 并假设所有输入都被填充到相同的长度，这样所有请求都变成形状为 [ 1, 5 ]，因此 Triton 可以将它们 batch 并作为一个 [ 3, 5 ] 张量发送到模型中。在这种情况下，会有padding的开销和对padding进行额外模型计算的开销。以下是输入配置：

max_batch_size: 16  
input [  
  {  
    name: "INPUT"  
    data_type: TYPE_FP32  
    dims: [ -1 ]  
  }  
]

那如果使用 Triton 的 ragged batching，模型将被实现为期望 INPUT 形状为 [ -1 ]，并且有一个额外的 batch 输入 INDEX，形状为 [ -1 ]，模型应该使用它来解释 INPUT 中的 batch 元素。对于这种模型，客户端请求不需要padding，可以按原样发送（形状为 [ 1, 3 ]、[ 1, 4 ]、[ 1, 5 ]）。上述后端将把输入 batch 成一个形状为 [ 12 ] 的张量，其中包含 3 + 4 + 5 个请求的连接。Triton 还创建了一个形状为 [ 3 ] 的 batch 输入张量，值为 [ 3, 7, 12 ]，它给出了每个 batch 元素在输入张量中的结束位置，即索引。以下是输入配置：

max_batch_size: 16  
input [  
  {  
    name: "INPUT"  
    data_type: TYPE_FP32  
    dims: [ -1 ]  
    allow_ragged_batch: true  
  }  
]  
batch_input [  
  {  
    kind: BATCH_ACCUMULATED_ELEMENT_COUNT  
    target_name: "INDEX"  
    data_type: TYPE_FP32  
    source_input: "INPUT"  
  }  
]

使用ragged batch需要实际的模型支持才行（如何处理ragged batch和index）。实际LLM推理中用到的比较多，比如TensorRT-LLM中kernel实现的时候已经考虑到了这种情况：

// TensorRT-LLM/cpp/tensorrt_llm/kernels/gptKernels.cu  
// This kernel also computes the padding offsets: Given the index (idx) of a token in a ragged tensor,  
// we need the index of the token in the corresponding tensor with padding. We compute an array  
// of numTokens elements, called the paddingOffsets, such that the position in the padded tensor  
// of the token "idx" in the ragged tensor is given by idx + paddingOffset[idx].  
//  
// That kernel uses a grid of batchSize blocks.

在图像中，ragged batch一般用不上。不过我们也可以自行设计一个使用ragged batch的策略，简单来说，我们可以把padding操作放在服务端（只是举个栗子，实际场景padding最好在客户端做）。

当客户端分别请求 [1,3,768,932] 和 [1,3,1024,768] 这两个shape的时候，我们可以在服务端将这两个请求 组batch ，然后进行某种预处理（padding或者resize），将input处理好再传入模型，需要我们设计模块处理这种ragged input并且利用起来：

    def execute(self, requests):  
  
        responses = []  
        num_request = len(requests)  
        input_tensors = []  
        time_start = time.time()  
          
        for idx, request in enumerate(requests):  
              
            input = pb_utils.get_input_tensor_by_name(request, "input")  
            input = input.as_numpy()  
            input = input.squeeze(0)  # NHWC -> HWC  
            input = Image.fromarray(input)  
            # 这里可以将不同shape的input 搞成一个shape  
            input_tensor = self.prepare_input(input).unsqueeze(0)  
            input_tensors.append(input_tensor)  
   
        input_tensor = torch.cat(input_tensors, dim=0).to("cuda")      
        self.logger.log_info("input_tensor: {}".format(input_tensor.shape))  
  
        output_texts = self.model.run(input_tensor, num_request, 4096)  
              
        for output_text in output_texts:  
              
            output = pb_utils.Tensor(  
                'output',  
                np.array(output_text).astype(self.output_dtype))  
  
            inference_response = pb_utils.InferenceResponse(  
                output_tensors=[output]  
            )  
  
            responses.append(inference_response)  
  
        # You should return a list of pb_utils.InferenceResponse. Length  
        # of this list must match the length of `requests` list.  
        return responses

使用场景还有很多，需要我们自行探索了。

Custom batching

自定义的一种batch策略，一般就是有 具体使用场景 才会有目的性的去设计。

举个实际的例子，比如LLM中多模态推理中，有个nougat模型 ^[3] ，会识别一幅图中所有的单词并且一个一个吐出来，这个模型是由两部分组成的：

encoder（普通的cv模型，传入图像传出特征）
decoder（可以理解为和llama一样的decoder模型，带有cross attention结构）

当decoder暂时不支持inflight batching的时候，我们只能使用static batching，但是显然在组batching的时候，字儿少的要等字儿多的都吐完才能一起返回。