llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

wip deepseek train

Browse Source
This commit is contained in:
Niels Geens
2025-03-03 16:57:04 +01:00
parent b64f8d7d33
commit 3ea3239a9b
21 changed files with 175 additions and 7058 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

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

14
Dockerfile
View File

@@ -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 && \
echo 'export PATH' >> /etc/profile
RUN for venv in /root/venvs/*/bin; do \
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"]

6906
collect.txt
View File

File diff suppressed because it is too large Load Diff

10
requirements.txt
View File

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

2
scripts/container.sh
View File

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

28
scripts/restart.sh
View File

@@ -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
echo "SIA exited with code 42. Restarting."
EXIT_CODE=$?
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
done
exit $EXIT_CODE

2
setup.py
View File

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

2
sia/__main__.py
View File

@@ -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:

19
sia/llm_engine/deepseek_llm_engine.py
View File

@@ -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
@@ -43,17 +43,28 @@ class DeepSeekLlmEngine(LlmEngine):
# Set padding token to avoid warnings
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,
)

11
tools/itb/itb.egg-info/PKG-INFO
View File

@@ -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

16
tools/itb/itb.egg-info/SOURCES.txt
View File

@@ -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

1
tools/itb/itb.egg-info/dependency_links.txt
View File

@@ -1 +0,0 @@

10
tools/itb/itb.egg-info/requires.txt
View File

@@ -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

1
tools/itb/itb.egg-info/top_level.txt
View File

@@ -1 +0,0 @@

8
tools/itb/requirements.txt
View File

@@ -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

13
tools/train/requirements.txt
View File

@@ -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

2
tools/train/setup.py
View File

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

2
tools/train/train.sh
View File

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

140
tools/train/train/unsloth_deepseek.py
View File

@@ -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
@@ -59,8 +65,12 @@ class Config:
@property
def api_key(self) -> str:
return self.args.api_key
@property
def device(self) -> str:
return self.args.device
def train_model(config: Config, training_data, train_params, commit_hash):
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
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
# Configure device and dtype
if train_params.mixed_precision == "bf16" and torch.cuda.is_bf16_supported():
dtype = torch.bfloat16
else:
dtype = torch.float16
# Configure quantization settings
load_in_4bit = train_params.quantization == "4bit"
load_in_8bit = train_params.quantization == "8bit"
# Load the model - always from a base model (no incremental updates)
# 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
output_dir.mkdir(parents=True, exist_ok=True)
config.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,
gradient_accumulation_steps=4,
per_device_train_batch_size=train_params.per_device_batch_size,
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(),
bf16=torch.cuda.is_bf16_supported(),
fp16=fp16,
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",
response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
instruction_part="<|im_start|>user",
response_part="<|im_start|>assistant",
)
# Train the model
trainer.train()
# Enable inference mode for the model
model = FastLanguageModel.for_inference(model)
if config.device != "cpu":
model = FastLanguageModel.for_inference(model)
# Save the model
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
model.save_pretrained(config.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,21 +267,23 @@ 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"
if current_link.exists() or current_link.is_symlink():
current_link.unlink()
os.symlink(model_dir, current_link, target_is_directory=True)
current_link = Path("/root/models/current")
if os.path.exists(current_link) or os.path.islink(current_link):
os.unlink(current_link)
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__":
exit(main())
exit(main())

39
tools/train/train/util.py
View File

@@ -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(
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)
)
train_params = TrainingParams.from_dict(config_data['params'])
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"""

5
training/config.yaml
View File

@@ -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/"