[Feature] triton backend optimization

[zhyncs]

Motivation

As described in #1280, there are some performance issues with the Triton backend currently implemented in the repo, and its throughput is not as good as that of the API server. In some companies, Triton Server is closely integrated with the model control platform. Therefore, it is necessary to implement an efficient Triton backend. Considering that @lvhan028 previously replied in the issue that there is currently no extra time available for this optimization.

@ispobock and I plan to re-implement an efficient triton backend, the details are as follows:

1. stream infer
   We will use the Triton Python backend and also employ the decoupled mode, which means starting only one instance, similar to this: https://github.com/triton-inference-server/vllm_backend/blob/main/src/model.py
   This is how we did it internally before open sourcing vllm-backend, and both throughput and latency meet expectations. NVIDIA Triton support vllm-project/vllm#541 (comment)
   It requires client support for stream infer.

2. normal infer
   We will implement it using the method of BLS + stream infer. It will start multiple BLS instances and one model instance. It requires Triton Server >= 23.04.
   It has been discussed here before NVIDIA Triton support vllm-project/vllm#541 (comment)
   It does not require client support for stream inference, normal synchronous inference is sufficient.

After completing these tasks, we will provide examples for the Python client and Java client, as these two clients are more commonly used in the enterprise. At the same time, performance test data will be provided to ensure that it is comparable to the performance of the API server. Cheers. Stay tuned.

@lvhan028 @lzhangzz @AllentDan @grimoire @irexyc Do you have any suggestions? Thanks.

Related resources

No response

Additional context

No response

[zhyncs]

Hi all @lvhan028 @lzhangzz @grimoire @irexyc @AllentDan

model: llama 2 13b chat
lmdeploy: 0.2.6

--------------------------------------------------
concurrency: 128
elapsed_time: 146.570s

number of prompt tokens: 244330
number of completion tokens: 239421
RPS (request per second): 6.823 req/s
RPM (request per minute): 409.361 req/min
--------------------------------------------------

@ispobock has implemented a simple version, and tested the throughput. It is basically on par with the API server, even slightly higher than the API server. Next, we will improve this program by adding error handling, necessary logging, configuration settings, as well as parsing and processing other sampling parameters. Do you have any suggestions? Thanks.

import json
import asyncio
import threading
import numpy as np

from lmdeploy import pipeline, GenerationConfig, TurbomindEngineConfig

import triton_python_backend_utils as pb_utils


class TritonPythonModel:

    def initialize(self, args):
        self.logger = pb_utils.Logger

        # parse model configs
        self.model_config = json.loads(args['model_config'])

        # make sure use decoupled mode
        self.using_decoupled = pb_utils.using_decoupled_model_transaction_policy(
            self.model_config)
        assert (
            self.using_decoupled
        ), "LMDeploy model should be configured to decoupled model transaction policy"

        # start lmdeploy engine
        engine_config = TurbomindEngineConfig()
        self.engine = pipeline('/workdir/llama2_13b_chat',
                               backend_config=engine_config)

        self.request_id = 0

        # create event loop to process requests asynchronously
        self.event_loop = asyncio.get_event_loop()
        self.loop_thread = threading.Thread(target=self.engine_loop,
                                            args=(self.event_loop, ))
        self.shutdown = asyncio.Event()
        self.loop_thread.start()

    def engine_loop(self, loop):
        asyncio.set_event_loop(loop)
        self.event_loop.run_until_complete(self.await_shutdown())

    async def await_shutdown(self):
        # first await the shutdown signal
        while self.shutdown.is_set() is False:
            await asyncio.sleep(5)

        for task in asyncio.all_tasks(loop=self.event_loop):
            if task is not asyncio.current_task():
                task.cancel()

    async def process_request(self, request_id, request):
        response_sender = request.get_response_sender()

        # parse request
        prompt = pb_utils.get_input_tensor_by_name(
            request, "prompt").as_numpy().item().decode()
        max_tokens = pb_utils.get_input_tensor_by_name(
            request, "max_tokens").as_numpy().item()
        temperature = pb_utils.get_input_tensor_by_name(
            request, "temperature").as_numpy().item()
        top_p = pb_utils.get_input_tensor_by_name(request,
                                                  "top_p").as_numpy().item()
        top_k = pb_utils.get_input_tensor_by_name(request,
                                                  "top_k").as_numpy().item()
        ignore_eos = pb_utils.get_input_tensor_by_name(
            request, "ignore_eos").as_numpy().item()
        stream = pb_utils.get_input_tensor_by_name(request,
                                                   "stream").as_numpy().item()

        gen_config = GenerationConfig(max_new_tokens=max_tokens,
                                      top_k=top_k,
                                      top_p=top_p,
                                      temperature=temperature,
                                      ignore_eos=ignore_eos)

        async for output in self.engine.generate(messages=prompt,
                                                 session_id=request_id,
                                                 stream_response=stream,
                                                 gen_config=gen_config,
                                                 do_preprocess=False):
            triton_output_tensor = pb_utils.Tensor(
                "response", np.asarray(output.response, dtype=np.object_))
            resp = pb_utils.InferenceResponse(
                output_tensors=[triton_output_tensor])
            if output.finish_reason is not None:
                response_sender.send(
                    resp, flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL)
            else:
                response_sender.send(resp)

    def execute(self, requests):
        for request in requests:
            asyncio.run_coroutine_threadsafe(
                self.process_request(self.request_id, request),
                self.event_loop)
            self.request_id += 1
        return None

    def finalize(self):
        self.shutdown.set()
        if self.loop_thread is not None:
            self.loop_thread.join()
            self.loop_thread = None

[lvhan028]

what's the evaluation result?

[zhyncs]

what's the evaluation result?

--------------------------------------------------
concurrency: 128
elapsed_time: 146.570s

number of prompt tokens: 244330
number of completion tokens: 239421
RPS (request per second): 6.823 req/s
RPM (request per minute): 409.361 req/min
--------------------------------------------------

[zhyncs]

And API Server RPS nearly 6.8 req/s.

[lvhan028]

I mean the evaluation on datasets, such as gms8k, mmlu and so on.

[zhyncs]

I mean the evaluation on datasets, such as gms8k, mmlu and so on.

Ok, we will conduct the test and share the conclusions with you asap. Please stay tuned.

[lvhan028]

Can the triton python backend support streaming output?

[zhyncs]

Can the triton python backend support streaming output?

sure

[zhyncs]

config.pbtxt

name: "lmdeploy_model"
backend: "python"
max_batch_size: 0

model_transaction_policy {
  decoupled: True
}

instance_group [
  {
    count: 1
    kind: KIND_CPU
  }
]

input [
  {
    name: "prompt"
    data_type: TYPE_STRING
    dims: [ -1 ]
  },
  {
    name: "max_tokens"
    data_type: TYPE_INT32
    dims: [ 1 ]
  },
  {
    name: "temperature"
    data_type: TYPE_FP64
    dims: [ 1 ]
  },
  {
    name: "top_p"
    data_type: TYPE_FP64
    dims: [ 1 ]
  },
  {
    name: "top_k"
    data_type: TYPE_INT32
    dims: [ 1 ]
  },
  {
    name: "ignore_eos"
    data_type: TYPE_BOOL
    dims: [ 1 ]
  },
  {
    name: "stream"
    data_type: TYPE_BOOL
    dims: [ 1 ]
  }
]

output [
  {
    name: "response"
    data_type: TYPE_STRING
    dims: [ -1 ]
  }
]

[ispobock]

@lvhan028 Here is the evaluation result for this TIS Python backend:

dataset                                 version    metric         mode    llama2-13b-chat-lmdeploy-tis
--------------------------------------  ---------  -------------  ------  ----------------------------
--------- 考试 Exam ---------           -          -              -       -
ceval                                   -          naive_average  gen     33.36
agieval                                 -          -              -       -
mmlu                                    -          naive_average  gen     42.90
GaokaoBench                             -          -              -       -
ARC-c                                   -          -              -       -
--------- 语言 Language ---------       -          -              -       -
WiC                                     d06864     accuracy       gen     13.01
summedits                               -          -              -       -
chid-dev                                -          -              -       -
afqmc-dev                               -          -              -       -
bustm-dev                               -          -              -       -
cluewsc-dev                             -          -              -       -
WSC                                     -          -              -       -
winogrande                              -          -              -       -
flores_100                              -          -              -       -
--------- 知识 Knowledge ---------      -          -              -       -
BoolQ                                   -          -              -       -
commonsense_qa                          -          -              -       -
nq                                      -          -              -       -
triviaqa                                2121ce     score          gen     64.30
--------- 推理 Reasoning ---------      -          -              -       -
cmnli                                   -          -              -       -
ocnli                                   -          -              -       -
ocnli_fc-dev                            -          -              -       -
AX_b                                    -          -              -       -
AX_g                                    -          -              -       -
CB                                      -          -              -       -
RTE                                     -          -              -       -
story_cloze                             -          -              -       -
COPA                                    -          -              -       -
ReCoRD                                  -          -              -       -
hellaswag                               -          -              -       -
piqa                                    -          -              -       -
siqa                                    -          -              -       -
strategyqa                              -          -              -       -
math                                    -          -              -       -
gsm8k                                   1d7fe4     accuracy       gen     42.23
TheoremQA                               -          -              -       -
openai_humaneval                        -          -              -       -
mbpp                                    -          -              -       -
bbh                                     -          -              -       -
--------- 理解 Understanding ---------  -          -              -       -
C3                                      -          -              -       -
CMRC_dev                                -          -              -       -
DRCD_dev                                -          -              -       -
MultiRC                                 -          -              -       -
race-middle                             9a54b6     accuracy       gen     74.58
race-high                               9a54b6     accuracy       gen     67.84
openbookqa_fact                         -          -              -       -
csl_dev                                 -          -              -       -
lcsts                                   -          -              -       -
Xsum                                    -          -              -       -
eprstmt-dev                             -          -              -       -
lambada                                 -          -              -       -
tnews-dev                               -          -              -       -

The result seems fine.

[zhyncs]

Motivation

As described in #1280, there are some performance issues with the Triton backend currently implemented in the repo, and its throughput is not as good as that of the API server. In some companies, Triton Server is closely integrated with the model control platform. Therefore, it is necessary to implement an efficient Triton backend. Considering that @lvhan028 previously replied in the issue that there is currently no extra time available for this optimization.

@ispobock and I plan to re-implement an efficient triton backend, the details are as follows:

1. stream infer
   We will use the Triton Python backend and also employ the decoupled mode, which means starting only one instance, similar to this: https://github.com/triton-inference-server/vllm_backend/blob/main/src/model.py
   This is how we did it internally before open sourcing vllm-backend, and both throughput and latency meet expectations. NVIDIA Triton support vllm-project/vllm#541 (comment)
   It requires client support for stream infer.
2. normal infer
   We will implement it using the method of BLS + stream infer. It will start multiple BLS instances and one model instance. It requires Triton Server >= 23.04.
   It has been discussed here before NVIDIA Triton support vllm-project/vllm#541 (comment)
   It does not require client support for stream inference, normal synchronous inference is sufficient.

After completing these tasks, we will provide examples for the Python client and Java client, as these two clients are more commonly used in the enterprise. At the same time, performance test data will be provided to ensure that it is comparable to the performance of the API server. Cheers. Stay tuned.

@lvhan028 @lzhangzz @AllentDan @grimoire @irexyc Do you have any suggestions? Thanks.

Related resources

No response

Additional context

No response

Hi all. As implemented in #1329, stream infer is currently supported. The original plan was to also support normal infer through BLS, but this requires Triton Server >= 23.04. Therefore, the priority of supporting this feature here is not very high at the moment. We plan to hold on for now and will consider supporting this feature when we have more time and resources available later on.

[Time-Limit]

what's the evaluation result?

--------------------------------------------------
concurrency: 128
elapsed_time: 146.570s

number of prompt tokens: 244330
number of completion tokens: 239421
RPS (request per second): 6.823 req/s
RPM (request per minute): 409.361 req/min
--------------------------------------------------

@zhyncs Can you provide TTFT and TPOT? Thank you very much!