From 8975a54d58af6a551679a69998a90eaade138bb1 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Mon, 24 Mar 2025 15:46:54 +0000
Subject: [PATCH] WIP QwQ train

---
 tools/train/train/dataset.py | 24 +++++++++++----
 tools/train/train/qwq.py     | 58 +++++++++++++++++++++++++++++++++---
 2 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/tools/train/train/dataset.py b/tools/train/train/dataset.py
index 1f6bbbc..0562eec 100644
--- a/tools/train/train/dataset.py
+++ b/tools/train/train/dataset.py
@@ -1,13 +1,13 @@
-from dataclasses import dataclass
+from datasets import Dataset as TransformersDataset
+from transformers import PreTrainedTokenizer
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Any, Iterator
+from typing import Dict, List, Iterator
 import hashlib
-import json
-import yaml
+import torch
 import xml.etree.ElementTree as ET
+import yaml
 
-
-class Dataset:
+class Dataset(torch.utils.data.Dataset):
     """Training dataset from XML iteration files"""
     
     def __init__(self, config_filename: str):
@@ -90,6 +90,18 @@ class Dataset:
             results.append(self[i])
         return results
     
+    def to_transformers_dataset(self, tokenizer: PreTrainedTokenizer) -> TransformersDataset:
+        def generator():
+            for item in self:
+                messages = item["messages"]
+                formatted_text = tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    add_generation_prompt=False
+                )
+                yield {"messages": formatted_text}
+        return TransformersDataset.from_generator(generator)
+    
     def validate(self) -> None:
         """Validate XML files"""
         print(f"Validating {len(self.files)} XML files...")
diff --git a/tools/train/train/qwq.py b/tools/train/train/qwq.py
index 4157e28..6fee924 100644
--- a/tools/train/train/qwq.py
+++ b/tools/train/train/qwq.py
@@ -1,12 +1,16 @@
 #!/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
 """
+
+# Unsloth should be imported before transformers to ensure all optimizations are applied.
+from unsloth import FastLanguageModel, is_bfloat16_supported
+
 from .dataset import Dataset
 from dataclasses import dataclass
 from pathlib import Path
-from unsloth import FastLanguageModel, is_bfloat16_supported
+from transformers import TrainingArguments
+from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
 import argparse
 import os
 import torch
@@ -62,7 +66,7 @@ def main():
     dataset = Dataset(args.config_path)
     dataset.validate()
 
-    max_seq_length = 1024 # Can increase for longer reasoning traces
+    max_seq_length = 2048 # Can increase for longer reasoning traces
     lora_rank = 64 # Larger rank = smarter, but slower
 
     model, tokenizer = FastLanguageModel.from_pretrained(
@@ -71,7 +75,7 @@ def main():
         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
+        gpu_memory_utilization = 0.85, # Reduce if out of memory
     )
 
     model = FastLanguageModel.get_peft_model(
@@ -86,5 +90,51 @@ def main():
         random_state = 3407,
     )
 
+    response_template = tokenizer.apply_chat_template(
+        [{"role": "assistant", "content": ""}],
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    training_args = TrainingArguments(
+        output_dir=str(args.output_dir),
+        num_train_epochs=3,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=16,
+        gradient_checkpointing=True,
+        learning_rate=2e-5,
+        lr_scheduler_type="cosine",
+        warmup_ratio=0.05,
+        weight_decay=0.01,
+        fp16=not is_bfloat16_supported(),
+        bf16=is_bfloat16_supported(),
+        logging_steps=10,
+        save_steps=200,
+        save_total_limit=3,
+        report_to="none",
+        optim="adamw_8bit",
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        args=training_args,
+        train_dataset=dataset.to_transformers_dataset(tokenizer),
+        dataset_text_field="messages",
+        max_seq_length=max_seq_length,
+        data_collator=DataCollatorForCompletionOnlyLM(
+            response_template=response_template, 
+            tokenizer=tokenizer
+        ),
+    )
+
+    trainer.train()
+    
+    model.save_pretrained_merged(
+        str(args.output_dir), 
+        tokenizer=tokenizer,
+        save_method="merged_16bit"
+    )
+
 if __name__ == "__main__":
     main()
\ No newline at end of file