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

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