[bnb] Let's make serialization of 4bit models possible 
Author: poedator
Pull request: yes

## What does this PR do?
Purpose: enable saving and loading transformers models in 4bit formats.

tested with Bloom-560 and Llama-2-7b. Save - Load - match tensors and quant_states - match inference results.

## connection with bitsandbytes
Requires this PR in bitsandbytes: https://github.com/TimDettmers/bitsandbytes/pull/753 to be able to save/load models

## testing:
the functionality was tested doing this series of commands:
```
model = m4 = transformers.AutoModelForCausalLM.from_pretrained(..., [with 4bit quant config])
model.save_pretrained(SAVE_PATH, safe_serialization= [False / True] )
model2 = transformers.AutoModelForCausalLM.from_pretrained( SAVE_PATH)
# then matching all params and quant_state items between the models, matching inference results
```
Specific tests will be added to this PR once the [bitsandbytes PR](https://github.com/TimDettmers/bitsandbytes/pull/753) merges.

## Open questions to the maintainers:
1. I added and tested code necessary for straightforward save/load operations with LLMs. Yet there may be other kinds of function calls that may need to be updated to handle 4bit save/load - pls suggest if/where to expand this PR.
specifically: this PR covers `load_state_dict_into_meta_model()` but not `load_state_dict_into_model()` for I did not find an  example that uses it. 

2. Some of my edits may not fit with style / refactoring plans of the maintainers - pls give guidance if needed.

Comments:
cc @younesbelkada @SunMarc 

## On storage format for the quantized parameters.
Right now there is no clear separation of code to load a pretrained model in fp16/fp32 with subsequent quantisation and loading already quantized model. I noticed that for 8-bit it comes down to detecting `SCB` in state_dict.keys(). In this PR the 4bit quantized param storage requires not one but about 10 components per weight to be loaded from the state_dict - it may become unpractical.
Also saving all quantization params as _buffers into state_dict() may be cluttering the state_dict, and be impractical for some data types.
This matter is important right now, because of multitude of quantization methods (and more to come): int8, nd4, fp4 in bnb, SpQR, GPTQ, etc. Some guidelines here will simplify transformers code and facilitate adding new quantizations later. 
One possibility is this:
Packing all of the quantisation components into single serializeable object per parameter (protobuf?) and then saving it alongside the main parameter data. Then the loading code may look like:
```
for k in state_dict:
    if k + '_quant_state' in loaded_state_dict:
        new_param_weight = load_quantized_weight(loaded_state_dict[k], loaded_state_dict[k+'_quant_state'], )
    else:
        new_param_weight = loaded_state_dict[k]
...
def unpack_quantized(data, packed_quant_state):
    quant_type, quant_state = unpack_q_state(packed_quant_state)
    if quant_type == 'bnb_nf4':
        new_param_weight = bnb.nn.Param4bit.from_prequantized(data, quant_state) # call specific unpacker
    ...
    return new_param_weight  # fully unpacked and loaded unto cuda

def unpack_q_state(packed_quant_state: protobuf) -> Tuple(str, SomeQuantStateStru):
    """from serialized format to dict(or some data structure) of quantization state"""
    ...
    return quant_type, quant_state
```
Alternatively, quant_type per parameter may be stored in `config.json`, but it becomes too fat then.
I suggest to allow various quantization types per parameter.
What do you think? I'll be glad to refactor this PR to make it future-proof. 
@younesbelkada @SunMarc @TimDettmers 

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

Thanks a lot @poedator for your great effort here! Let me know once this is ready for review!

> Thanks a lot @poedator for your great effort here! Let me know once this is ready for review!

@younesbelkada , thanks for your friendly attention to my work!

I updated the PR code to address your comments and also simplify/clarify some areas. Also stole the PR name from your #22177 ).

For saving format for 4bit params I changed scheme (see bnb my PR for more details) to this:
- save main weight tensor as Parameter using standard state_dict() code
- save tensor components from quant_state into state_dict() with custom code
- pack all other quant_state components into json->bytestring->int8_tensor and save to state dict using name like `...weight.quant_state.bitsandbytes__nf4`
- no extra per-weight metadata stored to safetensors
As a result, we have 6 items stored per weight (with double quant) vs ~12 in earlier draft.

Such scheme can be easily adopted for other quantizers: when loading, `_load_state_dict_into_meta_model()` checks for presence of `[param_name].quant_state.[QUANT_TYPE]` key in loaded dict and use QUANT_TYPE to call appropriate function to such weight in quantized form. It allows to load models with custom choice of quantizers on per-weight basis. Once this PR goes thru, I want to propose how to refactor model_utils to have a uniform interface for quantied models. 

I added tests covering saving / loading and matching by component of tensors and quant_states, + forward() and generate().

Now we need to see the enabling [PR in bnb](https://github.com/TimDettmers/bitsandbytes/pull/753) going thru. Other than that, this PR looks ready for your review.

Thanks a lot for you work, could you tell me how the 4bit quantized model should be merged with the adapter trained by lora?

Also I would like @SunMarc to have a look on the PR whenever possible, as he is also maintaining the quantization feature in transformers, thanks!

> Thanks a lot for you work, could you tell me how the 4bit quantized model should be merged with the adapter trained by lora?

Hi, @Yu-Yuqing 
I believe, such merge is done by [peft](https://github.com/huggingface/peft) package. It was implemented recently by https://github.com/huggingface/peft/pull/851.
This PR is not affecting such merges. But if(when) it gets merged, it would be possible to save 4-bit quantized models, including ones that had adapters merged into them. 


@younesbelkada, 
thank you for the review - I believe that I addressed your comments. Below is update on the tests and `bnb`:

- the slow 8 bit tests generally work, except 5 of them in `MixedInt8GPT2Test` group (starting from test_generate_quality) failed with `AssertionError: "Hello my name is John Doe, and I'm a fan of the" != "Hello my name is John Doe, and I'm a big fan of"` However this error is also present in `main` branch. I believe that it is still a legit generation result. Perhaps it was affected by version change somewhere.
- the slow 4bit tests generally work. My new test need some extra testing in multi-GPU setup. I will work on this soon.
- meanwhile [at bnb](https://github.com/TimDettmers/bitsandbytes/pull/753) we had the first round of comments.

Update: I am still optimistic that the enabling PR in BitsAndBytes will get merged soon, but can't bet on a specific date. 

As an alternative, I can try to bring 4-bit serialization code to here. But that would require subclassing `bnb.Linear4bit` with risk of compatibility issues.  Does it make sense?  @younesbelkada 

Thanks a lot @poedator ! 
I think it is better to have the bnb serialization live directly in bnb, there is no rush, we can wait until your PR gets merged!

rebased to main as of Nov 8. 
tested with recent `bitsandbytes==0.41.2.post2` - the 4bit serialization test works OK. 

@SunMarc @younesbelkada, you may want to review 

known issues:
- errors with `safetensors` serialization  - need to restore shared tensors (`quant_map`, rather small) somehow - need a hint. 

- error in CI with ExamplesTestsNoTrainer.test_run_clm_no_trainer - not sure how it is related to this PR?

@younesbelkada , thank you for  picking this up so fast. Indeed, `bnb_4bit_use_double_quant=False` causes error in serialization, it will be fixed on the bnb side. 
Meanwhile, let me know if any other issues pop up when testing with bnb_4bit_use_double_quant=True.
Also I need advice on shared tensors and on pytorch test - see my previous message.

UPDATE: fix PR is open now: https://github.com/TimDettmers/bitsandbytes/pull/868

Thanks @poedator ! 
For the safetensors issue please refer to this PR: https://github.com/huggingface/transformers/pull/27234 I had to fix something similar for llm.int8 weights. Let me know if this PR helps you digging a bit!

> Contributor

the example shows how to identify shared pointers. But where is that collection stored then and how it is used in unpacking?

I got the same issue for Mistral, too. 

> I got the same issue for Mistral, too.

hi, @KnutJaegersberg , thank you for the interest to bnb 4bit serialization. While most of the changes are already made on bnb side, for successful 4bit serialization one may need to wait for this PR to merge. Hopefully it is a matter of days.

The error that you mention comes from code that restores shared tensors. It is still in development. Try saving without safetensors `safe_serialization=False`.

Thanks for your answer, @poedator I started thinking I just made a stupid mistake. I'll try what you said as soon as my local compute is available again. 

Same exception with save_serialization=False. 
hmm
Loading checkpoint shards:   0%|                                                                                                                          | 0/4 [00:01<?, ?it/s]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/knut/Downloads/transformers-save4/src/transformers/models/auto/auto_factory.py", line 566, in from_pretrained
    return model_class.from_pretrained(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 3532, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 3927, in _load_pretrained_model
    new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 761, in _load_state_dict_into_meta_model
    recover_shared_tensors_values(state_dict, quantized_stats, param_name + ".", likely_shared_key_endings)
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 799, in recover_shared_tensors_values
    target_dict[target_prefix + ending] = replacement_tensors[0]
IndexError: list index out of range


Hi @poedator 

> the example shows how to identify shared pointers. But where is that collection stored then and how it is used in unpacking?

It is used directly by storing them on the state dict, that is what we do for 8bit layers, let me know if you want me to have a deeper look and propose something

@younesbelkada 
I decided to un-share these nested_quant_map tensors, which are only 256 elements each. So there is no need to deal with restoring them for 4bit quantization. 
This change makes earlier saved 4-bit models slightly incompatible. Please discard them if saved any before today. My test models at HF-hub are updated. [pt](https://huggingface.co/poedator/opt-125m-bnb-4bit_pt), [safetensors](https://huggingface.co/poedator/opt-125m-bnb-4bit)

I completed the needed changes on bnb side ([PR pending](https://github.com/TimDettmers/bitsandbytes/pull/868)). Also added a set of 4-bit serialization tests in a separate file. They are more detailed than you suggested, but this gives me mode peace of mind. 

Please resume the review.
---------------
@KnutJaegersberg - you may want to re-test this PR now. That problem should be gone.
the code that gave `IndexError: list index out of range` is removed entirely - no need to restore shared tensors cause there is none. Both torch and safetensors formats should work OK.

@poedator thank you :) 
I tried again after reinstalling tf, now I get another exception: 

Loading checkpoint shards:   0%|                                                                                                                          | 0/4 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/knut/Downloads/transformers-save4/src/transformers/models/auto/auto_factory.py", line 566, in from_pretrained
    return model_class.from_pretrained(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 3512, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 3907, in _load_pretrained_model
    new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 760, in _load_state_dict_into_meta_model
    set_module_quantized_tensor_to_device(
  File "/home/knut/Downloads/transformers-save4/src/transformers/integrations/bitsandbytes.py", line 121, in set_module_quantized_tensor_to_device
    new_value = bnb.nn.Params4bit.from_prequantized(
  File "/home/knut/transformers/lib/python3.9/site-packages/bitsandbytes/nn/modules.py", line 161, in from_prequantized
    self.quant_state = QuantState.from_dict(qs_dict=quantized_stats, device=device)
  File "/home/knut/transformers/lib/python3.9/site-packages/bitsandbytes/functional.py", line 645, in from_dict
    shape=torch.Size(qs_dict['shape']),
TypeError: 'NoneType' object is not iterable


@KnutJaegersberg ,
I am not seeing this error in my tests, but to be on the safe side, I edited the line 645 in the BNB PR.
Please help me understand what triggered it? Could it be that you try to load a model created with earlier version of either PRs? Try to run the whole cycle "quantize-save-load" using the latest PR commits. Or try loading [my model form HF](https://huggingface.co/poedator/opt-125m-bnb-4bit).

@poedator I I did not use the latest PR of BNB because it gave me a cuda setup exception. It might be a local issue then. Thanks! 

I'm trying to do something weird, I guess. I can load your opt model @poedator but not my llamafied yi-34b, though I tried it with the same config, double quantization, no safetensors. 

I'll upload my weights here, it might take a few hours:
https://huggingface.co/KnutJaegersberg/Deacon-34b-4bit


  File "/home/knut/Downloads/transformers-save4/src/transformers/models/auto/auto_factory.py", line 566, in from_pretrained
    return model_class.from_pretrained(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 3512, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 3907, in _load_pretrained_model
    new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
  File "/home/knut/Downloads/transformers-save4/src/transformers/modeling_utils.py", line 760, in _load_state_dict_into_meta_model
    set_module_quantized_tensor_to_device(
  File "/home/knut/Downloads/transformers-save4/src/transformers/integrations/bitsandbytes.py", line 121, in set_module_quantized_tensor_to_device
    new_value = bnb.nn.Params4bit.from_prequantized(
  File "/home/knut/miniconda3/envs/textgen/lib/python3.10/site-packages/bitsandbytes/nn/modules.py", line 161, in from_prequantized
    self.quant_state = QuantState.from_dict(qs_dict=quantized_stats, device=device)
  File "/home/knut/miniconda3/envs/textgen/lib/python3.10/site-packages/bitsandbytes/functional.py", line 633, in from_dict
    code=qs_dict['nested_quant_map'].to(device),
KeyError: 'nested_quant_map'


@KnutJaegersberg , 
Thank you for your persistence! I was able to reproduce the error with `llamafied yi-34b` with earlier commit.
Then I retested loading and saving with fresh commits and everything worked OK:
`# bnb commit 079d7afe3468a9f9ad0c8214d5c5055bdedaccbf`
`# transformers commit 8e44fd35712720222b5dcce2be9860f2bebd8f7d`
Testing code: https://hastebin.com/share/uyurinixus.python 

My guess is that either you used older commit when saving or loading the model, or there is some incompatibility in out environments. Please try my code with fresh save. 

@poedator I had issues building bnb for your newest commit on my local environment, have not resolved those so far. I tried to manually do the changes of your bnb commit to my local files, my transformers library should be exactly the same. It's properly my environment which is the cause. 
I'll just wait till the next bnb release, thanks :) 

cc @Titus-von-Koeller 

Hello, any update on this issue? Or is there any other way to save merged 4 bit model?

> Hello, any update on this issue? Or is there any other way to save merged 4 bit model?

Hi, @estibi 
it is 99% ready, and just need one last(hopefully) fix in bnb: https://github.com/TimDettmers/bitsandbytes/pull/868.  Please upvote it.

I'll look into #868 this week. Thanks @poedator for your continued work on
this!

On Mon, Dec 4, 2023 at 6:03 AM Poedator ***@***.***> wrote:

> Hello, any update on this issue? Or is there any other way to save merged
> 4 bit model?
>
> Hi, @estibi <https://github.com/estibi>
> it is 99% ready, and just need one last(hopefully) fix in bnb:
> TimDettmers/bitsandbytes#868
> <https://github.com/TimDettmers/bitsandbytes/pull/868>. Please upvote it.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/26037#issuecomment-1837848330>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ACFBEOYAYU5ETUDIJY6CXFTYHVKT7AVCNFSM6AAAAAA4PLZTBGVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQMZXHA2DQMZTGA>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>


With release of `bitsandbytes==0.41.3` ([PR 868 there](https://github.com/TimDettmers/bitsandbytes/pull/868)) the path to this PR seems clear. (Big thanks to @Titus-von-Koeller !)
I rebased the code to the latest main and retested.
All relevant long tests in quantization section pass.
@younesbelkada @SunMarc - please review.

> Thanks a lot @poedator for your work on this - I left a couple of comments regarding the bnb minimal version and about tests, can you please have a look? 🙏 Thanks!

@younesbelkada , thank you for pointing out the things that I overlooked.
I left one comment and fixed the rest. Pls check if it can fly now.

Hey all, progress on this PR is looking good. Thanks for all your work!

Just wanted to quickly make you aware, that due to failing PEFT tests, we had to [revert this change to Params4bit](https://github.com/TimDettmers/bitsandbytes/commit/f2e4b6d01f243f08c191bf6aaceb7ba281337354) and [just released a postfix to PyPi](https://pypi.org/project/bitsandbytes/0.41.3.post1/).

Not sure if this affects the PR here, but just wanted to be sure this final review takes this most recent change into account.

> Hey all, progress on this PR is looking good. Thanks for all your work!
> 
> Just wanted to quickly make you aware, that due to failing PEFT tests, we had to [revert this change to Params4bit](https://github.com/TimDettmers/bitsandbytes/commit/f2e4b6d01f243f08c191bf6aaceb7ba281337354) and [just released a postfix to PyPi](https://pypi.org/project/bitsandbytes/0.41.3.post1/).
> 
> Not sure if this affects the PR here, but just wanted to be sure this final review takes this most recent change into account.

@Titus-von-Koeller thanks for the notice. I tried to re-test with 0.41.3.post1, [but got RTE, details here](https://github.com/TimDettmers/bitsandbytes/issues/907) 0.41.3. was OK.

> Awesome work ! Left a couple of nits, can you confirm 4bit and 8bit slow tests pass with these changes?

something is not right between this PR and accelerate now. I get `ValueError: weight is on the meta device, we need a `value` to put in on cuda.`
Give me some time to debug. cc @Titus-von-Koeller just to watch in case we need changes on bnb side (I hope not).


We managed to fix the issue @poedator had with the bnb install. My GCC was too new and produces binaries that are not backward compatible enough. We now reissued the last release to be more backward compatible and could confirm that it works with an older GCC dependency.

The cuda library is packaged with a c++ standard library. If the user does not fulfill at least the c++ standard of the linked library, it cannot run and throws an error.

Awesome work, thanks everyone ! So this PR is ready for final review?

TLDR: found a weird 👾 bug in `accelerate`, only apparent with `device_map="cuda"` (as pure string).

Well, for a moment I thought that we are done, prepared to see this PR merge and celebrate. But then in the last minute the `tests/quantization/bnb/test_4bit.py::ExtendedSerializationTest` tests started failing with odd error `ValueError: weight is on the meta device, we need a value to put in on cuda`. This was especially odd, since almost same code worked flawlessly in my own tests which looked very similar. 

During debug, I found that the initially loaded model's weights were missing `quant_state` components, and therefore they were not saved and not available for subsequent loading. 
The `quant_state` components were created OK in `_load_state_dict_into_meta_model()` but later stripped out in `dispatch_model()`. With deeper dive into accelerate I found that magic happens in `accelerate.modeling.py::set_module_tensor_to_device()`:

1) [in this line the code took wrong branch](https://github.com/huggingface/accelerate/blob/ad3a5bc92084b9c5c6d5e141eb5ff0029f3fcc78/src/accelerate/utils/modeling.py#L327) because apparently `device(type='cuda') != device(type='cuda', index=0)`
`torch.device(device) != module._parameters[tensor_name].device`
`torch.device(device)` comes from device_map, based on `torch_device`, equal to `"cuda"` string
`module._parameters[tensor_name].device` is the result of torch `.cuda()` call in `bnb.Params4bit.cuda()`.
note that `torch.tensor([42]).to('cuda').device` returns `device(type='cuda', index=0)`

2) once in the wrong branch, the [code gets to line](https://github.com/huggingface/accelerate/blob/ad3a5bc92084b9c5c6d5e141eb5ff0029f3fcc78/src/accelerate/utils/modeling.py#L342) `new_value = param_cls(new_value, requires_grad=old_value.requires_grad)` and happily creates a `Param4bit` parameter with just 4bit weights and none of the `quant_state` components, which causes error much later.

3) Also found an issue 👾 with class name https://github.com/huggingface/accelerate/issues/2250 which is not relevant to this problem

It was not apparent until we switched from `"cuda:0"` to `"cuda"` in this commit https://github.com/huggingface/transformers/pull/26037/commits/6834f084e488737bc7252ffff093eb2a470258b1

Ideas for how to deal with this (pick any number of them):
1. compare not devices, but device types in [this accelerate line](https://github.com/huggingface/accelerate/blob/ad3a5bc92084b9c5c6d5e141eb5ff0029f3fcc78/src/accelerate/utils/modeling.py#L327) 
1. convert `device(type='cuda')` into `device(type='cuda', index=0)` along the way, for instance in `accelerate::dispatch_model()` 
2. simplify overall quantized weights handling, avoiding additional tinkering in accelerate, i.e move all logic from `set_module_tensor_to_device()` to `transformers`.
1. make bnb.Param4bit constructor more robust, so it would test data at init phase (and fail in case it gets just the qunatized weights like above).
1. ask `pytorch` to make `device(type='cuda')` equal `device(type='cuda', index=0)`  (makes sense but takes time and may cause other issues)
1. in transformers testing consider testing with both `"cuda"` and `device(type='cuda', index=0)`
1. re-test the rest of HF code to see where `(device(type='cuda') != device(type='cuda', index=0)` problem may arise.

While the code works and all slow tests would pass with empty `device map` or `device_map=device("cuda")` or  `"cuda:0"`, I hesitate to recommend merging it until the above is resolved. 

@younesbelkada , please, your guidance and/or make changes in `accelerate`  🙏.
cc @SunMarc @Titus-von-Koeller 


Hi @poedator 
Again thank you very much for all your efforts to push this forward! 
Thanks also for providing a detailed description of the new problem. The fix IMO should go into accelerate: https://github.com/huggingface/accelerate/pull/2254 
I suggest to wait until we get feedback from accelerate maintainers and see if this approach gets validated, if that's the case,  we'll merge this PR together with the accelerate PR and tell users to install accelerate from source if they face into a similar issue. Can you also meanwhile confirm https://github.com/huggingface/accelerate/pull/2254 fixes the problem? 


The PR is merged ! LMK @poedator if this fixes the issue 

@younesbelkada, @SunMarc ! thank you for the quick reaction. I re-tested with https://github.com/huggingface/accelerate/pull/2254 and the problem is gone now. Hope that the tests are all green on your side too, looking forward to the merge.

> Thanks a lot for iterating! Is `is_4bit_serializable` set into the model somewhere in `from_pretrained` ? Inspecting your branch I can't find it
you were correct about that. Fixed. Slow tests all pass. Pls review.

> Thanks for working on this @poedator for such a long time ! This is a very exciting feature that will enable many applications ! Left a few comments but nothing blocking.

Hi, @SunMarc thank you for the comments. I implemented most of them and left comments for the rest. Pls review.

Great work @poedator .

Where will docs be available to read about these updates?

> Where will docs be available to read about these updates?

@RonanKMcGovern, thanks for the interest to this feature. It should work with standard transformers' `.save_pretrained()` and `.from_pretrained()` methods. Check out [how it is done in the tests](https://github.com/huggingface/transformers/blob/f9a98c476c3a92beaba3b20e51f1ff49417231a6/tests/quantization/bnb/test_4bit.py#L521) or how 8bit quantized models did that.
Meanwhile  I am checking with the maintainers about needs to update documentation if any.

Hi @poedator, thank you for your wonderful work!

If I understand this correctly, this should work with all 4-bit models right? However if I run this directly,
```
from transformers import AutoModelForCausalLM 

model=AutoModelForCausalLM.from_pretrained("facebook/opt-125m",load_in_4bit=True)
model.save_pretrained("opt_cache")
```
The error
```
NotImplementedError: You are calling `save_pretrained` on a 4-bit converted model. This is currently not supported
```
will be raised.

I'm not sure whether this `load_in_4bit` makes any difference (because 8bit saving works well) and we need to load the 4bit model in another fashion, or is it just because this `NotImplementeError` is not removed?

I'm using `bitsandbytes 0.42.0` and `transformers 4.36.2`, which are already the latest. Thanks a lot!

Hi @hjenryin, 4.36.2 is the latest version release on pypi, but it doesn't contain all of the most recent commits (including this one). If you want to use the bleeding edge (not guaranteed stable) version of transformers, you can [install from source](https://huggingface.co/docs/transformers/installation#install-from-source). 

Alternatively, you can wait until v4.37 is released, which will include this feature. We release a new version roughly once every month, and are due to release one at the end of this week. 

> Hi @hjenryin, 4.36.2 is the latest version release on pypi, but it doesn't contain all of the most recent commits (including this one). If you want to use the bleeding edge (not guaranteed stable) version of transformers, you can [install from source](https://huggingface.co/docs/transformers/installation#install-from-source).
> 
> Alternatively, you can wait until v4.37 is released, which will include this feature. We release a new version roughly once every month, and are due to release one at the end of this week.

Hi @amyeroberts, thanks for putting this. I am having the same issue. I am curious when v4.37 would be published with these changes. Please let me know if you have any estimated timeline? Thanks for all the work.

Hi @hkhan9, we typically have new releases every month. The release for v4.37 is a bit later than that because of the Christmas and new year period. It will likely be the end of this week or the next. 

> Hi @hkhan9, we typically have new releases every month. The release for v4.37 is a bit later than that because of the Christmas and new year period. It will likely be the end of this week or the next.

Thank you @amyeroberts 

Thanks @poedator for this very helpful contribution and your (everyone
elses) persistence in getting it across the finishing line! I think this
was a significant step in a good direction.

On Wed, Jan 17, 2024, 16:05 hkhan9 ***@***.***> wrote:

> Hi @hkhan9 <https://github.com/hkhan9>, we typically have new releases
> every month. The release for v4.37 is a bit later than that because of the
> Christmas and new year period. It will likely be the end of this week or
> the next.
>
> Thank you @amyeroberts <https://github.com/amyeroberts>
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/26037#issuecomment-1896000621>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ACFBEO2YJMG6YTGLCARJMETYO7SC5AVCNFSM6AAAAAA4PLZTBGVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTQOJWGAYDANRSGE>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>


Hi

I first uninstalled **transformers** and **bitsandbytes**. I then install transformers from source, hoping it'll include the proper bnb version:

```sh
pip install git+https://github.com/huggingface/transformers.git
```

But, when I load a 4bits model I previously saved using **poedator**'s fork (_which works perfectly using that fork together with poedator's bnb fork_):

```python
        model = AutoModelForCausalLM.from_pretrained(
            model_name_or_path,
            low_cpu_mem_usage=True,
            local_files_only=model_name_or_path[0] == '.' or model_name_or_path[0] == '/',
            device_map=device_map if device_map else None,
            # torch_dtype=torch.float32
        )
```

It fails with following error:

```python
File ~/anaconda3/envs/musai/lib/python3.11/site-packages/transformers/integrations/bitsandbytes.py:58, in set_module_q
uantized_tensor_to_device(module, tensor_name, device, value, quantized_stats)
     55 is_buffer = tensor_name in module._buffers
     56 old_value = getattr(module, tensor_name)
---> 58 if old_value.device == torch.device("meta") and device not in ["meta", torch.device("meta")] and value is None
:
     59     raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
     61 prequantized_loading = quantized_stats is not None

NameError: name 'torch' is not defined
```

Then, I restart the environment and install latest **bitsandbytes** from source:

```sh
pip install git+https://github.com/TimDettmers/bitsandbytes.git
```

But, when loading the model again (in a fresh env), it fails:

```python
File ~/anaconda3/envs/musai/lib/python3.11/site-packages/transformers/integrations/bitsandbytes.py:113, in set_module_
quantized_tensor_to_device(module, tensor_name, device, value, quantized_stats)
    108     if new_value.dtype in (torch.int8, torch.uint8) and not is_4bit_serializable:
    109         raise ValueError(
    110             "Detected 4-bit weights but the version of bitsandbytes is not compatible with 4-bit serialization
. "
    111             "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
    112         )
--> 113     new_value = bnb.nn.Params4bit.from_prequantized(
    114         data=new_value,
    115         quantized_stats=quantized_stats,
    116         requires_grad=False,
    117         device=device,
    118         **kwargs,
    119     )
    120 else:
    121     new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(device)

File ~/anaconda3/envs/musai/lib/python3.11/site-packages/bitsandbytes/nn/modules.py:174, in Params4bit.from_prequantiz
ed(cls, data, quantized_stats, requires_grad, device, **kwargs)
    172 self = torch.Tensor._make_subclass(cls, data.to(device))
    173 self.requires_grad = requires_grad
--> 174 self.quant_state = QuantState.from_dict(qs_dict=quantized_stats, device=device)
    175 self.blocksize = self.quant_state.blocksize
    176 self.compress_statistics = self.quant_state.nested

File ~/anaconda3/envs/musai/lib/python3.11/site-packages/bitsandbytes/functional.py:626, in QuantState.from_dict(cls, 
qs_dict, device)
    623     qs_dict.update(unpack_tensor_to_dict(qs_dict.pop(qs_key)))
    625 qs_dict = {k.split('.')[-1]: v for k, v in qs_dict.items()}  # strip prefixes
--> 626 assert set(qs_dict.keys()).issubset(cls.valid_qs_keys)
    628 if 'nested_absmax' in qs_dict:
    629     offset = torch.tensor(float(qs_dict['nested_offset'])).to(device)

AssertionError: 
```

Any clue?



Hi, @webpolis! Thank you for interest to this serialization method.
Apparently, there are some keys in `qs_dict.keys()` which are not expected by the model. 
I have 2 guesses:
- (most likely) your model used some draft naming scheme for the quantization components, that was changed later as this PR progressed. In this case, try to re-save the quntized model using the latest transformers (from github until 4.37 is out) and bnb>=0.41.3.
- something changed on bnb side (try with 0.41.3 or whatever worked)

You may try to extract the mismatched keys from `qs_dict.keys()` to diagnose this more precisely, if none of the above helps.

Thanks @poedator .. that sounds reasonable. I'll try that.

Hi @poedator @amyeroberts thanks for releasing v4.37. I just installed it and saved the model using 4-bit. This was a huge achievement. Congratulations to the team ! 

> Hi @poedator @amyeroberts thanks for releasing v4.37. I just installed it and saved the model using 4-bit. This was a huge achievement. Congratulations to the team !

May I have your script for 4-bit model export? Just used bnb_config and did not work. Model still as large as original.

```py
from transformers import LlamaForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import torch

import argparse


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base_model_name_or_path", type=str)
    parser.add_argument("--output_dir", type=str)
    parser.add_argument("--device", type=str, default="cpu")

    return parser.parse_args()


def main():
    args = get_args()

    if args.device == "auto":
        device_arg = {"device_map": "auto"}
    else:
        device_arg = {"device_map": {"": args.device}}

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    print(f"Loading base model: {args.base_model_name_or_path}")

    model_4bit = LlamaForCausalLM.from_pretrained(
        args.base_model_name_or_path,
        low_cpu_mem_usage=True,
        return_dict=True,
        quantization_config=bnb_config,
        torch_dtype=torch.float16,
        **device_arg,
    )

    print(f"Running 4bit")

    tokenizer = AutoTokenizer.from_pretrained(args.base_model_name_or_path)

    model_4bit.save_pretrained(f"{args.output_dir}")
    tokenizer.save_pretrained(f"{args.output_dir}")
    print(f"Model saved to {args.output_dir}")


if __name__ == "__main__":
    main()
```

And the exported model is still big