[core] Fix offload issue (#248)

* fix offload dir

* remove offload index

* safety checker

* forward contrib credits from previous PR

---------

Co-authored-by: cosimoiaia <cosimoiaia@users.noreply.github.com>

src/peft/peft_model.py

@@ -164,6 +164,15 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
         if getattr(model, "hf_device_map", None) is not None:
             device_map = kwargs.get("device_map", "auto")
             max_memory = kwargs.get("max_memory", None)
+            offload_dir = kwargs.get("offload_dir", None)
+            offload_index = kwargs.get("offload_index", None)
+
+            dispatch_model_kwargs = {}
+            # Safety checker for previous `accelerate` versions
+            # `offload_index` was introduced in https://github.com/huggingface/accelerate/pull/873/
+            if "offload_index" in inspect.signature(dispatch_model).parameters:
+                dispatch_model_kwargs["offload_index"] = offload_index
+
             no_split_module_classes = model._no_split_modules
             if device_map != "sequential":
                 max_memory = get_balanced_memory(
@@ -176,7 +185,13 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
                 device_map = infer_auto_device_map(
                     model, max_memory=max_memory, no_split_module_classes=no_split_module_classes
                 )
-            model = dispatch_model(model, device_map=device_map)
+
+            model = dispatch_model(
+                model,
+                device_map=device_map,
+                offload_dir=offload_dir,
+                **dispatch_model_kwargs,
+            )
             hook = AlignDevicesHook(io_same_device=True)
             if model.peft_config.peft_type == PeftType.LORA:
                 add_hook_to_module(model.base_model.model, hook)