File size: 2,879 Bytes
b3d682a f4b0dc7 b3d682a 908b4bd b3d682a 15685a3 b3d682a 908b4bd b3d682a f4b0dc7 908b4bd f4b0dc7 5773257 908b4bd f4b0dc7 772f77b b3d682a d7ff6b6 b3d682a 15685a3 f4b0dc7 b3d682a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | # /// script
# dependencies = [
# "unsloth",
# "trl>=0.12.0",
# "peft>=0.7.0",
# "trackio",
# "datasets",
# "transformers>=4.40.0",
# "torch>=2.0.0",
# "huggingface-hub",
# ]
# ///
import os
# Import unsloth FIRST before anything else
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,
)
# Load tokenizer directly
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
# Explicitly set special tokens to avoid TokenizersBackend defaults
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()
|