Update TensorRT-LLM backend (#647)

kaiyux · web-flow · commit 869c2e008b29 · 2024-11-19T18:31:48.000+08:00
diff --git a/all_models/inflight_batcher_llm/ensemble/config.pbtxt b/all_models/inflight_batcher_llm/ensemble/config.pbtxt
@@ -62,6 +62,12 @@ input [
    dims: [ -1 ]
    optional: true
   },
+  {
+    name: "exclude_input_in_output"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  },
   {
     name: "end_id"
     data_type: TYPE_INT32
@@ -376,6 +382,10 @@ ensemble_scheduling {
         key: "decoder_input_lengths"
         value: "_REQUEST_DECODER_INPUT_LEN"
       }
+      input_map {
+        key: "exclude_input_in_output"
+        value: "exclude_input_in_output"
+      }
       input_map {
         key: "request_output_len"
         value: "_REQUEST_OUTPUT_LEN"
diff --git a/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py b/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py
@@ -809,11 +809,6 @@ def initialize(self, args):
         self.stats_check_period_ms = get_parameter(
             model_config, "stats_check_period_ms", int) or 100
 
-        if not self.decoupled:
-            raise pb_utils.TritonModelException(
-                "Please enable decoupled transaction policy in the model configuration to serve this model"
-            )
-
         self.create_metrics(args["model_name"],
                             args["model_version"],
                             is_v1_model=executor_config.batching_type ==
diff --git a/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py b/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py
@@ -93,6 +93,7 @@ class Request:
     lora_task_id: Optional[np.ndarray] = None
     lora_weights: Optional[np.ndarray] = None
     lora_config: Optional[np.ndarray] = None
+    exclude_input_in_output: Optional[np.ndarray] = None
 
     def validate(self):
         _validate_non_empty(self.text_input, "text_input is required")
diff --git a/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py b/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py
@@ -79,61 +79,27 @@ def __init__(self,
         ]
 
         self.input_names = [
-            "text_input",
-            "decoder_text_input",
-            "image_input",
-            "max_tokens",
-            "bad_words",
-            "stop_words",
-            "end_id",
-            "pad_id",
-            "top_k",
-            "top_p",
-            "temperature",
-            "length_penalty",
-            "repetition_penalty",
-            "min_length",
-            "presence_penalty",
-            "frequency_penalty",
-            "random_seed",
-            "return_log_probs",
-            "return_context_logits",
-            "return_generation_logits",
-            "beam_width",
-            "stream",
-            "prompt_embedding_table",
-            "prompt_vocab_size",
-            "prompt_table_extra_id",
-            "embedding_bias_words",
-            "embedding_bias_weights",
-            "num_draft_tokens",
-            "use_draft_logits",
-            "lora_task_id",
-            "lora_weights",
-            "lora_config",
+            "text_input", "decoder_text_input", "image_input", "max_tokens",
+            "bad_words", "stop_words", "end_id", "pad_id", "top_k", "top_p",
+            "temperature", "length_penalty", "repetition_penalty",
+            "min_length", "presence_penalty", "frequency_penalty",
+            "random_seed", "return_log_probs", "return_context_logits",
+            "return_generation_logits", "beam_width", "stream",
+            "prompt_embedding_table", "prompt_vocab_size",
+            "prompt_table_extra_id", "embedding_bias_words",
+            "embedding_bias_weights", "num_draft_tokens", "use_draft_logits",
+            "lora_task_id", "lora_weights", "lora_config",
+            "exclude_input_in_output"
         ]
 
         self.__undo_reshape_whitelist = {
-            "max_tokens",
-            "end_id",
-            "pad_id",
-            "top_k",
-            "top_p",
-            "temperature",
-            "length_penalty",
-            "repetition_penalty",
-            "min_length",
-            "presence_penalty",
-            "frequency_penalty",
-            "random_seed",
-            "return_log_probs",
-            "return_context_logits",
-            "return_generation_logits",
-            "beam_width",
-            "stream",
-            "prompt_vocab_size",
-            "num_draft_tokens",
-            "use_draft_logits",
+            "max_tokens", "end_id", "pad_id", "top_k", "top_p", "temperature",
+            "length_penalty", "repetition_penalty", "min_length",
+            "presence_penalty", "frequency_penalty", "random_seed",
+            "return_log_probs", "return_context_logits",
+            "return_generation_logits", "beam_width", "stream",
+            "prompt_vocab_size", "num_draft_tokens", "use_draft_logits",
+            "exclude_input_in_output"
         }
 
     def _exec_triton_request(self, request):
@@ -415,6 +381,7 @@ def _get_llm_tensors_from_request(
             "lora_task_id": "lora_task_id",
             "lora_weights": "lora_weights",
             "lora_config": "lora_config",
+            "exclude_input_in_output": "exclude_input_in_output",
         }
         batch_size = request.text_input.shape[0]
         tensors = self.create_triton_tensors(request, name_map)
diff --git a/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt b/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt
@@ -67,6 +67,12 @@ input [
    dims: [ -1 ]
    optional: true
   },
+  {
+    name: "exclude_input_in_output"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  },
   {
     name: "end_id"
     data_type: TYPE_INT32
diff --git a/inflight_batcher_llm/client/end_to_end_grpc_client.py b/inflight_batcher_llm/client/end_to_end_grpc_client.py
@@ -57,7 +57,8 @@ def prepare_inputs(prompt,
                    use_draft_logits=None,
                    num_return_sequences=1,
                    lora_dir=None,
-                   lora_task_id=None):
+                   lora_task_id=None,
+                   exclude_input_in_output=False):
 
     input0 = [[prompt]]
     input0_data = np.array(input0).astype(object)
@@ -140,6 +141,10 @@ def prepare_inputs(prompt,
     if pad_id is not None:
         pad_id_data = np.array([[pad_id]], dtype=np.int32)
         inputs["pad_id"] = pad_id_data
+    if exclude_input_in_output is not None:
+        exclude_input_in_output_data = np.array([[exclude_input_in_output]],
+                                                dtype=bool)
+        inputs["exclude_input_in_output"] = exclude_input_in_output_data
 
     if lora_dir and lora_task_id:
         inputs["lora_weights"] = np.load(
@@ -183,7 +188,8 @@ def run_inference(triton_client,
                   use_draft_logits=None,
                   num_return_sequences=None,
                   lora_dir=None,
-                  lora_task_id=None):
+                  lora_task_id=None,
+                  exclude_input_in_output=False):
 
     try:
         prompts = json.loads(prompt)
@@ -200,7 +206,8 @@ def run_inference(triton_client,
                            return_log_probs_data, return_context_logits_data,
                            return_generation_logits_data, end_id, pad_id,
                            num_draft_tokens, use_draft_logits,
-                           num_return_sequences, lora_dir, lora_task_id))
+                           num_return_sequences, lora_dir, lora_task_id,
+                           exclude_input_in_output))
 
     if batch_inputs:
         multiple_inputs = []
@@ -321,7 +328,7 @@ def run_inference(triton_client,
         for output_text in batch_output_text:
             output_texts.extend(output_text)
 
-    return output_texts
+    return prompts, output_texts
 
 
 if __name__ == '__main__':
@@ -524,6 +531,11 @@ def run_inference(triton_client,
                         required=False,
                         help="LoRA task ID")
 
+    parser.add_argument('--exclude-input-in-output',
+                        action="store_true",
+                        required=False,
+                        help='Option to exclude prompt in output text.')
+
     FLAGS = parser.parse_args()
     if FLAGS.url is None:
         FLAGS.url = "localhost:8001"
@@ -555,7 +567,7 @@ def run_inference(triton_client,
         return_generation_logits_data = np.array(
             [[FLAGS.return_generation_logits]], dtype=bool)
 
-    output_texts = run_inference(
+    prompts, output_texts = run_inference(
         client,
         FLAGS.prompt,
         FLAGS.output_len,
@@ -579,7 +591,8 @@ def run_inference(triton_client,
         FLAGS.pad_id,
         FLAGS.batch_inputs,
         True,
-        num_return_sequences=FLAGS.num_return_sequences)
+        num_return_sequences=FLAGS.num_return_sequences,
+        exclude_input_in_output=FLAGS.exclude_input_in_output)
 
     if FLAGS.check_outputs:
         expected_outputs = json.loads(FLAGS.expected_outputs)
@@ -589,5 +602,8 @@ def run_inference(triton_client,
         batched_output_texts = [
             output_texts[i:i + n] for i in range(0, len(output_texts), n)
         ]
-        for out_texts, expected in zip(batched_output_texts, expected_outputs):
+        for out_texts, prompt, expected in zip(batched_output_texts, prompts,
+                                               expected_outputs):
+            if not FLAGS.streaming and not FLAGS.exclude_input_in_output:
+                expected = prompt + expected
             assert all([out == expected for out in out_texts])
diff --git a/inflight_batcher_llm/src/utils.cc b/inflight_batcher_llm/src/utils.cc
@@ -776,8 +776,8 @@ std::vector<executor::Request> createRequestsFromInputTensors(std::vector<InputT
             = utils::getExternalDraftTokensConfigFromTensors(inputTensors, specDecFastLogits);
 
         auto request = executor::Request(inputTokens, maxNewTokens, streaming, samplingConfig, outConfig, endId, padId,
-            std::nullopt, badWords, stopWords, embeddingBias, externalDraftTokensConfig, pTuningConfig, loraConfig,
-            std::nullopt, std::nullopt, std::nullopt, encoderInputTokens);
+            std::nullopt, badWords, stopWords, embeddingBias, externalDraftTokensConfig, pTuningConfig, std::nullopt,
+            loraConfig, std::nullopt, std::nullopt, std::nullopt, encoderInputTokens);
 
         if (encoderInputFeatures.has_value())
         {
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,6 @@
 regex
 fire
 tritonclient[all]
-transformers==4.36.1
+transformers==4.45.1
 pandas
 tabulate
diff --git a/tensorrt_llm b/tensorrt_llm
@@ -1 +1 @@
-Subproject commit c629546ce429623c8a163633095230154a6f0574
+Subproject commit 535c9cc6730f5ac999e4b1cb621402b58138f819
diff --git a/tools/inflight_batcher_llm/speculative_decoding_test.py b/tools/inflight_batcher_llm/speculative_decoding_test.py
@@ -259,7 +259,7 @@
             # Calling control model only
             if FLAGS.verbose:
                 print(f"Calling control model", flush=True)
-            output_control = end_to_end_grpc_client.run_inference(
+            processed_prompt, output_control = end_to_end_grpc_client.run_inference(
                 client_control, prompt, output_len, str(request_id),
                 FLAGS.repetition_penalty, FLAGS.presence_penalty,
                 FLAGS.frequency_penalty, FLAGS.temperature, FLAGS.stop_words,
@@ -281,7 +281,7 @@
                     return_generation_logits_data = np.array(
                         [[FLAGS.return_generation_logits]], dtype=bool)
 
-                output_speculative = end_to_end_grpc_client.run_inference(
+                processed_prompt, output_speculative = end_to_end_grpc_client.run_inference(
                     client_target, prompt, output_len, str(request_id),
                     FLAGS.repetition_penalty, FLAGS.presence_penalty,
                     FLAGS.frequency_penalty, FLAGS.temperature,
diff --git a/tools/version.txt b/tools/version.txt
@@ -1 +1 @@
-72b706617a62c45bab6c76c87db9960745b73c4e
+0ce11a50f3523c4810be701fddf5cac664b37237

Original file line number	Diff line number	Diff line change
`@@ -776,8 +776,8 @@ std::vector<executor::Request> createRequestsFromInputTensors(std::vector<InputT`
`776`	`776`	`= utils::getExternalDraftTokensConfigFromTensors(inputTensors, specDecFastLogits);`
`777`	`777`
`778`	`778`	`auto request = executor::Request(inputTokens, maxNewTokens, streaming, samplingConfig, outConfig, endId, padId,`
`779`		`- std::nullopt, badWords, stopWords, embeddingBias, externalDraftTokensConfig, pTuningConfig, loraConfig,`
`780`		`- std::nullopt, std::nullopt, std::nullopt, encoderInputTokens);`
	`779`	`+ std::nullopt, badWords, stopWords, embeddingBias, externalDraftTokensConfig, pTuningConfig, std::nullopt,`
	`780`	`+ loraConfig, std::nullopt, std::nullopt, std::nullopt, encoderInputTokens);`
`781`	`781`
`782`	`782`	`if (encoderInputFeatures.has_value())`
`783`	`783`	`{`
Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-72b706617a62c45bab6c76c87db9960745b73c4e`
	`1`	`+0ce11a50f3523c4810be701fddf5cac664b37237`