Replaced deepseek with qwq
This commit is contained in:
Niels Geens
@@ -12,4 +12,4 @@ fi
|
||||
|
||||
mkdir -p "$OUTPUT_DIR"
|
||||
|
||||
train_deepseek --output-dir "$OUTPUT_DIR" --device cpu
|
||||
python -m train.qwq "$@"
|
||||
@@ -1,10 +0,0 @@
|
||||
#!/root/venvs/train/bin/python
|
||||
"""
|
||||
Command-line utility for fine-tuning DeepSeek models using Unsloth.
|
||||
Always trains from a base model to create a new fine-tuned model.
|
||||
"""
|
||||
import sys
|
||||
from train.unsloth_deepseek import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,9 +0,0 @@
|
||||
#!/root/venvs/train/bin/python
|
||||
"""
|
||||
Command-line utility for fine-tuning Mistral models using Mistral API.
|
||||
"""
|
||||
import sys
|
||||
from train.mistral_api import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -5,26 +5,27 @@ setup(
|
||||
version="0.1.0",
|
||||
packages=find_packages(),
|
||||
scripts=[
|
||||
'bin/train_deepseek',
|
||||
'bin/train_mistral'
|
||||
'bin/train'
|
||||
],
|
||||
|
||||
install_requires=[
|
||||
'pyyaml>=6.0',
|
||||
'requests>=2.28.0',
|
||||
'torch>=2.0.0',
|
||||
'transformers>=4.30.0',
|
||||
'accelerate>=0.25.0',
|
||||
'bitsandbytes>=0.41.1',
|
||||
'einops>=0.7.0',
|
||||
'sentencepiece>=0.1.99',
|
||||
'unsloth>=2025.2',
|
||||
'trl>=0.7.8',
|
||||
'datasets>=2.14.6',
|
||||
'peft>=0.8.0',
|
||||
'pytest>=7.0.0',
|
||||
'pytest-cov>=4.0.0',
|
||||
'black>=22.0.0',
|
||||
'flake8>=4.0.0'
|
||||
'datasets>=2.14.6',
|
||||
'einops>=0.7.0',
|
||||
'flake8>=4.0.0',
|
||||
'peft>=0.8.0',
|
||||
'peft>=0.8.0',
|
||||
'pytest-cov>=4.0.0',
|
||||
'pytest>=7.0.0',
|
||||
'pyyaml>=6.0',
|
||||
'requests>=2.28.0',
|
||||
'sentencepiece>=0.1.99',
|
||||
'torch>=2.0.0',
|
||||
'transformers>=4.30.0',
|
||||
'trl>=0.7.8',
|
||||
'unsloth>=2025.2',
|
||||
],
|
||||
classifiers=[
|
||||
'Development Status :: 3 - Alpha',
|
||||
|
||||
334
tools/train/train/qwq.py
Normal file
334
tools/train/train/qwq.py
Normal file
@@ -0,0 +1,334 @@
|
||||
#!/root/venvs/train/bin/python
|
||||
"""
|
||||
Fine-tuning script for QwQ models to support SIA's action schema.
|
||||
Supports both full and LoRA finetuning methods.
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
import json
|
||||
import logging
|
||||
import gc
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Import from shared library
|
||||
from .util import prepare_training_data
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
def __init__(self):
|
||||
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='Qwen/QwQ-32B',
|
||||
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'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--method',
|
||||
type=str,
|
||||
choices=['lora', 'qlora', 'full'],
|
||||
default='qlora',
|
||||
help='Finetuning method: LoRA, QLoRA (quantized LoRA), or full-model'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--device',
|
||||
type=str,
|
||||
default='auto',
|
||||
help='Override device (cpu, cuda, auto) from config'
|
||||
)
|
||||
self.args = parser.parse_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
|
||||
|
||||
@property
|
||||
def device(self) -> str:
|
||||
return self.args.device
|
||||
|
||||
@property
|
||||
def method(self) -> str:
|
||||
return self.args.method
|
||||
|
||||
def format_data_for_qwq(training_data):
|
||||
"""
|
||||
Format training data for QwQ model focusing on action schema formats.
|
||||
Ensures each example shows the model how to directly use action elements.
|
||||
"""
|
||||
formatted_data = []
|
||||
|
||||
for sample in training_data:
|
||||
# Get the system prompt, context, and response
|
||||
system_content = ""
|
||||
context_content = ""
|
||||
response_content = ""
|
||||
|
||||
for message in sample.get("messages", []):
|
||||
if message["role"] == "system":
|
||||
system_content = message["content"]
|
||||
elif message["role"] == "user":
|
||||
context_content = message["content"]
|
||||
elif message["role"] == "assistant":
|
||||
response_content = message["content"]
|
||||
|
||||
# Create conversations with explicit instruction to use action schema
|
||||
formatted_data.append({
|
||||
"conversations": [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": context_content},
|
||||
{"role": "assistant", "content": response_content}
|
||||
]
|
||||
})
|
||||
|
||||
logger.info(f"Formatted {len(formatted_data)} examples for QwQ training")
|
||||
return formatted_data
|
||||
|
||||
def train_model_lora(config, training_data, train_params):
|
||||
"""
|
||||
Train QwQ model using LoRA or QLoRA for parameter-efficient fine-tuning.
|
||||
This is the recommended approach for most use cases.
|
||||
"""
|
||||
try:
|
||||
# Import required libraries
|
||||
from transformers import (
|
||||
AutoModelForCausalLM, AutoTokenizer,
|
||||
TrainingArguments, DataCollatorForSeq2Seq
|
||||
)
|
||||
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
||||
from datasets import Dataset
|
||||
from trl import SFTTrainer
|
||||
|
||||
except ImportError as e:
|
||||
logger.error(f"Error importing required libraries: {e}")
|
||||
logger.error("Please ensure transformers, peft, and trl are installed.")
|
||||
sys.exit(1)
|
||||
|
||||
# Format data specifically for QwQ
|
||||
formatted_data = format_data_for_qwq(training_data)
|
||||
dataset = Dataset.from_list(formatted_data)
|
||||
|
||||
logger.info(f"Starting QwQ fine-tuning using {config.method}")
|
||||
logger.info(f"Base model: {config.base_model}")
|
||||
logger.info(f"Device: {config.device}")
|
||||
|
||||
# Configure device mapping and precision
|
||||
if torch.cuda.is_available() and config.device != "cpu":
|
||||
logger.info("Using GPU for training")
|
||||
device_map = "auto"
|
||||
|
||||
# Configure precision based on method
|
||||
if config.method == "qlora":
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
load_in_4bit = True
|
||||
load_in_8bit = False
|
||||
else:
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
load_in_4bit = False
|
||||
load_in_8bit = False
|
||||
else:
|
||||
logger.info("Using CPU for training")
|
||||
device_map = "cpu"
|
||||
dtype = torch.float32
|
||||
load_in_4bit = False
|
||||
load_in_8bit = False
|
||||
|
||||
# Configure quantization for QLoRA
|
||||
if config.method == "qlora":
|
||||
from transformers import BitsAndBytesConfig
|
||||
logger.info("Setting up 4-bit quantization for QLoRA")
|
||||
compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=compute_dtype,
|
||||
bnb_4bit_use_double_quant=True
|
||||
)
|
||||
else:
|
||||
bnb_config = None
|
||||
|
||||
# Load tokenizer
|
||||
logger.info(f"Loading tokenizer from {config.base_model}")
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
config.base_model,
|
||||
token=config.api_key,
|
||||
trust_remote_code=True
|
||||
)
|
||||
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
# Load model
|
||||
logger.info(f"Loading model from {config.base_model}")
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
config.base_model,
|
||||
torch_dtype=dtype,
|
||||
device_map=device_map,
|
||||
quantization_config=bnb_config,
|
||||
token=config.api_key,
|
||||
trust_remote_code=True
|
||||
)
|
||||
|
||||
# Configure LoRA
|
||||
if config.method in ["lora", "qlora"]:
|
||||
if config.method == "qlora":
|
||||
model = prepare_model_for_kbit_training(model)
|
||||
|
||||
logger.info("Setting up LoRA configuration")
|
||||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=16,
|
||||
lora_alpha=32,
|
||||
target_modules=target_modules,
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM"
|
||||
)
|
||||
|
||||
model = get_peft_model(model, lora_config)
|
||||
|
||||
# Create output directory
|
||||
config.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Configure training arguments
|
||||
batch_size = train_params.per_device_batch_size
|
||||
gradient_accumulation = train_params.gradient_accumulation_steps
|
||||
|
||||
# Scale down batch size based on model
|
||||
if "32B" in config.base_model and batch_size > 1:
|
||||
batch_size = 1
|
||||
gradient_accumulation *= 2
|
||||
|
||||
training_args = TrainingArguments(
|
||||
output_dir=str(config.output_dir),
|
||||
per_device_train_batch_size=batch_size,
|
||||
gradient_accumulation_steps=gradient_accumulation,
|
||||
learning_rate=train_params.learning_rate,
|
||||
num_train_epochs=train_params.epochs,
|
||||
logging_steps=10,
|
||||
save_strategy="epoch",
|
||||
save_total_limit=2,
|
||||
fp16=dtype == torch.float16,
|
||||
bf16=dtype == torch.bfloat16,
|
||||
report_to="none",
|
||||
remove_unused_columns=False,
|
||||
optim="adamw_torch",
|
||||
weight_decay=0.01,
|
||||
max_grad_norm=0.3,
|
||||
warmup_ratio=0.03,
|
||||
lr_scheduler_type="cosine",
|
||||
seed=42
|
||||
)
|
||||
|
||||
# Set up trainer
|
||||
logger.info("Setting up trainer")
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
args=training_args,
|
||||
train_dataset=dataset,
|
||||
tokenizer=tokenizer,
|
||||
max_seq_length=train_params.max_seq_length,
|
||||
dataset_text_field="conversations",
|
||||
packing=False
|
||||
)
|
||||
|
||||
# Start training
|
||||
logger.info("Starting training")
|
||||
trainer.train()
|
||||
|
||||
# Save the final model
|
||||
logger.info(f"Saving model to {config.output_dir}")
|
||||
trainer.save_model(config.output_dir)
|
||||
tokenizer.save_pretrained(config.output_dir)
|
||||
|
||||
# Create metadata file
|
||||
with open(config.output_dir / "training_info.json", "w") as f:
|
||||
json.dump({
|
||||
"base_model": config.base_model,
|
||||
"method": config.method,
|
||||
"learning_rate": train_params.learning_rate,
|
||||
"epochs": train_params.epochs,
|
||||
"dataset_size": len(dataset),
|
||||
"batch_size": batch_size,
|
||||
"gradient_accumulation": gradient_accumulation
|
||||
}, f, indent=2)
|
||||
|
||||
logger.info("Training complete!")
|
||||
return True
|
||||
|
||||
def main():
|
||||
# Initialize configuration
|
||||
config = Config()
|
||||
|
||||
# Prepare training data from config
|
||||
training_data, train_params = prepare_training_data(config.config_path)
|
||||
|
||||
if not training_data:
|
||||
logger.error("No valid training data found. Exiting.")
|
||||
return 1
|
||||
|
||||
# Force garbage collection
|
||||
gc.collect()
|
||||
|
||||
# Train using appropriate method
|
||||
if config.method in ["lora", "qlora"]:
|
||||
success = train_model_lora(config, training_data, train_params)
|
||||
else:
|
||||
logger.error(f"Training method '{config.method}' not yet implemented")
|
||||
return 1
|
||||
|
||||
if success:
|
||||
# Create symlink to current
|
||||
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)
|
||||
|
||||
logger.info(f"Training complete. Model saved to {config.output_dir}")
|
||||
logger.info(f"Symlink created at {current_link}")
|
||||
|
||||
return 0
|
||||
else:
|
||||
logger.error("Training failed")
|
||||
return 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,289 +0,0 @@
|
||||
#!/root/venvs/train/bin/python
|
||||
"""
|
||||
Script for fine-tuning DeepSeek models for SIA using Unsloth.
|
||||
Training always starts from a base model and creates a new fine-tuned model.
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
# Import from shared library
|
||||
from .util import prepare_training_data
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
def __init__(self):
|
||||
parser = argparse.ArgumentParser(description='Train SIA model using Unsloth')
|
||||
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='deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B',
|
||||
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'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--device',
|
||||
type=str,
|
||||
default='auto',
|
||||
help='Override device (cpu, cuda, auto) from config'
|
||||
)
|
||||
self.args = parser.parse_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
|
||||
|
||||
@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
|
||||
from transformers import TrainingArguments, DataCollatorForSeq2Seq
|
||||
from trl import SFTTrainer
|
||||
from datasets import Dataset
|
||||
from unsloth.chat_templates import get_chat_template, train_on_responses_only
|
||||
except ImportError as e:
|
||||
print(f"Error importing required libraries: {e}")
|
||||
print("Please ensure Unsloth and its dependencies are installed.")
|
||||
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)
|
||||
|
||||
# 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"
|
||||
|
||||
# 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=train_params.max_seq_length,
|
||||
dtype=dtype,
|
||||
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)
|
||||
|
||||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
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",
|
||||
# Only use gradient checkpointing for GPU
|
||||
use_gradient_checkpointing="unsloth" if config.device != "cpu" else None,
|
||||
random_state=3407,
|
||||
)
|
||||
|
||||
# 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}
|
||||
|
||||
# Standardize dataset and format
|
||||
from unsloth.chat_templates import standardize_sharegpt
|
||||
|
||||
# Add conversations field if not present
|
||||
if "conversations" not in dataset.column_names:
|
||||
if "messages" in dataset.column_names:
|
||||
dataset = dataset.rename_column("messages", "conversations")
|
||||
else:
|
||||
dataset = dataset.map(lambda x: {"conversations": [{"role": "system", "content": x.get("system_prompt", "")},
|
||||
{"role": "user", "content": x.get("prompt", "")},
|
||||
{"role": "assistant", "content": x.get("response", "")}]})
|
||||
|
||||
# Standardize format
|
||||
dataset = standardize_sharegpt(dataset)
|
||||
|
||||
# Apply formatting
|
||||
dataset = dataset.map(formatting_prompts_func, batched=True)
|
||||
|
||||
# Configure the trainer
|
||||
config.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Determine steps or epochs based on dataset size
|
||||
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 = -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=train_params.max_seq_length,
|
||||
data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
|
||||
dataset_num_proc=1 if config.device == "cpu" else 2,
|
||||
packing=False,
|
||||
args=TrainingArguments(
|
||||
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=fp16,
|
||||
bf16=bf16,
|
||||
logging_steps=10,
|
||||
optim="adamw_torch" if config.device == "cpu" else "adamw_8bit",
|
||||
weight_decay=0.01,
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
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="<|im_start|>user",
|
||||
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(config.output_dir)
|
||||
tokenizer.save_pretrained(config.output_dir)
|
||||
|
||||
# Create a metadata file with training information
|
||||
with open(config.output_dir / "training_info.json", "w") as f:
|
||||
json.dump({
|
||||
"base_model": config.base_model,
|
||||
"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)
|
||||
|
||||
def main():
|
||||
config = Config()
|
||||
|
||||
# Prepare training data
|
||||
training_data, train_params = prepare_training_data(config.config_path)
|
||||
|
||||
if not training_data:
|
||||
print("No valid training data found. Exiting.")
|
||||
return 1
|
||||
|
||||
# Train the model
|
||||
try:
|
||||
train_model(config, training_data, train_params)
|
||||
|
||||
# Create symlink to current
|
||||
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 {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())
|
||||
Reference in New Issue
Block a user