New web interface, move llm engine to separate process

tools/qwq_train/train/qwq.py

@@ -0,0 +1,135 @@
+#!/root/venvs/train/bin/python
+"""
+Fine-tuning for QwQ model
+"""
+
+# Unsloth should be imported before transformers to ensure all optimizations are applied.
+from unsloth import FastLanguageModel, is_bfloat16_supported
+
+from dataclasses import dataclass
+from pathlib import Path
+from transformers import AutoTokenizer, TrainingArguments
+from trl import SFTTrainer
+from typing import Optional, List
+import argparse
+import json
+import os
+
+from .dataset import Dataset
+
+@dataclass
+class Args:
+    def __init__(self, args: Optional[List[str]] = None):
+        parser = argparse.ArgumentParser(description='Train SIA model using QwQ')
+        parser.add_argument(
+            '--config',
+            type=Path,
+            default=Path('/root/sia/training/config.yaml'),
+            help='Path to config file'
+        )
+        parser.add_argument(
+            '--base-model',
+            type=str,
+            default='unsloth/QwQ-32B-bnb-4bit',
+            help='HuggingFace model ID for base model'
+        )
+        parser.add_argument(
+            '--output-dir',
+            type=Path,
+            required=True,
+            help='Directory to save the trained model'
+        )
+        parser.add_argument(
+            '--api-key',
+            type=str,
+            default=os.environ.get('SIA_HF_API_KEY'),
+            help='HuggingFace API key'
+        )
+        if args is None:
+            self.args = parser.parse_args()
+        else:
+            self.args = parser.parse_args(args)
+
+    @property
+    def config_path(self) -> Path:
+        return self.args.config
+
+    @property
+    def base_model(self) -> str:
+        return self.args.base_model
+    
+    @property
+    def output_dir(self) -> Path:
+        return self.args.output_dir
+    
+    @property
+    def api_key(self) -> str:
+        return self.args.api_key
+
+def main():
+    args = Args()
+    dataset = Dataset(args.config_path)
+    dataset.validate()
+
+    with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
+        tokenizer_config = json.load(f)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.base_model,
+        **tokenizer_config,
+    )
+
+    model, _returned_tokenizer = FastLanguageModel.from_pretrained(
+        model_name = args.base_model,
+        gpu_memory_utilization = 0.5,
+    )
+
+    model = FastLanguageModel.get_peft_model(
+        model,
+        target_modules = [
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ], # Remove QKVO if out of memory
+        lora_dropout = 0, # Supports any, but = 0 is optimized
+        bias = "none",    # Supports any, but = "none" is optimized
+        use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
+        random_state = 3407,
+    )
+
+    training_args = TrainingArguments(
+        output_dir=str(args.output_dir) + "_train",
+        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",
+    )
+
+    trainer.train()
+    
+    model.save_pretrained_merged(
+        str(args.output_dir), 
+        tokenizer=tokenizer,
+        save_method="merged_4bit_forced"
+    )
+
+if __name__ == "__main__":
+    main()