| |
| |
| |
| """Fine-tune Qwen3-0.6B on CodeForces-CoTS (100 examples)""" |
|
|
| import os |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" |
|
|
| from datasets import load_dataset |
| from peft import LoraConfig |
| from trl import SFTTrainer, SFTConfig |
| import torch |
|
|
| print(f"CUDA available: {torch.cuda.is_available()}") |
| if torch.cuda.is_available(): |
| print(f"GPU: {torch.cuda.get_device_name(0)}") |
|
|
| dataset = load_dataset("open-r1/codeforces-cots", "solutions", split="train").select(range(100)) |
| print(f"Dataset: {len(dataset)} examples") |
|
|
| peft_config = LoraConfig( |
| r=16, lora_alpha=32, lora_dropout=0.05, |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], |
| bias="none", task_type="CAUSAL_LM" |
| ) |
|
|
| training_args = SFTConfig( |
| output_dir="./qwen3-0.6b-codeforces-cots", |
| num_train_epochs=1, |
| per_device_train_batch_size=1, |
| gradient_accumulation_steps=8, |
| learning_rate=2e-4, |
| warmup_ratio=0.1, |
| logging_steps=5, |
| save_strategy="no", |
| eval_strategy="no", |
| max_length=2048, |
| push_to_hub=True, |
| hub_model_id="gilbaes/qwen3-0.6b-codeforces-cots", |
| report_to="none", |
| bf16=True, |
| gradient_checkpointing=True, |
| optim="adamw_torch_fused", |
| ) |
|
|
| trainer = SFTTrainer( |
| model="Qwen/Qwen3-0.6B", |
| train_dataset=dataset, |
| peft_config=peft_config, |
| args=training_args, |
| ) |
|
|
| print(f"Trainable params: {trainer.model.num_parameters(only_trainable=True):,}") |
| trainer.train() |
| trainer.push_to_hub() |
| print("Done! Model at gilbaes/qwen3-0.6b-codeforces-cots") |
|
|