[NLP] Evaluate batched inference calls singularly (#2538)

One at a time inference uses much less memory
and throughput is significantly diminished.

Author: davidkyle

bin/pytorch_inference/CCommandParser.cc

@@ -245,12 +245,6 @@ bool CCommandParser::checkArrayContainsUInts(const rapidjson::Value::ConstArray&
            }) == arr.End();
 }
 
-bool CCommandParser::checkArrayContainsDoubles(const rapidjson::Value::ConstArray& arr) {
-    return std::find_if(arr.Begin(), arr.End(), [](const auto& i) {
-               return i.IsDouble() == false;
-           }) == arr.End();
-}
-
 CCommandParser::SRequest
 CCommandParser::jsonToInferenceRequest(const rapidjson::Document& doc) {
     SRequest request;

bin/pytorch_inference/CCommandParser.h

@@ -183,7 +183,6 @@ class CCommandParser {
     static EMessageType validateControlMessageJson(const rapidjson::Document& doc,
                                                    const TErrorHandlerFunc& errorHandler);
     static bool checkArrayContainsUInts(const rapidjson::Value::ConstArray& arr);
-    static bool checkArrayContainsDoubles(const rapidjson::Value::ConstArray& arr);
     static SRequest jsonToInferenceRequest(const rapidjson::Document& doc);
     static SControlMessage jsonToControlMessage(const rapidjson::Document& doc);
 

bin/pytorch_inference/Main.cc

@@ -31,6 +31,7 @@
 #include "CThreadSettings.h"
 
 #include <ATen/Parallel.h>
+#include <ATen/ops/cat.h>
 #include <torch/csrc/api/include/torch/types.h>
 #include <torch/script.h>
 
@@ -44,26 +45,39 @@ torch::Tensor infer(torch::jit::script::Module& module_,
     std::vector<torch::jit::IValue> inputs;
     inputs.reserve(1 + request.s_SecondaryArguments.size());
 
-    std::array<std::int64_t, 2> dimensions = {request.s_NumberInferences,
-                                              request.s_NumberInputTokens};
+    std::array<std::int64_t, 2> dimensions = {1, request.s_NumberInputTokens};
     at::IntArrayRef inputSize{dimensions};
 
-    // Sequence tokens.
-    inputs.emplace_back(torch::from_blob(static_cast<void*>(request.s_Tokens.data()),
-                                         inputSize, at::dtype(torch::kInt64)));
-    // Attention mask.
-    for (auto& args : request.s_SecondaryArguments) {
-        inputs.emplace_back(torch::from_blob(static_cast<void*>(args.data()),
-                                             inputSize, at::dtype(torch::kInt64)));
-    }
+    std::vector<at::Tensor> all;
 
     torch::InferenceMode inferenceModeGuard;
-    auto result = module_.forward(inputs);
-    if (result.isTuple()) {
-        // For transformers the result tensor is the first element in a tuple.
-        return result.toTuple()->elements()[0].toTensor();
+
+    for (int i = 0; i < request.s_NumberInferences; i++) {
+
+        std::size_t offset = i * request.s_NumberInputTokens;
+
+        // Sequence tokens.
+        inputs.emplace_back(
+            torch::from_blob(static_cast<void*>(request.s_Tokens.data() + offset),
+                             inputSize, at::dtype(torch::kInt64)));
+        // Attention mask etc
+        for (auto& args : request.s_SecondaryArguments) {
+            inputs.emplace_back(torch::from_blob(static_cast<void*>(args.data() + offset),
+                                                 inputSize, at::dtype(torch::kInt64)));
+        }
+
+        auto output = module_.forward(inputs);
+        if (output.isTuple()) {
+            // For transformers the result tensor is the first element in a tuple.
+            all.push_back(output.toTuple()->elements()[0].toTensor());
+        } else {
+            all.push_back(output.toTensor());
+        }
+
+        inputs.clear();
     }
-    return result.toTensor();
+
+    return at::cat(all, 0);
 }
 
 bool handleRequest(ml::torch::CCommandParser::CRequestCacheInterface& cache,

docs/CHANGELOG.asciidoc

@@ -36,6 +36,9 @@
 * Improve detection of time shifts, for example for day light saving. (See {ml-pull}2479[#2479].)
 * Improve detection of calendar cyclic components with long bucket lengths. (See {ml-pull}2493[#2493].)
 
+=== Bug Fixes
+* Prevent high memory usage by evaluating batch inference singularly. (See {ml-pull}2538[#2538].)
+
 == {es} version 8.8.0
 
 === Enhancements