kintopp commited on
Commit
eb0b63b
·
verified ·
1 Parent(s): 1e0027e

Upload train_qwen_codeforces_v2.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train_qwen_codeforces_v2.py +114 -0
train_qwen_codeforces_v2.py ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = [
3
+ # "trl>=0.12.0",
4
+ # "peft>=0.7.0",
5
+ # "transformers>=4.45.0",
6
+ # "accelerate>=0.24.0",
7
+ # "trackio",
8
+ # "datasets",
9
+ # ]
10
+ # ///
11
+
12
+ """
13
+ Fine-tune Qwen3-0.6B on open-r1/codeforces-cots for competitive programming.
14
+ """
15
+
16
+ import trackio
17
+ from datasets import load_dataset
18
+ from peft import LoraConfig
19
+ from trl import SFTTrainer, SFTConfig
20
+ from transformers import AutoTokenizer
21
+
22
+ # Load dataset - using the solutions config with messages column
23
+ print("Loading dataset...")
24
+ dataset = load_dataset("open-r1/codeforces-cots", "solutions", split="train")
25
+ print(f"Dataset loaded: {len(dataset)} examples")
26
+
27
+ # The dataset has a 'messages' column in chat format
28
+ # We need to keep only the 'messages' column for SFT training
29
+ print("Preparing dataset - keeping only messages column...")
30
+ dataset = dataset.select_columns(["messages"])
31
+
32
+ # Create train/eval split
33
+ print("Creating train/eval split...")
34
+ dataset_split = dataset.train_test_split(test_size=0.05, seed=42)
35
+ train_dataset = dataset_split["train"]
36
+ eval_dataset = dataset_split["test"]
37
+ print(f"Train: {len(train_dataset)} examples")
38
+ print(f"Eval: {len(eval_dataset)} examples")
39
+
40
+ # Load tokenizer for chat template
41
+ print("Loading tokenizer...")
42
+ tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
43
+ if tokenizer.pad_token is None:
44
+ tokenizer.pad_token = tokenizer.eos_token
45
+
46
+ # Processing function to convert messages to text using chat template
47
+ def formatting_func(example):
48
+ return tokenizer.apply_chat_template(example["messages"], tokenize=False)
49
+
50
+ # Training configuration
51
+ config = SFTConfig(
52
+ # Hub settings - CRITICAL for saving results
53
+ output_dir="qwen3-codeforces-sft",
54
+ push_to_hub=True,
55
+ hub_model_id="kintopp/qwen3-0.6b-codeforces-cots",
56
+ hub_strategy="every_save",
57
+
58
+ # Training parameters
59
+ num_train_epochs=1,
60
+ per_device_train_batch_size=2,
61
+ gradient_accumulation_steps=8,
62
+ learning_rate=2e-5,
63
+ max_length=2048,
64
+
65
+ # Logging & checkpointing
66
+ logging_steps=25,
67
+ save_strategy="steps",
68
+ save_steps=500,
69
+ save_total_limit=2,
70
+
71
+ # Evaluation
72
+ eval_strategy="steps",
73
+ eval_steps=500,
74
+
75
+ # Optimization
76
+ warmup_ratio=0.1,
77
+ lr_scheduler_type="cosine",
78
+ gradient_checkpointing=True,
79
+ bf16=True,
80
+
81
+ # Monitoring with Trackio
82
+ report_to="trackio",
83
+ project="qwen3-codeforces",
84
+ run_name="qwen3-0.6b-codeforces-sft",
85
+ )
86
+
87
+ # LoRA configuration for efficient training
88
+ peft_config = LoraConfig(
89
+ r=16,
90
+ lora_alpha=32,
91
+ lora_dropout=0.05,
92
+ bias="none",
93
+ task_type="CAUSAL_LM",
94
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
95
+ )
96
+
97
+ # Initialize trainer
98
+ print("Initializing trainer...")
99
+ trainer = SFTTrainer(
100
+ model="Qwen/Qwen3-0.6B",
101
+ train_dataset=train_dataset,
102
+ eval_dataset=eval_dataset,
103
+ args=config,
104
+ peft_config=peft_config,
105
+ formatting_func=formatting_func,
106
+ )
107
+
108
+ print("Starting training...")
109
+ trainer.train()
110
+
111
+ print("Pushing final model to Hub...")
112
+ trainer.push_to_hub()
113
+
114
+ print("Complete! Model at: https://huggingface.co/kintopp/qwen3-0.6b-codeforces-cots")