Update src/train.py

src/train.py

@@ -6,7 +6,7 @@ import networkx as nx
 from tqdm import tqdm
 from peft import (LoraConfig, get_peft_model,
                   prepare_model_for_kbit_training)
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 
 
@@ -30,11 +30,17 @@ class QloraTrainer_CS:
         model_id = self.config['inference']["base_model"]
         print(model_id)
 
+        bnb_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+            bnb_8bit_use_double_quant=True,
+            bnb_8bit_quant_type="nf8",
+            bnb_8bit_compute_dtype=torch.bfloat16
+        )
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         tokenizer.model_max_length = self.config['training']['tokenizer']["max_length"]
         if not tokenizer.pad_token:
             tokenizer.pad_token = tokenizer.eos_token
-        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+        model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, torch_dtype=torch.bfloat16)
         if model.device.type != 'cuda':
             model.to('cuda')