WIP manual implementation of QwQ finetune

.gitignore

@@ -1,6 +1,7 @@
+**.egg-info/
 .env
 __pycache__/
+collect.txt
 data/
 model/
-**.egg-info/
-collect.txt
+unsloth_compiled_cache/

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',

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()