Quantization on multi-node GPUs about aqlm

[justheuristic]

Hi!

TL;DR alas, we do not support multi-node calibration. We support multi-gpu nodes, but not the other way around.

This is because our calibration is implemented with threading, not multiprocessing. We currently have no bandwidth for that, but if anyone intends to implement it, we'd gladly accept a pull request. However, we expect that this will need a lot of manual use of torch.distributed to work efficiently.
For instance, we rely on threading so that we can reuse shared memory, but that won't work multi-node, so it will need a different workaround. In other words, this may be tricky to get right.

If your model does not fit into a single node, the simplest workaround is to modify model loading so that it only loads one layer at a time. This is because main.py processes one transformer layer at a time, as can be seen in this for loop. Thus, if you load the appropriate layer (e.g. from disk) at the beginning of that loop, and deleting it once it after quantized layer has been saved, you can quantize very large models in due time.

If time is what you're after, we'd recommend running AQLM on a single node with higher mse tolerances (i.e. much faster), then running better/longer PV fine-tuning , which does support multi-node HW https://github.com/Vahe1994/AQLM/tree/pv-tuning?tab=readme-ov-file#3-refining-quantized-model-with-pv-tuning

[Devy99]

Thanks for the suggestion. I'll proceed as suggested then! 😄