Using torch_qaic gradScaler and making lora_dropout=0.05 (#320)

1. In case of finetuning on qaic, torch_qaic gradScaler will be used
2.  Moving back to lora_dropout = 0.05 on ML Framework team's ask.

Signed-off-by: Swati Allabadi <[email protected]>
Co-authored-by: Swati Allabadi <[email protected]>

Author: Swati Allabadi

QEfficient/finetune/configs/peft_config.py

@@ -19,7 +19,7 @@ class lora_config:
     target_modules: List[str] = field(default_factory=lambda: ["q_proj", "v_proj"])
     bias = "none"
     task_type: str = "CAUSAL_LM"
-    lora_dropout: float = 0.0
+    lora_dropout: float = 0.05
     inference_mode: bool = False  # should be False for finetuning
 
 

QEfficient/finetune/utils/train_utils.py

@@ -24,6 +24,7 @@
     import torch_qaic.debug as qaic_debug  # noqa: F401
     import torch_qaic.profile as qaic_profile  # noqa: F401
     import torch_qaic.utils as qaic_utils  # noqa: F401
+    from torch.qaic.amp import GradScaler as QAicGradScaler
 except ImportError as e:
     print(f"Warning: {e}. Moving ahead without these qaic modules.")
 
@@ -60,7 +61,6 @@ def train(
 
     Returns: results dictionary containing average training and validation perplexity and loss
     """
-
     train_prep = []
     train_loss = []
     val_prep = []
@@ -92,7 +92,10 @@ def train(
         tensorboard_updates = SummaryWriter()
 
     if train_config.grad_scaler:
-        scaler = GradScaler()
+        if device.startswith("qaic"):
+            scaler = QAicGradScaler()
+        else:
+            scaler = GradScaler()
 
     loss_0_counter = torch.tensor([0]).to(device)