Upload train_glm_qlora.py with huggingface_hub

train_glm_qlora.py

@@ -1,7 +1,11 @@
 # /// script
 # requires-python = ">=3.10"
 # dependencies = [
-#     "unsloth[cu124-ampere]",
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "transformers @ git+https://github.com/huggingface/transformers.git",
+#     "accelerate>=0.24.0",
+#     "bitsandbytes>=0.41.0",
 #     "trackio",
 #     "datasets",
 # ]
@@ -9,13 +13,16 @@
 
 """
 Agent Zero SFT: zai-org/GLM-4.7-Flash (30B MoE)
-QLoRA (4-bit) fine-tuning with Unsloth optimizations.
+QLoRA (4-bit) fine-tuning on agent-zero-sft-v1 dataset.
 Router layers frozen - only attention layers trained.
 """
 
+import torch
 import trackio
 from datasets import load_dataset
-from unsloth import FastLanguageModel
+from peft import LoraConfig
+from transformers import BitsAndBytesConfig
+from trl import SFTTrainer, SFTConfig
 
 # Load dataset
 print("Loading dataset...")
@@ -23,28 +30,14 @@ 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)}")
 
-# Load model in 4-bit with Unsloth
-print("Loading model with Unsloth (4-bit QLoRA)...")
-model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name="zai-org/GLM-4.7-Flash",
-    max_seq_length=2048,
+# 4-bit quantization config
+bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
-    dtype=None,  # auto-detect
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
 )
 
-# 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",  # Unsloth optimized
-)
-
-from trl import SFTTrainer, SFTConfig
-
 config = SFTConfig(
     output_dir="agent-zero-glm-4.7-v1",
     push_to_hub=True,
@@ -57,6 +50,7 @@ config = SFTConfig(
     gradient_accumulation_steps=16,
     learning_rate=1e-4,
     bf16=True,
+    gradient_checkpointing=True,
 
     logging_steps=10,
     save_strategy="steps",
@@ -74,13 +68,24 @@ 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=model,
-    tokenizer=tokenizer,
+    model="zai-org/GLM-4.7-Flash",
     train_dataset=train_ds,
     eval_dataset=val_ds,
     args=config,
+    peft_config=peft_config,
+    model_init_kwargs={"quantization_config": bnb_config},
 )
 
 print("Starting training...")