TensorRT LLM--In-Flight Batching

2023-11-21 11:57:33 浏览数 (2)

TensorRT LLM依赖于一个名为Batch Manager的组件来支持请求的即时批处理,该技术旨在减少队列中的等待时间,达到更高的GPU利用率。

更详细地说,该功能允许当一个请求再处理中,同时开始处理另一个新请求。

批处理管理器API

客户端可以使用两个主要的回调与批处理管理器交互,它们的签名在callbacks.h文件中定义。

添加新请求、获取处理结果

向批处理管理器传递新请求的入口点是GetInferenceRequestsCallback,该回调的实现返回需要批处理管理器处理的请求列表(std::list<std::shared_ptr<InferenceRequest>),有一个参数表示可以接受的最大请求数(负值表示可以接受无限数量的请求)。

传递新请求的实现如下:

using GetInferenceRequestsCallback = std::function<std::list<std::shared_ptr<InferenceRequest>>(int32_t)>;

对于每个新请求,客户端必须向批处理管理器提供输入张量和一个64位无符号数字(uint64_t),该数字将唯一标识该请求。该标识符称为请求ID。输入张量被收集在一个map(std::map<std::string,Tensor>)中,该map将输入名称与张量相关联。

结果通过SendResponseCallback传递到客户端,一致回调必须接受唯一标识请求的64位请求ID、输出张量列表、布尔值(设置为true时标识请求的最后一个响应)和潜在的非空错误消息。非空错误消息表示遇到错误。在这种情况下,指示这是最后一个响应的布尔值将设置为true,回调必须正确处理错误。

获取请求结果的实现如下:

using SendResponseCallback = std::function<void(uint64_t, std::list<std::shared_ptr<Tensor>> const&, bool, const std::string&)>;

需要注意的是,如果客户端传递的请求ID与批次管理器正在处理的请求的请求ID相同,则批次管理器将拒绝使用GetInferenceRequestsCallback回调发送的任何请求。请求ID出现在对标记为final(第三个参数设置为true)的SendResponseCallback回调的调用中后,可以重用。

中断请求

批处理管理器允许用户停止执行当前正在运行的请求,需要停止的一组请求ID可以通过回调传递给批处理管理器:

using PollStopSignalCallback = std::function<std::unordered_set<uint64_t>()>;

当一个正在处理中的请求出现在要中断的请求集中时,批处理管理器将确保它被正确停止。

处理中请求统计

当调用以下接口时,批处理管理器可以返回统计信息:

using ReturnBatchManagerStatsCallback = std::function<void(const std::string&)>;

统计信息被打包为JSON字符串。该字符串包含三个字段: 时间戳,请求的时间戳(使用std::put_time(&tm,“%m-%d-%Y%H:%m:%S”)获得), 迭代计数器(Iteration Counter)是与给定请求的执行相对应的计数器值, Active Request Count,正在处理中的请求数。

GptManager设计

GptManager管理活跃工作项池,活跃工作项池由服务器主动接收的请求流组成。GptManager在其构造函数中生成一个工作线程,然后一直循环获取新请求。工作线程在每个循环迭代开始时调用GetInferenceRequestsCallback,用于读取新请求。当一个或多个请求处理结束了要返回时,它会在每次迭代结束时调用SendResponseCallback。在启用流模式的请求的情况下,此响应可以是单个令牌,或者在禁用流模式时,此响应是完整响应。PollStopSignalCallback和ReturnBatchManagerStatsCallback(如果提供)都会在每次迭代循环结束时调用。当系统没有活动请求时,不调用ReturnBatchManagerStatsCallback。当批处理管理器在SendResponseCallback中通知(通过final_response布尔参数)完成时,服务器可以安全地从其工作项池中撤回请求。在此之前,与该请求相关的所有TensorRTLLM内部状态都将被释放。可以如下创建批处理管理器的实例以服务于像GPT这样的自回归模型:

#include <tensorrt_llm/batch_manager/GptManager.h> using namespace tensorrt_llm::batch_manager; GptManager batchManager(pathToTrtEngine, // Path to the TensorRT engine of the model, TrtGptModelType::InflightBatching, // Use in-flight batching, maxBeamWidth, // Maximum beam width (must be >= 1), schedulerPolicy, // Scheduling policy (see below), maxNumRequests, // Maximum number of requests, getInferenceRequestsCb, // The Get callback (see above), sendResponseCb); // The Send callback (see above).

调度程序策略帮助批处理管理器调整请求的执行调度方式。批处理管理器可以尝试通过积极地调度请求(schedulerPolicy设置为MAX_utilization)来最大限度地提高GPU的利用率,如果KV缓存的内存不足,则可能不得不暂停请求。请注意,任何暂停的请求都将自动恢复,用户唯一可见的影响可能是延迟增加。它还可以采用更保守的方法,只有在知道内存分配足以处理所有活动请求时才调度请求,即使在KV缓存消耗的最坏情况下也是如此。该模式对应于设置为GUARANTED_NO_EVICT的schedulerPolicy。

多GPU计算

当使用张量并行或流水线并行在多个GPU上运行时,需要服务器启动的进程数量与GPU排列的进程数量一样多,并且每个进程都运行自己的GptManager副本。

给定节点上可见的GPU数量可以使用CUDA_visible_DEVICES环境变量进行控制。

必须注意确保所有列在生成循环的每次迭代中都能看到相同的输入,在TensorRT LLM Triton后端,在GetInferenceRequestsCallback中执行MPI广播,以确保每个MPI列都能看到相同的请求集。ReturnBatchManagerStatsCallback只需要从单个列组中调用;所有等级都持有最终结果的相同副本。

0 人点赞