luiscosio commited on
Commit
2f199c0
·
verified ·
1 Parent(s): 3e3fc0e

Upload train_qwen3_codeforces_v2.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train_qwen3_codeforces_v2.py +96 -0
train_qwen3_codeforces_v2.py ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # requires-python = ">=3.10"
3
+ # dependencies = [
4
+ # "trl>=0.12.0",
5
+ # "peft>=0.7.0",
6
+ # "transformers>=4.45.0",
7
+ # "datasets>=2.18.0",
8
+ # "accelerate>=0.30.0",
9
+ # "torch>=2.0.0",
10
+ # ]
11
+ # ///
12
+
13
+ import os
14
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
15
+
16
+ from datasets import load_dataset
17
+ from peft import LoraConfig
18
+ from trl import SFTTrainer, SFTConfig
19
+ from transformers import AutoModelForCausalLM, AutoTokenizer
20
+
21
+ print("Loading tokenizer...")
22
+ tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)
23
+ if tokenizer.pad_token is None:
24
+ tokenizer.pad_token = tokenizer.eos_token
25
+
26
+ print("Loading model...")
27
+ model = AutoModelForCausalLM.from_pretrained(
28
+ "Qwen/Qwen3-0.6B",
29
+ torch_dtype="auto",
30
+ trust_remote_code=True,
31
+ )
32
+
33
+ print("Loading dataset...")
34
+ dataset = load_dataset("open-r1/codeforces-cots", "solutions_py_decontaminated", split="train")
35
+ print(f"Dataset size: {len(dataset)}")
36
+
37
+ # Take a subset for faster training (full dataset is large)
38
+ dataset = dataset.shuffle(seed=42).select(range(min(10000, len(dataset))))
39
+ print(f"Using {len(dataset)} examples")
40
+
41
+ # Split
42
+ split = dataset.train_test_split(test_size=0.05, seed=42)
43
+ train_dataset = split["train"]
44
+ eval_dataset = split["test"]
45
+
46
+ print("Setting up LoRA...")
47
+ peft_config = LoraConfig(
48
+ r=16,
49
+ lora_alpha=32,
50
+ lora_dropout=0.05,
51
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
52
+ bias="none",
53
+ task_type="CAUSAL_LM",
54
+ )
55
+
56
+ print("Setting up training...")
57
+ training_args = SFTConfig(
58
+ output_dir="qwen3-0.6b-codeforces-sft",
59
+ push_to_hub=True,
60
+ hub_model_id="luiscosio/qwen3-0.6b-codeforces-sft",
61
+ num_train_epochs=3,
62
+ per_device_train_batch_size=2,
63
+ gradient_accumulation_steps=8,
64
+ gradient_checkpointing=True,
65
+ learning_rate=2e-4,
66
+ lr_scheduler_type="cosine",
67
+ warmup_ratio=0.1,
68
+ eval_strategy="steps",
69
+ eval_steps=100,
70
+ save_strategy="steps",
71
+ save_steps=100,
72
+ save_total_limit=3,
73
+ logging_steps=10,
74
+ bf16=True,
75
+ max_length=2048,
76
+ report_to="none",
77
+ )
78
+
79
+ print("Creating trainer...")
80
+ trainer = SFTTrainer(
81
+ model=model,
82
+ tokenizer=tokenizer,
83
+ train_dataset=train_dataset,
84
+ eval_dataset=eval_dataset,
85
+ peft_config=peft_config,
86
+ args=training_args,
87
+ )
88
+
89
+ print("Starting training...")
90
+ trainer.train()
91
+
92
+ print("Saving and pushing to Hub...")
93
+ trainer.save_model()
94
+ trainer.push_to_hub()
95
+
96
+ print("Done!")