gilbaes commited on
Commit
be41874
·
verified ·
1 Parent(s): 413b449

Upload train_qwen3_hf_v3.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. train_qwen3_hf_v3.py +87 -0
train_qwen3_hf_v3.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # /// script
2
+ # dependencies = ["trl>=0.12.0", "peft>=0.7.0", "transformers>=4.45.0", "datasets", "accelerate", "torch"]
3
+ # ///
4
+ """Fine-tune Qwen3-0.6B on CodeForces-CoTS (100 examples)"""
5
+
6
+ import os
7
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
8
+
9
+ from datasets import load_dataset
10
+ from peft import LoraConfig
11
+ from trl import SFTTrainer, SFTConfig
12
+ import torch
13
+
14
+ print(f"CUDA available: {torch.cuda.is_available()}")
15
+ if torch.cuda.is_available():
16
+ print(f"GPU: {torch.cuda.get_device_name(0)}")
17
+ print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
18
+
19
+ # Load 100 examples
20
+ print("\nLoading dataset...")
21
+ dataset = load_dataset("open-r1/codeforces-cots", "solutions", split="train").select(range(100))
22
+ print(f"Dataset: {len(dataset)} examples")
23
+
24
+ # Format for SFT: use the 'messages' column directly (already in chat format)
25
+ # The dataset has 'messages' column with list of {'role': ..., 'content': ...}
26
+ print(f"Sample messages: {dataset[0]['messages'][:1]}")
27
+
28
+ # Split: 90 train, 10 val
29
+ splits = dataset.train_test_split(test_size=0.1, seed=42)
30
+ train_ds, val_ds = splits["train"], splits["test"]
31
+ print(f"Train: {len(train_ds)}, Val: {len(val_ds)}")
32
+
33
+ peft_config = LoraConfig(
34
+ r=8,
35
+ lora_alpha=16,
36
+ lora_dropout=0.05,
37
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
38
+ bias="none",
39
+ task_type="CAUSAL_LM"
40
+ )
41
+
42
+ # 90 examples, batch=1, accum=4 -> ~22 steps/epoch
43
+ training_args = SFTConfig(
44
+ output_dir="./qwen3-0.6b-codeforces-cots",
45
+ num_train_epochs=1,
46
+ per_device_train_batch_size=1,
47
+ gradient_accumulation_steps=4,
48
+ learning_rate=2e-4,
49
+ warmup_ratio=0.1,
50
+ logging_steps=2,
51
+ logging_first_step=True,
52
+ save_strategy="no",
53
+ eval_strategy="steps",
54
+ eval_steps=5,
55
+ max_length=1024,
56
+ push_to_hub=True,
57
+ hub_model_id="gilbaes/qwen3-0.6b-codeforces-cots",
58
+ report_to="none",
59
+ bf16=True,
60
+ gradient_checkpointing=True,
61
+ optim="adamw_torch_fused",
62
+ dataset_text_field=None, # Use messages format
63
+ )
64
+
65
+ print("\nInitializing trainer...")
66
+ trainer = SFTTrainer(
67
+ model="Qwen/Qwen3-0.6B",
68
+ train_dataset=train_ds,
69
+ eval_dataset=val_ds,
70
+ peft_config=peft_config,
71
+ args=training_args,
72
+ )
73
+
74
+ print(f"Trainable params: {trainer.model.num_parameters(only_trainable=True):,}")
75
+ print(f"Total params: {trainer.model.num_parameters():,}")
76
+
77
+ print("\n" + "="*50)
78
+ print("TRAINING START")
79
+ print("="*50 + "\n")
80
+
81
+ trainer.train()
82
+
83
+ print("\n" + "="*50)
84
+ print("PUSHING TO HUB")
85
+ print("="*50)
86
+ trainer.push_to_hub()
87
+ print("\nDone! Model: https://huggingface.co/gilbaes/qwen3-0.6b-codeforces-cots")