Upload train_design_sft.py with huggingface_hub

train_design_sft.py

@@ -1,5 +1,5 @@
 # /// script
-# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "accelerate", "torch"]
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "accelerate", "torch", "bitsandbytes"]
 # ///
 
 from datasets import load_dataset, concatenate_datasets
@@ -83,8 +83,8 @@ trainer = SFTTrainer(
     train_dataset=dataset_split["train"],
     eval_dataset=dataset_split["test"],
     peft_config=LoraConfig(
-        r=32,
-        lora_alpha=64,
+        r=16,  # Reduced from 32 to save memory
+        lora_alpha=32,
         target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
         lora_dropout=0.05,
         bias="none",
@@ -95,21 +95,22 @@ trainer = SFTTrainer(
         push_to_hub=True,
         hub_model_id="chaddy81/qwen3-0.6b-design-sft",
         num_train_epochs=3,
-        per_device_train_batch_size=2,
-        gradient_accumulation_steps=8,
+        per_device_train_batch_size=1,  # Reduced from 2
+        gradient_accumulation_steps=16,  # Increased to maintain effective batch size
         learning_rate=2e-4,
         lr_scheduler_type="cosine",
         warmup_ratio=0.1,
-        max_length=4096,
+        max_length=2048,  # Reduced from 4096 to save memory
         logging_steps=10,
-        eval_strategy="steps",
-        eval_steps=100,
+        eval_strategy="no",  # Skip eval during training to save memory
         save_strategy="steps",
-        save_steps=200,
+        save_steps=300,
         bf16=True,
         gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},  # More memory efficient
+        optim="adamw_8bit",  # 8-bit optimizer to save memory
         report_to="trackio",
-        run_name="qwen3-design-sft-v1",
+        run_name="qwen3-design-sft-v2",
     )
 )