Increase grad_accum to 8, add gc.collect and empty_cache before training

train.py

@@ -1,9 +1,7 @@
 """
-Full fine-tuning script:
+Full fine-tuning script with aggressive memory optimizations:
   Model: google/gemma-2-2b-it
   Dataset: talkmap/telecom-conversation-corpus
-  Converts turn-based telecom dialogues into alternating-role conversational messages for SFT.
-  Manually tokenizes with truncation to control sequence length and avoid OOM.
 """
 import os
 from collections import defaultdict
@@ -40,7 +38,6 @@ print("Loading dataset...")
 ds = load_dataset(DATASET_ID, split="train")
 print(f"Rows: {len(ds)}, Columns: {ds.column_names}")
 
-# Group rows by conversation_id and sort by date_time
 print("Grouping conversations...")
 conversations = defaultdict(list)
 for row in ds:
@@ -54,7 +51,6 @@ for row in ds:
 for conv_id in conversations:
     conversations[conv_id].sort(key=lambda x: x["date_time"])
 
-# Convert each conversation into messages format with merged consecutive same-role turns
 print("Converting to messages format...")
 messages_data = []
 for conv_id, turns in conversations.items():
@@ -107,7 +103,7 @@ if tokenizer.pad_token is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
 # ------------------------------------------------------------------
-# Pre-tokenize dataset with truncation (avoids SFTTrainer auto-tokenization)
+# Pre-tokenize dataset with truncation
 # ------------------------------------------------------------------
 print("Pre-tokenizing dataset...")
 
@@ -125,11 +121,10 @@ def apply_and_tokenize(example):
 raw_dataset = Dataset.from_list(messages_data)
 raw_dataset = raw_dataset.map(apply_and_tokenize, remove_columns=["messages"])
 raw_dataset = raw_dataset.filter(lambda x: len(x["text"]) > 0)
-
 print(f"Dataset after filtering: {len(raw_dataset)}")
 
 # ------------------------------------------------------------------
-# Model
+# Model - load on CPU first to control placement
 # ------------------------------------------------------------------
 print("Loading model...")
 model = AutoModelForCausalLM.from_pretrained(
@@ -140,6 +135,11 @@ model = AutoModelForCausalLM.from_pretrained(
 
 model.gradient_checkpointing_enable()
 
+# Empty cache to free up fragmentation
+import gc
+gc.collect()
+torch.cuda.empty_cache()
+
 # ------------------------------------------------------------------
 # Training arguments
 # ------------------------------------------------------------------
@@ -149,7 +149,7 @@ args = SFTConfig(
     push_to_hub=True,
     num_train_epochs=1,
     per_device_train_batch_size=1,
-    gradient_accumulation_steps=4,
+    gradient_accumulation_steps=8,
     learning_rate=2e-5,
     max_seq_length=MAX_SEQ_LENGTH,
     logging_strategy="steps",