140 lines
3.9 KiB
Python
140 lines
3.9 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 .dataset import Dataset
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from transformers import TrainingArguments
|
|
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
|
|
import argparse
|
|
import os
|
|
import torch
|
|
|
|
@dataclass
|
|
class Args:
|
|
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'
|
|
)
|
|
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
|
|
|
|
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
|
|
|
|
model, 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.85, # Reduce if out of memory
|
|
)
|
|
|
|
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,
|
|
use_gradient_checkpointing = "unsloth", # Enable long context finetuning
|
|
random_state = 3407,
|
|
)
|
|
|
|
response_template = tokenizer.apply_chat_template(
|
|
[{"role": "assistant", "content": ""}],
|
|
tokenize=False,
|
|
add_generation_prompt=True
|
|
)
|
|
|
|
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,
|
|
data_collator=DataCollatorForCompletionOnlyLM(
|
|
response_template=response_template,
|
|
tokenizer=tokenizer
|
|
),
|
|
)
|
|
|
|
trainer.train()
|
|
|
|
model.save_pretrained_merged(
|
|
str(args.output_dir),
|
|
tokenizer=tokenizer,
|
|
save_method="merged_16bit"
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
main() |