112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
import torch
|
|
import os
|
|
from datasets import load_dataset
|
|
from transformers import (
|
|
AutoModelForCausalLM,
|
|
AutoTokenizer,
|
|
BitsAndBytesConfig,
|
|
TrainingArguments
|
|
)
|
|
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
|
from trl import SFTTrainer, SFTConfig
|
|
|
|
def train():
|
|
"""
|
|
Main training function to fine-tune Qwen2.5-0.5B using QLoRA.
|
|
Includes debug prints for step-by-step monitoring on Windows.
|
|
"""
|
|
print(">>> DEBUG: Entering train() function...")
|
|
|
|
# Clear CUDA memory before starting to prevent fragmentation
|
|
if torch.cuda.is_available():
|
|
print(">>> DEBUG: Clearing CUDA cache...")
|
|
torch.cuda.empty_cache()
|
|
|
|
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
|
|
data_file = "arxiv_cs_2000.jsonl"
|
|
output_dir = "./qwen-resercher-checkpoints"
|
|
|
|
# 1. Load Tokenizer
|
|
print(f">>> DEBUG: Loading tokenizer for {model_id}...")
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.padding_side = "right"
|
|
|
|
# 2. BitsAndBytes Configuration (QLoRA)
|
|
print(">>> DEBUG: Configuring BitsAndBytes for 4-bit quantization...")
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_use_double_quant=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16
|
|
)
|
|
|
|
# 3. Load Base Model
|
|
print(f">>> DEBUG: Loading base model from {model_id}...")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_id,
|
|
quantization_config=bnb_config,
|
|
device_map="auto",
|
|
trust_remote_code=True
|
|
)
|
|
|
|
print(">>> DEBUG: Preparing model for kbit training...")
|
|
model = prepare_model_for_kbit_training(model)
|
|
|
|
# 4. LoRA Configuration
|
|
print(">>> DEBUG: Setting up LoRA configuration (Rank=16)...")
|
|
peft_config = LoraConfig(
|
|
r=16,
|
|
lora_alpha=32,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
|
lora_dropout=0.05,
|
|
bias="none",
|
|
task_type="CAUSAL_LM"
|
|
)
|
|
|
|
# 5. Load Dataset
|
|
print(f">>> DEBUG: Loading dataset from {data_file}...")
|
|
dataset = load_dataset("json", data_files=data_file, split="train")
|
|
|
|
# 6. Training Arguments
|
|
print(">>> DEBUG: Defining training arguments...")
|
|
training_args = SFTConfig(
|
|
output_dir=output_dir,
|
|
per_device_train_batch_size=1,
|
|
gradient_accumulation_steps=16,
|
|
learning_rate=2e-4,
|
|
logging_steps=10,
|
|
max_steps=200,
|
|
save_steps=100,
|
|
optim="paged_adamw_8bit",
|
|
bf16=True,
|
|
fp16=False,
|
|
dataset_text_field="text",
|
|
max_length=1024,
|
|
gradient_checkpointing=True,
|
|
packing=False
|
|
)
|
|
|
|
# 7. Initialize Trainer
|
|
print(">>> DEBUG: Initializing SFTTrainer...")
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
train_dataset=dataset,
|
|
peft_config=peft_config,
|
|
args=training_args,
|
|
processing_class=tokenizer,
|
|
)
|
|
|
|
# 8. Start Training
|
|
print(">>> DEBUG: Starting training loop...")
|
|
trainer.train()
|
|
|
|
# 9. Save the Adapter
|
|
print(">>> DEBUG: Saving the fine-tuned adapter to ./qwen-resercher...")
|
|
trainer.model.save_pretrained("qwen-resercher")
|
|
tokenizer.save_pretrained("qwen-resercher")
|
|
print(">>> DEBUG: Training complete.")
|
|
|
|
if __name__ == "__main__":
|
|
train()
|