Question about quantized model with zero3
Author: mxjmtxrm
Pull request: no
Labels: name: DeepSpeed

name: Quantization

### System Info

- `transformers` version: 4.41.0.dev0
- Platform: Linux-5.15.0-92-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.21.4
- Safetensors version: 0.4.2
- Accelerate version: 0.28.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.2.0a0+81ea7a4 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@pacman100 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype="bfloat16",
            bnb_4bit_quant_storage="bfloat16",
        )
model = AutoModelForCausalLM.from_pretrained(
            args.model_name_or_path,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            quantization_config=bnb_config,
            attn_implementation="flash_attention_2" if args.use_flash_attn else "eager",
        )
```
I found that when fine-tune a quantized model using trainer with **Zero3**, the quantized model will  be loaded all to the GPU first, and then partitioning the parameters across data-parallel processes. **What if there is not enough memory to load the whole quantized model?**
The code that load all quantized model is in the deepspeed/runtime/engine.py about line262:
```
 self._configure_distributed_model(model)
```
It was entered from transformers trainer: inner_train_loop: about line1082:
```                   
 model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
```


### Expected behavior

How to partitioning the parameters during load from_pretrain instead of in trainer? like load float model.

Comments:
cc @younesbelkada 

Hi, is there any updates?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @mxjmtxrm 
Apologies for the delay in replying, I am not super familiar with DeepSpeed, I assume the process will be OOM-ed if you face such scenario, however, linear layers are quantized one by one and the model is first loaded on the meta device (0 RAM). Unless you finetune a >100B model (which will need around 50-60GB GPU VRAM), recent hardwares should be able to properly load most models in the ecosystem

Same error! When I fine tuning llama2 with zero3 and qlora, parameters will be fully loaded to each GPU and then partitioned. I think they should be partitioned first and then loaded to GPUs. So can your team fix it? @pacman100 

cc @muellerzr @SunMarc 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.