From 1ef32ed33e90a26ece694002ab37181463451e44 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Mon, 24 Mar 2025 12:05:43 +0000
Subject: [PATCH] WIP manual implementation of QwQ finetune

---
 .gitignore               |  5 +++--
 tools/train/setup.py     |  3 ++-
 tools/train/train/qwq.py | 30 ++++++++++++++++++++++++++++--
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 4b8d1db..9bf1836 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
+**.egg-info/
 .env
 __pycache__/
+collect.txt
 data/
 model/
-**.egg-info/
-collect.txt
\ No newline at end of file
+unsloth_compiled_cache/
\ No newline at end of file
diff --git a/tools/train/setup.py b/tools/train/setup.py
index eebde0d..38cadc5 100644
--- a/tools/train/setup.py
+++ b/tools/train/setup.py
@@ -25,7 +25,8 @@ setup(
         'torch>=2.0.0',
         'transformers>=4.30.0',
         'trl>=0.7.8',
-        'unsloth>=2025.2',
+        'unsloth>=2025.3',
+        'vllm>=0.8',
     ],
     classifiers=[
         'Development Status :: 3 - Alpha',
diff --git a/tools/train/train/qwq.py b/tools/train/train/qwq.py
index ea685ab..4157e28 100644
--- a/tools/train/train/qwq.py
+++ b/tools/train/train/qwq.py
@@ -1,12 +1,15 @@
 #!/root/venvs/train/bin/python
 """
 Fine-tuning for QwQ model
+Based on: https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen2.5_(3B)-GRPO.ipynb
 """
 from .dataset import Dataset
 from dataclasses import dataclass
-import argparse
 from pathlib import Path
+from unsloth import FastLanguageModel, is_bfloat16_supported
+import argparse
 import os
+import torch
 
 @dataclass
 class Args:
@@ -58,7 +61,30 @@ def main():
     args = Args()
     dataset = Dataset(args.config_path)
     dataset.validate()
-    print(dataset[3])
+
+    max_seq_length = 1024 # Can increase for longer reasoning traces
+    lora_rank = 64 # Larger rank = smarter, but slower
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = args.base_model,
+        max_seq_length = max_seq_length,
+        load_in_4bit = True, # False for LoRA 16bit
+        fast_inference = True, # Enable vLLM fast inference
+        max_lora_rank = lora_rank,
+        gpu_memory_utilization = 0.5, # Reduce if out of memory
+    )
+
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+        target_modules = [
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ], # Remove QKVO if out of memory
+        lora_alpha = lora_rank,
+        use_gradient_checkpointing = "unsloth", # Enable long context finetuning
+        random_state = 3407,
+    )
 
 if __name__ == "__main__":
     main()
\ No newline at end of file