Revisiting AirLLM and now RabbitLLM functionality... #20045
Replies: 3 comments 2 replies
|
Today Qwen announced their small 3.5 models. A quick look suggests that the 9B parameter model will be capable of some great things but my own small 6GB GPU wont run it, and while it will run the 4GB model that is substantially less capable. But if we could run the 9B model layer by layer, it would fit and be very capable. This functionality really mighty make a difference IMO. |
|
I am interested in working on this. Can we have some irc/discord chat to discuss implementing this? |
|
I have posted some additional ideas on the RabbitLLM repo, because it has become clear that swapping out layer by layer is way slower than CPU inferencing or mixed GPU / CPU inferencing. So I have suggested a much less ambitious idea - i.e. for models where most of the layers fit into vRAM, instead of having a handful of layers inferencing on CPU and having to store the context and KV caches in shared RAM (which is way slower), instead we keep the context and KV caches entirely in vRAM and arrange to e.g. swap out every 5th layer and keep the rest in memory - so we run e.g. permanent layers 1-4, and then a swapped in layer 5, then permanent layers 6-9 and whilst these are executing we swap in layer 10, and hopefully it is ready to run by the point layers 6-9 have executed - and because we are also talking smaller layer sizes the time to swap in a layer will be less. The idea is that if we get the proportions right, we can keep the GPU executing and achieve faster execution than running a few layers using CPU inference. RabbitLLM is therefore going to start off being an experimental platform looking at the metrics and trying to calculate the sweet spot. |
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
References (issues & discussions, numeric order)
Concept
Break a large model into layers and load and run layers individually so that you are no longer limited by fitting the model and context into memory, but instead by fitting largest layer and context into memory.
Previous experiences with AirLLM
The primary use case is to load the model layer by layer into GPU in order to get GPU speeds, but you only get these whilst the layer is executing and it stops when waiting for the next layer to be loaded.
Experience of AirLLM a few years ago appears to have concluded that even whilst starting to pre-fetch the next layer into memory whilst the current layer is executing, you end up waiting substantially for the next layer to be read from disk into normal memory and then into GPU.
Additionally, because newer models attempt to improve quality and speed using novel techniques (such as dual attention), some models run fine, some need special handling and some probably won't run at all.
This may be why AirLLM development ground to a halt. Also, AirLLM does not provide an HTTP server with OpenAI (or Anthropic) API, so adoption was likely to be limited.
RabbitLLM is a recent (c. 1 week ago) hard fork and update of AirLLM. But still no HTTP server. Aside from some code restructuring, the main improvement is to support Qwen3 (but not yet 3.5). OpenAI API is still not supported.
Use Cases
I believe that there are actually several use cases for this general technological approach not just the single AirLLM use case:
The AirLLM use case - Model too big for vRAM and GPU inference - instead of loading a fixed selection of layers into GPU with the rest executing on CPU, swap layers into GPU so that they all execute there.
Model too big for even RAM and CPU inference - the same layer by layer technique but applied to CPU inference in main memory.
I also believe that the following technological changes over the past few years since AirLLM was tried will also result in there being a bigger "market" for this functionality:
Emergence of unified memory systems (e.g. Apple silicon, AMD Ryzen AI) with inference built into the CPU - there is no reason why a technique like this wouldn't enable running e.g. Kimi 2.5 with a significantly smaller memory providing that reasonable TPS can be achieved.
Widespread adoption of NVMe SSDs which are substantially faster than SATA / SAS SSDs - bandwidth for loading layers is now much better
Models keep growing and growing - so hardware that would run Qwen2.5 won't run Qwen 3, and the quality gains from a significantly better (but b=significantly bigger) model might still make a lower TPS worthwhile.
Manufacturing supply constraints have driven up hardware prices substantially in recent months and there is no sign of this abating. Techniques like this can help reduce the demand and allow existing hardware to continue to be used despite larger models.
Performance techniques
It may also be possible to improve performance through other techniques:
An improved pipeline - have the currently executing layer and the next to execute layer both in GPU - needs layers to be < 1/2 size of GPU memory, and still only really useful if you are not constrained by bandwidth from disk to normal memory.
Holding all layers in normal memory and not having to read them from disk. If you have (say) 32GB of vRAM and 192GB of main memory, you could easily hold a 160GB model entirely in main memory and swap out the GPU layers quickly enough to keep the GPU running all of the time (2 layers in vRAM) or most of the time (1 layer in VRAM). I do not believe that this use case has ever been evaluated for its performance characteristics.
Even if you only have enough memory for the majority of layers but not all, you could still get pretty good performance if you held most layers permanently in memory and read in a few layers from disk when they were needed.
As an example, suppose that you know that layers execute strictly in sequence, and you have enough memory to hold 16 of 20 layers in memory. If you held 1-4, 6-9, 11-14 and 16-19 permanently in memory (i.e. 15 of the 20 layers), then you could start reading layer 5 as soon as layer 20 starts to execute, and you would have until layer 4 finishes to have read it from disk into memory i.e. 5x as long as AirLLM would have allowed.
Similarly, if you can already hold (say) 16 of 20 layers in vRAM, then you could put the samed numbered layers into vRAM and hold the other 5 layers in RAM to swap in as needed. I suspect that this use case could be quite common, and performance might be pretty good.
AirLLM seems to have used Weight quantization to reduce the layer sizes - I can foresee the community experimenting and finding ways to reduce the size of models and layers in a way that specifically takes advantage of these capabilities whilst minimising issues and loss of quality.
I have no idea what impact MoE models would have on this approach - I am simply lacking in anywhere near enough knowledge even to make a guess.
Technical constraints
TL;DR Summary
I believe that the potential benefits of this type of approach are worth revisiting, taking into account all of the above points.
But it is still possible that performance won't be good enough in any use case for it to be worthwhile, or that performance gains will be limited to small use cases that make it not worth the effort. Or that it is simply too difficult to integrate it with all the existing functionality.
All reactions