1 整体架构
三部分组成:
- Triton backend
tensorRT_backend、onnx_backend、tfs_backend、torch_backend
- **Triton model **
不同的模型
- **Triton model instance **
模型实例
![P2}5X%2ULV(2OAC$_`OKOP.png
2 设计思路
需要实现七个接口:
- TRITONBACKEND_Initialize:
- 初始化 Triton backend。
- TRITONBACKEND_ModelInitialize:
- 初始化模型配置,包括在 model.config 中指定的配置。
- TRITONBACKEND_ModelInstanceInitialize:
- 初始化模型实例,指定模型在哪个 GPU 上执行。
- TRITONBACKEND_ModelInstanceExecute:
- 执行模型。
剩下的三个 TRITONBACKEND_ModelFinalize、TRITONBACKEND_ModelInstanceFinalize 和 TRITONBACKEND_Finalize 负责清理工作,不需要关注。
重点是TRITONBACKEND_ModelInstanceExecute,其成员变量ModelState、ModelInstanceState负责实际的模型执行,其具有两方面作用:
- 维护状态
- ModelState:
- 维护模型配置信息,例如模型名称、模型输入输出信息、batching size等。
- ModelInstanceState:
- 记录模型在哪个机器上运行。
- 模型执行
- modelState:
- 加载模型。
- modelInstanceState:
- 模型执行。
3 源码分析
七个接口+两个具体执行类
3.1 TRITONBACKEND_Initialize
检查backend API version
3.2 TRITONBACKEND_ModelInitialize
设置模型名称和版本,创建model_state
model_state具有以下字段,BackendModel的字段是不同模型共有的,ModelState私有字段是torch模型专属的。
3.3 TRITONBACKEND_ModelInstanceInitialize
指定模型示例名称、设备id、以及创建ModelInstanceState
3.4 TRITONBACKEND_ModelInstanceExecute
通过执行ProcessRequests处理Sheduler传过来的批次 requests
3.4.1 max_batch_size_校验
如果max_batch_size > 1 , 输入的total_batch_size <= max_batch_size,shape的第一个纬度是batch大小,对每个请求的batch请和获取total_batch_size。
如果max_batch_size = 0, 则每个请求的batch默认为1。
3.4.2 setInputTensor
- 所有request具有相同数量的输入tensor
- 所有request汇聚到连续的CPU或GPU内存中
3.4.3 Execute
调用具体模型执行推理
3.4.4 ReadOutputTensors
处理推理结果
4 backend调优
backend模型配置分为两类,一类是所有backend通用的,一类是不同backend自有的。
4.1 通用配置
可以调整的max_batch_size_、enable_pinned_input_、enable_pinned_output_、ragged_inputs_
std::string name_;
uint64_t version_;
std::string repository_path_;
bool allow_optional_;
common::TritonJson::Value model_config_;
int max_batch_size_;
bool enable_pinned_input_;
bool enable_pinned_output_;
std::vector<BatchInput> batch_inputs_;
std::vector<BatchOutput> batch_outputs_;
std::map<std::string, const BatchOutput*> batch_output_map_;
std::set<std::string> ragged_inputs_;
std::set<std::string> optional_inputs_;
4.1.1 max_batch_size_
不设置默认批次为1
4.1.2 enable_pinned_input_
TritonMemoryManager是否使用Pinned
Pinned内存优化技术优势包括:
- 加速数据传输:减少CPU与GPU间数据交换延迟,提升计算效率。
- 高性能:保证数据传输速度,适合实时和高性能计算场景。
- 性能稳定:内存访问时间可预测,利于优化应用响应时间。
存在的局限性有:
- 内存占用:固定占用物理RAM,可能影响系统其他部分的内存需求。
- 管理难度:配置复杂,需精确评估内存需求以避免资源浪费。
- 系统影响:不合理使用可能削弱系统整体稳定性和其他应用表现。
4.1.3 ragged_inputs_
Triton 提供了动态批处理功能,该功能可以将多个针对同一模型执行的请求合并,以实现更高的吞吐量。默认情况下,只有当每个请求的输入形状相同时,这些请求才能被动态批处理。为了在输入形状经常变化的情况下利用动态批处理,客户端需要对请求中的输入张量进行填充以达到相同的形状。
4.2 torch_backend自有配置
// Flag to indicate whether optimized execution is enabled. Defaults to true.
bool enable_optimized_execution_;
// Flag to indicate whether inference mode is enabled. Defaults to false.
bool enable_inference_mode_;
// Flag to indicate whether cache cleaning after each run is enabled.
// Defaults to false.
bool enable_cache_cleaning_;
// Flag to indicate whether weight sharing is enabled. Defaults to false.
bool enable_weight_sharing_;
// Flag pairs to indicate if various JIT settings are set and
// enabled respectively. Defaults to (false, true). Default behavior
// is to do nothing if not explicitly set.
std::pair<bool, bool> enable_tensor_fuser_pair_;
std::pair<bool, bool> enable_jit_profiling_pair_;
std::pair<bool, bool> enable_jit_executor_pair_;
4.2.1 enable_optimized_execution_
通启用或禁用图执行器的优化模式。这影响了模型执行时的性能。
4.2.2 enable_inference_mode_
启用或禁用推理模式。这通常会禁用梯度计算,提高推理速度。
4.2.3 enable_cache_cleaning_
是否在每次运行后清理缓存。默认为false,意味着通常不自动清理缓存。
4.2.4 enable_weight_sharing_
是否启用权重共享。这个设置在某些特定结构中(如循环神经网络RNN的多个时间步共享权重)非常有用,可以减少内存占用。默认为false,表示不启用共享,每个权重副本独立存储。
4.2.5 enable_tensor_fuser_pair_
允许根据配置开关来决定是否启用Tensor表达式融合器,这是优化图执行效率的一个重要手段,特别是在GPU上。
4.2.6 enable_jit_profiling_pair_
如果在model_state_中指定了JIT profiling的启用状态,则通过torch::jit::getProfilingMode来设定是否开启profiling,这对于性能分析和优化很有帮助。
4.2.4 enable_jit_executor_pair_
JIT executor模式决定了模型图是否会被即时编译(JIT compiled),通常用于提升执行速度。