145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
from unsloth import FastLanguageModel
|
|
import torch
|
|
from datasets import load_dataset
|
|
|
|
|
|
fourbit_models = [
|
|
"unsloth/Qwen3-1.7B-unsloth-bnb-4bit", # Qwen 14B 2x faster
|
|
"unsloth/Qwen3-4B-unsloth-bnb-4bit",
|
|
"unsloth/Qwen3-8B-unsloth-bnb-4bit",
|
|
"unsloth/Qwen3-14B-unsloth-bnb-4bit",
|
|
"unsloth/Qwen3-32B-unsloth-bnb-4bit",
|
|
|
|
# 4bit dynamic quants for superior accuracy and low memory use
|
|
"unsloth/gemma-3-12b-it-unsloth-bnb-4bit",
|
|
"unsloth/Phi-4",
|
|
"unsloth/Llama-3.1-8B",
|
|
"unsloth/Llama-3.2-3B",
|
|
"unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit" # [NEW] We support TTS models!
|
|
]
|
|
|
|
|
|
# 加载预训练模型
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name = "unsloth/Qwen3-1.7B-unsloth-bnb-4bit",
|
|
max_seq_length = 2048, # Context length - can be longer, but uses more memory
|
|
load_in_4bit = True, # 4bit uses much less memory , 启用QLoRA
|
|
load_in_8bit = False, # A bit more accurate, uses 2x memory
|
|
full_finetuning = False, # We have full finetuning now!
|
|
# token = "hf_...", # use one if using gated models
|
|
)
|
|
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r = 16, # 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",],
|
|
lora_alpha = 16, # Best to choose alpha = rank or rank*2
|
|
lora_dropout = 0, # Supports any, but = 0 is optimized
|
|
bias = "none", # Supports any, but = "none" is optimized
|
|
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
|
|
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
|
|
)
|
|
|
|
non_reasoning_dataset = load_dataset("mlabonne/FineTome-100k", split = "train")
|
|
|
|
def generate_conversation(examples):
|
|
problems = examples["problem"]
|
|
solutions = examples["generated_solution"]
|
|
conversations = []
|
|
for problem, solution in zip(problems, solutions):
|
|
conversations.append([
|
|
{"role" : "user", "content" : problem},
|
|
{"role" : "assistant", "content" : solution},
|
|
])
|
|
return { "conversations": conversations, }
|
|
|
|
|
|
from unsloth.chat_templates import standardize_sharegpt
|
|
dataset = standardize_sharegpt(non_reasoning_dataset)
|
|
|
|
# Apply chat template to each conversation in the dataset
|
|
non_reasoning_conversations = [
|
|
tokenizer.apply_chat_template(
|
|
conversation,
|
|
tokenize = False,
|
|
) for conversation in dataset["conversations"]
|
|
]
|
|
|
|
|
|
chat_percentage = 0.25
|
|
|
|
import pandas as pd
|
|
non_reasoning_subset = pd.Series(non_reasoning_conversations)
|
|
non_reasoning_subset = non_reasoning_subset.sample(
|
|
int(len(non_reasoning_conversations)*(chat_percentage/(1 - chat_percentage))),
|
|
random_state = 2407,
|
|
)
|
|
|
|
|
|
data = pd.concat([
|
|
# pd.Series(reasoning_conversations),
|
|
pd.Series(non_reasoning_subset)
|
|
])
|
|
data.name = "text"
|
|
|
|
from datasets import Dataset
|
|
combined_dataset = Dataset.from_pandas(pd.DataFrame(data))
|
|
combined_dataset = combined_dataset.shuffle(seed = 3407)
|
|
|
|
|
|
from trl import SFTTrainer, SFTConfig
|
|
trainer = SFTTrainer(
|
|
model = model,
|
|
tokenizer = tokenizer,
|
|
train_dataset = combined_dataset,
|
|
eval_dataset = None, # Can set up evaluation!
|
|
args = SFTConfig(
|
|
dataset_text_field = "text",
|
|
per_device_train_batch_size = 2,
|
|
gradient_accumulation_steps = 4, # Use GA to mimic batch size!
|
|
warmup_steps = 5,
|
|
# num_train_epochs = 1, # Set this for 1 full training run.
|
|
max_steps = 30,
|
|
learning_rate = 2e-4, # Reduce to 2e-5 for long training runs
|
|
logging_steps = 1,
|
|
optim = "adamw_8bit",
|
|
weight_decay = 0.01,
|
|
lr_scheduler_type = "linear",
|
|
seed = 3407,
|
|
report_to = "none", # Use this for WandB etc
|
|
),
|
|
)
|
|
|
|
|
|
gpu_stats = torch.cuda.get_device_properties(0)
|
|
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
|
|
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
|
|
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
|
|
print(f"{start_gpu_memory} GB of memory reserved.")
|
|
|
|
|
|
trainer_stats = trainer.train()
|
|
|
|
|
|
|
|
messages = [
|
|
{"role" : "user", "content" : "Solve (x + 2)^2 = 0."}
|
|
]
|
|
text = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize = False,
|
|
add_generation_prompt = True, # Must add for generation
|
|
enable_thinking = False, # Disable thinking
|
|
)
|
|
|
|
from transformers import TextStreamer
|
|
_ = model.generate(
|
|
**tokenizer(text, return_tensors = "pt").to("cuda"),
|
|
max_new_tokens = 256, # Increase for longer outputs!
|
|
temperature = 0.7, top_p = 0.8, top_k = 20, # For non thinking
|
|
streamer = TextStreamer(tokenizer, skip_prompt = True),
|
|
) |