WIP

1 files changed, 14 insertions, 10 deletions

diff --git a/finetuning.py b/finetuning.py
index a46b544..3d7b011 100755
--- a/finetuning.py
+++ b/finetuning.py
@@ -67,23 +67,27 @@ def main():
     parser.add_argument("--gradient-checkpointing", action="store_true")
     parser.add_argument("--batch-size", default=1, type=int)
     parser.add_argument("--gradient-accumulation", default=4, type=int)
-    parser.add_argument("--learning-rate", default=1e-4, type=float)
+    parser.add_argument("--learning-rate", default=2e-4, type=float)
     parser.add_argument("--epochs", default=1, type=int)
     parser.add_argument("--warmup-ratio", default=0.03, type=float)
     parser.add_argument("--scheduler-type", default="constant", type=str)
-    parser.add_argument("--logging-steps", default=3, type=int)
+    parser.add_argument("--logging-steps", default=10, type=int)
     parser.add_argument("--lora-alpha", default=32, type=int)
     parser.add_argument("--lora-dropout", default=0.05, type=float)
     parser.add_argument("--lora-r", default=16, type=int)
+    parser.add_argument("--4bit", action="store_true")
     args = parser.parse_args()
-
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.bfloat16,
-        bnb_4bit_quant_storage=torch.bfloat16,
-    )
+    
+    if args.4bit:
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_quant_storage=torch.bfloat16,
+        )
+	else:
+		bnb_config = None
 
     processor = AutoProcessor.from_pretrained(args.model, use_fast=True)