[Cieric]

I was actually just working on the same thing as this, but I went down the route of mmapping the entire model into memory to avoid the extra ram usage. I also had Claude implement Medusa[1] on the model to try and avoid loading an additional model into memory but still get the benefits of MTP. Currently at a stop light so I can't list everything and I didn't get to read your full post either yet.

To expand since I just got home, I'm making all of my modifications to llama.cpp, the goal was to eventually put this on a SBC of some kind with an nvme to handle the mmapped files. I think the theoretical limit of my current setup is about 1.8 tok/s based on prior testing but that is also with the additional medusa heads not fully trained (I honestly don't know if the counting it's generated tokens or not.)

In the end it seems like the idea we had is similar, I just don't know how to write an llm parser/runner from scratch yet and instead of specifying what needed to stay in memory I just let the linux kernel handle it.

Oh last note, I also capped llama.cpp usage to 16GB of my 32GB, so it might be possible to get it down even lower.

[1] https://arxiv.org/abs/2401.10774

[Roxxik]

Not sure if mmapping is the right way. In my own tests I noticed that simple mmapping will produce many small reads and not keep the SSD queue saturated. So if RAM is large enough to cache most experts, that is a rounding error. But if the base weigths without experts fill more than half or RAM and you basically need to load in a few experts for each layer of each token, the latency gets important and mmap sadly blocks until the data is loaded. You can't do concurrent requests for multiple experts with mmap (but you know all the ones you need right after the router ran). And even going one step further, depending on the arch / the tensors you need, you could eagerly load some, start computing with them and load the rest of the expert tensors in the background (extra thread or async io) parallel to the compute. This is not really possible with mmap, even with madvise.

One further step is predicting which experts will be needed next token / next layer. LRU does this okish. But a learned projection from the hidden state can do better. Or even a simple correlation from past activated experts. Expert usage is heavily skewed.

[Cieric]

On the technical details of mmap I agree (at least while single threaded which is how I believe I'm running it), but making it async does sound like an interesting method for speeding it up. My only goal with my testing was get as large of a model as possible running on a computer that "can't run it." I'm going to have to actually read through the code and figure out how it really works to really make any good optimizations, but as before this was just experiments with using and LLM (claude + codex) to just get it running. Since if they couldn't get it running I'm not sure if I would have wanted to spend time trying to get it working myself.

I also know I did some things that would actually make the perf worse to, like I believe I also had AI mmap the KV Cache to make sure to runs under any circumstance. For actual optimizations based on what I currently know, I'm probably going to try and get the llm running under my igpu on my laptop with persistent shader that has some kind of inbuilt request mechanism. That way the weights that are loaded can be used as fast as possible.

For the expert prediction, I assume I could use the medusa paper as kind of a kick off point for that since I'm already using it to try and predict the next 4 tokens. Doing verification on those 4 tokens is about as much as I can do though since it started to thrash on loading the experts. So some method of predicting even more tokens, but then batching together those with the same experts would probably yield slightly better results in this weird case.

Note: All of my tests have been around programming since that's the use case I'm interested in. I don't actually know if this would preform well in other cases (and anything more broad than that I assume would be slower.)

[Roxxik]

My recommendation if you are willing to push it further is to keep KV and always needed weights pinned in memory, so they do not get evicted. This is likely already the case, as they are touched on each token. mmap is slow on evicted pages, as it does not load the whole tensor, but only the touched pages. And it does this through a page fault, thus blocking your code. So it loads a page hands control back to you and the code goes on to touch another evicted page, repeating the loop. Now on an HDD that is not a big problem (yes reads can be coalesced, but a HDD is fundamentally serial in reading) while an SSD can overlay reads better and it is good to keep a few reads in flight at all times to keep its queue fed.

One option is locking the pages. But for that size you need extra privileges.

I experimented with some options. For example: one problem with io_uring is that it still reads to page cache so your reads gets copied in memory after they landed. Now if you pass O_DIRECT that does not happen, but it has its own can of worms.

For full transparency: I had opus write the io_uring layer into llama.cpp for me. And it yielded something slightly short of a 2x tok/s speedup vs simple mmaping. Also I noticed that disabling the warmup and initial test dramatically increases startup time.

The medusa paper looks interesting. My work was a few months ago, multi-token decoding was not a thing then.

[adrian_b]

I agree that the maximum SSD performance can be achieved with carefully planned asynchronous reads, using either the modern liburing (io_uring) or the older libaio (Linux asynchronous I/O).

Nevertheless, I do not think that "this is not really possible with mmap, even with madvise".

If the kernel is not eager enough to prefetch pages from the SSD when you use madvise with MADV_WILLNEED, then you can use madvise with MADV_POPULATE_READ, which should force the reading of the pages that you request.

Using madvise with MADV_POPULATE_READ for appropriate page ranges at the right moments should be able to provide a performance not much lower than when using explicit asynchronous reads.

[BigTTYGothGF]

> Currently at a stop light so I can't list everything and I didn't get to read your full post either yet.

Posting on the Hacker News forum is not worth risking your, or anybody else's, life.

[Cieric]

Being fully stopped at a stop light isn't putting my or anyone else's life in danger. If I have a phone in my hand the car is not moving. I also finished the comment before the light turned green so the phone also was no longer in my hand when I needed to start moving.

[embedding-shape]

Paying attention to your surroundings tend to be a good idea always in traffic, even if you're standing still. More than once I've beeped at someone who didn't notice a red light turning green, and their first reaction is to accelerate, then look up. I'm guessing if they were actually still paying attention even though they were stopped, they would have noticed better what's going on around them.

[vforno]

if you like, colibrì always needs to improve so if you have ideas or anything else you are welcome for pull request issues and also benchmarks!

[Cieric]

Yeah I'll see what I can transfer over from my llama.cpp work. As before I'm not too experienced with llm work, but I have a lot of experiments I'm trying out. So I'll make a PR if I get any interesting results.

[dicroce]

This is the approach I was wondering about.

[Cieric]

Let me know if you want to hear anything specific about it. It kind of works, so it's not something I recommend doing if it can be avoided, but as Roxxik pointed out there is much room for improvement since this was just a naive just get it to run experiment.