2024-06-30 15:37:20 -06:00
|
|
|
#include <fmt/std.h>
|
2024-07-08 16:32:41 -06:00
|
|
|
#include <nvml.h>
|
2024-07-08 16:08:49 -06:00
|
|
|
#include <spdlog/spdlog.h>
|
2024-06-30 15:37:20 -06:00
|
|
|
|
|
|
|
#include "backend.h"
|
|
|
|
|
2024-07-08 16:08:49 -06:00
|
|
|
void huggingface::tgi::backends::InitializeBackend() {
|
|
|
|
SPDLOG_INFO("Initializing Backend...");
|
2024-07-08 16:32:41 -06:00
|
|
|
nvmlInit_v2();
|
2024-07-08 16:08:49 -06:00
|
|
|
initTrtLlmPlugins();
|
|
|
|
}
|
2024-07-03 15:12:24 -06:00
|
|
|
|
2024-07-08 16:08:49 -06:00
|
|
|
tle::ExecutorConfig huggingface::tgi::backends::GetExecutorConfig(const json &config, const std::string &workerPath) {
|
|
|
|
tle::ExecutorConfig execConfig(1);
|
2024-07-03 15:38:17 -06:00
|
|
|
|
2024-07-09 06:15:41 -06:00
|
|
|
// Get the compute capabilities of the current hardware
|
2024-07-08 16:32:41 -06:00
|
|
|
nvmlDevice_t device;
|
2024-07-09 06:15:41 -06:00
|
|
|
int32_t cudaComputeCapabilitiesMajor = 0, cudaComputeCapabilitiesMinor = 0;
|
2024-07-08 16:32:41 -06:00
|
|
|
if(nvmlDeviceGetHandleByIndex_v2(0, &device) == NVML_SUCCESS) {
|
2024-07-09 06:15:41 -06:00
|
|
|
SPDLOG_DEBUG("Successfully acquired nvmlDevice_t = 0");
|
2024-07-08 16:32:41 -06:00
|
|
|
if(nvmlDeviceGetCudaComputeCapability(device, &cudaComputeCapabilitiesMajor, &cudaComputeCapabilitiesMinor) == NVML_SUCCESS) {
|
|
|
|
SPDLOG_INFO(FMT_STRING("Detected sm_{:d}{:d} compute capabilities"), cudaComputeCapabilitiesMajor, cudaComputeCapabilitiesMinor);
|
|
|
|
}
|
|
|
|
}
|
2024-07-08 16:08:49 -06:00
|
|
|
|
2024-07-09 06:15:41 -06:00
|
|
|
// Single engine (TP = PP = 1) -> using leader mode (no MPI involved)
|
2024-07-08 16:08:49 -06:00
|
|
|
if(config["/pretrained_config/mapping/world_size"_json_pointer].get<uint8_t>() == 1){
|
|
|
|
SPDLOG_INFO("Detected single engine deployment, using leader mode");
|
|
|
|
execConfig.setParallelConfig(tle::ParallelConfig(
|
|
|
|
tle::CommunicationType::kMPI,
|
|
|
|
tle::CommunicationMode::kLEADER,
|
|
|
|
std::nullopt,
|
|
|
|
std::nullopt,
|
|
|
|
std::nullopt
|
|
|
|
));
|
2024-07-09 06:15:41 -06:00
|
|
|
} else { // Multiple engines -> using orchestrator mode (MPI involved)
|
2024-07-08 16:08:49 -06:00
|
|
|
SPDLOG_INFO("Detected sharded engine deployment, using orchestrator mode");
|
|
|
|
execConfig.setParallelConfig(tle::ParallelConfig(
|
|
|
|
tle::CommunicationType::kMPI,
|
|
|
|
tle::CommunicationMode::kORCHESTRATOR,
|
|
|
|
std::nullopt,
|
|
|
|
std::nullopt,
|
|
|
|
tle::OrchestratorConfig(true, workerPath)
|
|
|
|
));
|
|
|
|
}
|
2024-07-09 06:15:41 -06:00
|
|
|
|
|
|
|
// Define some configuration variables
|
|
|
|
execConfig.setKvCacheConfig(tle::KvCacheConfig(true));
|
|
|
|
execConfig.setEnableChunkedContext(cudaComputeCapabilitiesMajor >= 8);
|
2024-07-03 15:12:24 -06:00
|
|
|
return execConfig;
|
|
|
|
}
|
|
|
|
|
|
|
|
huggingface::tgi::backends::TensorRtLlmBackend::TensorRtLlmBackend(
|
2024-07-08 16:08:49 -06:00
|
|
|
const std::filesystem::path &enginesFolder,
|
2024-07-03 15:12:24 -06:00
|
|
|
const std::filesystem::path &executorWorker
|
|
|
|
):
|
2024-07-08 16:08:49 -06:00
|
|
|
config(json::parse(std::ifstream(enginesFolder / "config.json"))),
|
|
|
|
executor(
|
|
|
|
enginesFolder,
|
|
|
|
tensorrt_llm::executor::ModelType::kDECODER_ONLY,
|
|
|
|
GetExecutorConfig(config, executorWorker.string()
|
|
|
|
))
|
2024-07-03 15:12:24 -06:00
|
|
|
{
|
2024-07-08 16:08:49 -06:00
|
|
|
SPDLOG_INFO(FMT_STRING("Engine (version={})"), config["/version"_json_pointer].get_ref<const std::string&>());
|
2024-06-30 15:37:20 -06:00
|
|
|
}
|
2024-07-03 02:27:53 -06:00
|
|
|
|
|
|
|
tle::IdType huggingface::tgi::backends::TensorRtLlmBackend::Submit(
|
2024-07-08 16:08:49 -06:00
|
|
|
const std::vector<tle::TokenIdType> &tokens,
|
2024-07-03 02:27:53 -06:00
|
|
|
const int32_t maxNewTokens,
|
2024-07-08 16:08:49 -06:00
|
|
|
const int32_t topK,
|
2024-07-03 02:27:53 -06:00
|
|
|
const float_t topP,
|
|
|
|
const float_t temperature,
|
|
|
|
const int32_t minLength,
|
2024-07-08 16:08:49 -06:00
|
|
|
std::optional<float_t> repetitionPenalty,
|
|
|
|
std::optional<float_t> frequencyPenalty,
|
|
|
|
std::optional<uint32_t> seed,
|
|
|
|
std::optional<uint32_t> nTopTokens
|
2024-07-03 02:27:53 -06:00
|
|
|
) {
|
2024-07-08 16:08:49 -06:00
|
|
|
spdlog::debug(
|
2024-07-08 16:32:41 -06:00
|
|
|
FMT_STRING("Submitting inference over {:d} tokens to the executor {:d}"),
|
2024-07-08 16:08:49 -06:00
|
|
|
tokens.size(),
|
|
|
|
executor.getLatestIterationStats().back().numActiveRequests
|
|
|
|
);
|
|
|
|
|
|
|
|
const auto sampling = tle::SamplingConfig{
|
|
|
|
1,
|
|
|
|
topK,
|
|
|
|
topP,
|
|
|
|
std::nullopt,
|
|
|
|
std::nullopt,
|
|
|
|
std::nullopt,
|
|
|
|
seed,
|
|
|
|
temperature,
|
|
|
|
minLength,
|
|
|
|
std::nullopt,
|
|
|
|
repetitionPenalty,
|
|
|
|
std::nullopt,
|
|
|
|
frequencyPenalty,
|
|
|
|
};
|
|
|
|
const auto output = tle::OutputConfig{false, false, nTopTokens.value_or(1) > 1};
|
|
|
|
const auto request = tle::Request{tokens, maxNewTokens, true, sampling, output};
|
|
|
|
|
|
|
|
return executor.enqueueRequest(request);
|
2024-07-03 02:27:53 -06:00
|
|
|
}
|
2024-07-08 16:08:49 -06:00
|
|
|
|
|
|
|
std::vector<tle::Response> huggingface::tgi::backends::TensorRtLlmBackend::Poll(const tle::IdType reqId) {
|
2024-07-08 16:32:41 -06:00
|
|
|
SPDLOG_DEBUG(FMT_STRING("Polling request {:d}"), reqId);
|
2024-07-08 16:08:49 -06:00
|
|
|
const auto responses = executor.awaitResponses(reqId);
|
|
|
|
return responses;
|
|
|
|
}
|