Upload train_unsloth.py with huggingface_hub

train_unsloth.py

@@ -0,0 +1,92 @@
+# /// 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
+from datasets import load_dataset
+from trl import SFTTrainer, SFTConfig
+from unsloth import FastLanguageModel
+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 Unsloth...")
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=BASE_MODEL,
+    max_seq_length=1024,
+    dtype=None,
+    load_in_4bit=True,
+    trust_remote_code=True,
+)
+
+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,
+    tokenizer=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,
+    ),
+)
+
+print("Starting training...")
+trainer.train()
+print("Training complete, pushing to hub...")
+trainer.push_to_hub()
+print("Done!")
+trackio.finish()