SFT script

train_sft.py

@@ -0,0 +1,64 @@
+# /// script
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "torch"]
+# ///
+"""SFT on compile-verified C++ curriculum. Designed for Hugging Face Jobs (uv)."""
+
+import os
+
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTConfig, SFTTrainer
+
+DATASET_ID = os.environ.get("DATASET_ID", "gonzalolinares/cpp-compiler-curriculum")
+MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
+HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "gonzalolinares/qwen25-1.5b-cpp-sft")
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "qwen25-1.5b-cpp-sft")
+
+
+def main() -> None:
+    ds = load_dataset(DATASET_ID, split="train")
+    # Keep messages format for SFTTrainer chat templates
+    if "messages" not in ds.column_names:
+        raise SystemExit(f"Dataset must have 'messages'; got {ds.column_names}")
+
+    split = ds.train_test_split(test_size=0.1, seed=42)
+
+    trainer = SFTTrainer(
+        model=MODEL_ID,
+        train_dataset=split["train"],
+        eval_dataset=split["test"],
+        peft_config=LoraConfig(
+            r=16,
+            lora_alpha=32,
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        ),
+        args=SFTConfig(
+            output_dir=OUTPUT_DIR,
+            num_train_epochs=3,
+            per_device_train_batch_size=2,
+            per_device_eval_batch_size=2,
+            gradient_accumulation_steps=8,
+            learning_rate=2e-4,
+            logging_steps=5,
+            eval_strategy="steps",
+            eval_steps=20,
+            save_steps=50,
+            max_length=1024,
+            bf16=True,
+            push_to_hub=True,
+            hub_model_id=HUB_MODEL_ID,
+            report_to="trackio",
+            project="cpp-compiler-rl",
+            run_name="sft-qwen25-1.5b",
+        ),
+    )
+    trainer.train()
+    trainer.push_to_hub()
+    print(f"Pushed to {HUB_MODEL_ID}")
+
+
+if __name__ == "__main__":
+    main()