d0287bc596
* tests for additional merge fix unsloth zoo pr 163 * fixed load_dataset indent in mistral perplexity test file
299 lines
8.9 KiB
Python
299 lines
8.9 KiB
Python
from unsloth import FastLanguageModel, FastVisionModel, UnslothVisionDataCollator
|
|
from unsloth.chat_templates import get_chat_template
|
|
from trl import SFTTrainer, SFTConfig
|
|
from transformers import DataCollatorForLanguageModeling, DataCollatorForSeq2Seq, TrainingArguments
|
|
from datasets import load_dataset, Dataset
|
|
import torch
|
|
from tqdm import tqdm
|
|
import pandas as pd
|
|
import multiprocessing as mp
|
|
from multiprocessing import Process, Queue
|
|
import gc
|
|
|
|
# ruff: noqa
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
REPO_ROOT = Path(__file__).parents[3]
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from tests.utils.cleanup_utils import safe_remove_directory
|
|
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def load_and_compute_8bit_ppl(result_queue, load_in_4bit=False, load_in_8bit=False):
|
|
"""Load model and compute perplexity in subprocess"""
|
|
from unsloth import FastLanguageModel
|
|
from tests.utils.perplexity_eval import ppl_model
|
|
|
|
# Load model
|
|
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name="./unsloth_out/merged_mistral_text_model",
|
|
max_seq_length=2048,
|
|
load_in_4bit=load_in_4bit,
|
|
load_in_8bit=load_in_8bit,
|
|
)
|
|
# Set up tokenizer
|
|
# merged_tokenizer = get_chat_template(
|
|
# merged_tokenizer,
|
|
# chat_template="llama-3.1",
|
|
# )
|
|
|
|
# Load dataset fresh in subprocess
|
|
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
|
|
|
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
|
|
|
|
### Instruction:
|
|
{}
|
|
|
|
### Input:
|
|
{}
|
|
|
|
### Response:
|
|
{}"""
|
|
|
|
EOS_TOKEN = merged_tokenizer.eos_token
|
|
|
|
def formatting_prompts_func(examples):
|
|
instructions = []
|
|
inputs = []
|
|
outputs = []
|
|
texts = []
|
|
|
|
for conversation in examples["messages"]:
|
|
# Extract user message and assistant response
|
|
user_message = ""
|
|
assistant_message = ""
|
|
|
|
for turn in conversation:
|
|
if turn["role"] == "user":
|
|
user_message = turn["content"]
|
|
elif turn["role"] == "assistant":
|
|
assistant_message = turn["content"]
|
|
|
|
# Store intermediate format
|
|
instruction = "Complete the statement"
|
|
instructions.append(instruction)
|
|
inputs.append(user_message)
|
|
outputs.append(assistant_message)
|
|
|
|
# Create formatted text
|
|
text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
|
|
texts.append(text)
|
|
|
|
return {
|
|
"instruction": instructions,
|
|
"input": inputs,
|
|
"output": outputs,
|
|
"text": texts
|
|
}
|
|
|
|
|
|
|
|
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
|
|
|
# Compute perplexity using the passed dataset
|
|
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
|
|
|
|
|
|
# IMPORTANT: Convert to Python float if it's a tensor
|
|
if torch.is_tensor(ppl_value):
|
|
ppl_value = ppl_value.cpu().item() # Move to CPU and convert to Python scalar
|
|
elif hasattr(ppl_value, 'item'):
|
|
ppl_value = ppl_value.item() # Convert numpy or other array types
|
|
else:
|
|
ppl_value = float(ppl_value) # Ensure it's a float
|
|
|
|
# Return only the perplexity value
|
|
result_queue.put(ppl_value)
|
|
|
|
# Clean up
|
|
del merged_model
|
|
del merged_tokenizer
|
|
del dataset_ppl
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
# Main execution code should be wrapped in this guard
|
|
if __name__ == "__main__":
|
|
mp.set_start_method('spawn', force=True)
|
|
|
|
if torch.cuda.is_bf16_supported():
|
|
compute_dtype = torch.bfloat16
|
|
attn_implementation = 'flash_attention_2'
|
|
else:
|
|
compute_dtype = torch.float16
|
|
attn_implementation = 'sdpa'
|
|
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name="unsloth/mistral-7b-v0.3",
|
|
max_seq_length=2048,
|
|
dtype=compute_dtype,
|
|
load_in_4bit=True,
|
|
load_in_8bit=False,
|
|
full_finetuning=False,
|
|
attn_implementation=attn_implementation
|
|
)
|
|
|
|
|
|
EOS_TOKEN = tokenizer.eos_token
|
|
|
|
|
|
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
|
|
|
|
### Instruction:
|
|
{}
|
|
|
|
### Input:
|
|
{}
|
|
|
|
### Response:
|
|
{}"""
|
|
|
|
|
|
# Define helper functions outside of main
|
|
def formatting_prompts_func(examples):
|
|
instructions = []
|
|
inputs = []
|
|
outputs = []
|
|
texts = []
|
|
|
|
for conversation in examples["messages"]:
|
|
# Extract user message and assistant response
|
|
user_message = ""
|
|
assistant_message = ""
|
|
|
|
for turn in conversation:
|
|
if turn["role"] == "user":
|
|
user_message = turn["content"]
|
|
elif turn["role"] == "assistant":
|
|
assistant_message = turn["content"]
|
|
|
|
# Store intermediate format
|
|
instruction = "Complete the statement"
|
|
instructions.append(instruction)
|
|
inputs.append(user_message)
|
|
outputs.append(assistant_message)
|
|
|
|
# Create formatted text
|
|
text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
|
|
texts.append(text)
|
|
|
|
|
|
return {
|
|
"instruction": instructions,
|
|
"input": inputs,
|
|
"output": outputs,
|
|
"text": texts
|
|
}
|
|
|
|
|
|
|
|
dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split="train")
|
|
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
|
|
|
dataset_train = dataset_train.map(formatting_prompts_func, batched=True)
|
|
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched=True)
|
|
|
|
add_to_comparison("Base model 4 bits", ppl_model(model, tokenizer, dataset_ppl))
|
|
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r=16,
|
|
target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
|
lora_alpha=16,
|
|
lora_dropout=0,
|
|
bias="none",
|
|
use_gradient_checkpointing="unsloth",
|
|
random_state=3407,
|
|
use_rslora=False,
|
|
loftq_config=None,
|
|
)
|
|
|
|
from unsloth import is_bfloat16_supported
|
|
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
train_dataset=dataset_train,
|
|
dataset_text_field="text",
|
|
max_seq_length=2048,
|
|
dataset_num_proc=2,
|
|
packing=False,
|
|
args=TrainingArguments(
|
|
per_device_train_batch_size=2,
|
|
gradient_accumulation_steps=4,
|
|
warmup_ratio=0.1,
|
|
max_steps=200,
|
|
learning_rate=2e-4,
|
|
fp16=not is_bfloat16_supported(),
|
|
bf16=is_bfloat16_supported(),
|
|
logging_steps=50,
|
|
optim="adamw_8bit",
|
|
lr_scheduler_type="linear",
|
|
seed=3407,
|
|
output_dir="outputs",
|
|
report_to="none",
|
|
),
|
|
)
|
|
|
|
# run training
|
|
trainer_stats = trainer.train()
|
|
|
|
add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
|
|
|
|
# saving and merging the model to local disk
|
|
print("merge and save to local disk")
|
|
model.save_pretrained_merged(
|
|
save_directory='./unsloth_out/merged_mistral_text_model',
|
|
tokenizer=tokenizer
|
|
)
|
|
|
|
# print("cleaning")
|
|
# del model
|
|
# del tokenizer
|
|
# torch.cuda.empty_cache()
|
|
# gc.collect()
|
|
|
|
# load model from local disk and test
|
|
print("Loading merged model in 4 bit for perplexity test")
|
|
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name="./unsloth_out/merged_mistral_text_model",
|
|
max_seq_length=2048,
|
|
load_in_4bit=True,
|
|
load_in_8bit=False,
|
|
)
|
|
|
|
add_to_comparison("merged model load 4bit", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
|
|
|
|
|
print("Computing 8-bit model perplexity in subprocess...")
|
|
result_queue = mp.Queue()
|
|
p = mp.Process(target=load_and_compute_8bit_ppl, args=(result_queue, False, True))
|
|
p.start()
|
|
p.join()
|
|
|
|
ppl_8bit = result_queue.get()
|
|
add_to_comparison("merged model loaded 8bits", ppl_8bit)
|
|
|
|
print("Loading merged model in 16 bit for perplexity test")
|
|
merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name="./unsloth_out/merged_mistral_text_model",
|
|
max_seq_length=2048,
|
|
load_in_4bit=False,
|
|
load_in_8bit=False,
|
|
)
|
|
|
|
add_to_comparison("merged model loaded 16bits", ppl_model(merged_model, merged_tokenizer, dataset_ppl))
|
|
|
|
print_model_comparison()
|
|
|
|
safe_remove_directory("./outputs")
|
|
safe_remove_directory("./unsloth_compiled_cache")
|
|
safe_remove_directory("./unsloth_out")
|