Add --prefetch-experts to stream mmap'd MoE experts into page cache#2101
Conversation
| // register with the MoE expert prefetch engine, which needs (fd, offset) | ||
| // to stream expert weights into the page cache with pread() workers. | ||
| // The mapping starts at file offset 0, so VA -> offset is direct. | ||
| prefetch_fd = dup(fd); |
There was a problem hiding this comment.
You want to do this unconditionally? Even when not using the prefetch option?
There was a problem hiding this comment.
I don't think it affected anything when not using the prefetch option, but regardless, it's not the right spot for it so I'll move this logic elsewhere.
|
I don't have enough disk space left to download a model that is larger than the RAM I have. Hence, would appreciate if at least one more user tested and confirmed that it works for them. |
|
Notes:
box/proc/cpuinfo: lspci | grep -i nvme: Host ram: Disk benchmark:
GPU stats (note, I rented three GPUs by accident. I only enabled visibility of one for the test.) cmdline: Bench (with --defer-experts accidentally disabled)with patchwithout patchtldr: it works, even with prefilled host ram |
|
I did not repeat with |
| } else { | ||
| size_t done = 0; | ||
| while (done < j.len) { | ||
| ssize_t n = pread(j.fd, buf.data(), std::min<size_t>(j.len - done, buf.size()), (off_t)(j.off + done)); |
There was a problem hiding this comment.
would madvise(MADV_POPULATE_READ) work here? I think it would avoid copying out of the page cache info buf, only for buf to be immediately dropped?
There was a problem hiding this comment.
I did some testing and MADV_POPULATE_READ does work here; I think its a better call over using pread and we don't have to dance around the descriptors
|
shower thought (and fantasy request). I think this PR does the following host2host memcopies:
This can be reduced to two with MADV_POPULATE_WRITE:
The fantasy request is: This PR contains a generic prefetch engine that already teleports weights into the application ahead of usage. With that capability in mind, is it foreseeable to be able to prefetch directly into pinned memory, and then from there, directly DMA the expert straight into the GPU? that would reduce the host2host memcopies to 0,as weights would go straight from NVME into pinned memory, and straight from pinned memory into GPU. asking because I observed on the rented box that:
So with all those observations in mind, I wonder if it's possible to imagine an extension to this PR where the weights are prefetched into a vector of pinned buffers (eg 4 buffers * sizeof(expert_tensor)), and from there immediately DMAd into the GPU. |
|
Thanks for testing; I'll have to do some more experimenting to see how much further I can push this beyond swapping out |
|
Don't let my silly fantasies stop get in the way of what you want to merge, doubling performance for disk streaming prefill is already an achievement! |
|
Ah, no worries. I have been doing some experiments regarding IO, but I don't think it's the right move for this PR. I'm not seeing too much more improvement, at least for my hardware, so I think it would just be adding unnecessary complexity. However, I do realize that batch and ubatch sizes are also helping greatly (I think with models closer to the RAM capacity). With Step 3.7 Flash I'm seeing PP increase from ~284 t/s at -b/-ub 4096 to ~419 t/s at 8192 and ~788 t/s at 16384. |
|
Overall LGTM, but I would like to propose that we do not expose the As it stands, The usage in The usage in
With that the code in if (params.prefetch_experts) {
if (ggml_backend_prefetch_init()) {
LLAMA_LOG_INFO("%s: enabling MoE expert read-ahead (prefetch_experts) via threaded populate engine\n", __func__);
} else {
LLAMA_LOG_INFO("%s: failed to initialize MoE expert read-ahead -> using madvise fallback\n");
}
for (const auto & mapping : model->mappings) {
ggml_backend_prefetch_register_mapping(mapping->addr(), mapping->size());
}
}With that, the As it stands, there is a singleton |
|
Moved the header into ggml/src and added the ggml_backend_prefetch_* wrappers, so llama.cpp/llama-mmap.cpp only go through ggml-backend.h now. I also made threads an argument ( |
If this were possible, then running a quad x4 nvme pcie card should be as fast as the link to the gpu, so maybe it's a way to both gain speed and use an awful lot less RAM? Predictive expert prefetch would definitely then become the new hotness... |
|
This is 100% my fantasy! Although it totally depends on your motherboard's exact allocation of PCIE lanes and speeds, my home desktop board (relatively cheap AM5 board, but not uncommon) can't actually saturate its GPU slot with bandwidth direct from disk (I believe there are desktop class boards that are less stingy where this is possible), and totally doable with worksation and server class hardware. 
And even with an awful lot less RAM, whatever you do have can still be filled to the brim with weights to take pressure off the disk streaming setup. This is obviously great for prefill (with a moderate GPU and huge batch sizes, GLM5.2 would prefill at hundreds of tokens per second 🤤), but as with any weight streaming, doesn't really move the needle at all for decode. |
I am on older kit, so limited to pcie3 :-( but I can get full 4 x4 NVME and have 512GB RAM but only DDR4 - it's my dream too to get larger models running on my v100 32gb or 4080, but I think it really needs that predictive expert prefetch to work well. But it seems to me if younhave enough PCIE lanes, large amounts of RAM might not be necessary at all, I mean for GPU you're always capped at what you can send over pcie anyway, yes? |
|
For decode/token-generation, yes. WIth a dense model, you're forced to stream the entire model for every token, but for MoE, where you might only need to load 2-5% of the model per token, it becomes really hard to predict which parts you'd need next. And you you ahve to fetch a different set of experts at every layer of every tokens. It's really hard! But for prefill, the trick is to set the batch size insanely high (16k, 32k, whatever), as this amortizes the amount of ewight transfer that you actually need to do. This is where this PR hits the sweet spot -- the faster you can get weights into your GPU, the less ridiculous your batch size needs to be to get good prefill performance! |
4 participants
This PR introduces
--prefetch-experts; it allows for ~2x gains in PP performance for those running models that do not fully fit in RAM (as such, get paged in from storage).When expert weights live in a host mmap (e.g.
-ot exps=CPUwith a model larger than RAM), the consumers of those weights (the CPU MoE matmuls at decode, the scheduler's host-to-device expert copies for offloaded batches) stall on page faults, as they are serviced one miss at a time. This PR adds a small pool ofpread()workers that warm the page cache ahead of use. This gives us two behaviors:icomputes, the next few MoE-bearing splits' expert tensors are streamed in full (batch graphs activate most experts anyway). Freshly streamed pages areMADV_COLD-ed after use so this one-shot traffic is reclaimed ahead of the resident working set.Also added env vars for tuning, though the current defaults gave me the best performance:
GGML_MOE_PREFETCH_THREADS: default min(8, hw)GGML_MOE_PREFETCH_AHEAD: splits to stream ahead, default 3GGML_MOE_PREFETCH_DEBUG: print counters at shutdownI tested on a system that has i9-14900F, RTX 4080 SUPER, 64 GB RAM, NVMe (kyber scheduler). As for models, I tested these two:
Step 3.7 Flash just slightly spills out of memory, while GLM 5.2 is double my RAM capacity.
Command:
All numbers are cold-start. The prompts I used were just varied natural-language text.
Step-3.7-Flash UD-IQ3_XXS @ ~73.6 GB:
GLM-5.2 1.630bpw @ ~143 GB:
Changes in TG speeds are likely noise.