Upload train_code_reasoning.py

train_code_reasoning.py

@@ -13,6 +13,7 @@
 import os
 import random
 from datasets import load_dataset, concatenate_datasets
+from transformers import AutoTokenizer
 from trl import SFTTrainer, SFTConfig
 import trackio
 
@@ -24,6 +25,9 @@ OUTPUT_DIR = "./code-reasoning-1.5b"
 # Initialize Trackio
 trackio.init(project="code-reasoning-ft", name="qwen2.5-coder-1.5b-code-reasoning")
 
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
 print("Loading and preparing datasets...")
 
 all_datasets = []
@@ -159,8 +163,8 @@ training_args = SFTConfig(
     hub_model_id=HUB_MODEL_ID,
     push_to_hub=True,
     num_train_epochs=2,
-    per_device_train_batch_size=4,
-    gradient_accumulation_steps=4,
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=8,
     learning_rate=5e-5,
     warmup_steps=300,
     lr_scheduler_type="cosine",
@@ -171,7 +175,7 @@ training_args = SFTConfig(
     logging_first_step=True,
     save_strategy="steps",
     save_steps=10,
-    packing=True,
+    packing=False,
     dataset_num_proc=4,
     disable_tqdm=True,
     report_to=["trackio"],
@@ -184,6 +188,7 @@ trainer = SFTTrainer(
     model=MODEL_ID,
     train_dataset=train_dataset,
     args=training_args,
+    processing_class=tokenizer,
 )
 
 print("Starting training...")