| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import os |
|
|
| |
| from unsloth import FastLanguageModel |
| from datasets import load_dataset |
| from trl import SFTTrainer, SFTConfig |
| from transformers import AutoTokenizer |
| 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-text") |
| 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 Unsloth...") |
| model, _ = FastLanguageModel.from_pretrained( |
| model_name=BASE_MODEL, |
| max_seq_length=1024, |
| dtype=None, |
| load_in_4bit=True, |
| trust_remote_code=True, |
| ) |
|
|
| |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) |
| |
| tokenizer.eos_token = "<|im_end|>" |
| tokenizer.pad_token = tokenizer.eos_token |
| tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token) |
| tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids(tokenizer.eos_token) |
|
|
| model = FastLanguageModel.get_peft_model( |
| model, |
| r=16, |
| lora_alpha=32, |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], |
| lora_dropout=0.05, |
| bias="none", |
| use_gradient_checkpointing="unsloth", |
| random_state=3407, |
| ) |
|
|
| trackio.init( |
| project="qwen-reasoning-finetune", |
| space_id="oExz/trackio", |
| ) |
|
|
| trainer = SFTTrainer( |
| model=model, |
| processing_class=tokenizer, |
| train_dataset=train_ds, |
| eval_dataset=eval_ds, |
| 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=2, |
| gradient_accumulation_steps=4, |
| 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, |
| dataset_text_field="text", |
| packing=False, |
| eos_token="<|im_end|>", |
| ), |
| ) |
|
|
| print("Starting training...") |
| trainer.train() |
| print("Training complete, pushing to hub...") |
| trainer.push_to_hub() |
| print("Done!") |
| trackio.finish() |
|
|
