add deepspeed support for adalora finetune

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

Author: sywangyi

src/peft/tuners/adalora/layer.py

@@ -15,19 +15,21 @@
 import warnings
 from typing import Any, List, Optional
 
+import packaging
 import torch
+import transformers
 from torch import nn
 
 from peft.tuners.lora import LoraLayer
 from peft.tuners.tuners_utils import check_adapters_to_merge
 from peft.utils import transpose
-import transformers
-import packaging
+
 
 if packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.33.0"):
-    from transformers.integrations import is_deepspeed_zero3_enabled, deepspeed_config
+    from transformers.integrations import deepspeed_config
 else:
-    from transformers.deepspeed import is_deepspeed_zero3_enabled, deepspeed_config
+    from transformers.deepspeed import deepspeed_config
+
 
 class AdaLoraLayer(LoraLayer):
     # List all names of layers that may contain adapter weights
@@ -262,6 +264,7 @@ def update_ipt(self, model):
                 with torch.no_grad():
                     if deepspeed_config() is not None:
                         import deepspeed
+
                         grad = deepspeed.utils.safe_get_full_grad(p)
                         self.ipt[n] = (p * grad).abs().detach()
                     else:

src/peft/tuners/adalora/model.py

@@ -14,7 +14,9 @@
 
 import warnings
 
+import packaging
 import torch
+import transformers
 from transformers.pytorch_utils import Conv1D
 
 from peft.import_utils import is_bnb_4bit_available, is_bnb_available
@@ -30,14 +32,12 @@
 
 from .gptq import SVDQuantLinear
 from .layer import AdaLoraLayer, RankAllocator, SVDLinear
-import transformers
-import packaging
+
 
 if packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.33.0"):
-    from transformers.integrations import is_deepspeed_zero3_enabled, deepspeed_config
+    from transformers.integrations import is_deepspeed_zero3_enabled
 else:
-    from transformers.deepspeed import is_deepspeed_zero3_enabled, deepspeed_config
-
+    from transformers.deepspeed import is_deepspeed_zero3_enabled
 
 
 class AdaLoraModel(LoraModel):
@@ -253,9 +253,13 @@ def forward(self, *args, **kwargs):
             for n, p in self.model.named_parameters():
                 if ("lora_A" in n or "lora_B" in n) and self.trainable_adapter_name in n:
                     if is_deepspeed_zero3_enabled():
-                        import deepspeed
                         import contextlib
-                        with deepspeed.zero.GatheredParameters(p, modifier_rank=0, fwd_module=self) if p.shape==torch.Size([0]) else contextlib.nullcontext() :
+
+                        import deepspeed
+
+                        with deepspeed.zero.GatheredParameters(
+                            p, modifier_rank=0, fwd_module=self
+                        ) if p.shape == torch.Size([0]) else contextlib.nullcontext():
                             para_cov = p @ p.T if "lora_A" in n else p.T @ p
                     else:
                         para_cov = p @ p.T if "lora_A" in n else p.T @ p