Add support for memory maps with --load-in-8bit

[mayaeary]

Related to #193

Now it's possible to use --gpu-memory or --auto-devices with --load-in-8bits.
It uses new transformers library, so maybe you will need an upgrade

I've tested it on pygmalion-6b, gpt-j-6b and OPT-6.7B with my 8GB RTX2070, and it worked at least twice as fast as same models without 8bit or with any other accelerators.

[oobabooga]

It seems to work. Here is what I have tested:

• The results with the Debug-deterministic preset are the same in the old and new implementations with python server.py --load-in-8bit (without --gpu-memory).
• When --gpu-memory is used, the offload works: the GPU memory usage is reduced and the program doesn't crash with a bitsandbytes error like it used to.
• The results with offloading using Debug-deterministic are somewhat different. Probably an upstream bitsandbytes issue.

I will leave the PR open so that more people can test it before merging. Thanks for the refactor, by the way :)

[ghost]

This is a great feature.

When offloading, perhaps outputs are different because the modules that don't end up in VRAM do not get converted to 8-bit, as explained in the Transformers documentation.

Having an explicit use for the --auto-devices flag alongside --load-in-8bit would allow the user a choice whether to offload or not.

[ghost]

Suggested change

	if shared.args.load_in_8bit:
	params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True, llm_int8_enable_fp32_cpu_offload=True)
	if shared.args.load_in_8bit and shared.args.auto_devices:
	params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True, llm_int8_enable_fp32_cpu_offload=True)
	elif shared.args.load_in_8bit:
	params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True)

[oobabooga]

Oops, I merged at the same time that you commented, @81300. As it is, --load-in-8bit does not do any memory map unless --gpu-memory is also specified, so my understanding is that llm_int8_enable_fp32_cpu_offload=True will just be ignored in this case.

perhaps outputs are different because the modules that don't end up in VRAM do not get converted to 8-bit

Yes, that's probably the cause of the difference. This issue also appeared here: #198

[ghost]

@oobabooga,

This part of the code can add a GPU memory limit by default.

my understanding is that llm_int8_enable_fp32_cpu_offload=True will just be ignored in this case.

In my testing it was not. With quantization if the model doesn't fit in VRAM and offloading is requested by the automatically computed map, the graceful error message in the current version of Transformers should be:

ValueError: If you want to offload some keys to `cpu` or `disk`, you need to set `load_in_8bit_fp32_cpu_offload=True`. Note that these modules will not be  converted to 8-bit but kept in 32-bit.

There's a simple way to test this: change llm_int8_enable_fp32_cpu_offload to False which is the default and try loading a model that wouldn't normally fit on your video card via python server.py --load-in-8bit. Then attempt something like python server.py --load-in-8bit --gpu-memory 1 and you should get the message above.

Edit: also note that --load-in-8bit actually does a memory map, even before this PR, as load_in_8bit=True requires device_map="auto". This is the case in the current implementation as well.

[oobabooga]

You are right, the layers were being offloaded by default. See if this settles it

ee164d1

[ghost]

It works, thank you. For example:

python server.py --load-in-8bit --model llama-65b

ValueError: 
                        Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit
                        the quantized model. If you want to dispatch the model on the CPU or the disk while keeping
                        these modules in 32-bit, you need to set `load_in_8bit_fp32_cpu_offload=True` and pass a custom
                        `device_map` to `from_pretrained`. Check
                        https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu
                        for more details.