| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import os |
| import torch |
| from datasets import load_dataset |
| from peft import LoraConfig |
| from trl import SFTTrainer, SFTConfig |
| from transformers import AutoModelForCausalLM, BitsAndBytesConfig |
| import trackio |
|
|
| hf_token = os.getenv("HF_TOKEN") |
| if hf_token: |
| os.environ.setdefault("HF_HUB_TOKEN", hf_token) |
|
|
| DATASET = os.getenv("DATASET", "oExz/reasoning-distill-claude-opus-4-7-max-qwen-format") |
| BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen3.5-4B") |
| OUTPUT_REPO = os.getenv("OUTPUT_REPO", "oExz/qwen3.5-4b-reasoning-lora") |
|
|
| print(f"Loading dataset {DATASET}...") |
| dataset = load_dataset(DATASET) |
| train_ds = dataset["train"] |
| eval_ds = dataset["eval"] |
| print(f"Train: {len(train_ds)}, Eval: {len(eval_ds)}") |
|
|
| print(f"Loading base model {BASE_MODEL} with 4-bit quantization...") |
|
|
| bnb_config = BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| bnb_4bit_use_double_quant=True, |
| ) |
|
|
| model = AutoModelForCausalLM.from_pretrained( |
| BASE_MODEL, |
| trust_remote_code=True, |
| quantization_config=bnb_config, |
| device_map="auto", |
| ) |
|
|
| peft_config = LoraConfig( |
| r=16, |
| lora_alpha=32, |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], |
| lora_dropout=0.05, |
| bias="none", |
| task_type="CAUSAL_LM", |
| ) |
|
|
| trackio.init( |
| project="qwen-reasoning-finetune", |
| space_id="oExz/trackio", |
| ) |
|
|
| trainer = SFTTrainer( |
| model=model, |
| train_dataset=train_ds, |
| eval_dataset=eval_ds, |
| peft_config=peft_config, |
| args=SFTConfig( |
| output_dir="/tmp/qwen-lora", |
| push_to_hub=True, |
| hub_model_id=OUTPUT_REPO, |
| hub_strategy="every_save", |
| num_train_epochs=2, |
| per_device_train_batch_size=1, |
| gradient_accumulation_steps=8, |
| learning_rate=2e-4, |
| lr_scheduler_type="cosine", |
| warmup_ratio=0.1, |
| eval_strategy="steps", |
| eval_steps=50, |
| save_strategy="steps", |
| save_steps=50, |
| logging_steps=10, |
| report_to="trackio", |
| run_name="qwen3.5-4b-reasoning-qlora", |
| max_length=1024, |
| bf16=True, |
| gradient_checkpointing=True, |
| ), |
| ) |
|
|
| print("Starting training...") |
| trainer.train() |
| print("Training complete, pushing to hub...") |
| trainer.push_to_hub() |
| print("Done!") |
| trackio.finish() |
|
|
