147 lines
4.2 KiB
Python
147 lines
4.2 KiB
Python
#!/root/venvs/train/bin/python
|
|
"""
|
|
Fine-tuning for QwQ model
|
|
"""
|
|
|
|
# Unsloth should be imported before transformers to ensure all optimizations are applied.
|
|
from unsloth import FastLanguageModel, is_bfloat16_supported
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from transformers import AutoTokenizer, TrainingArguments
|
|
from trl import SFTTrainer
|
|
from typing import Optional, List
|
|
import argparse
|
|
import json
|
|
import os
|
|
|
|
from .dataset import Dataset
|
|
|
|
@dataclass
|
|
class Args:
|
|
def __init__(self, args: Optional[List[str]] = None):
|
|
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'
|
|
)
|
|
if args is None:
|
|
self.args = parser.parse_args()
|
|
else:
|
|
self.args = parser.parse_args(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
|
|
|
|
def main():
|
|
args = Args()
|
|
dataset = Dataset(args.config_path)
|
|
dataset.validate()
|
|
|
|
max_seq_length = 2048 # Can increase for longer reasoning traces
|
|
lora_rank = 64 # Larger rank = smarter, but slower
|
|
|
|
with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
|
|
tokenizer_config = json.load(f)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
args.base_model,
|
|
tokenizer_config=tokenizer_config,
|
|
)
|
|
|
|
model, _returned_tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name = args.base_model,
|
|
max_seq_length = max_seq_length,
|
|
load_in_4bit = True, # False for LoRA 16bit
|
|
fast_inference = True, # Enable vLLM fast inference
|
|
max_lora_rank = lora_rank,
|
|
gpu_memory_utilization = 0.5, # Reduce if out of memory
|
|
tokenizer = tokenizer,
|
|
)
|
|
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
|
target_modules = [
|
|
"q_proj", "k_proj", "v_proj", "o_proj",
|
|
"gate_proj", "up_proj", "down_proj",
|
|
], # Remove QKVO if out of memory
|
|
lora_alpha = lora_rank,
|
|
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
|
|
)
|
|
|
|
training_args = TrainingArguments(
|
|
output_dir=str(args.output_dir),
|
|
num_train_epochs=3,
|
|
per_device_train_batch_size=1,
|
|
gradient_accumulation_steps=16,
|
|
gradient_checkpointing=True,
|
|
learning_rate=2e-5,
|
|
lr_scheduler_type="cosine",
|
|
warmup_ratio=0.05,
|
|
weight_decay=0.01,
|
|
fp16=not is_bfloat16_supported(),
|
|
bf16=is_bfloat16_supported(),
|
|
logging_steps=10,
|
|
save_steps=200,
|
|
save_total_limit=3,
|
|
report_to="none",
|
|
optim="adamw_8bit",
|
|
)
|
|
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
args=training_args,
|
|
train_dataset=dataset.to_transformers_dataset(tokenizer),
|
|
dataset_text_field="messages",
|
|
max_seq_length=max_seq_length,
|
|
)
|
|
|
|
trainer.train()
|
|
|
|
model.save_pretrained_merged(
|
|
str(args.output_dir),
|
|
tokenizer=tokenizer,
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
main() |