wip deepseek train

.gitignore

@@ -2,4 +2,4 @@
 __pycache__/
 data/
 model/
-sia.egg-info/
+**.egg-info/

Dockerfile

@@ -64,19 +64,15 @@ RUN npm run build
 # Final image
 FROM base
 
-# Copy virtual environments (these layers only change if setup.py files change)
 COPY --from=itb-env /root/venvs/itb /root/venvs/itb
 COPY --from=train-env /root/venvs/train /root/venvs/train
 COPY --from=sia-env /root/venvs/sia /root/venvs/sia
-
-# Copy source code and scripts (these change frequently but don't affect venv layers)
-COPY --from=itb-env /root/sia/tools/itb /root/sia/tools/itb
-COPY --from=train-env /root/sia/tools/train /root/sia/tools/train
-COPY --from=sia-env /root/sia /root/sia
 COPY --from=web-build /app/dist /root/static/
 
-RUN echo 'for venv in /root/venvs/*/bin; do PATH="$venv:$PATH"; done' >> /etc/profile && \
+RUN for venv in /root/venvs/*/bin; do \
-    echo 'export PATH' >> /etc/profile
+    echo "export PATH=\"$venv:\$PATH\"" >> /etc/profile.d/sia.sh ; \
+done
 WORKDIR /root/desktop
 
-CMD ["/bin/bash", "-l", "-c", "/root/sia/scripts/restart.sh"]
+ENTRYPOINT ["/bin/bash", "-l", "-c"]
+CMD ["/root/sia/scripts/restart.sh"]

requirements.txt

@@ -1,10 +0,0 @@
-accelerate
-aiohttp
-bs4
-mistral_common
-mistralai
-openai
-python-dotenv
-tiktoken
-torch
-transformers

scripts/container.sh

@@ -1,4 +1,4 @@
 #!/bin/bash
 
 container_id=$(docker ps -q)
-docker exec -it $container_id bash
+docker exec -it $container_id bash -l

scripts/restart.sh

@@ -1,11 +1,31 @@
 #!/bin/bash
 
+set -e
+
+echo "=== Preparing SIA environment ==="
+
+echo "Installing ITB tool..."
+/root/venvs/itb/bin/pip install -e /root/sia/tools/itb/ || echo "Warning: Failed to install ITB tool"
+
+echo "Installing Train tool..."
+/root/venvs/train/bin/pip install -e /root/sia/tools/train/ || echo "Warning: Failed to install Train tool"
+
+echo "Installing SIA core..."
+/root/venvs/sia/bin/pip install -e /root/sia/ || { echo "Error: Failed to install SIA core"; exit 1; }
+
+echo "=== Starting SIA ==="
+
 while true; do
     sia
-    if [ $? -eq 42 ]; then
+    EXIT_CODE=$?
-        echo "SIA exited with code 42. Restarting."
+    
+    if [ $EXIT_CODE -eq 42 ]; then
+        echo "SIA exited with code 42. Restarting in 2 seconds..."
+        sleep 2
     else
-        echo "SIA exited with code $?. Not restarting."
+        echo "SIA exited with code $EXIT_CODE. Not restarting."
         break
     fi
 done
+
+exit $EXIT_CODE

setup.py

@@ -5,7 +5,9 @@ setup(
     version="0.1.0",
     packages=find_packages(),
     install_requires=[
+        'accelerate>=0.26.0',
         'aiohttp>=3.8.0',
+        'bitsandbytes>=0.41.0',
         'dotenv-python>=0.0.1',
         'huggingface_hub>=0.16.0',
         'lxml>=4.9.0',

sia/__main__.py

@@ -67,7 +67,7 @@ class Main:
                 config.deepseek_model,
                 config.deepseek_temperature,
                 config.deepseek_token_limit,
-                config.hf_api_key,  # Use the existing HF API key
+                config.hf_api_key,
             )
 
         if not self._llms:

sia/llm_engine/deepseek_llm_engine.py

@@ -1,6 +1,6 @@
 from typing import Callable, Iterator, Optional
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, BitsAndBytesConfig
 from threading import Thread
 from pathlib import Path
 
@@ -44,16 +44,27 @@ class DeepSeekLlmEngine(LlmEngine):
         if self._tokenizer.pad_token is None:
             self._tokenizer.pad_token = self._tokenizer.eos_token
         
-        # Load model with 4-bit quantization by default
+        # Configure 4-bit quantization with CPU offloading
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            llm_int8_enable_fp32_cpu_offload=True
+        )
+        
+        # Configure device map for efficient memory usage
+        # "auto" with the proper quantization config will handle the memory constraints
         self._device_map = "auto"
         
+        # Load model with quantization config
         self._model = AutoModelForCausalLM.from_pretrained(
             self._model_path,
             return_dict=True,
             low_cpu_mem_usage=True,
             trust_remote_code=True,
             device_map=self._device_map,
-            load_in_4bit=True, 
+            quantization_config=quantization_config,
             torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
             token=api_key,
         )

tools/itb/itb.egg-info/PKG-INFO

@@ -1,11 +0,0 @@
-Metadata-Version: 2.1
-Name: itb
-Version: 0.1.0
-Summary: UNKNOWN
-Home-page: UNKNOWN
-License: UNKNOWN
-Platform: UNKNOWN
-Provides-Extra: dev
-
-UNKNOWN
-

tools/itb/itb.egg-info/SOURCES.txt

@@ -1,16 +0,0 @@
-README.md
-setup.py
-bin/itb_click
-bin/itb_forms
-bin/itb_input
-bin/itb_links
-bin/itb_navigate
-bin/itb_refresh
-bin/itb_screenshot
-bin/itb_scroll
-bin/itb_start
-itb.egg-info/PKG-INFO
-itb.egg-info/SOURCES.txt
-itb.egg-info/dependency_links.txt
-itb.egg-info/requires.txt
-itb.egg-info/top_level.txt

tools/itb/itb.egg-info/dependency_links.txt

@@ -1 +0,0 @@
-

tools/itb/itb.egg-info/requires.txt

@@ -1,10 +0,0 @@
-beautifulsoup4>=4.9.0
-click>=8.0.0
-selenium>=4.0.0
-webdriver-manager>=3.8.0
-
-[dev]
-black>=22.0.0
-flake8>=4.0.0
-pytest-cov>=4.0.0
-pytest>=7.0.0

tools/itb/itb.egg-info/top_level.txt

@@ -1 +0,0 @@
-

tools/itb/requirements.txt

@@ -1,8 +0,0 @@
-selenium>=4.0.0
-webdriver-manager>=3.8.0
-click>=8.0.0
-beautifulsoup4>=4.9.0
-pytest>=7.0.0
-pytest-cov>=4.0.0
-black>=22.0.0
-flake8>=4.0.0

tools/train/requirements.txt

@@ -1,13 +0,0 @@
-pyyaml>=6.0
-requests>=2.28.0
-torch>=2.0.0
-transformers>=4.30.0
-# DeepSeek support
-accelerate>=0.25.0
-bitsandbytes>=0.41.1
-einops>=0.7.0
-sentencepiece>=0.1.99
-unsloth>=2024.3
-trl>=0.7.8
-datasets>=2.14.6
-peft>=0.8.0

tools/train/setup.py

@@ -17,7 +17,7 @@ setup(
         'bitsandbytes>=0.41.1',
         'einops>=0.7.0',
         'sentencepiece>=0.1.99',
-        'unsloth>=2024.3',
+        'unsloth>=2025.2',
         'trl>=0.7.8',
         'datasets>=2.14.6',
         'peft>=0.8.0',

tools/train/train.sh

@@ -12,4 +12,4 @@ fi
 
 mkdir -p "$OUTPUT_DIR"
 
-train_deepseek --output-dir "$OUTPUT_DIR"
+train_deepseek --output-dir "$OUTPUT_DIR" --device cpu

tools/train/train/unsloth_deepseek.py

@@ -42,6 +42,12 @@ class Config:
             default=os.environ.get('SIA_HF_API_KEY'),
             help='HuggingFace API key'
         )
+        parser.add_argument(
+            '--device',
+            type=str,
+            default='auto',
+            help='Override device (cpu, cuda, auto) from config'
+        )
         self.args = parser.parse_args()
 
     @property
@@ -60,7 +66,11 @@ class Config:
     def api_key(self) -> str:
         return self.args.api_key
         
-def train_model(config: Config, training_data, train_params, commit_hash):
+    @property
+    def device(self) -> str:
+        return self.args.device
+
+def train_model(config: Config, training_data, train_params):
     """Train the model using Unsloth"""
     try:
         from unsloth import FastLanguageModel
@@ -74,52 +84,83 @@ def train_model(config: Config, training_data, train_params, commit_hash):
         sys.exit(1)
 
     print(f"Starting training from base model: {config.base_model}")
+    print(f"Using device: {config.device}")
+    print(f"Training configuration:")
+    print(f"  Max sequence length: {train_params.max_seq_length}")
+    print(f"  Quantization: {train_params.quantization}")
+    print(f"  Batch size: {train_params.per_device_batch_size}")
+    print(f"  Gradient accumulation: {train_params.gradient_accumulation_steps}")
+    print(f"  Mixed precision: {train_params.mixed_precision}")
     
     # Convert to datasets format
     dataset = Dataset.from_list(training_data)
     
-    # Determine if bfloat16 is supported
+    # Configure device and dtype
-    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    if train_params.mixed_precision == "bf16" and torch.cuda.is_bf16_supported():
+        dtype = torch.bfloat16
+    else:
+        dtype = torch.float16
         
-    # Load the model - always from a base model (no incremental updates)
+    # Configure quantization settings
+    load_in_4bit = train_params.quantization == "4bit"
+    load_in_8bit = train_params.quantization == "8bit"
+    
+    # Configure device mapping
+    device_map = config.device
+    if config.device == "cpu":
+        # Force CPU even for quantized model
+        bnb_config = None
+        # When on CPU, we should disable quantization
+        load_in_4bit = False
+        load_in_8bit = False
+        dtype = torch.float32
+        print("CPU-only mode: Disabling quantization and using float32")
+    else:
+        # Setup quantization config for GPU
+        from transformers import BitsAndBytesConfig
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=load_in_4bit,
+            load_in_8bit=load_in_8bit,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=dtype,
+            llm_int8_enable_fp32_cpu_offload=True
+        ) if (load_in_4bit or load_in_8bit) else None
+    
+    # Load the model with appropriate settings
     try:
         model, tokenizer = FastLanguageModel.from_pretrained(
             model_name=config.base_model,
-            max_seq_length=2048,
+            max_seq_length=train_params.max_seq_length,
             dtype=dtype,
-            load_in_4bit=True,
+            quantization_config=bnb_config,
+            device_map=device_map,
             token=config.api_key,
         )
     except Exception as e:
         print(f"Error loading base model: {e}")
         sys.exit(1)
     
-    # Apply LoRA
     target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
     model = FastLanguageModel.get_peft_model(
         model,
-        r=16,
+        r=8 if config.device == "cpu" else 16,  # Lower rank for CPU to save memory
         target_modules=target_modules,
         lora_alpha=16,
         lora_dropout=0,
         bias="none",
-        use_gradient_checkpointing="unsloth",
+        # Only use gradient checkpointing for GPU
+        use_gradient_checkpointing="unsloth" if config.device != "cpu" else None,
         random_state=3407,
     )
     
-    # Apply chat template
-    tokenizer = get_chat_template(
-        tokenizer, 
-        chat_template="llama-3.1",  # Compatible with DeepSeek
-    )
-    
     # Function to format conversations
     def formatting_prompts_func(examples):
         convos = examples["conversations"]
         texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
         return {"text": texts}
     
-    # Standarize dataset and format
+    # Standardize dataset and format
     from unsloth.chat_templates import standardize_sharegpt
     
     # Add conversations field if not present
@@ -138,80 +179,87 @@ def train_model(config: Config, training_data, train_params, commit_hash):
     dataset = dataset.map(formatting_prompts_func, batched=True)
     
     # Configure the trainer
-    output_dir = config.output_dir / commit_hash
+    config.output_dir.mkdir(parents=True, exist_ok=True)
-    output_dir.mkdir(parents=True, exist_ok=True)
     
     # Determine steps or epochs based on dataset size
-    max_steps = None
+    max_steps = -1
     num_train_epochs = train_params.epochs
     if len(dataset) < 100:  # Small dataset
         # Aim for at least 500 steps for small datasets
         max_steps = 500
-        num_train_epochs = None
+        num_train_epochs = -1
+    
+    # Configure mixed precision settings
+    fp16 = train_params.mixed_precision == "fp16"
+    bf16 = train_params.mixed_precision == "bf16" and torch.cuda.is_bf16_supported()
     
     trainer = SFTTrainer(
         model=model,
         tokenizer=tokenizer,
         train_dataset=dataset,
         dataset_text_field="text",
-        max_seq_length=2048,
+        max_seq_length=train_params.max_seq_length,
         data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
-        dataset_num_proc=2,
+        dataset_num_proc=1 if config.device == "cpu" else 2,
         packing=False,
         args=TrainingArguments(
-            per_device_train_batch_size=2,
+            per_device_train_batch_size=train_params.per_device_batch_size,
-            gradient_accumulation_steps=4,
+            gradient_accumulation_steps=train_params.gradient_accumulation_steps,
             warmup_steps=5,
             max_steps=max_steps,
             num_train_epochs=num_train_epochs,
             learning_rate=train_params.learning_rate,
-            fp16=not torch.cuda.is_bf16_supported(),
+            fp16=fp16,
-            bf16=torch.cuda.is_bf16_supported(),
+            bf16=bf16,
             logging_steps=10,
-            optim="adamw_8bit",
+            optim="adamw_torch" if config.device == "cpu" else "adamw_8bit",
             weight_decay=0.01,
             lr_scheduler_type="linear",
             seed=3407,
-            output_dir=str(output_dir),
+            output_dir=str(config.output_dir),
             report_to="none",
+            dataloader_num_workers=0 if config.device == "cpu" else 2,
+            gradient_checkpointing=config.device != "cpu",
+            max_grad_norm=0.3,
         ),
     )
     
     # Train only on responses
     trainer = train_on_responses_only(
         trainer,
-        instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
+        instruction_part="<|im_start|>user",
-        response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
+        response_part="<|im_start|>assistant",
     )
     
     # Train the model
     trainer.train()
     
     # Enable inference mode for the model
+    if config.device != "cpu":
         model = FastLanguageModel.for_inference(model)
     
     # Save the model
-    model.save_pretrained(output_dir)
+    model.save_pretrained(config.output_dir)
-    tokenizer.save_pretrained(output_dir)
+    tokenizer.save_pretrained(config.output_dir)
     
     # Create a metadata file with training information
-    with open(output_dir / "training_info.json", "w") as f:
+    with open(config.output_dir / "training_info.json", "w") as f:
         json.dump({
             "base_model": config.base_model,
-            "commit_hash": commit_hash,
             "learning_rate": train_params.learning_rate,
             "epochs": train_params.epochs,
             "dataset_size": len(dataset),
+            "device": config.device,
             "training_method": "unsloth",
+            "max_seq_length": train_params.max_seq_length,
+            "quantization": train_params.quantization,
         }, f, indent=2)
     
-    return output_dir
-
 def main():
     config = Config()
     
     # Prepare training data
-    training_data, train_params, commit_hash = prepare_training_data(config.config_path)
+    training_data, train_params = prepare_training_data(config.config_path)
     
     if not training_data:
         print("No valid training data found. Exiting.")
@@ -219,20 +267,22 @@ def main():
     
     # Train the model
     try:
-        model_dir = train_model(config, training_data, train_params, commit_hash)
+        train_model(config, training_data, train_params)
         
         # Create symlink to current
-        current_link = config.output_dir / "current"
+        current_link = Path("/root/models/current")
-        if current_link.exists() or current_link.is_symlink():
+        if os.path.exists(current_link) or os.path.islink(current_link):
-            current_link.unlink()
+            os.unlink(current_link)
-        os.symlink(model_dir, current_link, target_is_directory=True)
+        os.symlink(config.output_dir, current_link, target_is_directory=True)
         
-        print(f"Training complete. Model saved to {model_dir}")
+        print(f"Training complete. Model saved to {config.output_dir}")
         print(f"Symlink created at {current_link}")
         
         return 0
     except Exception as e:
         print(f"Error during training: {e}")
+        import traceback
+        traceback.print_exc()
         return 1
 
 if __name__ == "__main__":

tools/train/train/util.py

@@ -4,7 +4,7 @@ Contains common code for both API-based and local training.
 """
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Set, Tuple
+from typing import Dict, List, Optional, Set, Tuple, Any
 import hashlib
 import json
 import subprocess
@@ -17,7 +17,26 @@ class TrainingParams:
     """Parameters for model training"""
     learning_rate: float
     epochs: int
-    batch_size: int = 1
+    batch_size: int
+    max_seq_length: int
+    quantization: str
+    per_device_batch_size: int
+    gradient_accumulation_steps: int
+    mixed_precision: str
+    
+    @classmethod
+    def from_dict(cls, config_dict: Dict[str, Any]) -> 'TrainingParams':
+        """Create from config dictionary with defaults"""
+        return cls(
+            learning_rate=float(config_dict.get('learning_rate', 1e-5)),
+            epochs=int(config_dict.get('epochs', 1)),
+            batch_size=int(config_dict.get('batch_size', 1)),
+            max_seq_length=int(config_dict.get('max_seq_length', 1024)),
+            quantization=config_dict.get('quantization', '4bit'),
+            per_device_batch_size=int(config_dict.get('per_device_batch_size', 1)),
+            gradient_accumulation_steps=int(config_dict.get('gradient_accumulation_steps', 8)),
+            mixed_precision=config_dict.get('mixed_precision', 'no')
+        )
 
 class DatasetCreator:
     """Creates training datasets from XML iteration files"""
@@ -147,7 +166,7 @@ def format_chat_for_mistral(messages):
     # Format according to Mistral chat template
     return f"<s>[INST] {instruction} [/INST] {assistant_content} </s>"
 
-def prepare_training_data(config_path: Path) -> Tuple[List[Dict], TrainingParams, str]:
+def prepare_training_data(config_path: Path) -> Tuple[List[Dict], TrainingParams]:
     """Prepare training data from config and XML files"""
     with open(config_path) as f:
         config_data = yaml.safe_load(f)
@@ -155,12 +174,6 @@ def prepare_training_data(config_path: Path) -> Tuple[List[Dict], TrainingParams
     data_paths = [Path(p) for p in config_data['data']]
     xml_files = find_xml_files(data_paths)
     
-    paths = list(xml_files)
-    paths.append(config_path)
-    paths.append(Path(config_data['model']['system_prompt_path']))
-    paths.append(Path(config_data['model']['action_schema']))
-    commit_hash = check_git_status(paths)
-    
     creator = DatasetCreator(
         xml_files=xml_files,
         system_prompt_file=config_data['model']['system_prompt_path'],
@@ -169,13 +182,9 @@ def prepare_training_data(config_path: Path) -> Tuple[List[Dict], TrainingParams
     
     training_data = creator.create_dataset()
     
-    train_params = TrainingParams(
+    train_params = TrainingParams.from_dict(config_data['params'])
-        learning_rate=config_data['params'].get('learning_rate', 1e-5),
-        epochs=config_data['params'].get('epochs', 3),
-        batch_size=config_data['params'].get('batch_size', 1)
-    )
     
-    return training_data, train_params, commit_hash
+    return training_data, train_params
 
 def save_jsonl_dataset(data: List[Dict], output_path: Path) -> None:
     """Save dataset in JSONL format"""

training/config.yaml

@@ -4,6 +4,11 @@ model:
 params:
   learning_rate: 1e-5
   epochs: 3
+  max_seq_length: 1024
+  quantization: "4bit"  # Options: "none", "4bit", "8bit"
+  per_device_batch_size: 1
+  gradient_accumulation_steps: 8
+  mixed_precision: "no"  # Options: "no", "fp16", "bf16"
 data:
   - "/root/sia/training/clean_start/"
   - "/root/sia/training/delete_indicated_entries/"