gilbaes commited on
Commit
6605eeb
·
verified ·
1 Parent(s): 1dfdb43

Upload train_sft.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train_sft.py +106 -0
train_sft.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = [
3
+ # "trl>=0.12.0",
4
+ # "peft>=0.7.0",
5
+ # "transformers>=4.51.0",
6
+ # "accelerate>=0.24.0",
7
+ # "trackio",
8
+ # "datasets",
9
+ # ]
10
+ # ///
11
+
12
+ """
13
+ SFT training: Qwen3-0.6B on open-r1/codeforces-cots for instruction following.
14
+ Quick test run with 100 examples.
15
+ """
16
+
17
+ import trackio
18
+ from datasets import load_dataset
19
+ from peft import LoraConfig
20
+ from trl import SFTTrainer, SFTConfig
21
+
22
+ # Load dataset - using solutions_w_editorials config for instruction following
23
+ print("Loading dataset...")
24
+ dataset = load_dataset(
25
+ "open-r1/codeforces-cots",
26
+ "solutions_w_editorials",
27
+ split="train"
28
+ )
29
+ print(f"Full dataset: {len(dataset)} examples")
30
+
31
+ # Take 100 examples for quick test
32
+ dataset = dataset.select(range(min(100, len(dataset))))
33
+ print(f"Using {len(dataset)} examples for training")
34
+
35
+ # Create train/eval split (90/10)
36
+ dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
37
+ train_dataset = dataset_split["train"]
38
+ eval_dataset = dataset_split["test"]
39
+ print(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
40
+
41
+ # Training configuration
42
+ config = SFTConfig(
43
+ output_dir="qwen3-codeforces-sft",
44
+ push_to_hub=True,
45
+ hub_model_id="gilbaes/qwen3-0.6b-codeforces-sft",
46
+ hub_strategy="every_save",
47
+
48
+ # Training parameters
49
+ num_train_epochs=1,
50
+ per_device_train_batch_size=2,
51
+ gradient_accumulation_steps=4,
52
+ learning_rate=2e-4,
53
+ max_length=2048,
54
+
55
+ # Logging & checkpoints
56
+ logging_steps=5,
57
+ save_strategy="steps",
58
+ save_steps=50,
59
+ save_total_limit=2,
60
+
61
+ # Evaluation
62
+ eval_strategy="steps",
63
+ eval_steps=25,
64
+
65
+ # Optimization
66
+ warmup_ratio=0.1,
67
+ lr_scheduler_type="cosine",
68
+ bf16=True,
69
+
70
+ # Monitoring
71
+ report_to="trackio",
72
+ project="qwen3-codeforces",
73
+ run_name="sft-test-100examples",
74
+ )
75
+
76
+ # LoRA configuration
77
+ peft_config = LoraConfig(
78
+ r=16,
79
+ lora_alpha=32,
80
+ lora_dropout=0.05,
81
+ bias="none",
82
+ task_type="CAUSAL_LM",
83
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
84
+ )
85
+
86
+ # Initialize trainer
87
+ print("Initializing trainer with Qwen3-0.6B...")
88
+ trainer = SFTTrainer(
89
+ model="Qwen/Qwen3-0.6B",
90
+ train_dataset=train_dataset,
91
+ eval_dataset=eval_dataset,
92
+ args=config,
93
+ peft_config=peft_config,
94
+ )
95
+
96
+ # Train
97
+ print("Starting training...")
98
+ trainer.train()
99
+
100
+ # Push to Hub
101
+ print("Pushing model to Hub...")
102
+ trainer.push_to_hub()
103
+
104
+ print("Training complete!")
105
+ print(f"Model: https://huggingface.co/gilbaes/qwen3-0.6b-codeforces-sft")
106
+ print(f"Trackio: https://huggingface.co/spaces/gilbaes/trackio")