Upload train_glm_qlora.py with huggingface_hub

train_glm_qlora.py

@@ -1,11 +1,8 @@
 # /// script
 # requires-python = ">=3.10"
 # dependencies = [
-#     "trl>=0.12.0",
-#     "peft>=0.7.0",
-#     "transformers @ git+https://github.com/huggingface/transformers.git",
-#     "accelerate>=0.24.0",
-#     "bitsandbytes>=0.41.0",
+#     "unsloth",
+#     "transformers>=4.49.0",
 #     "trackio",
 #     "datasets",
 # ]
@@ -13,16 +10,13 @@
 
 """
 Agent Zero SFT: zai-org/GLM-4.7-Flash (30B MoE)
-QLoRA (4-bit) fine-tuning on agent-zero-sft-v1 dataset.
+QLoRA (4-bit) fine-tuning with Unsloth on agent-zero-sft-v1 dataset.
 Router layers frozen - only attention layers trained.
 """
 
-import torch
 import trackio
 from datasets import load_dataset
-from peft import LoraConfig
-from transformers import BitsAndBytesConfig
-from trl import SFTTrainer, SFTConfig
+from unsloth import FastLanguageModel
 
 # Load dataset
 print("Loading dataset...")
@@ -30,14 +24,28 @@ train_ds = load_dataset("wheattoast11/agent-zero-sft-v1", data_files="data/train
 val_ds = load_dataset("wheattoast11/agent-zero-sft-v1", data_files="data/validation.jsonl", split="train")
 print(f"Train: {len(train_ds)}, Val: {len(val_ds)}")
 
-# 4-bit quantization config
-bnb_config = BitsAndBytesConfig(
+# Load model in 4-bit with Unsloth
+print("Loading model with Unsloth (4-bit QLoRA)...")
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name="unsloth/GLM-4.7-Flash",
+    max_seq_length=2048,
     load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
+    dtype=None,
 )
 
+# Apply LoRA adapters via Unsloth
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=16,
+    lora_alpha=32,
+    lora_dropout=0.05,
+    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+    bias="none",
+    use_gradient_checkpointing="unsloth",
+)
+
+from trl import SFTTrainer, SFTConfig
+
 config = SFTConfig(
     output_dir="agent-zero-glm-4.7-v1",
     push_to_hub=True,
@@ -50,7 +58,6 @@ config = SFTConfig(
     gradient_accumulation_steps=16,
     learning_rate=1e-4,
     bf16=True,
-    gradient_checkpointing=True,
 
     logging_steps=10,
     save_strategy="steps",
@@ -68,24 +75,13 @@ config = SFTConfig(
     run_name="glm-4.7-flash-qlora-v1",
 )
 
-# LoRA targeting attention layers only (router layers frozen)
-peft_config = LoraConfig(
-    r=16,
-    lora_alpha=32,
-    lora_dropout=0.05,
-    bias="none",
-    task_type="CAUSAL_LM",
-    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
-)
-
 print("Initializing trainer...")
 trainer = SFTTrainer(
-    model="zai-org/GLM-4.7-Flash",
+    model=model,
+    tokenizer=tokenizer,
     train_dataset=train_ds,
     eval_dataset=val_ds,
     args=config,
-    peft_config=peft_config,
-    model_init_kwargs={"quantization_config": bnb_config},
 )
 
 print("Starting training...")