Show HN: Zipslicer, a library for loading LLM checkpoints on consumer hardware (github.com)
This is a low-level opensource library I developed for my own use and decided to share, as it makes it possible to process large checkpoints of neural networks without renting high-RAM instances, on a regular PC. It replaces torch.load() with a custom function that produces a dictionary that materializes tensors on the fly. Compared to other solutions it doesn't require sharding or re-encoding checkpoints and uses them completely as-is.
It is a foundation to make it possible to run inference and compress language models and other large models one layer at a time - in principle, even one tensor at a time.
I describe the rationale and technical details of the library's design in the blogpost: https://kir-gadjello.github.io/posts/zipslicer/
8 comments
[ 3.8 ms ] story [ 30.5 ms ] threadTo support the inference for the model of your choice, someone should write code that uses zipslicer to instantiate layers of the model while threading activations through active layers - and it will make economic sense if you use this code with large batch size in long-term computation. There is no magic trick yet to compute interactive sampling (think ChatGPT) without fitting the whole model in, at least, RAM+VRAM, better in VRAM.
You could go slightly different route and use this non-interactive mode just one time to compress the weights down to int8, and this can, in turn, make the model runnable on your hardware in interactive mode without layer streaming from disk. Say, your model is 12B params and you have 16G RAM - it can work perfectly well if each weight takes 1 byte, but the weights are distributed in bf16/fp16 and your PC doesn't even support these so it instantiates the weights in fp32 - and you are out of RAM. Zipslicer can help here by allowing you to execute weight compression.
While this "offline execution engine" or "parameter compression engine" is slightly beyond the scope of this library, I can write it as an example or as a separate library, and describe this in a blogpost. For the purpose of making an example I would prefer to work with permissively licensed models such as this one https://huggingface.co/Salesforce/codegen-6B-nl - but someone could easily fork the repo and adapt the code later to whatever model they need.
It would be nice if people proposed desirable features in project's Github issues and I could see which features get votes.
Back of the envelope calculation for GPT-like models is simple: just divide the number of parameters by number of layers (ignore the embeddings and lm head size - you will overestimate layer size and add some safety margin that way) and multiply by size of your datatype. Also leave some RAM for activations and maybe KV-cache, but we will ignore this here. Say, for OPT-175B which has 96 layers at fp16 precision we have 2 * 175e9 / 96 ~= 3.65GB VRAM. It should fit into a 8GB GPU, I think.
If you want a precise way to calculate expected RAM consumption of your GPT-like model's layer, simply subtract the embedding and lm head size and divide the remainder by number of layers. Or you can instantiate one layer in python REPL and measure parameter count with this function https://stackoverflow.com/a/62508086
About performance. From my early behchmarks I see that:
1. You obviously need to store at least one layer (could be a dense sublayer - basically one matrix - and we could imagine a scheme to extend this library to load these dense layers as shards, but I don't think it's necessary right now) on the GPU or CPU+GPU at the precision of your model (most commonly bf16 or fp16, so 2 bytes per parameter).
2. In small batch regime, your inference performance is bottlenecked by disk read bandwidth (I could imagine it starts being bottlenecked by tensor materializing code for very fast SSDs - we might need a native extension here, but only after a good benchmark) - you can mask some, but not all of it by clever interleaving of tensor materialization and layer computation. As you grow batch size, you start being bottlenecked by computation throughput of your main computation engine (could be CPU, GPU or some exotic accelerator). At some point you can also become bottlenecked by memory bandwidth if your batch size is too small, but it is more about the case of running from RAM. Personally I don't yet see performance wins from stuffing less than all and more than one layer in RAM, but being able to fit all layers in RAM obviously makes a great difference.