Upload train_qwen.py with huggingface_hub

train_qwen.py

@@ -2,6 +2,7 @@
 # dependencies = [
 #     "trl>=0.12.0",
 #     "peft>=0.7.0",
+#     "bitsandbytes",
 #     "trackio",
 #     "datasets",
 #     "transformers>=4.40.0",
@@ -13,9 +14,11 @@
 # ///
 
 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")
@@ -32,7 +35,21 @@ 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}...")
+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,
@@ -49,7 +66,7 @@ trackio.init(
 )
 
 trainer = SFTTrainer(
-    model=BASE_MODEL,
+    model=model,
     train_dataset=train_ds,
     eval_dataset=eval_ds,
     peft_config=peft_config,
@@ -71,7 +88,7 @@ trainer = SFTTrainer(
         logging_steps=10,
         report_to="trackio",
         run_name="qwen3.5-4b-reasoning-qlora",
-        max_length=2048,
+        max_length=1024,
         bf16=True,
         gradient_checkpointing=True,
     ),