backend.cpp 7.79 KB
Newer Older
1
#include <cstdlib>
Nicolas Patry's avatar
Nicolas Patry committed
2
3
4
5
6
7
8
9
10
#include <fstream>

#include <fmt/ranges.h>
#include <spdlog/spdlog.h>
#include <nvml.h>

#include "backend.h"
#include "hardware.h"

11
12
13

void huggingface::tgi::backends::InitializeLogging() {
#ifdef NDEBUG
14
15
16
17
18
19
20
21
22
23
24
    if (const auto TRTLLM_LOG_LEVEL_CSTR = std::getenv("TRTLLM_LOG_LEVEL")) {
        std::string log_level(TRTLLM_LOG_LEVEL_CSTR);
        std::transform(log_level.begin(), log_level.end(), log_level.begin(), [](unsigned char c) {
            return std::tolower(c);
        });

        if (log_level == "debug")
            spdlog::set_level(spdlog::level::debug);
        else
            spdlog::set_level(spdlog::level::info);
    }
25
26
27
28
#else
    spdlog::set_level(spdlog::level::debug);
#endif
}
29

30
void huggingface::tgi::backends::InitializeBackend() {
Nicolas Patry's avatar
Nicolas Patry committed
31
32
33
34
    SPDLOG_INFO("Initializing Backend...");
    nvmlInit_v2();
    initTrtLlmPlugins();

35
36
    InitializeLogging();

37
    SPDLOG_INFO("Backend Executor Version: {}", tle::version());
Nicolas Patry's avatar
Nicolas Patry committed
38
39
40
41
42
43
44
45
    const auto numGpus = huggingface::hardware::cuda::GetNumDevices();
    if (numGpus.has_value()) {
        SPDLOG_INFO("Detected {:d} Nvidia GPU(s)", numGpus.value());
    } else {
        SPDLOG_WARN("Failed to detected Nvidia GPU(s) on the system");
    }
}

46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
[[nodiscard]]
tle::ParallelConfig
huggingface::tgi::backends::GetParallelConfig(const size_t worldSize, const std::string workerPath) noexcept {
    auto mode = tle::CommunicationMode::kLEADER;
    std::optional<tle::OrchestratorConfig> orchestratorConfig = std::nullopt;

    if (worldSize > 1) {
        SPDLOG_INFO("Detected sharded engine deployment, using orchestrator mode");
        mode = tle::CommunicationMode::kORCHESTRATOR;
        orchestratorConfig = std::make_optional<tle::OrchestratorConfig>(true, workerPath, nullptr, true);
    } else {
        SPDLOG_INFO("Detected single engine deployment, using leader mode");
    }

    return tle::ParallelConfig(tle::CommunicationType::kMPI, mode, std::nullopt, std::nullopt, orchestratorConfig);
}

Nicolas Patry's avatar
Nicolas Patry committed
63
64
[[nodiscard]]
tle::ExecutorConfig huggingface::tgi::backends::GetExecutorConfig(const json &config, const std::string &workerPath) {
65
    tle::ExecutorConfig execConfig(/* maxBeamWidth = */ 1);
Nicolas Patry's avatar
Nicolas Patry committed
66
67
68
69
70

    // Retrieve the compute capabilities to enable some options at runtime
    const auto computeCapabilities = huggingface::hardware::cuda::GetCudaComputeCapabilities();

    // Single engine (TP = PP = 1) -> using leader mode (no MPI involved)
71
72
    const auto worldSize = config["/pretrained_config/mapping/world_size"_json_pointer].get<size_t>();
    execConfig.setParallelConfig(GetParallelConfig(worldSize, workerPath));
Nicolas Patry's avatar
Nicolas Patry committed
73
74
75

    // Define some configuration variables
    execConfig.setKvCacheConfig(tle::KvCacheConfig(true));
76
77
    execConfig.setEnableChunkedContext(computeCapabilities.IsPostAmpere());
    execConfig.setSchedulerConfig(tle::SchedulerConfig(tle::CapacitySchedulerPolicy::kMAX_UTILIZATION));
Nicolas Patry's avatar
Nicolas Patry committed
78
79
80
81
    return execConfig;
}

tle::SamplingConfig huggingface::tgi::backends::GetSamplingConfig(
82
83
84
85
86
87
88
        const uint32_t topK,
        const float_t topP,
        const float_t temperature,
        const float_t repetition_penalty,
        const float_t frequency_penalty,
        const uint64_t seed) noexcept {

Nicolas Patry's avatar
Nicolas Patry committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
    return tle::SamplingConfig(
            1,  // TGI only use a single beam
            topK,
            topP,
            std::nullopt,
            std::nullopt,
            std::nullopt,
            seed,
            temperature,
            temperature,
            std::nullopt,
            repetition_penalty,
            std::nullopt,
            frequency_penalty
    );
}

106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
std::optional<std::list<std::vector<huggingface::tgi::backends::TokenId>>>
huggingface::tgi::backends::GetStopWordsFromConfig(
        const std::filesystem::path &generationConfigPath) noexcept {
    if (exists(generationConfigPath)) {
        const auto generationConfig = json::parse(std::ifstream(generationConfigPath));
        if (const auto eosTokenIds = generationConfig["/eos_token_id"_json_pointer]; eosTokenIds.is_array()) {
            SPDLOG_INFO(FMT_STRING("Found {:d} EOS tokens"), eosTokenIds.size());
            std::list<std::vector<huggingface::tgi::backends::TokenId>> stopWords(eosTokenIds.size());

            const auto to_single_token = [](const auto tokenIdObj) -> decltype(stopWords)::value_type {
                return {tokenIdObj.template get<tle::TokenIdType>()};
            };

            std::transform(eosTokenIds.cbegin(), eosTokenIds.cend(), stopWords.begin(), to_single_token);
            return stopWords;
        } else {
            SPDLOG_INFO("Invalid EOS tokens entry found (not an array)");
        }
    } else {
        SPDLOG_INFO("No EOS tokens found, generation_config.json doesn't exist");
    }

    return std::nullopt;
}

Nicolas Patry's avatar
Nicolas Patry committed
131
132
133
134
135
huggingface::tgi::backends::TensorRtLlmBackend::TensorRtLlmBackend(
        const std::filesystem::path &enginesFolder,
        const std::filesystem::path &executorWorker
) :
        config(json::parse(std::ifstream(enginesFolder / "config.json"))),
136
137
        executor(enginesFolder, tensorrt_llm::executor::ModelType::kDECODER_ONLY,
                 GetExecutorConfig(config, executorWorker.string())) {
138
139
140
141
142
143
144
145
146
147

    SPDLOG_INFO(FMT_STRING("Engine (version={})"), config["/version"_json_pointer].get<std::string_view>());

    // Ensure we have enough GPUs on the system
    const auto worldSize = config["/pretrained_config/mapping/world_size"_json_pointer].get<size_t>();
    const auto numGpus = huggingface::hardware::cuda::GetNumDevices().value_or(0);
    if (numGpus < worldSize) {
        SPDLOG_CRITICAL(FMT_NOT_ENOUGH_GPUS, numGpus, worldSize);
        // todo : raise exception to catch on rust side
    }
Nicolas Patry's avatar
Nicolas Patry committed
148

149
150
    // Cache variables
    maxNumTokens = config["/build_config/max_num_tokens"_json_pointer].get<uint32_t>();
151
152
153
154

    // Attempt to discover stopWords from the generation_config.json
    const auto generationConfigPath = enginesFolder / "generation_config.json";
    stopWords = GetStopWordsFromConfig(generationConfigPath).value_or(std::list<std::vector<TokenId>>());
Nicolas Patry's avatar
Nicolas Patry committed
155
156
157
158
}

[[nodiscard("Returned number of requests needs to be consumed")]]
size_t huggingface::tgi::backends::TensorRtLlmBackend::NumResponsesReady() const {
159
160
161
#ifdef NDEBUG
    return executor.getNumResponsesReady();
#else
162
    const auto numResponses = executor.getNumResponsesReady();
163
    if (numResponses > 0) SPDLOG_INFO(FMT_STRING("Num responses ready: {:d}"), numResponses);
164
    return numResponses;
165
#endif
Nicolas Patry's avatar
Nicolas Patry committed
166
167
168
169
170
}

[[nodiscard("Returned request id needs to be provided back to gather generated tokens")]]
tle::IdType huggingface::tgi::backends::TensorRtLlmBackend::Submit(
        const std::vector<tle::TokenIdType> &tokens,
171
        const uint32_t maxNewTokens,
Nicolas Patry's avatar
Nicolas Patry committed
172
173
174
        const int32_t topK,
        const float_t topP,
        const float_t temperature,
175
176
        const float_t repetitionPenalty,
        const float_t frequencyPenalty,
Nicolas Patry's avatar
Nicolas Patry committed
177
178
        const uint64_t seed
) {
179
180
181
182
183
184
185
    const auto maxNewTokensChecked = std::min(maxNewTokens, static_cast<uint32_t>(maxNumTokens - tokens.size()));
#ifndef NDEBUG
    {
        const auto &iterations = executor.getLatestIterationStats();
        const auto &lastIteration = iterations.front();

        SPDLOG_DEBUG(FMT_EXECUTOR_STATS, fmt::join(tokens, ", "), lastIteration.numActiveRequests);
186
        SPDLOG_DEBUG(FMT_SAMPLING_CONFIG, topK, topP, temperature, repetitionPenalty, frequencyPenalty, seed);
187
188
        SPDLOG_DEBUG(FMT_STRING("Asking for max_new_tokens={:d}"), maxNewTokensChecked);
    }
Nicolas Patry's avatar
Nicolas Patry committed
189
190
#endif

191
192
193
194
195
196
197
198
    const auto sampling = GetSamplingConfig(topK, topP, temperature, repetitionPenalty, frequencyPenalty, seed);

    // Build the request
    auto request = tle::Request{tokens, CAST_SIZETYPE(maxNewTokensChecked), true, sampling, OUTPUT_CONFIG};
    request.setStopWords(stopWords);

    // Submit to the executor for batching
    return executor.enqueueRequest(request);
Nicolas Patry's avatar
Nicolas Patry committed
199
200
}

201
202
std::vector<tle::Response> huggingface::tgi::backends::TensorRtLlmBackend::PullNewTokens() {
    return executor.awaitResponses();
Nicolas Patry's avatar
Nicolas Patry committed
203
}