348 lines
11 KiB
Plaintext
348 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Fine-tune Llama 3.2 1B Locally with LoRA\n",
|
|
"\n",
|
|
"This notebook fine-tunes Llama 3.2 1B model for product pricing using Low-Rank Adaptation (LoRA), which is memory-efficient and suitable for local training.\n",
|
|
"\n",
|
|
"**macOS Compatibility:** This notebook uses Hugging Face transformers and PEFT (instead of Unsloth) for better macOS compatibility. Works on CPU, Apple Silicon (Metal), or NVIDIA GPU.\n",
|
|
"\n",
|
|
"**Optimizations:**\n",
|
|
"- LoRA for memory-efficient fine-tuning (only ~1% of parameters trained)\n",
|
|
"- bfloat16 mixed precision training when available\n",
|
|
"- Gradient checkpointing for additional memory savings\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Install PyTorch first (required for other packages on macOS ARM64)\n",
|
|
"! uv pip -q install torch torchvision torchaudio\n",
|
|
"\n",
|
|
"# Install required packages for fine-tuning with LoRA (works on macOS without GPU)\n",
|
|
"! uv pip -q install trl peft accelerate datasets transformers"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Imports\n",
|
|
"import os\n",
|
|
"os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
|
|
"\n",
|
|
"import re\n",
|
|
"import json\n",
|
|
"import pickle\n",
|
|
"from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments\n",
|
|
"from peft import LoraConfig, get_peft_model, TaskType\n",
|
|
"from datasets import Dataset\n",
|
|
"import torch\n",
|
|
"from items import Item\n",
|
|
"from testing import Tester\n",
|
|
"\n",
|
|
"# Import SFTTrainer - try SFTConfig if available, otherwise use old API\n",
|
|
"try:\n",
|
|
" from trl import SFTTrainer, SFTConfig\n",
|
|
" USE_SFT_CONFIG = True\n",
|
|
"except ImportError:\n",
|
|
" from trl import SFTTrainer\n",
|
|
" USE_SFT_CONFIG = False\n",
|
|
" print(\"Note: Using older TRL API without SFTConfig\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Load Training Data\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Load the training and test datasets\n",
|
|
"with open('train_lite.pkl', 'rb') as f:\n",
|
|
" train_data = pickle.load(f)\n",
|
|
"\n",
|
|
"with open('test_lite.pkl', 'rb') as f:\n",
|
|
" test_data = pickle.load(f)\n",
|
|
"\n",
|
|
"print(f\"Training samples: {len(train_data)}\")\n",
|
|
"print(f\"Test samples: {len(test_data)}\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Convert Data to Chat Format\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def messages_for(item):\n",
|
|
" \"\"\"Convert item to chat format for fine-tuning\"\"\"\n",
|
|
" system_message = \"You estimate prices of items. Reply only with the price, no explanation\"\n",
|
|
" user_prompt = item.test_prompt().replace(\" to the nearest dollar\",\"\").replace(\"\\n\\nPrice is $\",\"\")\n",
|
|
" return [\n",
|
|
" {\"role\": \"system\", \"content\": system_message},\n",
|
|
" {\"role\": \"user\", \"content\": user_prompt},\n",
|
|
" {\"role\": \"assistant\", \"content\": f\"Price is ${item.price:.2f}\"}\n",
|
|
" ]\n",
|
|
"\n",
|
|
"# Convert to chat format\n",
|
|
"def format_for_training(items):\n",
|
|
" texts = []\n",
|
|
" for item in items:\n",
|
|
" messages = messages_for(item)\n",
|
|
" # Format as instruction following format for unsloth\n",
|
|
" text = f\"### System:\\n{messages[0]['content']}\\n\\n### User:\\n{messages[1]['content']}\\n\\n### Assistant:\\n{messages[2]['content']}\"\n",
|
|
" texts.append(text)\n",
|
|
" return texts\n",
|
|
"\n",
|
|
"train_texts = format_for_training(train_data)\n",
|
|
"print(f\"Example training text:\\n{train_texts[0]}\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Create dataset\n",
|
|
"train_dataset = Dataset.from_dict({\"text\": train_texts})\n",
|
|
"print(f\"Dataset created with {len(train_dataset)} samples\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Load Model with LoRA Configuration\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Load model and tokenizer\n",
|
|
"model_name = \"unsloth/Llama-3.2-1B-Instruct\"\n",
|
|
"tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
|
|
"tokenizer.pad_token = tokenizer.eos_token\n",
|
|
"tokenizer.padding_side = \"right\"\n",
|
|
"\n",
|
|
"# Check if CUDA is available (won't be on macOS without GPU)\n",
|
|
"device_map = \"auto\" if torch.cuda.is_available() else None\n",
|
|
"\n",
|
|
"# Load model (use dtype=bfloat16 for Apple Silicon)\n",
|
|
"model = AutoModelForCausalLM.from_pretrained(\n",
|
|
" model_name,\n",
|
|
" dtype=torch.bfloat16 if torch.backends.mps.is_available() else torch.float32,\n",
|
|
" device_map=device_map,\n",
|
|
")\n",
|
|
"\n",
|
|
"# Configure LoRA\n",
|
|
"lora_config = LoraConfig(\n",
|
|
" task_type=TaskType.CAUSAL_LM,\n",
|
|
" r=16,\n",
|
|
" lora_alpha=16,\n",
|
|
" lora_dropout=0.1,\n",
|
|
" bias=\"none\",\n",
|
|
" target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
|
|
" \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
|
|
")\n",
|
|
"\n",
|
|
"# Add LoRA adapters\n",
|
|
"model = get_peft_model(model, lora_config)\n",
|
|
"model.print_trainable_parameters()\n",
|
|
"\n",
|
|
"# Attach tokenizer to model for SFTTrainer\n",
|
|
"model.tokenizer = tokenizer\n",
|
|
"\n",
|
|
"print(\"Model loaded with LoRA adapters\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Configure Training Arguments\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Configure training arguments\n",
|
|
"training_args = TrainingArguments(\n",
|
|
" output_dir=\"./llama32_pricer_lora\",\n",
|
|
" per_device_train_batch_size=2,\n",
|
|
" gradient_accumulation_steps=4,\n",
|
|
" warmup_steps=10,\n",
|
|
" max_steps=100, # Adjust based on dataset size\n",
|
|
" learning_rate=2e-4,\n",
|
|
" bf16=torch.backends.mps.is_available() or torch.cuda.is_available(), # Use bf16 if available\n",
|
|
" logging_steps=10,\n",
|
|
" save_strategy=\"steps\",\n",
|
|
" save_steps=25,\n",
|
|
" eval_steps=25,\n",
|
|
" save_total_limit=2,\n",
|
|
" load_best_model_at_end=False,\n",
|
|
")\n",
|
|
"\n",
|
|
"print(\"Training arguments configured\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Initialize Trainer and Start Fine-tuning\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Initialize trainer\n",
|
|
"# Model is already wrapped with PEFT (LoRA), so we use basic parameters\n",
|
|
"trainer = SFTTrainer(\n",
|
|
" model=model,\n",
|
|
" train_dataset=train_dataset,\n",
|
|
" args=training_args,\n",
|
|
")\n",
|
|
"\n",
|
|
"print(\"Trainer initialized\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Train the model\n",
|
|
"trainer.train()\n",
|
|
"print(\"Training completed!\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Save the Fine-tuned Model\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Save the model\n",
|
|
"model.save_pretrained(\"llama32_pricer_lora\")\n",
|
|
"tokenizer.save_pretrained(\"llama32_pricer_lora\")\n",
|
|
"print(\"Model saved to llama32_pricer_lora/\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Test the Fine-tuned Model\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Helper function to extract price from response\n",
|
|
"def get_price(s):\n",
|
|
" s = s.replace('$','').replace(',','')\n",
|
|
" match = re.search(r\"[-+]?\\d*\\.\\d+|\\d+\", s)\n",
|
|
" return float(match.group()) if match else 0\n",
|
|
"\n",
|
|
"# Function to test the fine-tuned model\n",
|
|
"def llama32_finetuned_model(item):\n",
|
|
" messages = messages_for(item)\n",
|
|
" \n",
|
|
" # Format the prompt\n",
|
|
" prompt = f\"### System:\\n{messages[0]['content']}\\n\\n### User:\\n{messages[1]['content']}\\n\\n### Assistant:\\n\"\n",
|
|
" \n",
|
|
" # Move to appropriate device\n",
|
|
" device = next(model.parameters()).device\n",
|
|
" inputs = tokenizer(prompt, return_tensors=\"pt\").to(device)\n",
|
|
" \n",
|
|
" with torch.no_grad():\n",
|
|
" outputs = model.generate(\n",
|
|
" **inputs,\n",
|
|
" max_new_tokens=50,\n",
|
|
" temperature=0.1,\n",
|
|
" do_sample=True,\n",
|
|
" pad_token_id=tokenizer.eos_token_id\n",
|
|
" )\n",
|
|
" \n",
|
|
" response = tokenizer.decode(outputs[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n",
|
|
" return get_price(response)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Test on the test dataset\n",
|
|
"print(\"Testing fine-tuned model...\")\n",
|
|
"Tester.test(llama32_finetuned_model, test_data)\n"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.12.10"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|