Add support for jupyter notebooks for training

scripts/bootstrap.sh

@@ -18,6 +18,11 @@ if [ -z "$SIA_REPO_PAT" ]; then
     exit 1
 fi
 
+# Install required packages
+apt-get update
+apt-get install -y \
+    vim
+
 # Create directory structure
 echo "Creating directory structure..."
 mkdir -p "/root/data/iterations"
@@ -33,6 +38,7 @@ if [ ! -d "/root/sia" ]; then
     cd "/root/sia"
     git config --global user.name "Niels Geens"
     git config --global user.email "niels.geens@gmail.com"
+    git config --global core.editor vim
 fi
 
 # Fixing permissions
@@ -58,10 +64,10 @@ ln -s "/root/sia/web/dist" "/root/static"
 # Install SIA dependencies
 source "/root/sia/scripts/install.sh"
 
-# Finetune model
+## Finetune model
-echo "Finetuning model..."
+#echo "Finetuning model..."
-train
+#train
-
+#
-# Start SIA using restart script
+## Start SIA using restart script
-echo "Run restart script..."
+#echo "Run restart script..."
-"/root/sia/scripts/restart.sh"
+#"/root/sia/scripts/restart.sh"

scripts/install.sh

@@ -9,6 +9,7 @@ source "/etc/profile.d/venv_itb.sh"
 echo "Installing Train tool..."
 python3 -m venv "/root/venvs/train"
 /root/venvs/train/bin/pip install -e /root/sia/tools/train/
+/root/venvs/train/bin/ipython kernel install --name=train
 echo "PATH=\"/root/venvs/train/bin/:\$PATH\"" > "/etc/profile.d/venv_train.sh"
 source "/etc/profile.d/venv_train.sh"
 

tools/train/setup.py

@@ -15,6 +15,8 @@ setup(
         'datasets>=2.14.6',
         'einops>=0.7.0',
         'flake8>=4.0.0',
+        'ipykernel>=6.0.0', 
+        'ipywidgets>=8.0.0',
         'peft>=0.8.0',
         'peft>=0.8.0',
         'pytest-cov>=4.0.0',

tools/train/train/qwq.py

@@ -11,13 +11,13 @@ from dataclasses import dataclass
 from pathlib import Path
 from transformers import TrainingArguments
 from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
+from typing import Optional, List
 import argparse
 import os
-import torch
 
 @dataclass
 class Args:
-    def __init__(self):
+    def __init__(self, args: Optional[List[str]]):
         parser = argparse.ArgumentParser(description='Train SIA model using QwQ')
         parser.add_argument(
             '--config',
@@ -43,7 +43,10 @@ class Args:
             default=os.environ.get('SIA_HF_API_KEY'),
             help='HuggingFace API key'
         )
-        self.args = parser.parse_args()
+        if args is None:
+            self.args = parser.parse_args()
+        else:
+            self.args = parser.parse_args(args)
 
     @property
     def config_path(self) -> Path:
@@ -86,8 +89,12 @@ def main():
             "gate_proj", "up_proj", "down_proj",
         ], # Remove QKVO if out of memory
         lora_alpha = lora_rank,
-        use_gradient_checkpointing = "unsloth", # Enable long context finetuning
+        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,
+        use_rslora = False,  # We support rank stabilized LoRA
+        loftq_config = None, # And LoftQ
     )
 
     response_template = tokenizer.apply_chat_template(