Repository: huggingface/peft

Issue:
[FSDP+QLoRA] ValueError: Expected a cuda device, but got: cpu
### System Info

pip list
```
accelerate                0.29.3
bitsandbytes              0.43.1
datasets                  2.14.6
huggingface-hub           0.20.3
llama-recipes             0.0.1
peft                      0.10.0
safetensors               0.4.2      
tokenizers                0.19.1
torch                     2.1.2
transformers              4.40.0
cupy-cuda12x              12.1.0
nvidia-cuda-cupti-cu12    12.1.105
nvidia-cuda-nvrtc-cu12    12.1.105
nvidia-cuda-runtime-cu12  12.1.105
```
8xA6000 48G, CUDA Version: 12.2

### Who can help?

_No response_

### Information

- [x] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder
- [x] My own task or dataset (give details below)

### Reproduction

Code from https://github.com/huggingface/alignment-handbook/tree/main/recipes/zephyr-141b-A35b
Set `use_dora=True` in LoRAConfig
Running with my modified command from the following
```
ACCELERATE_LOG_LEVEL=info TRANSFORMERS_VERBOSITY=info accelerate launch --config_file recipes/accelerate_configs/fsdp.yaml scripts/run_orpo.py recipes/zephyr-141b-A35b/orpo/config_qlora.yaml
```

Raise ValueError
```
Traceback (most recent call last):
  File "/root/kyzhang/llms/UltraMedical/llm_dpo/run_sft.py", line 209, in <module>
    main()
  File "/root/kyzhang/llms/UltraMedical/llm_dpo/run_sft.py", line 141, in main
    trainer = SFTTrainer(
  File "/root/miniconda3/lib/python3.10/site-packages/trl/trainer/sft_trainer.py", line 228, in __init__
    model = get_peft_model(model, peft_config)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/mapping.py", line 136, in get_peft_model
    return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config, adapter_name=adapter_name)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/peft_model.py", line 1094, in __init__
    super().__init__(model, peft_config, adapter_name)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/peft_model.py", line 129, in __init__
    self.base_model = cls(model, {adapter_name: peft_config}, adapter_name)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/lora/model.py", line 136, in __init__
    super().__init__(model, config, adapter_name)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/tuners_utils.py", line 148, in __init__
    self.inject_adapter(self.model, adapter_name)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/tuners_utils.py", line 325, in inject_adapter
    self._create_and_replace(peft_config, adapter_name, target, target_name, parent, current_key=key)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/lora/model.py", line 220, in _create_and_replace
    new_module = self._create_new_module(lora_config, adapter_name, target, **kwargs)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/lora/model.py", line 295, in _create_new_module
    new_module = dispatcher(target, adapter_name, lora_config=lora_config, **kwargs)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/lora/bnb.py", line 506, in dispatch_bnb_4bit
    new_module = Linear4bit(target, adapter_name, **fourbit_kwargs)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/lora/bnb.py", line 293, in __init__
    self.update_layer(
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/lora/layer.py", line 126, in update_layer
    self.dora_init(adapter_name)
  File "/root/miniconda3/lib/python3.10/site-packages/peft/tuners/lora/layer.py", line 186, in dora_init
    weight = dequantize_bnb_weight(weight, state=quant_state)  # no-op if not bnb
  File "/root/miniconda3/lib/python3.10/site-packages/peft/utils/integrations.py", line 58, in dequantize_bnb_weight
    return bnb.functional.dequantize_4bit(weight.data, weight.quant_state)
  File "/root/miniconda3/lib/python3.10/site-packages/bitsandbytes/functional.py", line 1353, in dequantize_4bit
    device = pre_call(A.device)
  File "/root/miniconda3/lib/python3.10/site-packages/bitsandbytes/functional.py", line 459, in pre_call
    torch.cuda.set_device(device)
  File "/root/miniconda3/lib/python3.10/site-packages/torch/cuda/__init__.py", line 402, in set_device
    device = _get_device_index(device)
  File "/root/miniconda3/lib/python3.10/site-packages/torch/cuda/_utils.py", line 35, in _get_device_index
    raise ValueError(f"Expected a cuda device, but got: {device}")
ValueError: Expected a cuda device, but got: cpu
```

### Expected behavior

-


Discussion:
I successfully trained the LLaMA-3-70B model using the script from the official PEFT example: [run_peft_qlora_fsdp.sh](https://github.com/huggingface/peft/blob/main/examples/sft/run_peft_qlora_fsdp.sh).

However, I'm still encountering this problem when I set `use_dora=True` in the code.
Thanks for reporting. It looks like at initialization time, the model is still on CPU. As initializing DoRA requires us to dequantize the bnb weights, which is not supported on CPU, we see this error. This should hopefully not be that hard to fix on our side. Meanwhile, perhaps you can adjust your scripts so that the base model is sent to GPU before calling `get_peft_model` and check if that works.

**Edit**: Honestly not sure how the weights can be on CPU here, maybe some form of offloading? In that case, the problem probably runs deeper. Are you aware if any offloading goes on here?
I have this same issue.  I can do Lora/Dora, DDP Lora/Dora, QLora/QDora, DDP QLora/QDora,  FSDP Lora/Dora, and FSDP QLora but FSDP QDora does not seem to be working.

Resolving patch:
diff --git a/src/peft/tuners/lora/layer.py b/src/peft/tuners/lora/layer.py
--- a/src/peft/tuners/lora/layer.py
+++ b/src/peft/tuners/lora/layer.py
@@ -23,7 +23,7 @@
 from transformers.pytorch_utils import Conv1D
 
 from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge
-from peft.utils.integrations import dequantize_bnb_weight, gather_params_ctx
+from peft.utils.integrations import dequantize_module_weight, gather_params_ctx
 from peft.utils.other import transpose
 
 from .config import LoraConfig
@@ -195,12 +195,7 @@ def dora_init(self, adapter_name: str) -> None:
         scaling = self.scaling[adapter_name]
         with gather_params_ctx(self.get_base_layer().parameters()):
             base_layer = self.get_base_layer()
-            if hasattr(base_layer, "W_q"):  # For handling HQQ quantized weight
-                weight = base_layer.dequantize()
-            else:
-                weight = base_layer.weight
-                quant_state = getattr(base_layer, "state", None)
-                weight = dequantize_bnb_weight(weight, state=quant_state)  # no-op if not bnb
+            weight = dequantize_module_weight(base_layer)
             if weight.data.ndim == 4:  # For handling LoRAs applied to Conv2Ds.
                 lora_weight = torch.mm(lora_B.flatten(start_dim=1), lora_A.flatten(start_dim=1))
                 lora_weight = lora_weight.reshape(weight.shape)
@@ -231,12 +226,7 @@ def _apply_dora(self, x, lora_A, lora_B, scaling, active_adapter):
         lora_weight = lora_B.weight @ lora_A.weight
         magnitude = self.lora_magnitude_vector[active_adapter]
         base_layer = self.get_base_layer()
-        if hasattr(base_layer, "W_q"):  # For handling HQQ quantized weight
-            weight = base_layer.dequantize()
-        else:
-            weight = base_layer.weight
-            quant_state = getattr(base_layer, "state", None)
-            weight = dequantize_bnb_weight(weight, state=quant_state)  # no-op if not bnb
+        weight = dequantize_module_weight(base_layer)
         weight = weight.to(x.dtype)
         weight_norm = self._get_weight_norm(weight, lora_weight, scaling)
         # see section 4.3 of DoRA (https://arxiv.org/abs/2402.09353)
diff --git a/src/peft/utils/integrations.py b/src/peft/utils/integrations.py
--- a/src/peft/utils/integrations.py
+++ b/src/peft/utils/integrations.py
@@ -38,12 +38,19 @@ def gather_params_ctx(param, modifier_rank: int = 0, fwd_module: torch.nn.Module
     return
 
 
-def dequantize_bnb_weight(weight: torch.nn.Parameter, state=None):
+def dequantize_module_weight(module: torch.nn.Module) -> torch.nn.Parameter:
     """
-    Helper function to dequantize 4bit or 8bit bnb weights.
+    Helper function to dequantize a quantized weight.
+
+    This function should be extended if more quantization schemes are added to the library.
 
-    If the weight is not a bnb quantized weight, it will be returned as is.
+    If the weight is not quantized, it will be returned as is.
     """
+    if hasattr(module, "W_q"):  # For handling HQQ quantized weight
+        weight = module.dequantize()
+        return weight
+
+    weight = module.weight
     if not isinstance(weight, torch.nn.Parameter):
         raise TypeError(f"Input weight should be of type nn.Parameter, got {type(weight)} instead")
 
@@ -51,10 +58,35 @@ def dequantize_bnb_weight(weight: torch.nn.Parameter, state=None):
     if cls_name not in ("Params4bit", "Int8Params"):
         return weight
 
+    quant_state = getattr(module, "state", None)
+    device = weight.device
+    is_cpu = device.type == torch.device("cpu").type
+    weight = dequantize_bnb_weight(weight, state=quant_state)  # no-op if not bnb
+    if is_cpu:
+        # dequantize_bnb_weight for 8bit moves the device in-place, thus we need to move it back to CPU if necessary
+        module.weight = module.weight.to(device)
+    return weight
+
+
+def dequantize_bnb_weight(weight: torch.nn.Parameter, state=None):
+    """Helper function to dequantize 4bit or 8bit bnb weights.
+
+    Since dequantization is not supported on CPU, the weight will be temporarily moved to CUDA if necessary.
+    """
     import bitsandbytes as bnb
 
+    # BNB requires CUDA weights
+    device = weight.device
+    is_cpu = device.type == torch.device("cpu").type
+    if is_cpu:
+        weight = weight.to(torch.device("cuda"))
+
+    cls_name = weight.__class__.__name__
     if cls_name == "Params4bit":
-        return bnb.functional.dequantize_4bit(weight.data, weight.quant_state)
+        dequantized = bnb.functional.dequantize_4bit(weight.data, weight.quant_state)
+        if is_cpu:
+            dequantized = dequantized.to(device)
+        return dequantized
 
     if state.SCB is None:
         state.SCB = weight.SCB
@@ -65,4 +97,7 @@ def dequantize_bnb_weight(weight: torch.nn.Parameter, state=None):
     if state.CxB is None:
         state.CxB, state.SB = bnb.functional.transform(weight.data, to_order=state.formatB)
     out32, Sout32 = bnb.functional.igemmlt(im, state.CxB, Sim, state.SB)
-    return bnb.functional.mm_dequant(out32, Sout32, SCim, state.SCB, bias=None).t()
+    dequantized = bnb.functional.mm_dequant(out32, Sout32, SCim, state.SCB, bias=None).t()
+    if is_cpu:
+        dequantized = dequantized.to(device)
+    return dequantized


Test patch:
diff --git a/tests/test_gpu_examples.py b/tests/test_gpu_examples.py
--- a/tests/test_gpu_examples.py
+++ b/tests/test_gpu_examples.py
@@ -1093,6 +1093,33 @@ def test_causal_lm_training_gpt2_dora(self):
             # assert loss is not None
             assert trainer.state.log_history[-1]["train_loss"] is not None
 
+    @parameterized.expand(["4bit", "8bit"])
+    def test_initialize_dora_with_bnb_on_cpu(self, kbit):
+        # 1674
+        # The issue is that to initialize DoRA, we need to dequantize the weights. That only works on GPU for bnb.
+        # Therefore, intializing DoRA with bnb on CPU used to fail.
+        model_id = "facebook/opt-125m"
+        if kbit == "4bit":
+            bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4")
+        elif kbit == "8bit":
+            bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+        else:
+            raise ValueError("Only 4bit and 8bit bnb allowed")
+
+        model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)
+        model = model.cpu()  # ensure that we're on CPU
+        # sanity check that all weights are on CPU
+        weights_not_cpu = [name for name, p in model.named_parameters() if p.device != torch.device("cpu")]
+        assert not weights_not_cpu
+
+        lora_config = LoraConfig(use_dora=True)
+
+        # should not raise
+        peft_model = get_peft_model(model, lora_config)
+        # check that the weights are still on CPU
+        weights_not_cpu = [name for name, p in peft_model.named_parameters() if p.device != torch.device("cpu")]
+        assert not weights_not_cpu
+
 
 @require_torch_gpu
 @require_auto_gptq