Questions about serving PyTorch LLM in Python backend with token streaming using "Decoupled Mode"

[jackylu0124]

I am planning to use Triton's Python backend to serve a LLM model in Pytorch; and more specifically, I want to implement token streaming and hence based on the suggestions I read here #5913, I think I am going to use the "Decoupled Mode" to achieve this. I went over the square_model.pyhttps://github.com/triton-inference-server/python_backend/blob/main/examples/decoupled/square_model.py example and have some questions about the "Decoupled Mode" and also the Python backend in general, and I would really appreciate your help and insights! Thanks for your time in advance!

Questions:

1. For the Python backend in general, if I use multiple model instances (e.g. 4 in the example config below), since the documentation (https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/python_backend/README.html#multiple-model-instance-support) mentions that the multi-model instance support for the Python backend is achieved via multi-processing instead of multi-threading, does that mean the model's weights would be loaded 4 times on the GPU, and each instance of the model's weights on the GPU would correspond to one process? For example, if each model is 4 GB, then with the following multiple model instance config, the server will take up 16 GB of the GPU memory in total?

  instance_group [
    {
      count: 4
      kind: KIND_GPU
    }
  ]

2. For the following config where I have max_batch_size: 8, and count: 4 in the instance group config, does it mean the max batch size for each one of the model instances is 8 (e.g. batch size of 8 for instance 0, batch size of 8 for instance 1, batch size of 8 for instance 2, batch size of 8 for instance 3), or does it mean the max batch size for all 4 model instances in combination is 8 (e.g. batch size of 2 for instance 0, batch size of 3 for instance 1, batch size of 1 for instance 2, batch size of 2 for instance 3)?

  max_batch_size: 8
  instance_group [
    {
      count: 4
      kind: KIND_GPU
    }
  ]

3. For the square_model.pyhttps://github.com/triton-inference-server/python_backend/blob/main/examples/decoupled/square_model.py example in particular, I see that we are spawning a new thread to handle each request in the batch of requests; my question is: is threading in decoupled mode required or is it meant for demonstration purposes only? In other words, is it valid to handle each request in the batch of requests in a loop sequentially without using threading? This is because I am not sure if the LLM PyTorch model that I am using is thread-safe during inference (I could put a lock for each thread's token generation, but I am not sure if that would impact the server performance).
4. As a follow-up to question 3, if I do decide to use multi-threading to handle each request in the batch of requests with a Pytorch model that's set as a member variable of the TritonPythonModel class, and assuming I only use 1 model instance, my question is: if I use the Pytorch model to perform inference on each thread, the model's weights are only loaded once on the GPU right? In other words, if my model is 4 GB and I have 2 threads serving a batch of requests of size 2, the model is still only loaded once and hence only takes up 4 GB of the GPU memory right?
5. For the square_model.pyhttps://github.com/triton-inference-server/python_backend/blob/main/examples/decoupled/square_model.py example in particular, I see that we are setting the spawned thread to daemon thread with the line thread.daemon = True. Is this line always necessary/required because I also read in that file's comments that In real-world models, the developer should be mindful of when to return from execute and be willing to accept next request batch.? So if I only want the server to accept the next request batch after all requests in the current batch have been completed, should I use join() to wait for the spawned threads to complete before calling return None in execute()?
6. In general, if my model's weights are 4 GB in size, and my GPU has 16 GB of memory, which of the following designs do you recommend based on your expertise and past experiences?
   1. Use 4 instances of the model (e.g. I think in practice maybe 3 is more realistic since not every last bit of GPU memory can be used), and using 1 thread to sequentially serve the 2 requests in each batch (of size 2).
   2. Use 1 instance of the model, and set the request batch size to a high number like 8, and spawn 8 threads to serve each request concurrently.
   3. Use a mix and match of the first and second approach, for example, use 2 instances of the model, and then set the batch size to 4, and in each of the 2 processes, there will be 4 threads serving the batch of 4 requests concurrently.

[Tabrizian]

Thanks for your questions!

1. Correct.

2. The max batch size for each model instance will be 8.

3. It is for demonstration purposes only. It is mainly meant to show that you can send responses after returning from execute function. It could be a workaround to the memory usage issue you brought up in (1) (i.e. you can have a single model instance with multiple threads inside to avoid extra memory usage).

4. That's correct.

5. The warning is mainly about not accepting too many requests and spawning a lot of threads (i.e. having a limit on the number of requests that you want to accept). I think it is not required to set the daemon flag. You don't need to join the threads before returning None. As shown in the example, the threads continue to send responses after returning None from the execute function.

6. I think it is hard to have a general recommendation for your scenario since it is dependent on a lot of parameters. It would be good to experiment with these different options to see which one performs best for your model. Would love to learn more which route worked best for you.

Let us know if you have any additional questions!

[jackylu0124]

Hi @Tabrizian , thank you very much for your detailed reply and insights! I have a few follow-up questions:

1. Regarding quetsion 5, if I don't join the threads before returning None in execute(), then in order to prevent too many requests from being accepted and to prevent too many threads from being spawned, I would then have to keep track of the number of requests that are still on-going/the number of threads that are in still in flight, is that correct? Is the self.inflight_thread_count member variable in the square_model.py file (https://github.com/triton-inference-server/python_backend/blob/main/examples/decoupled/square_model.py) a way to keep track of this?
2. In general, regarding the amount of memory a model needs during inference, the amount of memory inference takes is not just the size of the weights but also the size of the intermediate layer tensor output/activation right? In other words, do I have to leave some additional space outside of the raw weights memory size for the intermediate layer tensor output/activation during inference?

Thanks a lot for your time and help again!

[Tabrizian]

1. correct. You can also probably use semaphores so that you block until you have acquired enough resources: https://docs.python.org/3/library/threading.html#semaphore-objects

2. Yes, it would be good to run your model thorough model analyzer to better understand the memory requirements.

[jackylu0124]

1. correct. You can also probably use semaphores so that you block until you have acquired enough resources: https://docs.python.org/3/library/threading.html#semaphore-objects
2. Yes, it would be good to run your model thorough model analyzer to better understand the memory requirements.

I see, thank you very much for your answers and insights!