Loading models with mmap

[polyrand]

Hi, first, thank you so much for the wonderful work on mlx!

Are there any plans to support memory-mapping models from disk?

When I try to load a model that's larger than the memory I have available, it uses swap memory anyway, so maybe allowing to mmap the model directly is more convenient and allows for faster loading (i.e: faster script startup and faster time-to-first-token)? It would also make it more efficient to run multiple processes that read the same model file?

I'm putting those as questions because I'm not super familiar with the unified memory model of the new Macs and how it interacts with the disk / GPU memory / etc. So maybe what I'm saying doesn't make a lot of sense.

[awni]

I think it's quite possible to have a mmap mode, but we would still need to allocate memory to use the weights once they are needed and if the model is too big then it will still swap. There is no easy way to avoid that, as far as I know we can't mmap memory in a way that makes it available for the GPU as well.

So right now we have lazy loading from a file path which may effectively achieve the same goal: We don't mmap the model instead just read the weights into memory at the time they are needed.

[antmikinka]

Would a 8GB MacBook Pro be able to run a 4bit 7.79 GB MLX model? such as mlx-community/CodeLlama-13b-Instruct-hf-4bit-MLX. My understanding is that the weights are read in and correctly de allocated as the model is processing the prompt. This allows for something like above to be possible?

I first installed MLX and built when MLX was V.0.0.11, the last version. I seen in the new release, some memory issues were fixed. Would I just need to update my 'pip install mlx' or would I need to rebuild the entire c++ and python bindings for this new update?

Is there anything else I can do to run decent LLMs on my 8GB 2020 MacBook Pro M1, or is it time to upgrade?

I am trying to run the above model as well rn, it's using python bindings, so thats why it might be so slow, I am not sure. Only 10% - 15% of my CPU is being used, my RAM is staying steady at 3.2GB. Below is a photo of my ASITOP computer metrics.

Any help is appreciated.

[awni]

That model is too big to run on 8GB. You're machine needs some RAM for other things so it's not going to work to have just the model taking up all the RAM, it will swap and be extremely slow.

If you want to run that4-bit quantized 13B model you will need more RAM (16GB should be ok). Otherwise I would try a smaller model like the StableLM 1.6 B https://huggingface.co/stabilityai/stablelm-2-zephyr-1_6b or a 4-bit quantized 7B model

[polyrand]

Thanks for the response.

as far as I know we can't mmap memory in a way that makes it available for the GPU as well.

Regardless of the original question, this is good to know!

So right now we have lazy loading from a file path which may effectively achieve the same goal: We don't mmap the model instead just read the weights into memory at the time they are needed.

Is this done by default when loading a model? i.e: If I do:

model, tokenizer = load("mlx-community/Nous-Hermes-2-Mixtral-8x7B-DPO-4bit")
response = generate(model, tokenizer, prompt="hello", verbose=True)

it's already lazy-loading the weights from the file on-demand?

[awni]

it's already lazy-loading the weights from the file on-demand?

Unless you explicitly eval the arrays if you load from a string file then it should be lazy. In the case you point to specifically it is in fact evaluating the arrays so it's not lazy https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/utils.py#L262

[antbob]

tl;dr is its certainly possible but not very practical. i've made a prototype that enhances mlx metal allocator to use mmap and create metal buffers via newBufferWithBytesNoCopy via separate api which is then used when loading from safetensors (as their format is straightforward enough to do that) and using that for inference streaming.

the first problem is that you cant really reconcile mmap offsets constraints with metal buffer offset constraints if you want to maintain a zero copy approach as the former can provide vm page aligned offsets and byte offsets from those once mapped and the later has to use vm page aligned base and offsets which are data type aligned. so you either have to copy or convert safetensor files to match. i ended up creating temp files (one for each tensor) and mapping them to metal buffers without offsets as this was the most straightforward workaround.

the second problem is that relying on the kernel page cache eviction algorithms here is not a good fit for this. yes, you can run larger models but since you cannot control what is swapped out and when it turns out to be very inefficient as once you throw i/o in to the mix in a wrong place at the wrong time it really nose dives the performance. on my tests for stream inference mlx-community/Llama-3.3-70B-Instruct-8bit (which is about 70Gb model) on 64Gb m4 mini does about 0.025 tokens-per-sec (in comparison 4bit variant that fits does about 6 tokens-per-sec) which is not practical at all, even more so for even larger models.

while i reckon (theoretically) it should be possible to optimize this further for specific models and graphs (with the knowledge of their specifics), by mlocking/unlocking specific buffers at specific points of the computation and madvise on the rest according to priority ie by pinning as much as possible up to a barrier or some similar approach where you dont have to rely on kernel page cache eviction strategy, it is unlikely to bring it into some practical tps territory, especially for larger models and longer contexts, tho it might be a fun exercise to do.

anyways, hopefully this is useful to the next person looking into this and would save you some time digging.

[lennartvoelz]

Could you share that work?

[antbob]

i've just pushed https://github.com/antbob/mlx/tree/mmap-prototype this is what i could find. i think that was the latest revision at the time of my comment above. its been a minute so i'm not sure if that snapshot even compiles at this point so dunno if this would be useful to you at all but hopefully it is somehow. i think with faster, more modern h/w and maybe a sparse model this can get into somewhat usable tps territory if optimized for specific model which i've never had the time to look into and try. good luck!

[lennartvoelz]

Cool ty, might be useful!