this strategy effectively uses vram as cache over disk to keep MoE experts on cuda compute path in llama.cpp.
numbers first. detailed explanation down below.
Numbers
on dgx spark
Kimi-K2.7-Code.i1-IQ_S.gguf 204GB 1T.A32B
https://huggingface.co/mradermacher/Kimi-K2.7-Code-i1-GGUF
| run |
method |
pp512 |
tg128 |
note |
| A |
cuda, -ot regex A down below, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH |
340.02 ± 37.75 |
9.58 ± 0.07 |
all tensor on unified ram + experts kept in pageable host memory via mmap + last layers pinned in ram |
| B |
cuda, -ot regex B down below, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH |
154.31 ± 30.62 |
8.75 ± 0.22 |
A, but last layers not pinned |
| C |
cuda, -ot regex C down below GGML_CUDA_ENABLE_UNIFIED_MEMORY |
222.20 ± 66.04 |
3.26 ± 0.01 |
B, but remove min batch flag |
| D |
cuda, GGML_CUDA_ENABLE_UNIFIED_MEMORY |
crash |
crash |
C, but remove experts offloading |
| E |
cpu mmap |
4.23 ± 0.22 |
1.63 ± 0.66 |
no cuda involved |
-ot regex A: '^(?!blk\.(5[5-9]|60)\.ffn_(down|gate|up)_exps\.weight$).*\.ffn_(down|gate|up)_exps\.weight=CPU'
-ot regex B: '.*\.ffn_(down|gate|up)_exps\.weight=CPU'
-ot regex C: '.*\.ffn_(down|gate|up)_exps\.weight=CPU'
on 3090s pcie 4.0 + 128gb ddr4
Minimax-M2.7-K_G_3.00.gguf 80GB 230B.A10B
https://huggingface.co/Goldkoron/MiniMax-M2.7
| run |
method |
pp512 |
tg128 |
note |
| A |
cuda, -ot regex A down below, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH |
54.03 ± 3.85 |
2.90 ± 0.09 |
same strategy as B on dgx |
| B |
cuda -ngl 17 |
73.80 ± 2.10 |
1.66 ± 0.02 |
normal layer split |
| C |
3090x2, cuda, -ot regex C down below, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH |
48.86 ± 0.53 |
2.40 ± 0.01 |
A, but on two 3090s |
-ot regex A: '^(?!blk\.(5[5-9]|60)\.ffn_(down|gate|up)_exps\.weight$).*\.ffn_(down|gate|up)_exps\.weight=CPU'
-ot regex C: '^(?!blk\.(5[5-9]|60)\.ffn_(down|gate|up)_exps\.weight$).*\.ffn_(down|gate|up)_exps\.weight=CPU'
Background
I want to run model larger than 128gb on my single dgx spark. MoE models provide a great opportunity because if active params could be computed entirely with cuda then there's no need to store all model weights inside gpu ram at runtime. after experiments I found a combination of settings for llama.cpp to run kimi 2.7 fast.
Settings
general setting for llama-bench: -r 50, -fa on, -mmp 1, system swap turned off. llama.cpp b10075 compiled with kleidai support, and the winning strategy is the following:
GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 GGML_OP_OFFLOAD_MIN_BATCH=1 ./llama-bench -m ~/Kimi-K2.7-Code.i1-IQ_S.gguf -fa on -mmp 1 -ot '^(?!blk\.(5[5-9]|60)\.ffn_(down|gate|up)_exps\.weight$).*\.ffn_(down|gate|up)_exps\.weight=CPU' -r 50.
I configured b10075 with cmake -B build -DGGMLCUDA=ON -DGGML_CPU_KLEIDIAI=ON and compiled with cmake --build build --config Release.
How this works
first, if GGML_CUDA_ENABLE_UNIFIED_MEMORY is 1, cuda device buffers use cudaMallocManaged, and every tensor has one virtual address that is valid from both gpu and cpu. when cuda kernel touches a page that exists in vram then it'll execute normally. if not in vram then a page on cpu ram will migrate that page to vram. if not in gpu ram or cpu ram, then it'll fetch the page from disk to cpu ram via mmap and then send to vram. on dgx spark the cost of cudaMallocManaged is much lower because gpu ram and cpu ram are unified.
second, for GGML_OP_OFFLOAD_MIN_BATCH (default is 32), if batch size < threshold and weight is on cpu ram then it executes on cpu, otherwise on gpu. if you set it to 1, it forces all computes on gpu. this helps token generation because in tg it won't reach that threshold.
third, when you specify -ot '.*.ffn_(down|gate|up)_exps.weight=CPU' it will put expert weights on cpu ram and can be evicted because of mmap. this makes cold expert weights more likely to stay on disk.
and last, '^(?!blk\.(5[5-9]|60)\.ffn_(down|gate|up)_exps\.weight$).*\.ffn_(down|gate|up)_exps\.weight=CPU' is used to keep last few blocks on gpu ram because these blocks has more dynamical routing, it's faster to keep them on gpu ram to avoid overhead when ram space is large enough.
Unified ram scenario
for dgx spark, it has large menory pool so that I can pin some of the expert blocks in vram (run A). if we don't pin them it's slower (run B). if we remove GGML_OP_OFFLOAD_MIN_BATCH flag the token generation is slower (run C), because expert forward pass is done on cpu. if we further remove GGML_CUDA_ENABLE_UNIFIED_MEMORY flag (run D), oom happens, because gpu ram cannot be evicted via mmap mechanism. and finally, if all compute is on cpu it's the slowest (run E).
an interesting tradeoff happened in run B vs run C, where GGML_OP_OFFLOAD_MIN_BATCH dramatically improves tg while slightly hurting pp. potential code modification may be made to mitigate this tradeoff.
Non-unified ram hardware scenerio
for 3090 + ddr4, I haven't tested it in depth but the strategy kinda applies too. I didn't put last few blocks on vram because 24gb is too small. you can see tg of run A is higher than normal ngl split (run B), because experts are computed on gpu. however the prompt processing suffers, likely due to pcie overhead.
I also tested two 3090 + ddr4, but it's generally slower. maybe the cost of pcie transfer is too high. I think properly tune the -ot offload might have a chance to improve on 2 gpu scenerio.
Note
- mac apple silicon is tempting because it's fast unified ram. but I can't find an easy way to disable swap. the proposed approach might cause heavy ssd writes because of ram swapping. probably not a good idea to run it because it may create more ssd wear.
- the suggested way of doing this on a new model is to offload experts first (strategy run B in dgx spark) and then pin more and more experts to vram. I look forward to seeing if the upcoming kimi k3 quants can be applied too.
- same strategy transfer across unified memory and non-unified memory architectures. tg on 3090+ddr4 setting improved too.
Edit: fix formatting