Parallelize weight loading for weights targeted at anonymous host ram and also GPU

[Farmadupe]

• I have read the contributing guidelines
• Self-reported review complexity:
  • Medium

Another PR to speed up model loading. Another PR in a "spend my pocket money on runpod to spend less pocket money on runpod" series.

Observation

I'm playing with kimi k2.5/6/7 on cheap runpod nodes. Startup time is dominated by the "loading weights into buffer" phase. When I look at htop, cpu usage is low (25% of a single core), and disk read rate is moderate (0.7G/s -- big metal on runpod is often good for nearly 10G/s )

So, I observe:

• Reading model weights into host ram is currently single threaded, and pingponging between IO and CPU.
• The task should be trivially parallelizable. My expectation is that sequential disk read should be the limiting factor for performance.
• Obvious answer is to add threads.

Implementation

Existing loader function can load weights into 4 backend buffer kinds, all serialized loading. I am targeting the "load to host ram" path and my intention is to leave the rest alone.

Backend	Description	Execution
mmap	host page cache	Serial
host	system RAM	Parallel 🎯
cuda	GPU	Serial
rest	All other backends	Serial

Implementation is to extract main body of the loading function into a lambda, spawn a fixed number of threads for loading. To prevent changes in of behavior for the other three buffer "kinds" which I do not intent to parellelize, I added a mutex so that only one thread can progress at a time.

Design considerations

• Why do you spawn a fixed number of threads even for serialized loading and then make them fight over a mutex?
  • This is extreme coarse grained work, so the contention will never dominate.
  • This implementation is actually the most reduced code churn compared to my first attempt. My first draft had a separated "serial path" and "parallel path", which was hard to handle.
• Why did you pick 8 threads? This number nearly saturated the disk read speed on the runpod I am testing (actually 16 threads is 25% faster, but 8 threads is good). I claim all SSDs (including enterprise and consumer) do not suffer significant throughput penalty for saturating with large sequential reads, so a fixed number of threads is not harmful.
• In the existing code, there's a note that says that a tensor may have no weights with a continue statement. This check is still present but it is moved to make_next_tensor.

Benchmarks

• CPU: AMD EPYC 7742 64-Core Processor
• Storage 8x Samsung Electronics Co Ltd NVMe SSD Controller PM173X
• Model: https://huggingface.co/AesSedai/Kimi-K2.7-Code-GGUF
• Amount of model weights loaded to host: 479.57 GiB

Sequential loading	Parallel loading
955s	409s

todo

• I tested on the big metal, but I will test on my home rig too. There should be a similar speedup
• Cuda path seems unchanged.
• I did not test the mmap path is not broken or regressed yet. I will do this.
• I should test on windows. This will take some willpower to set up a build but I will do this. I don't have mac hardware so I can't test that.
• I did not check the rest path. (is this for vulkan? sycl?) Anyway I declare that I did not break it because I don't know how to activate it.

Thoughts

• On the test rig, This patch accelerated host load speed from 0.7G/s to 5G/s average. The Cuda load path is doing disk reads at 0.7G/s average, so maybe a similar speedup is available.
• 16 threads is faster than 8 on the big metal rig. I am scared to propose 16 because it sounds like a very high number, so this patch has 8

[Farmadupe]

Additional commit with cuda implementation using the same threadpool

Backend	Description	Execution
mmap	host page cache	Serial
host	system RAM	Parallel 🎯
cuda	GPU	Parallel 🎯
rest	All other backends	Serial

Big metal benchmark

Same system now loads kimi k2.7 in 383s (This is ~479G weights in host ram and ~65 in VRAM, so I think still a noticeable speedup)

Sequential loading	Parallel (host-part-only)	Parallel (host and GPU)
955s	409s	383s

Desktop Benchmarks

on my home PC, I benchmarked loading time for qwen3VL-8B at BF16 (9950x + 3090)

Config	Serial	Parallel	Speedup
cpu-mmap	2.30	2.37	−3.0%
cpu-no-mmap	4.27	4.09	+4.2%
gpu-mmap	3.32	3.40	−2.4%
gpu-no-mmap	3.03	3.13	−3.3%

The result on desktop is totally inconclusive. It's the entire startup time for llama-server so lots of noise.

I guess this is a speedup for large models on large machines, and no significant change with models for normal people.

[Farmadupe]

Can't help sharing 46.9G/s loading speed :)

[ikawrakow]

How can I test this? I observe zero difference in loading speed. Is it because my Threadripper 3995WX is a "desktop" rather than a server CPU?

[Farmadupe]

This is my current theory. My initial premise is that this work timing is neither dominated by CPU work (syscalls, memcpy), nor disk IO, so I determined that more threads is an easy win. On the big metal, more threads scales near linearly until disk saturation is reached at some ungodly speed (I could have taken a screenshot of 80GiB/sec with 16 threads). However the best consumer SSDs at Pcie4 claim 7GiB/sec but more realistically somewhere between 3-7 Gib/sec. So on any model that normal people can load on normal hardware, the proportion of startup time dedicated to weights loading is lower when limited to single thread, and you are much close to saturating your disk anyway, so multithreading doesn't scale so much.

On my desktop with slightly higher memcpy throughput that your threadripper, but likely similar disk throughput, the speedup is insignificant:

Config	Serial	Parallel	Speedup
cpu-no-mmap	4.27	4.09	+4.2%

The benchmark I took was just measuring process startup to "model loading complete" message on stdout, so lots of room for noise.

possible confounds are:

• If your host RAM is signfiicantly larger than the destination buffer, then the weights will be resident in apge cache, and the benchamark will fail to measure host IO performance
• If your weights are stored on sata or spinning rust, you are already limited by disk throughput.

---

TLDR: If you do not have a huge model and a big raid 0, then expect no benefit. Otherwise, the cost is

• spawning 8 threads (negligable)
• Mutex contention on 8 threads, on a per tensor basis during loading (seems to be a 3% loss on qwen3vl-8B dense.).

---

edit: my memcpy is almost defintiely lower han yours 2channel DDR5 does not beat 8 channel DDR4.

[Farmadupe]

So my mental model for who wins is:

• Are you CPU limited on loading? (big metal with bad single core perf): You win with multithreading 🚀
• Do you have unused IO availability during loading? (bit metal with > 80GiB/sec sequential throughput): You win with multithreading 🚀
• Are you using desktop or workstation class hardware (mine and your home PCs with fast CPUs, but 3-7 GiB/sec sequential throughput and no 8-way parallel raid-0)? No difference, due to IO limiting ☹️, or your model is small and it loads quickly anyway 🤷

[Farmadupe]

Sorry for the spam, but other meaningful change is that before, each thread had 4x1M buffers for async PCiE loading (2 for double buffering, and I think 2 for good luck (the extra buffers make no sense as-is)). I replaced them with 1x16M buffer per thread. This decreases cuda synchronization, and slightly drops pcie overhead. I did not benchmark, but it should be a win.

Double buffering removed on the basis of "now there are threads".

I suspect original author allocated 4 buffers to take advantage of more than one DMA queue, but afaik GGML will only use one anyway.

[Farmadupe]

@ikawrakow Also parallelized loading for -sm graph. If that's something that you use 😉 then hopefully you might get the speedup now

[Farmadupe]

Unbounded memory use. Not mergeable.

[Nexesenex]

@Farmadupe , loading my GLM 4.6 on my rig (265K, 192GB DDR5, 2x3090, 1xRTXA4000, SSD PCIE4, 155GB or so to load) went from 110-120s to 75-80s.

Much appreciated, thank you, I was unable to do anything meaningful while playing with DIO implementations.

[Farmadupe]

I have a plan up my sleeve which should hopefully help with the fact that -sm graph can only eat whole tensors at a time (if you have 16 threads each trying to load a 3G tensor then that's a 50G allocation -- which you may have already allocated and pinned for host weights)

It's a bit brute force but the answer is just to.... ask the OS how much memory is available for allocation and then never exceed it.

[Nexesenex]

Funny that you speak about that. Being on Win11, 2 days ago, I had to make myself a custom pinmem implementation to get.. 54 GiB of pinned memory, with a tool to test the different allocations methods, and to test the WDDM and TCC drivers.
Doesn't help much for big MOEs, but for mid sized ones (100 GB), it's all good.

But the transfer rate to RAM is still lame for now, even with your PR (2.5 GB/sec to host for now, while I now reach 3.5-4 GB/sec to the graphic cards, thanks to your implementation which tripled the speed of the Cuda tensors offloading (!).