Question about quantized model with zero3

[mxjmtxrm]

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?:
• Using distributed or parallel set-up in script?:

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.

[LysandreJik]

cc @younesbelkada

[mxjmtxrm]

Hi, is there any updates?

[github-actions]

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 are likely to be ignored.

[younesbelkada]

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

[CHNRyan]

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

[amyeroberts]

cc @muellerzr @SunMarc

[github-actions]

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 are likely to be ignored.