triton-server中的TensorRT backend

[imoldpan]

本篇文章分析下 triton inference server 中，TensorRT Backend 的具体实现细节，使用的版本为 r23.01。

整理流程

各种结构体：

// A struct to hold TensorRT execution context and its meta data, a
// backend context can have multiple of this struct if multiple
// optimization profiles is specified.
struct TensorRTContext {
  TensorRTContext(
      const std::string& profile_name, const int profile_idx,
      const int binding_cnts, const int event_set_cnts)
      : profile_name_(profile_name), profile_idx_(profile_idx),
        context_(nullptr), cuda_graph_execs_(event_set_cnts),
        min_dims_(binding_cnts), max_dims_(binding_cnts),
        opt_dims_(binding_cnts), min_shapes_(binding_cnts),
        max_shapes_(binding_cnts), opt_shapes_(binding_cnts),
        is_dynamic_per_binding_(binding_cnts)
  {}

GetMostOptimizedProfile

TRITONSERVER_Error*
ModelInstanceState::GetMostOptimizedProfile(
    size_t total_batch_size, TRITONBACKEND_Request** requests,
    uint32_t request_count,
    const std::map<int, std::vector<int32_t>>& request_shape_values,
    std::map<int, TensorRTContext>::iterator* citr)
{
  // Returns the TensorRT context that uses profile with shortest
  // Manhattan distance in terms of input dimensions [TODO] traverse
  // it with more efficient data structure (i.e. K-D tree)
  *citr = trt_contexts_.begin();
  if (trt_contexts_.size() != 1) {
    int64_t shortest_distance = LLONG_MAX;
    for (auto cit = trt_contexts_.begin(); cit != trt_contexts_.end(); cit++) {
      int64_t current_distance = 0;
      EvaluateTensorRTContext(
          cit, total_batch_size, requests, request_count, request_shape_values,
          &current_distance);
      if (current_distance < shortest_distance) {
        *citr = cit;
        shortest_distance = current_distance;
      }
    }
    if (shortest_distance == LLONG_MAX) {
      std::string profiles_str;
      for (const auto& trt_context : trt_contexts_) {
        profiles_str +=
            (" " + trt_context.second.profile_name_ + "[" +
             std::to_string(trt_context.first) + "]");
      }
      return TRITONSERVER_ErrorNew(
          TRITONSERVER_ERROR_INVALID_ARG,
          (std::string("failed to find any Optimization Profile among [") +
           profiles_str +
           "] to support the "
           "requested dimensions (or shape values), proceeding with "
           "first "
           "profile.")
              .c_str());
    }
  }

  LOG_MESSAGE(
      TRITONSERVER_LOG_VERBOSE,
      (std::string("Optimization profile ") + (*citr)->second.profile_name_ +
       " [" + std::to_string((*citr)->first) + "] is selected for " + Name())
          .c_str());
  return nullptr;
}

这个函数 GetMostOptimizedProfile 的目的是为给定的输入数据选择最优的 TensorRT 优化配置文件 (optimization profile)。

在 TensorRT 中，当模型具有动态输入尺寸时，可以定义多个优化配置文件，每个配置文件都有一组输入尺寸的范围（最小、最大、优化尺寸）。在运行时，需要选择一个与输入数据尺寸最匹配的配置文件来执行模型。

以下是此函数的详细步骤：

1. 初始化上下文迭代器:

   • 默认情况下，将迭代器 citr 设置为 trt_contexts_ 的第一个元素。

2. 选择最优的配置文件:

   • 如果 trt_contexts_ 中只有一个配置文件，那么就使用这个配置文件，否则：
     • 初始化一个变量 shortest_distance 来跟踪最短的曼哈顿距离。
     • 遍历所有的配置文件，并为每个配置文件计算与输入数据尺寸之间的曼哈顿距离。
     • 使用 EvaluateTensorRTContext 函数计算每个配置文件的曼哈顿距离。
     • 如果某个配置文件的曼哈顿距离比 shortest_distance 小，那么更新 citr 和 shortest_distance。
   • 如果所有的配置文件都不匹配，那么返回一个错误，并指出没有找到任何匹配的配置文件。

3. 输出选定的配置文件:

   • 输出一条消息，指明已选择哪个配置文件来执行模型。

最后，如果所有步骤都成功完成，函数将返回 nullptr（表示没有错误）。

从代码中的 TODO 注释可以看出，当前的实现使用线性搜索来查找最优的配置文件。这可能不是最高效的方法，特别是当有大量的配置文件时。未来的实现可能会考虑使用更高效的数据结构（如 K-D 树）来加速搜索过程。

InitOptimizationProfiles

TRITONSERVER_Error*
ModelInstanceState::InitOptimizationProfiles()
{
  total_bindings_ = engine_->getNbBindings();
  const int total_profiles = engine_->getNbOptimizationProfiles();

  // TRT sets the optimization profile index to be 0 implicitly with
  // the first context creation. As currently triton supports one
  // context per engine, in order to set the specified profile_index,
  // another context is created and the previous context is destroyed.
  std::shared_ptr<nvinfer1::IExecutionContext> default_trt_context(
      engine_->createExecutionContext());
  if (default_trt_context == nullptr) {
    return TRITONSERVER_ErrorNew(
        TRITONSERVER_ERROR_INTERNAL, "unable to create TensorRT context");
  }

  num_expected_bindings_ = total_bindings_ / total_profiles;

  std::vector<std::pair<std::string, int>> profile_name_index;
  // No optimization profile is set for this TensorRT plan
  if (ProfileNames().empty()) {
    profile_name_index.emplace_back("default", 0);
  } else {
    for (const auto& profile_name : ProfileNames()) {
      int profile_index = 0;
      RETURN_IF_ERROR(GetProfileIndex(profile_name, &profile_index));
      profile_name_index.emplace_back(profile_name, profile_index);
    }
  }

  // Create one TRT context for each specified profile
  for (const auto& name_index : profile_name_index) {
    const auto& profile_name = name_index.first;
    const int profile_index = name_index.second;
    auto res = trt_contexts_.emplace(
        profile_index, TensorRTContext(
                            profile_name, profile_index,
                            num_expected_bindings_, EVENT_SET_COUNT));
    if (!res.second) {
      LOG_MESSAGE(
          TRITONSERVER_LOG_WARN,
          (profile_name + " maps to profile index " +
            std::to_string(profile_index) + " which has been mapped by " +
            res.first->second.profile_name_ +
            ", existing optimization profile will be reused")
              .c_str());
      continue;
    }
    if (profile_index == 0) {
      res.first->second.context_ = std::move(default_trt_context);
    } else {
      res.first->second.context_.reset(engine_->createExecutionContext());
      if (res.first->second.context_ == nullptr) {
        return TRITONSERVER_ErrorNew(
            TRITONSERVER_ERROR_INTERNAL, "unable to create TensorRT context");
      }
      if (!res.first->second.context_->setOptimizationProfileAsync(
              profile_index, stream_)) {
        return TRITONSERVER_ErrorNew(
            TRITONSERVER_ERROR_INVALID_ARG,
            (std::string("Can not set the specified optimization "
                          "profile ") +
              profile_name + "[" + std::to_string(profile_index) + "] for " +
              name_ + ". Expected optimization profile index range 0-" +
              std::to_string(engine_->getNbOptimizationProfiles() - 1))
                .c_str());
      }
      cudaStreamSynchronize(CudaStream());
    }

    // Store the profile dimensions for later initializing the input bindings
    for (int io_index = 0; io_index < num_expected_bindings_; io_index++) {
      const auto binding_index =
          profile_index * num_expected_bindings_ + io_index;
      if (engine_->bindingIsInput(binding_index)) {
        RETURN_IF_ERROR(GetProfileDimensions(
            io_index, profile_index, &res.first->second));
      }
    }
  }

  return nullptr;
}

此代码段的主要功能是初始化 TensorRT 的优化配置文件 (optimization profiles)。在 TensorRT 中，优化配置文件用于动态张量尺寸和精度，它们定义了一组输入尺寸的范围（最小、最大、优化尺寸）。当有多个配置文件时，可以根据输入的实际尺寸选择合适的配置文件来执行模型。

以下是此代码段的详细步骤：

1. 获取绑定和配置文件数量:

   • total_bindings_ = engine_->getNbBindings(); 获取 TensorRT 引擎的总绑定数（输入和输出的总数）。
   • const int total_profiles = engine_->getNbOptimizationProfiles(); 获取 TensorRT 引擎的优化配置文件数量。

2. 创建默认的 TensorRT 上下文:

   • 默认情况下，TensorRT 将优化配置文件的索引设置为 0。
   • 创建默认的 TensorRT 执行上下文，并检查是否成功创建。

3. 计算每个配置文件的预期绑定数:

   • num_expected_bindings_ = total_bindings_ / total_profiles;

4. 确定要使用的配置文件:

   • 如果没有为此 TensorRT 计划设置优化配置文件，那么将使用默认配置文件。
   • 否则，它将为每个配置文件名称获取相应的索引。

5. 为每个指定的配置文件创建一个 TensorRT 上下文:

   • 遍历配置文件名称和索引，为每个配置文件创建一个上下文。
   • 如果配置文件的索引为 0（即默认配置文件），则使用先前创建的默认上下文。
   • 对于其他配置文件，创建一个新的 TensorRT 上下文，并为其设置相应的优化配置文件。
   • 如果创建新的上下文或设置优化配置文件失败，返回错误。

6. 存储配置文件的尺寸:

   • 对于每个配置文件，存储其输入绑定的尺寸，这将在后面初始化输入绑定时使用。

最后，如果所有步骤都成功完成，函数将返回 nullptr（表示没有错误）。

总的来说，此函数的目的是为每个优化配置文件创建和初始化一个 TensorRT 执行上下文。这允许模型在运行时根据输入数据的实际尺寸选择合适的配置文件来执行。

[imoldpan]

关于refit

TRT Refit API is not suitable for production level serving systems. Making a TensorRT engine refittable comes with a performance-cost. The engine can not be optimized to the same degree when running with refittable mode. This means higher inference latency.
Additionally, the engine can not be used to run inference when the model is being updated with new weights. This can lead to large tail latency.

Reloading the model with new weights is better suited solution for running serving. There is no service downtime as the requests can still run on the previous model till the new model with updated weights is ready. And the obtained engine is highly optimized.

参考：

• Llama 2 with LoRA · Issue #68 · NVIDIA/TensorRT-LLM · GitHub
• https://github.com/triton-inference-server/server/issues/5968