`HfQuantizer` class for quantization-related stuff in `modeling_utils.py`
Author: poedator
Pull request: yes

### What does this PR do?
Refactoring `modeling_utils.py` to move all quantization-related logic into new `HFQuantizer` class

### Reasons and benefits:
- easier to understand how quantization works during loading.
- easier to add new quantization methods like SPQR.
- much easier to implement 4-bit serialization (once BnB supports it)
- it was a rainy day outside )

### Things to be done:
- extend `HFQuantizer` functionality to cover `_load_pretrained_model()` and `_load_state_dict_into_model()`
- extend `HFQuantizer` functionality in BnB to cover `set_module_quantized_tensor_to_device()` from `integrations/bitsandbytes.py`
- review multiple TODOs left in the comments.
- undo temporary changes in the tests
- check for backward compatibility issues
- eliminate repeats and redundancies, if any
- consider fully absorbing `integrations/bitsandbytes.py` into new class
- move new classes / functions into proper project file. possibly separate folder.

### Current state:
- reworked code up to `_load_state_dict_into_model()`
- all tests in BnB and GPTQ still pass with `RUN_SLOW=1`.

summoning  @SunMarc and @younesbelkada to comment on the idea and current state.

Comments:
Thanks a lot for your huge work! I really like the idea! Let me know once this is ready for a first pass, I'll have a deep look

Hi, @younesbelkada!
By now I reworked `_load_pretrained_model()` and `_load_state_dict_into_model()` functions. 
`Integations.bitsandbytes.py::set_module_quantized_tensor_to_device()` is replicated in new code and no longer needed.
I like that by having separate `HFQuantizer` classes there is no need to have multiple ifs in the code.

The code works fine and passes all RUN_SLOW quantization tests: gptq, 4/8 bit bnb.
It is now ready for the first detailed review - please take a look.

It may be further optimized, but I tried to preserve some resemblance with the earlier code.
Three are quite a few TODOs left - some of them require your guidance.
Please recommend ways to simplify the `HFQuantizer` class, to make it easier to implement new quantizers. 
For instance, there are few small functions to validate and adjust environment. Could they be consolidated? Also, shall we separate fresh quantization code from loading pre-quantized?

@younesbelkada , I addressed most of your comments and left couple open. Take a look. All SLOW tests pass now.

@younesbelkada ,
I changed to enum, added typing. 
let's try these battle tests.

Thanks so much @poedator ! Will try that out by today and let you know

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26610). All of your documentation changes will be reflected on that endpoint.

hi @poedator , sorry did not had time to have a deep look today, will do it tomorrow first thing in the morning! Thanks again for your patience! 🙏 

> ...can you also help fix the merge conflicts?

I fixed the merge conflicts, however it cancelled some of the changes from https://github.com/huggingface/transformers/pull/27043 by @patrickvonplaten.  
As I see it, keeping changes to the config makes no effect, since `process_model_after_weight_loading()` only returns the model. Please recheck, whether and how to reapply changes from #27043. Perhaps you have some test to check if the problem solved by 27043 reappears.

Ideas for before and after merge:
- move quantization code to `quantize` folder - it will contain quantization_config.py, quantizers.py and files for future quant classes
- should there be more documentation/blog post on how to build a new quantizer?  The AWQ #27045 may be a good basis for this.
- the checks for `is_loaded_in_4bit` should now be replaced to checks based on model.config.quantization_config.

Thanks for all your work and for addressing all comments!

> I fixed the merge conflicts, however it cancelled some of the changes from https://github.com/huggingface/transformers/pull/27043 by @patrickvonplaten.
As I see it, keeping changes to the config makes no effect, since process_model_after_weight_loading() only returns the model. Please recheck, whether and how to reapply changes from https://github.com/huggingface/transformers/pull/27043. Perhaps you have some test to check if the problem solved by 27043 reappears.

Sure yes I will check that out myself and confirm you here

> should there be more documentation/blog post on how to build a new quantizer? The AWQ https://github.com/huggingface/transformers/pull/27045 may be a good basis for this.

Definitely 💯 ! we can always do it after this PR being merged as a follow up, so no worries I would say

@younesbelkada,
I see that #27045 is merged now. I guess it is time to update this PR again -  I can try it over the weekend. Please confirm that this makes sense, or is there any other way forward? 

hi @poedator , thanks for getting back , yes what you have said makes sense we still want to merge this PR. On my end the tests I performed were good, I can try again once you merge with main! 🙏  

@younesbelkada , @SunMarc , @ArthurZucker, 
The PR is rebased to `main` as of Nov 3,  including the AWQ changes. It is ready for review - please take a look.

major issues:
- seemingly unrelated test fails in `examples_torch` 
- some tests with multiple GPUs fail - need more time to debug
- the number of methods in HFQuantizer bothers me. Some could be combined if we pass context from above. See separate note
- there is still a number of TODOs left - need hints on how to approach them.
- get rid of `is_loaded_in_4[8]bit` use in `modeling_utils.py`. (see `.save_pretrained()` and `.num_parameters()`) - use quant_config refs.
- can we drop `is_loaded_in_4[8]bit` completely? or replace with @property with warning? 

Here are the calls to `HFQuantizer` from `modeling_utils.py`. Would be great to reduce their number. 
So far I only see it possible with contextlib. See `HFQuantizer.get_locals_from_above()` as prototype.


after initialization:
-	`quantizer.validate_environment(torch_dtype=torch_dtype, from_tf=from_tf, from_flax=from_flax)  # odd args selection`
-	`quantizer.set_torch_dtype(torch_dtype)`
-	`quantizer.update_device_map(device_map)`
the calls in this group could be combined with `__init__()` if we pass context from .from_pretrained()

before weight loading:
-	quantizer.process_model_before_weight_loading(model, device_map, torch_dtype, keep_in_fp32_modules)
-	quantizer.get_special_dtypes_update(model, torch_dtype)
-	quantizer.adjust_target_dtype(target_dtype)
-	quantizer.adjust_max_memory(max_memory)
-	quantizer.validate_device_map(device_map)
the calls in this group could be combined into `process_model_before_weight_loading` if we pass context from .from_pretrained()

in `_load_pretrained_model()`:
-	`quantizer.check_quantized_param(model, param_value=value, param_name=key, state_dict={})`
-	`quantizer.update_mismatched_keys(unexpected_keys, missing_keys)`

in `_load_state_dict_into_meta_model()`:
-	`quantizer.create_quantized_param` 

after weight loading:
-	`quantizer.process_model_after_weight_loading(model)`

@ArthurZucker , @SunMarc , could you comment on this PR please.
cc @younesbelkada 

Yeah sorry, it's a big PR and not super time sensitive so I've been pushing it back 😓 will try to do this week

cc @Titus-von-Koeller for visibility

Hi @poedator ! Happy new year and thanks for putting so much efforts on this PR! Would you be happy helping us finalising the PR ? Otherwise happy to take it over as well as we want to have this refactor ASAP on transformers!

Hi, @younesbelkada ,
I am ready to finalize this together with you and your team.  For now I see the following open issues:
- completing response to comments of @SunMarc 
- rebase to resolve conflicts (hope that this is trivial this time)
- updating `save_prequantized` (optional)
- another round of tests and reviews ?

Hi @poedator thanks for getting back ! I think it would be great if you could first resolve the merge conflicts, then complete the final comments from Marc, after that we'll do a final review and run the tests if that sounds good to you

The PR code has been rebased to the latest main (as of 2023.01.07).  It is pretty much ready for the next round of reviews @younesbelkada @SunMarc @ArthurZucker @amyeroberts 

Most of the comments of Marc are resolved, a few are addressed but left in the open form for Marc to revisit.
One open item is # TODO: consider removing used param_parts from state_dict before return. I want to do some more tests before committing it.

I left a few TODOs in the quantizers.py for the maintainers to consider choices - please comment on them or delete.

There are further optimisation opportunities in this refactoring - I did not try to pursue them actively. 
Examples:
- refactor the corresponding code in `accelerate:: ... ::set_module_tensor_to_device()`
- refactor save_pretrained() - DONE
- replace `is_loaded_in_4bit` and `is_4bit_serializable` with references to config - DONE
- integrations/bitsandbytes can be integrated into quantizers.py (set_module_quantized_tensor_to_device is not used in the main code)
- quantizers.py may be sliced into smaller files (one per method)

The Long tests pass with 2 exceptions
- FAILED tests/quantization/gptq/test_gptq.py::GPTQTest::test_change_loading_attributes - AssertionError: True != False
- FAILED tests/quantization/gptq/test_gptq.py::GPTQTest::test_serialization - ValueError: Found modules on cpu/disk. ...
Not sure yet what caused them, maybe they are specific to my setup. Will update here if I find the causes.

@ArthurZucker, thank you for the kind words and deep comments!

General reply first: I started this PR trying to repack existing quantization-related code to make it easier adding other quantization methods. What you see by today is effectively the same operations as before but packaged into `HFQuantizer` classes. This explains the numerous calls to quantizer - they remain where I found them. Combining calls to quantizers would be nice, but may require more massive refactoring (and understanding of the code history) and I am feeling a bit out of my depth here. More specific refactoring comments from the maintainers do help me moving forward. 

Another reason why I refrained from deeper refactoring was to keep code familiar to the reviewers.
With this context, would it be possible to postpone deeper refactoring to a separate PR? Although it would be ideal to set the class interface properly now.

Now more on some of the topics that you raised:

The QuantizationMethod enum and config classes like BitsAndBytesConfig were pre-existing.  QuantizationMethod can be removed, but I'd still need a way to keep track of chosen q-method. Maybe it is possible to use config class for that.

re "IMO We should not need an additional parser." - I packaged there the original code that selected `QuantizationMethod`. It should be possible to make it a class method instead. But the HF quantizer is already quite heavy. Currently the parser is a factory for quantizers. 

re "about potential uses outside ... community interests / use cases?" - I envisage possibility of people writing/importing custom quantizers as subclasses of HFQuantizer, and using them to quantize models as they like. Also this PR makes it easier to add more quantization methods, like SPQR which I coauthored recently :)
Perhaps it is time to raise it in Reddit/LocalLlama or other similar place where LLM quantization enthusiasts meet.  

UPD (16.01.23): stopped using quantization method in fac=vor of just type(quantization_config); refactored parser into simpler class with single call.

cc @ArthurZucker I just ported out the commits from https://github.com/huggingface/transformers/pull/28703 here

Merging! 

Thank you, @ArthurZucker  @younesbelkada @SunMarc @amyeroberts and the rest of HF team involved, for getting this to the merge! I was very impressed watching you work on this PR in the past month!