Update finetune.py

finetune.py

@@ -31,23 +31,36 @@ if torch.cuda.is_available():
 # 3. Dataset
 dataset = load_dataset("lamekemal/brvm_finetune")
 
+#bnb_config = BitsAndBytesConfig(
+#    load_in_8bit=True,              # ou load_in_4bit=True pour encore moins de VRAM
+#    bnb_8bit_use_fp32_cpu_offload=True,  # décharge certains modules sur CPU
+#)
+# Configuration pour 8-bit + offload CPU
 bnb_config = BitsAndBytesConfig(
-    load_in_8bit=True,              # ou load_in_4bit=True pour encore moins de VRAM
-    bnb_8bit_use_fp32_cpu_offload=True,  # décharge certains modules sur CPU
+    load_in_8bit=True,
+    llm_int8_enable_fp32_cpu_offload=True
 )
 
+
 # 4. Charger modèle + tokenizer (FP16 natif, pas de quantization)
+#model = AutoModelForCausalLM.from_pretrained(
+#    base_model,
+#    torch_dtype=torch.float16,
+#    device_map="auto",
+#    trust_remote_code=True,
+#    quantization_config=bnb_config
+#)
 model = AutoModelForCausalLM.from_pretrained(
-    base_model,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True,
-    quantization_config=bnb_config
+    "mistralai/Mistral-7B-Instruct-v0.3",
+    quantization_config=bnb_config,
+    device_map="auto"  # ou tu peux préciser {'': 'cpu'} pour offload complet
 )
 model.config.use_cache = False
 model.gradient_checkpointing_enable()
 
-tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
+#tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = "right"