Upload train_sft_qwen.py with huggingface_hub

train_sft_qwen.py

@@ -1,10 +1,17 @@
 # /// script
-# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "torch", "transformers>=4.40.0"]
+# dependencies = [
+#   "trl>=0.12.0",
+#   "peft>=0.7.0",
+#   "trackio",
+#   "transformers>=4.40.0",
+#   "datasets>=2.18.0",
+#   "torch>=2.0.0",
+# ]
 # ///
 
 """
-SFT Fine-tuning Script for Qwen/Qwen2.5-0.5B
-Optimized for Hugging Face Jobs with Trackio monitoring
+SFT (Supervised Fine-Tuning) training script for Qwen/Qwen2.5-0.5B
+Uses TRL with LoRA, Trackio monitoring, and automatic Hub push
 """
 
 from datasets import load_dataset
@@ -12,81 +19,83 @@ from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
 import trackio
 
-# Load dataset - using TRL-compatible Capybara dataset
-print("Loading dataset...")
+# Load a high-quality instruction dataset
 dataset = load_dataset("trl-lib/Capybara", split="train")
 
-# Create train/eval split for monitoring progress
+# Create train/eval split for monitoring training progress
 dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
-print(f"Train size: {len(dataset_split['train'])}, Eval size: {len(dataset_split['test'])}")
 
-# LoRA configuration for efficient fine-tuning
+# Configure LoRA for efficient fine-tuning
 peft_config = LoraConfig(
-    r=16,                    # LoRA rank
-    lora_alpha=32,           # LoRA alpha (scaling factor)
-    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Attention layers
-    lora_dropout=0.05,       # Dropout for regularization
-    bias="none",             # Don't train bias terms
-    task_type="CAUSAL_LM"    # Causal language modeling
+    r=16,                      # LoRA rank
+    lora_alpha=32,             # LoRA alpha scaling
+    lora_dropout=0.05,         # Dropout for regularization
+    bias="none",               # Don't train bias parameters
+    task_type="CAUSAL_LM",     # Causal language modeling
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Qwen attention layers
 )
 
-# Training configuration
-training_args = SFTConfig(
-    # Output and Hub settings
-    output_dir="qwen-0.5b-sft-capybara",
+# Configure trainer
+training_config = SFTConfig(
+    # Model and output
+    output_dir="qwen-sft-capybara",
+
+    # Hub configuration - CRITICAL for saving results
     push_to_hub=True,
-    hub_model_id="vgtomahawk/qwen-0.5b-sft-capybara",
-    hub_strategy="every_save",
+    hub_model_id="qwen-sft-capybara-demo",  # Will use format: username/qwen-sft-capybara-demo
+    hub_strategy="every_save",              # Push checkpoints during training
     hub_private_repo=False,
 
     # Training parameters
     num_train_epochs=3,
-    per_device_train_batch_size=4,
-    per_device_eval_batch_size=4,
-    gradient_accumulation_steps=4,  # Effective batch size: 4 * 4 = 16
-    gradient_checkpointing=True,
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=4,  # Effective batch size: 2 * 4 = 8
 
     # Optimization
     learning_rate=2e-4,
     lr_scheduler_type="cosine",
     warmup_ratio=0.1,
-    optim="paged_adamw_8bit",  # Memory-efficient optimizer
 
-    # Evaluation and logging
+    # Evaluation
     eval_strategy="steps",
     eval_steps=50,
-    logging_steps=10,
+    per_device_eval_batch_size=2,
+
+    # Checkpointing
     save_strategy="steps",
     save_steps=100,
     save_total_limit=3,  # Keep only last 3 checkpoints
 
-    # Trackio monitoring
+    # Logging - Trackio integration
+    logging_steps=10,
     report_to="trackio",
-    run_name="qwen-0.5b-sft-capybara-test",
+    run_name="qwen-0.5b-sft-demo",
+
+    # Performance optimization
+    bf16=True,  # Use bfloat16 for better performance on modern GPUs
+    gradient_checkpointing=True,  # Reduce memory usage
 
-    # Performance
-    bf16=True,  # Use bfloat16 for better numerical stability
+    # Misc
+    seed=42,
     dataloader_num_workers=4,
-    remove_unused_columns=True,
 )
 
 # Initialize trainer
-print("Initializing SFT Trainer...")
 trainer = SFTTrainer(
     model="Qwen/Qwen2.5-0.5B",
     train_dataset=dataset_split["train"],
     eval_dataset=dataset_split["test"],
     peft_config=peft_config,
-    args=training_args,
+    args=training_config,
 )
 
 # Train the model
 print("Starting training...")
 trainer.train()
 
-# Save final model to Hub
-print("Pushing final model to Hub...")
-trainer.push_to_hub(commit_message="Training completed")
+# Final push to Hub
+print("Training complete! Pushing final model to Hub...")
+trainer.push_to_hub()
 
-print("✅ Training completed successfully!")
-print(f"Model saved to: https://huggingface.co/{training_args.hub_model_id}")
+print("✅ Training complete and model saved to Hub!")
+print(f"Model available at: https://huggingface.co/{trainer.hub_model_id}")