Huggingface model quantization has odd behavior with BitsandbyesPrecisionPlugin

[ElleLeonne]

Huggingface model quantization has odd behavior with BitsandbyesPrecisionPlugin

[ElleLeonne]

Of note, if I bypass this and handle everything manually in setup, it appears that my GPU doesn't actually take advantage of the quantization benefits. I am told my quantized model is 2gb, the trainer says 6gb, but when I go to train, my 16gb GPU overflows immediately. It even happens on my 24gb GPU.

I'm unsure what's going wrong.

[ElleLeonne]

An update:

When I load the model using just quantization, it takes up 2gb.

Lightning says it takes up 6gb. I assume lightning does a sample backward pass and that the excess is stored gradients.

I can use the trainer.init_module() context manager to keep lightning faithful to its own stated size.

However as soon as my model receives any textual data at all, it goes OOM. My suspicion now, is that the optimizer is not properly handling the quantization or respecting the frozen layers. I can think of no other reason that my 2gb double quantized 4bit model would OOM on a single backwards pass.

[ElleLeonne]

Solved: Trainer defaults to mixed precision when handed tags that are not explicitly set to "32-true" or similar.

This produces duplicate tensor overhead.

Additionally, the trainer appears to use as a base of the weights that the model was saved in, rather than the weights that the model is currently in, for its size prediction. It can be safely ignored.