HfQuantizer class for quantization-related stuff in modeling_utils.py

[poedator]

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.

[younesbelkada]

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

[poedator]

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]

Thanks a lot for your huge work ! I really like the refactor and the cleanup, this makes things much cleaner in from_pretrained. I left few open questions before some final reviews, let me know what do you think!

[younesbelkada]

can't this be done inside update_device_map ? then we would change the method name to update_and_validate_device_map - what do you think?

[poedator]

I had the same thought but hesitated to combine these two methods because because between
quantizer.update_device_map() and quantizer. validate_device_map():

• changes are made to device_map, by infer_auto_device_map()
• modules_to_not_convert are updated during quantizer.process_model_before_weight_loading()
  Let me revisit this in the future, after some testing...

[younesbelkada]

I see, I think that it is fine if the method first validates the DM and then modify the device map

[poedator]

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

[younesbelkada]

Thanks a lot @poedator for your great work, my main comment is only: https://github.com/huggingface/transformers/pull/26610/files#r1364571404 it would be great if you can use an enum instead!
After that I will battle test your PR locally and let you know!

[younesbelkada]

I see, I think that it is fine if the method first validates the DM and then modify the device map

[poedator]

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

[younesbelkada]

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

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[younesbelkada]

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! 🙏

[younesbelkada]

Looks really nice! The refactor is amazing and makes things much easier to follow and understand @poedator I have battle tested evertyhing and looks great on my end!
I re-ran BNB tests and ran some quantization training with PEFT and all tests seem to pass so far!
I left 3 questions - can you also help fix the merge conflicts? I will also let @ArthurZucker & @SunMarc have a look! Before merging I'll re-run the tests a last time before merging
Would be great if we can merge this PR before: #27045 as well 🙏

[younesbelkada]

shouldn't we call validate_environment inside init as well? What do you think?

[poedator]

The present choice of arguments to validate_environment() looks awkward to me: torch_dtype, from_tf, from_flax.
It was driven by validation requirements of present quantizers. Passing them via quantization_config_parser.get_quantizer() would be even more awkward.

I would rather keep the functions separate, but pass the whole context to the validation fn. This would make life easier for the future quantizers.

[poedator]

@younesbelkada, How would you like the idea of passing the whole context of from_pretrained() to HFQuantizer methods? It is not very pythonic, but this way we'd be able to combine validate_environment(), set_torch_dtype(), update_device_map() into one method.

[younesbelkada]

I would rather keep the functions separate, but pass the whole context to the validation fn. This would make life easier for the future quantizers.

Yes OK got it, makes sense, let's keep it as it is i would say. I personally really like the fact that everything is separated and makes things very clear, validate dtype, device_map, etc. Let's wait for other maintainers to see their thoughts

[poedator]

...can you also help fix the merge conflicts?

I fixed the merge conflicts, however it cancelled some of the changes from #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.

[poedator]

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 [core / Quantization ] AWQ integration #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.

[younesbelkada]

Thanks for all your work and for addressing all comments!

I fixed the merge conflicts, however it cancelled some of the changes from #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.

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 #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

[poedator]

@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?

[younesbelkada]

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! 🙏

[poedator]

@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?

[poedator]

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)

[poedator]

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

[younesbelkada]

cc @ArthurZucker I just ported out the commits from #28703 here

[younesbelkada]

Thanks @poedator for your great work on this - the community will largely benefit from this refactor and adding new quantization methods natively in transformers will be much more easy and scalable thanks to your great efforts ! Thanks again for all your time, contribution and patience for this PR and I look forward to collaborating more with you
I can confirm all quantization slow tests pass with these changes
Thanks also @ArthurZucker for your patience and reviewing #28703 ! 🚀

[ArthurZucker]

🚀 Thanks both for the hard work! Looks really nice.
Let's keep an eye on the tests cc @amyeroberts and @ydshieh

[ArthurZucker]

Merging!

[poedator]

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!

[winglian]

Missed a guard clause here to check if hf_quantizer is not None and getattr(

[younesbelkada]

Thanks ! This should be fixed here: #28804 no?