nishant-research commited on
Commit
a7744d6
·
verified ·
1 Parent(s): 3ee61d6

Upload train_qwen_codeforces.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train_qwen_codeforces.py +85 -0
train_qwen_codeforces.py ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "transformers>=4.37.0", "datasets", "torch"]
3
+ # ///
4
+
5
+ from datasets import load_dataset
6
+ from peft import LoraConfig
7
+ from trl import SFTTrainer, SFTConfig
8
+ import trackio
9
+
10
+ print("Loading dataset: open-r1/codeforces-cots...")
11
+ dataset = load_dataset("open-r1/codeforces-cots", "solutions", split="train")
12
+
13
+ # Take a subset for quick training (t4-small is memory-constrained)
14
+ print(f"Original dataset size: {len(dataset)}")
15
+ dataset = dataset.select(range(min(500, len(dataset))))
16
+ print(f"Using subset: {len(dataset)} examples")
17
+
18
+ # Create small eval split for monitoring
19
+ dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
20
+ print(f"Train: {len(dataset_split['train'])}, Eval: {len(dataset_split['test'])}")
21
+
22
+ # Configure LoRA for efficient training
23
+ lora_config = LoraConfig(
24
+ r=16,
25
+ lora_alpha=32,
26
+ lora_dropout=0.05,
27
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
28
+ bias="none",
29
+ task_type="CAUSAL_LM"
30
+ )
31
+
32
+ print("Initializing trainer...")
33
+ trainer = SFTTrainer(
34
+ model="Qwen/Qwen2.5-0.5B-Instruct",
35
+ train_dataset=dataset_split["train"],
36
+ eval_dataset=dataset_split["test"],
37
+ peft_config=lora_config,
38
+ args=SFTConfig(
39
+ output_dir="sr-test-qwen-codeforces-ft",
40
+
41
+ # Training hyperparameters optimized for t4-small
42
+ num_train_epochs=1,
43
+ per_device_train_batch_size=1,
44
+ per_device_eval_batch_size=1,
45
+ gradient_accumulation_steps=8, # Effective batch size = 8
46
+ gradient_checkpointing=True,
47
+
48
+ # Learning rate
49
+ learning_rate=2e-4,
50
+ warmup_ratio=0.03,
51
+ lr_scheduler_type="cosine",
52
+
53
+ # Logging and evaluation
54
+ logging_steps=10,
55
+ eval_strategy="steps",
56
+ eval_steps=50,
57
+ save_strategy="steps",
58
+ save_steps=100,
59
+ save_total_limit=2,
60
+
61
+ # Memory optimization
62
+ optim="adamw_torch",
63
+ bf16=True, # Use bf16 if supported, else will fall back to fp32
64
+
65
+ # Hub configuration
66
+ push_to_hub=True,
67
+ hub_model_id="nishant-research/sr-test-qwen-codeforces-ft",
68
+ hub_strategy="every_save",
69
+ hub_private_repo=False,
70
+
71
+ # Trackio monitoring
72
+ report_to="trackio",
73
+ project="qwen-codeforces-training",
74
+ run_name="qwen2.5-0.5b-codeforces-ft-test",
75
+ )
76
+ )
77
+
78
+ print("Starting training...")
79
+ trainer.train()
80
+
81
+ print("Training complete! Pushing final model to Hub...")
82
+ trainer.push_to_hub()
83
+
84
+ print("✅ Training job completed successfully!")
85
+ print(f"Model saved to: https://huggingface.co/nishant-research/sr-test-qwen-codeforces-ft")