AI 文摘

借着tritoninferenceserver聊一下各种batching方法





作者: 吃果冻不吐果冻皮 来源: 吃果冻不吐果冻皮

####**【点击】加入大模型技术交流群**

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

不过当然有解决方案:

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送给模型:

image

当模型实际拿到的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

当然实际情况更为复杂,这里不进行详细讨论。另外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的时候,字儿少的要等字儿多的都吐完才能一起返回。

怎么办,我们可以利用nougat的特性,提前用一个检测模型检测当前图片中字儿的数量,然后根据数量将数量相近的请求组成batch传入模型 ,并且将那种数量特别多的请求单独推理,这样就可以避免一些等待的情况。

这种根据数量进行组batch就是一种custom batching策略。

参考

更多AI工具,参考Github-AiBard123国内AiBard123

可关注我们的公众号:每天AI新工具