ajustement to avoid oom err

train_with_unsloth.py

@@ -20,6 +20,7 @@ from datasets import load_dataset
 import os
 from transformers.integrations import WandbCallback
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
 PROJECT_NAME='SmolLM2-135M-Instruct-TaiwanChat'
 BASE_MODEL_ID="unsloth/SmolLM2-135M-Instruct"
@@ -27,6 +28,8 @@ DATASET_ID="yentinglin/TaiwanChat"
 N_SAMPLES=80000
 MAX_LEN=2048
 
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:128"
+
 # Tell wandb which project to use, and that you want to log your model
 os.environ["WANDB_PROJECT"]    = f"{PROJECT_NAME}_CLOUD"
 os.environ["WANDB_LOG_MODEL"]  = "end"
@@ -73,15 +76,15 @@ dataset = dataset.map(fmt, batched=True, remove_columns=["messages"])
 new_dataset = dataset.train_test_split(test_size = 0.01)
 
 training_args = SFTConfig(
-        fp16_full_eval = True,
-        per_device_eval_batch_size = 2,
+        fp16_full_eval = False,
+        per_device_eval_batch_size = 1,
         eval_accumulation_steps = 4,
-        eval_strategy = "steps",
+        eval_strategy = "epoch",
         eval_steps = 1,
         dataset_text_field="text",
         output_dir=PROJECT_NAME,
         max_seq_length = MAX_LEN,
-        per_device_train_batch_size = 2,
+        per_device_train_batch_size = 1,
         gradient_accumulation_steps = 4,
         warmup_steps = 10,
         max_steps = 60,
@@ -106,6 +109,7 @@ trainer = SFTTrainer(
     train_dataset = new_dataset["train"],
     eval_dataset = new_dataset["test"],
 )
+torch.cuda.empty_cache()
 trainer = train_on_responses_only(
   trainer,
   instruction_part = "<|im_start|>user\n",