Add support for CPU offloading for quantizing bigger models on smaller GPUs

[abhinavkulkarni]

Hi,

This PR has the following changes:

1. Further refinement to the branch dev/more_models to do quantization layer by layer on GPU
2. Using device_map="auto" and max_memory kwargs to do LM Evaluation on smaller GPUs if needed

Thanks!

[abhinavkulkarni]

low_cpu_mem_usage is needed here as loading HF models on CPU with fp16 dtype is slow, more details here: https://huggingface.co/mosaicml/mpt-7b-instruct/discussions/6#6470dec93df93fddece5fcde and

[abhinavkulkarni]

Let me know if you would want me to change device_map to balanced instead of auto.

[Sakits]

Hi @abhinavkulkarni,

Thank you for your work on this PR!
However, when I was reviewing and testing it, I encountered the following two issues:

1. I received the following error when running the original model directly: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
   This seems to be due to the code that moves the model to the GPU being inside the if branch of quantization. This should probably be outside the quantization if statement.

2. When running the evaluation with fake quantization, the results I obtained were identical to those of the original model. This raises a question about whether the state_dict=model.state_dict() on line 167 in entry.py is functioning as intended. Additionally, when applying AWQ to models using GeLU, we add an additional scaling node, which may not be well supported in this context.

Have you encountered these issues as well? If I've misunderstood or missed anything, I'd appreciate your correction!

[abhinavkulkarni]

Hey @Sakits,

Thanks for your reply. I have made the necessary changes. Here's the output I got on RTX 3060 (12GB VRAM):

# LM Eval: Original model
$ python -m awq.entry --model_path mosaicml/mpt-7b-instruct \
        --max_memory 0:9GiB cpu:99GiB \
	--tasks wikitext

Time: 05:51
Results:
|  Task  |Version|    Metric     | Value |   |Stderr|
|--------|------:|---------------|------:|---|------|
|wikitext|      1|word_perplexity|10.8864|   |      |
|        |       |byte_perplexity| 1.5628|   |      |
|        |       |bits_per_byte  | 0.6441|   |      |

# LM Eval: Fake quantization
$ python -m awq.entry --model_path mosaicml/mpt-7b-instruct \
        --max_memory 0:9GiB cpu:99GiB \
	--tasks wikitext \
	--w_bit 4 --q_group_size 128 \
	--load_awq awq_cache/mpt-7b-instruct-w4-g128.pt \
	--q_backend fake

Time: 05:53
Results:
|  Task  |Version|    Metric     | Value |   |Stderr|
|--------|------:|---------------|------:|---|------|
|wikitext|      1|word_perplexity|11.2684|   |      |
|        |       |byte_perplexity| 1.5729|   |      |
|        |       |bits_per_byte  | 0.6534|   |      |

# LM Eval: Real quantization
$ python -m awq.entry --model_path mosaicml/mpt-7b-instruct \
    --tasks wikitext \
    --w_bit 4 --q_group_size 128 \
    --load_quant quant_cache/mpt-7b-instruct-w4-g128-awq.pt

Time: 06:52
Results:
|  Task  |Version|    Metric     | Value |   |Stderr|
|--------|------:|---------------|------:|---|------|
|wikitext|      1|word_perplexity|11.2696|   |      |
|        |       |byte_perplexity| 1.5729|   |      |
|        |       |bits_per_byte  | 0.6535|   |      |

[Sakits]

Hi @abhinavkulkarni,

I've reviewed your changes and everything looks great!
Thank you very much for your valuable contributions to awq project!