High Scale Performance of LLMs -- Needed features?

Hey @mgwidmann! Many deployment considerations are built into Nx, specifically Nx.Serving. A couple thins it does for you:

• batching inference requests
• if you have multiple hosts in a connected cluster, the work is load balanced transparently
• if you have multiple GPUs on the given host, you can easily load balance the work between them too
• it supports streaming (for example, Bumblebee.Text.generation builds a serving instance and all you need to do is pass stream: true)
• encapsulates model pre/post processing (so it can be treated as an end-to-end pipeline in a way)

You can look at https://hexdocs.pm/bumblebee/llama.html for an example of using LLM with Bumblebee.

As for the other points:

• optimised attentions - we don’t have it yet, but that’s something we plan on supporting in Bumblebee
• sharding a single model across GPUs - for that we don’t have an abstraction yet
• quantization - there is an ongoing work in the EXLA compiler to use MLIR, which should enable quantization

Thanks a lot, the use case I was looking to compare with would be something like LLaMa2 70B sharded across 8 A100 GPUs since that model cannot fit within a single GPU, especially unquantized.