Fix Gemma3ForCausalLm does not have attribute self.llm
This commit is contained in:
Roland Tannous
parent
4a70f8e880
commit
17563176d9
8 changed files with 13 additions and 45376 deletions
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -1,240 +0,0 @@
|
|||
from unsloth import FastModel, 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
|
||||
|
||||
|
||||
|
||||
|
||||
from cleanup_utils import safe_remove_directory
|
||||
from perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False).removeprefix('<bos>') for convo in convos]
|
||||
return { "text" : texts, }
|
||||
|
||||
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 FastModel
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from perplexity_eval import ppl_model
|
||||
|
||||
# Load model
|
||||
merged_model, merged_tokenizer = FastModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_gemma3_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 = "gemma-3",
|
||||
)
|
||||
|
||||
# Load dataset fresh in subprocess
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
# Format the dataset
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False).removeprefix('<bos>') for convo in convos]
|
||||
return { "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 = FastModel.from_pretrained(
|
||||
model_name="unsloth/gemma-3-1b-it",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template = "gemma-3",
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import standardize_sharegpt
|
||||
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 = FastModel.get_peft_model(
|
||||
model,
|
||||
finetune_vision_layers = False, # Turn off for just text!
|
||||
finetune_language_layers = True, # Should leave on!
|
||||
finetune_attention_modules = True, # Attention good for GRPO
|
||||
finetune_mlp_modules = True, # SHould leave on always!
|
||||
|
||||
r=16,
|
||||
#target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
target_modules = "all-linear",
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
modules_to_save=[
|
||||
"lm_head",
|
||||
"embed_tokens"
|
||||
]
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
max_seq_length=2048,
|
||||
packing=False,
|
||||
args=SFTConfig(
|
||||
dataset_text_field="text",
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
gradient_checkpointing=True,
|
||||
gradient_checkpointing_kwargs={"use_reentrant":False},
|
||||
warmup_ratio=0.03,
|
||||
max_steps=40,
|
||||
learning_rate=3e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=5,
|
||||
#optim="adamw_8bit",
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
max_grad_norm=0.3,
|
||||
dataset_num_proc=2,
|
||||
),
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part = "<start_of_turn>user\n",
|
||||
response_part = "<start_of_turn>model\n",
|
||||
)
|
||||
|
||||
# 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_gemma3_text_model',
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
# Clean up
|
||||
del model
|
||||
del tokenizer
|
||||
del trainer
|
||||
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 = FastModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_gemma3_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))
|
||||
|
||||
# Clean up
|
||||
del merged_model
|
||||
del merged_tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
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 = FastModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_gemma3_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()
|
||||
|
||||
# final cleanup
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth_out")
|
||||
|
|
@ -1,228 +0,0 @@
|
|||
from unsloth import FastModel, 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[2]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
print(sys.path)
|
||||
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
||||
# Define helper functions outside of main
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False).removeprefix('<bos>') for convo in convos]
|
||||
return { "text" : texts, }
|
||||
|
||||
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 FastModel
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
from tests.utils.perplexity_eval import ppl_model
|
||||
|
||||
# Load model
|
||||
merged_model, merged_tokenizer = FastModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_gemma3_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 = "gemma-3",
|
||||
)
|
||||
|
||||
# Load dataset fresh in subprocess
|
||||
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split="eval")
|
||||
|
||||
# Format the dataset
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["messages"]
|
||||
texts = [merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False).removeprefix('<bos>') for convo in convos]
|
||||
return { "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 = FastModel.from_pretrained(
|
||||
model_name="unsloth/gemma-3-4b-it",
|
||||
max_seq_length=2048,
|
||||
dtype=compute_dtype,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
attn_implementation=attn_implementation
|
||||
)
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template = "gemma-3",
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import standardize_sharegpt
|
||||
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 = FastModel.get_peft_model(
|
||||
model,
|
||||
finetune_vision_layers = False, # Turn off for just text!
|
||||
finetune_language_layers = True, # Should leave on!
|
||||
finetune_attention_modules = True, # Attention good for GRPO
|
||||
finetune_mlp_modules = True, # SHould leave on always!
|
||||
|
||||
r=16,
|
||||
#target_modules=['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
|
||||
target_modules = "all-linear",
|
||||
lora_alpha=16,
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
use_rslora=False,
|
||||
loftq_config=None,
|
||||
modules_to_save=[
|
||||
"lm_head",
|
||||
"embed_tokens"
|
||||
]
|
||||
)
|
||||
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=dataset_train,
|
||||
max_seq_length=2048,
|
||||
packing=False,
|
||||
args=SFTConfig(
|
||||
dataset_text_field="text",
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4,
|
||||
gradient_checkpointing=True,
|
||||
gradient_checkpointing_kwargs={"use_reentrant":False},
|
||||
warmup_ratio=0.03,
|
||||
max_steps=40,
|
||||
learning_rate=3e-4,
|
||||
fp16=not is_bfloat16_supported(),
|
||||
bf16=is_bfloat16_supported(),
|
||||
logging_steps=5,
|
||||
#optim="adamw_8bit",
|
||||
optim="adamw_8bit",
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
output_dir="outputs",
|
||||
report_to="none",
|
||||
max_grad_norm=0.3,
|
||||
dataset_num_proc=2,
|
||||
),
|
||||
)
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part = "<start_of_turn>user\n",
|
||||
response_part = "<start_of_turn>model\n",
|
||||
)
|
||||
|
||||
# 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_gemma3_text_model',
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
|
||||
# load model from local disk and test
|
||||
print("Loading merged model in 4 bit for perplexity test")
|
||||
merged_model, merged_tokenizer = FastModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_gemma3_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 = FastModel.from_pretrained(
|
||||
model_name="./unsloth_out/merged_gemma3_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()
|
||||
|
||||
# final cleanup
|
||||
safe_remove_directory("./outputs")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./unsloth_out")
|
||||
|
|
@ -1,802 +0,0 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
"""test_Llama3_1_(3B)_GRPO_LoRA (1).ipynb
|
||||
|
||||
### Unsloth
|
||||
|
||||
"""
|
||||
# import os
|
||||
# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
|
||||
# os.environ['TORCH_USE_CUDA_DSA'] = '1'
|
||||
#
|
||||
# # Add at the beginning of your training script
|
||||
# import torch
|
||||
# torch.backends.cuda.matmul.allow_tf32 = False
|
||||
# torch.backends.cudnn.allow_tf32 = False
|
||||
|
||||
from unsloth import FastLanguageModel
|
||||
import torch
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import multiprocessing as mp
|
||||
import gc
|
||||
from multiprocessing import Queue
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[2]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.aime_eval import evaluate_model_aime, compare_aime_results
|
||||
|
||||
|
||||
max_seq_length = 2048 # Can increase for longer reasoning traces
|
||||
lora_rank = 64 # Larger rank = smarter, but slower
|
||||
|
||||
|
||||
def evaluate_merged_model(result_queue, load_in_4bit=False, load_in_8bit=False):
|
||||
from unsloth import FastLanguageModel
|
||||
from tests.utils.aime_eval import evaluate_model_aime
|
||||
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 = "./final_merged_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.8, # Reduce if out of memory
|
||||
)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
if load_in_4bit:
|
||||
print("🔍 EVALUATION Merged model: 4 bits load")
|
||||
model_type="merged_model_4bits"
|
||||
elif load_in_8bit:
|
||||
print("🔍 EVALUATION Merged model: 8 bits load")
|
||||
model_type="merged_model_8bits"
|
||||
else:
|
||||
print("🔍 EVALUATION Merged model: 16 bits load")
|
||||
model_type="merged_model_16bits"
|
||||
print(f"{'='*60}")
|
||||
|
||||
evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type=model_type,
|
||||
temperature=0.3,
|
||||
n_sampling=8,
|
||||
max_tokens=32768,
|
||||
top_p=0.95,
|
||||
seed=0
|
||||
)
|
||||
|
||||
result_queue.put(results)
|
||||
|
||||
del model
|
||||
del tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
|
||||
# Main execution code should be wrapped in this guard
|
||||
def training_run(result_queue):
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "unsloth/gemma-3-1b-it",
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = False, # False for LoRA 16bit
|
||||
fast_inference = True, # Enable vLLM fast inference
|
||||
max_lora_rank = lora_rank,
|
||||
gpu_memory_utilization = 0.8, # Reduce if out of memory
|
||||
)
|
||||
|
||||
"""### Helper Functions
|
||||
<a name="Data"></a>
|
||||
|
||||
#### Helper functions - Data Prep
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answer>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
def extract_hash_answer(text):
|
||||
"""Extract answer from GSM8K format"""
|
||||
if "####" not in text:
|
||||
return None
|
||||
return text.split("####")[1].strip()
|
||||
|
||||
def prepare_gsm8k_dataset(dataset):
|
||||
"""Format GSM8K dataset for training"""
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answer>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
system_prompt = (
|
||||
f"You are given a problem. Think about the problem and reason step by step. "
|
||||
f"Place your thinking process between {reasoning_start} and {reasoning_end}. "
|
||||
f"Then, provide your final numerical solution between {solution_start}{solution_end}"
|
||||
)
|
||||
|
||||
def format_gsm8k(example):
|
||||
return {
|
||||
"prompt": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": example["question"]},
|
||||
],
|
||||
"answer": extract_hash_answer(example["answer"]),
|
||||
}
|
||||
|
||||
return dataset.map(format_gsm8k)
|
||||
|
||||
def prepare_limo_dataset(dataset):
|
||||
"""Format LIMO dataset for SFT training"""
|
||||
if dataset is None:
|
||||
return None
|
||||
|
||||
system_prompt = """You are a helpful reasoning assistant. When given a problem, think through it step by step and provide your answer in the following format:
|
||||
|
||||
<reasoning>
|
||||
[Your detailed step-by-step reasoning and solution process]
|
||||
</reasoning>
|
||||
<answer>
|
||||
[Your final numerical answer]
|
||||
</answer>"""
|
||||
|
||||
def format_limo(example):
|
||||
# Create the assistant response
|
||||
assistant_response = f"<reasoning>\n{example['solution']}\n</reasoning>\n<answer>\n{example['answer']}\n</answer>"
|
||||
|
||||
# Return a DICTIONARY with the conversation in a field
|
||||
return {
|
||||
"prompt": [ # ← This is the key change - wrap in a dict
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": example["question"]},
|
||||
{"role": "assistant", "content": assistant_response}
|
||||
]
|
||||
}
|
||||
|
||||
return dataset.map(format_limo)
|
||||
|
||||
print("\n✅ Dataset preparation functions defined!")
|
||||
|
||||
"""#### Helper functions - Evaluation"""
|
||||
|
||||
def get_max_prompt_length(dataset, tokenizer):
|
||||
"""Calculate maximum and average prompt length in dataset"""
|
||||
print("Analyzing prompt lengths...")
|
||||
|
||||
lengths = dataset.map(
|
||||
lambda x: {
|
||||
"tokens": tokenizer.apply_chat_template(
|
||||
x["prompt"],
|
||||
add_generation_prompt=True,
|
||||
tokenize=True
|
||||
)
|
||||
},
|
||||
batched=True,
|
||||
).map(lambda x: {"length": len(x["tokens"])})["length"]
|
||||
|
||||
max_length = max(lengths)
|
||||
avg_length = sum(lengths) / len(lengths)
|
||||
min_length = min(lengths)
|
||||
|
||||
print(f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}")
|
||||
return max_length, avg_length
|
||||
|
||||
def extract_unsloth_answer(text, start_tag="<SOLUTION>", end_tag="</SOLUTION>"):
|
||||
"""Extract answer from Unsloth SOLUTION tags"""
|
||||
pattern = re.escape(start_tag) + r"(.*?)" + re.escape(end_tag)
|
||||
matches = re.findall(pattern, text, re.DOTALL)
|
||||
|
||||
if matches:
|
||||
answer = matches[-1] # Get the last match
|
||||
answer = re.sub(r"[%$,]", "", answer).strip()
|
||||
return answer
|
||||
return ""
|
||||
|
||||
def find_number(search_string):
|
||||
"""Find the last number in a string"""
|
||||
numbers = re.compile(
|
||||
r"-?[\d,]*\.?\d+",
|
||||
re.MULTILINE | re.DOTALL | re.IGNORECASE,
|
||||
).findall(search_string)
|
||||
|
||||
if numbers:
|
||||
return numbers[-1].replace(",", "").strip()
|
||||
return ""
|
||||
|
||||
def remove_symbols(x: str) -> str:
|
||||
"""Remove commas, percent and dollar symbols"""
|
||||
if not x:
|
||||
return ""
|
||||
return x.replace(",", "").replace("%", "").replace("$", "").strip()
|
||||
|
||||
def get_num_tokens(text, tokenizer_instance):
|
||||
"""Count tokens in text"""
|
||||
if not text:
|
||||
return 0
|
||||
encoding = tokenizer_instance(text, return_tensors="pt")
|
||||
return len(encoding["input_ids"][0])
|
||||
|
||||
def check_format_compliance(text, format_type="unsloth"):
|
||||
"""Check if response follows expected format"""
|
||||
if format_type == "unsloth":
|
||||
reasoning_start = "<start_reasoning>"
|
||||
reasoning_end = "<end_reasoning>"
|
||||
solution_start = "<SOLUTION>"
|
||||
solution_end = "</SOLUTION>"
|
||||
|
||||
pattern = (
|
||||
rf"^[\s]*{re.escape(reasoning_start)}.+?{re.escape(reasoning_end)}.*?"
|
||||
rf"{re.escape(solution_start)}.+?{re.escape(solution_end)}[\s]*$"
|
||||
)
|
||||
else:
|
||||
return False
|
||||
|
||||
return bool(re.match(pattern, text.strip(), re.DOTALL))
|
||||
|
||||
def normalize_answer(answer):
|
||||
"""Normalize answer for comparison"""
|
||||
if not answer:
|
||||
return ""
|
||||
|
||||
normalized = remove_symbols(str(answer))
|
||||
|
||||
try:
|
||||
float_val = float(normalized)
|
||||
if float_val.is_integer():
|
||||
return str(int(float_val))
|
||||
else:
|
||||
return str(float_val)
|
||||
except (ValueError, TypeError):
|
||||
return normalized
|
||||
|
||||
def evaluate_answer_correctness(extracted_answer, ground_truth):
|
||||
"""Evaluate answer correctness with multiple criteria"""
|
||||
if not extracted_answer or not ground_truth:
|
||||
return False, False, 0.0
|
||||
|
||||
norm_extracted = normalize_answer(extracted_answer)
|
||||
norm_ground_truth = normalize_answer(ground_truth)
|
||||
|
||||
if norm_extracted == norm_ground_truth:
|
||||
return True, True, 1.0
|
||||
|
||||
try:
|
||||
extracted_num = float(norm_extracted)
|
||||
ground_truth_num = float(norm_ground_truth)
|
||||
|
||||
if ground_truth_num != 0:
|
||||
relative_error = abs(extracted_num - ground_truth_num) / abs(ground_truth_num)
|
||||
|
||||
if relative_error < 0.01:
|
||||
return True, True, 0.9
|
||||
elif relative_error < 0.05:
|
||||
return False, True, 0.7
|
||||
elif relative_error < 0.10:
|
||||
return False, True, 0.5
|
||||
else:
|
||||
if extracted_num == 0:
|
||||
return True, True, 1.0
|
||||
elif abs(extracted_num) < 0.01:
|
||||
return False, True, 0.7
|
||||
|
||||
except (ValueError, TypeError):
|
||||
if norm_extracted.lower() == norm_ground_truth.lower():
|
||||
return True, True, 1.0
|
||||
|
||||
return False, False, 0.0
|
||||
|
||||
"""#### Reward Functions for GRPO"""
|
||||
|
||||
def match_format_exactly(completions, **kwargs):
|
||||
"""Reward function for exact format matching"""
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answer>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
pattern = (
|
||||
rf"^[\s]*{re.escape(reasoning_start)}.+?{re.escape(reasoning_end)}.*?"
|
||||
rf"{re.escape(solution_start)}.+?{re.escape(solution_end)}[\s]*$"
|
||||
)
|
||||
|
||||
responses = [completion[0]["content"] for completion in completions]
|
||||
rewards = [3.0 if re.match(pattern, response, re.DOTALL) else 0.0 for response in responses]
|
||||
return rewards
|
||||
|
||||
def match_format_approximately(completions, **kwargs):
|
||||
"""Reward function for approximate format matching"""
|
||||
reasoning_start = "<reasoning>"
|
||||
reasoning_end = "</reasoning>"
|
||||
solution_start = "<answerr>"
|
||||
solution_end = "</answer>"
|
||||
|
||||
scores = []
|
||||
for completion in completions:
|
||||
score = 0
|
||||
response = completion[0]["content"]
|
||||
score += 0.5 if response.count(reasoning_start) == 1 else -1.0
|
||||
score += 0.5 if response.count(reasoning_end) == 1 else -1.0
|
||||
score += 0.5 if response.count(solution_start) == 1 else -1.0
|
||||
score += 0.5 if response.count(solution_end) == 1 else -1.0
|
||||
scores.append(score)
|
||||
return scores
|
||||
|
||||
def check_answer_correctness(prompts, completions, answer, **kwargs):
|
||||
"""Reward function for answer correctness"""
|
||||
def extract_solution_answer(text):
|
||||
pattern = r"<answer>(.*?)</answer>"
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
if match:
|
||||
return re.sub(r"[%$,]", "", match.group(1)).strip()
|
||||
return ""
|
||||
|
||||
responses = [completion[0]["content"] for completion in completions]
|
||||
extracted_responses = [extract_solution_answer(r) for r in responses]
|
||||
|
||||
scores = []
|
||||
for guess, true_answer in zip(extracted_responses, answer):
|
||||
score = 0
|
||||
if not guess:
|
||||
scores.append(0)
|
||||
continue
|
||||
|
||||
if guess == true_answer:
|
||||
score += 3.0
|
||||
elif guess.strip() == true_answer.strip():
|
||||
score += 1.5
|
||||
else:
|
||||
try:
|
||||
ratio = float(guess) / float(true_answer)
|
||||
if 0.9 <= ratio <= 1.1:
|
||||
score += 1.0
|
||||
elif 0.8 <= ratio <= 1.2:
|
||||
score += 0.5
|
||||
else:
|
||||
score -= 1.5
|
||||
except:
|
||||
score -= 1.5
|
||||
scores.append(score)
|
||||
return scores
|
||||
|
||||
print("✅ Reward functions defined!")
|
||||
|
||||
"""#### Main Evaluation Function"""
|
||||
|
||||
import gc
|
||||
|
||||
|
||||
|
||||
"""#### Comparison and Memory Management"""
|
||||
|
||||
def compare_model_results(all_results):
|
||||
"""Generate comprehensive comparison of multiple model results"""
|
||||
print(f"\n{'='*80}")
|
||||
print("COMPREHENSIVE MODEL COMPARISON")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Main table
|
||||
print(f"{'Model':<15} {'Format %':<10} {'Exact %':<10} {'Plausible %':<12} {'Confidence':<12}")
|
||||
print("-" * 80)
|
||||
|
||||
for result in all_results:
|
||||
print(f"{result['model_type']:<15} "
|
||||
f"{result['correct_format_pct']:<10.1f} "
|
||||
f"{result['exact_match_pct']:<10.1f} "
|
||||
f"{result['plausible_match_pct']:<12.1f} "
|
||||
f"{result['avg_confidence']:<12.3f}")
|
||||
|
||||
# Improvement analysis
|
||||
if len(all_results) > 1:
|
||||
print(f"\n{'='*50}")
|
||||
print("IMPROVEMENT ANALYSIS")
|
||||
print(f"{'='*50}")
|
||||
|
||||
base_result = all_results[0]
|
||||
for result in all_results[1:]:
|
||||
print(f"\n{result['model_type']} vs {base_result['model_type']}:")
|
||||
format_improvement = result['correct_format_pct'] - base_result['correct_format_pct']
|
||||
exact_improvement = result['exact_match_pct'] - base_result['exact_match_pct']
|
||||
plausible_improvement = result['plausible_match_pct'] - base_result['plausible_match_pct']
|
||||
|
||||
print(f" Format compliance: {format_improvement:+.1f}%")
|
||||
print(f" Exact matches: {exact_improvement:+.1f}%")
|
||||
print(f" Plausible matches: {plausible_improvement:+.1f}%")
|
||||
|
||||
# Save comparison
|
||||
comparison_data = {
|
||||
"summary": all_results,
|
||||
"best_model": max(all_results, key=lambda x: x['exact_match_pct']),
|
||||
}
|
||||
|
||||
with open("model_comparison_comprehensive.json", "w") as f:
|
||||
json.dump(comparison_data, f, indent=4)
|
||||
|
||||
print(f"\nBest performing model: {comparison_data['best_model']['model_type']} "
|
||||
f"({comparison_data['best_model']['exact_match_pct']:.1f}% exact matches)")
|
||||
|
||||
def cleanup_memory():
|
||||
"""Comprehensive memory cleanup"""
|
||||
print("🧹 Cleaning up GPU memory...")
|
||||
for _ in range(10):
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
allocated = torch.cuda.memory_allocated() / 1024**3
|
||||
reserved = torch.cuda.memory_reserved() / 1024**3
|
||||
print(f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
|
||||
|
||||
"""#### Data Loading and Preparation"""
|
||||
|
||||
from datasets import load_dataset
|
||||
|
||||
|
||||
# Load GSM8K
|
||||
gsm8k_dataset = load_dataset("openai/gsm8k", "main", split="train")
|
||||
|
||||
# Load LIMO (adjust this based on your access method)
|
||||
limo_train = load_dataset("GAIR/LIMO", split="train")
|
||||
|
||||
# Prepare datasets
|
||||
gsm8k_train = prepare_gsm8k_dataset(gsm8k_dataset)
|
||||
limo_train = prepare_limo_dataset(limo_train)
|
||||
|
||||
|
||||
print(f" GSM8K train: {len(gsm8k_train)}")
|
||||
print(f" LIMO train: {len(limo_train) if limo_train else 0}")
|
||||
|
||||
# Store results
|
||||
all_results = []
|
||||
|
||||
# Single temperature evaluation on combined dataset
|
||||
# results = evaluate_model_aime(
|
||||
# model=model,
|
||||
# tokenizer=tokenizer,
|
||||
# model_type="base",
|
||||
# temperature=0.3,
|
||||
# n_sampling=8,
|
||||
# max_tokens=32768,
|
||||
# top_p=0.95,
|
||||
# seed=0
|
||||
# )
|
||||
#
|
||||
from unsloth.chat_templates import get_chat_template
|
||||
|
||||
tokenizer = get_chat_template(
|
||||
tokenizer,
|
||||
chat_template = "gemma-3",
|
||||
)
|
||||
|
||||
def formatting_prompts_func(examples):
|
||||
convos = examples["prompt"]
|
||||
texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False).removeprefix('<bos>') for convo in convos]
|
||||
return { "text" : texts, }
|
||||
|
||||
limo_train = limo_train.map(formatting_prompts_func, batched = True,)
|
||||
|
||||
from trl import SFTTrainer
|
||||
from transformers import DataCollatorForSeq2Seq, TrainingArguments
|
||||
from unsloth import is_bfloat16_supported
|
||||
|
||||
|
||||
print(f"\n{'*'*60}")
|
||||
print("🎯 STAGE 1: Qlora Fine-Tuning on LIMO")
|
||||
print(f"{'*'*60}")
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
if limo_train is not None:
|
||||
trainer = SFTTrainer(
|
||||
model = model,
|
||||
tokenizer = tokenizer,
|
||||
train_dataset = limo_train,
|
||||
dataset_text_field = "text",
|
||||
max_seq_length = max_seq_length,
|
||||
data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),
|
||||
dataset_num_proc = 2,
|
||||
packing = False, # Can make training 5x faster for short sequences.
|
||||
args = TrainingArguments(
|
||||
per_device_train_batch_size = 2,
|
||||
gradient_accumulation_steps = 4,
|
||||
warmup_steps = 5,
|
||||
num_train_epochs = 1, # Set this for 1 full training run.
|
||||
#max_steps = 60,
|
||||
learning_rate = 2e-4,
|
||||
fp16 = not is_bfloat16_supported(),
|
||||
bf16 = is_bfloat16_supported(),
|
||||
logging_steps = 1,
|
||||
optim = "adamw_8bit",
|
||||
weight_decay = 0.01,
|
||||
lr_scheduler_type = "linear",
|
||||
seed = 3407,
|
||||
output_dir = "outputs",
|
||||
report_to = "none", # Use this for WandB etc
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
from unsloth.chat_templates import train_on_responses_only
|
||||
trainer = train_on_responses_only(
|
||||
trainer,
|
||||
instruction_part = "<start_of_turn>user\n",
|
||||
response_part = "<start_of_turn>model\n",
|
||||
)
|
||||
|
||||
# Train
|
||||
#print(f"🚂 Starting SFT training on {len(limo_train)} examples...")
|
||||
#trainer.train()
|
||||
|
||||
# Save checkpoint
|
||||
#model.save_pretrained("qlora_checkpoint")
|
||||
#tokenizer.save_pretrained("qlora_checkpoint")
|
||||
#print("💾 Qlora checkpoint saved!")
|
||||
|
||||
# Cleanup
|
||||
del trainer
|
||||
cleanup_memory()
|
||||
|
||||
#print("✅ Qlora training completed!")
|
||||
else:
|
||||
print("⚠️ Skipping Qlora training - no LIMO dataset available")
|
||||
|
||||
# Cleanup
|
||||
cleanup_memory()
|
||||
|
||||
global PRINTED_TIMES
|
||||
PRINTED_TIMES = 0
|
||||
global PRINT_EVERY_STEPS
|
||||
PRINT_EVERY_STEPS = 5
|
||||
|
||||
match_numbers = re.compile(
|
||||
solution_start + r".*?([\d\.\,]{1,})",
|
||||
flags = re.MULTILINE | re.DOTALL
|
||||
)
|
||||
|
||||
def check_numbers(prompts, completions, answer, **kwargs):
|
||||
question = prompts[0][-1]["content"]
|
||||
responses = [completion[0]["content"] for completion in completions]
|
||||
|
||||
extracted_responses = [
|
||||
guess.group(1)
|
||||
if (guess := match_numbers.search(r)) is not None else None \
|
||||
for r in responses
|
||||
]
|
||||
|
||||
scores = []
|
||||
# Print only every few steps
|
||||
global PRINTED_TIMES
|
||||
global PRINT_EVERY_STEPS
|
||||
if PRINTED_TIMES % PRINT_EVERY_STEPS == 0:
|
||||
print('*'*20, f"Question:\n{question}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
|
||||
PRINTED_TIMES += 1
|
||||
|
||||
for guess, true_answer in zip(extracted_responses, answer):
|
||||
if guess is None:
|
||||
scores.append(0)
|
||||
continue
|
||||
# Convert to numbers
|
||||
try:
|
||||
true_answer = float(true_answer.strip())
|
||||
# Remove commas like in 123,456
|
||||
guess = float(guess.strip().replace(",", ""))
|
||||
scores.append(1.5 if guess == true_answer else -0.5)
|
||||
except:
|
||||
scores.append(0)
|
||||
continue
|
||||
return scores
|
||||
|
||||
print(f"\n{'*'*60}")
|
||||
print("🎯 STAGE 2: GRPO Fine-Tuning on GSM8K")
|
||||
print(f"{'*'*60}")
|
||||
|
||||
# Get max prompt length
|
||||
max_prompt_length, _ = get_max_prompt_length(gsm8k_train, tokenizer)
|
||||
max_prompt_length = min(max_prompt_length + 10, 512) # Add buffer, cap at 512
|
||||
|
||||
print(f"Using max_prompt_length: {max_prompt_length}")
|
||||
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
training_args = GRPOConfig(
|
||||
learning_rate = 5e-6,
|
||||
weight_decay = 0.1,
|
||||
warmup_ratio = 0.1,
|
||||
lr_scheduler_type = "cosine",
|
||||
optim = "adamw_torch_fused",
|
||||
logging_steps = 1,
|
||||
per_device_train_batch_size = 1,
|
||||
gradient_accumulation_steps = 4, # Increase to 4 for smoother training
|
||||
num_generations = 8, # Decrease if out of memory
|
||||
max_prompt_length = max_prompt_length,
|
||||
max_completion_length = max_seq_length - max_prompt_length,
|
||||
# num_train_epochs = 1, # Set to 1 for a full training run
|
||||
#max_steps = 250,
|
||||
max_steps = 1000,
|
||||
save_steps = 250,
|
||||
max_grad_norm = 0.1,
|
||||
report_to = "none", # Can use Weights & Biases
|
||||
output_dir = "outputs",
|
||||
)
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model = model,
|
||||
processing_class = tokenizer,
|
||||
reward_funcs = [
|
||||
match_format_exactly,
|
||||
match_format_approximately,
|
||||
check_answer_correctness,
|
||||
check_numbers,
|
||||
],
|
||||
args = training_args,
|
||||
train_dataset = gsm8k_train,
|
||||
)
|
||||
|
||||
|
||||
# Train
|
||||
print(f"🚂 Starting GRPO training on {len(gsm8k_train)} examples...")
|
||||
trainer.train()
|
||||
|
||||
# Save checkpoint
|
||||
model.save_pretrained("grpo_checkpoint")
|
||||
tokenizer.save_pretrained("grpo_checkpoint")
|
||||
print("💾 GRPO checkpoint saved!")
|
||||
|
||||
# Cleanup
|
||||
del trainer
|
||||
del training_args
|
||||
cleanup_memory()
|
||||
|
||||
print("✅ GRPO training completed!")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print("🔍 EVALUATION 3: Final GRPO Model")
|
||||
print(f"{'='*60}")
|
||||
|
||||
grpo_results = evaluate_model_aime(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
model_type="grpo",
|
||||
temperature=0.3,
|
||||
n_sampling=8,
|
||||
max_tokens=32768,
|
||||
top_p=0.95,
|
||||
seed=0
|
||||
)
|
||||
|
||||
all_results.append(grpo_results)
|
||||
print("✅ Final model evaluation complete!")
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print("💾 SAVING FINAL MODEL")
|
||||
print(f"{'='*60}")
|
||||
|
||||
# Save as merged model
|
||||
try:
|
||||
model.save_pretrained_merged("final_merged_model", tokenizer, save_method="merged_16bit")
|
||||
print("✅ Merged model saved to: final_merged_model/")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Could not save merged model: {e}")
|
||||
print("Final model saved as LoRA adapter only")
|
||||
|
||||
print("💾 Model saving complete!")
|
||||
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
result_queue.put(results)
|
||||
|
||||
# Clean up
|
||||
del model
|
||||
del tokenizer
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
# # Merged model load 16 bits model AIME eval
|
||||
# result_queue = mp.Queue()
|
||||
# p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, False))
|
||||
# p.start()
|
||||
# p.join()
|
||||
#
|
||||
# merged_16bits = result_queue.get()
|
||||
# all_results.append(merged_16bits)
|
||||
#
|
||||
# # Clean up
|
||||
# del merged_model
|
||||
# del merged_tokenizer
|
||||
# del dataset_ppl
|
||||
# torch.cuda.empty_cache()
|
||||
# gc.collect()
|
||||
#
|
||||
# safe_remove_directory("./unsloth_compiled_cache")
|
||||
#
|
||||
# # Merged model load 8 bits model AIME eval
|
||||
#
|
||||
# result_queue = mp.Queue()
|
||||
# p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, True))
|
||||
# p.start()
|
||||
# p.join()
|
||||
#
|
||||
# merged_16bits = result_queue.get()
|
||||
# all_results.append(merged_16bits)
|
||||
|
||||
|
||||
# Merged model load 4 bits AIME eval
|
||||
# result_queue = mp.Queue()
|
||||
# p = mp.Process(target=evaluate_merged_model, args=(result_queue, True, False))
|
||||
# p.start()
|
||||
# p.join()
|
||||
#
|
||||
# merged_16bits = result_queue.get()
|
||||
# all_results.append(merged_16bits)
|
||||
|
||||
if __name__ == "__main__":
|
||||
mp.set_start_method('spawn', force=True)
|
||||
result_queue = mp.Queue()
|
||||
all_results = []
|
||||
|
||||
|
||||
# run main finetuning and grpo loop
|
||||
p = mp.Process(target=training_run, args=(result_queue,))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
results = result_queue.get()
|
||||
all_results = results
|
||||
|
||||
# evaluate merged model loaded 16bits
|
||||
p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, False))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
merged_load_16bits = result_queue.get()
|
||||
all_results.append(merged_load_16bits)
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
# Merged model load 8 bits model AIME eval
|
||||
p = mp.Process(target=evaluate_merged_model, args=(result_queue, False, True))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
merged_load_8bits = result_queue.get()
|
||||
all_results.append(merged_load_8bits)
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
# Merged model load 4 bits model AIME eval
|
||||
p = mp.Process(target=evaluate_merged_model, args=(result_queue, True, False))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
merged_load_4bits = result_queue.get()
|
||||
all_results.append(merged_load_4bits)
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
|
||||
|
||||
# AIME-specific comparison function
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🏆 FINAL TRAINING PIPELINE RESULTS")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Use the AIME-specific comparison
|
||||
compare_aime_results(all_results)
|
||||
|
|
@ -184,8 +184,9 @@ def grpo_trainer__prepare_inputs(function_name, function):
|
|||
rest = re.sub(r"^[ \t]*free, total = torch.cuda.mem_get_info\(\)\s*\n", "", rest)
|
||||
rest = re.sub(r"^[ \t]*print\(f?\".*cuda.*\"\)\s*\n", "", rest)
|
||||
insert = (
|
||||
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
|
||||
" self.llm.wake_up()\n"
|
||||
" if hasattr(self, 'llm'):\n"
|
||||
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
|
||||
" self.llm.wake_up()\n"
|
||||
)
|
||||
function = function[:sig_end] + insert + rest
|
||||
else:
|
||||
|
|
@ -199,8 +200,9 @@ def grpo_trainer__prepare_inputs(function_name, function):
|
|||
rest = re.sub(r"^[ \t]*free, total = torch.cuda.mem_get_info\(\)\s*\n", "", rest)
|
||||
rest = re.sub(r"^[ \t]*print\(f?\".*cuda.*\"\)\s*\n", "", rest)
|
||||
insert = (
|
||||
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
|
||||
" self.llm.wake_up()\n"
|
||||
" if (hasattr(self, 'llm'):\n"
|
||||
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
|
||||
" self.llm.wake_up()\n"
|
||||
)
|
||||
function = header_and_comments + insert + rest
|
||||
|
||||
|
|
@ -218,8 +220,9 @@ def grpo_trainer__prepare_inputs(function_name, function):
|
|||
"self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False)",
|
||||
)
|
||||
sleep_and_cache = (
|
||||
"if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
|
||||
" self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1))\n"
|
||||
"if hasattr(self, 'llm'):\n"
|
||||
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
|
||||
" self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1))\n"
|
||||
" "
|
||||
)
|
||||
if re.search(r"\n\s*return ", function):
|
||||
|
|
@ -310,7 +313,7 @@ def grpo_trainer_compute_loss(function_name, function):
|
|||
logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens
|
||||
_input_ids = input_ids
|
||||
_logits_to_keep = logits_to_keep
|
||||
|
||||
|
||||
per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
|
||||
|
||||
# Compute the KL divergence between the model and the reference model
|
||||
|
|
@ -330,12 +333,12 @@ def grpo_trainer_compute_loss(function_name, function):
|
|||
# loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
|
||||
if "old_per_token_logps" in inputs.keys():
|
||||
old_hidden_states = inputs["old_per_token_logps"]
|
||||
else:
|
||||
else:
|
||||
old_hidden_states = None
|
||||
input_ids = input_ids[:, -logits_to_keep:]
|
||||
if per_token_logps is not None:
|
||||
loss, completion_length, mean_kl = grpo_compute_loss_slow(
|
||||
ref_per_token_logps, per_token_logps, old_hidden_states, input_ids, completion_mask, self.beta, advantages,
|
||||
ref_per_token_logps, per_token_logps, old_hidden_states, input_ids, completion_mask, self.beta, advantages,
|
||||
loss_type = self.args.loss_type,
|
||||
epsilon_low = self.epsilon_low, epsilon_high = self.epsilon_high,
|
||||
max_completion_length = self.args.max_completion_length,
|
||||
|
|
@ -356,7 +359,7 @@ def grpo_trainer_compute_loss(function_name, function):
|
|||
loss, completion_length, mean_kl = grpo_accumulated_loss(
|
||||
self, _input_ids, logits_to_keep, completion_mask, advantages, old_hidden_states,
|
||||
n_chunks = self.args.unsloth_num_chunks,
|
||||
)
|
||||
)
|
||||
|
||||
# Log the metrics
|
||||
# completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue