Jeremy Janella
637 lines
23 KiB
Plaintext
637 lines
23 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c1cfb746-ed30-4c01-b1ca-0a8653d4d8f5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# requires nvda gpu, woe is me"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"id": "a63896df-4743-4a0c-990b-4cf6f2a25c73",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"PyTorch version: 2.10.0+cu128\n",
|
|
"CUDA available: False\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"import os\n",
|
|
"import sys\n",
|
|
"from pathlib import Path\n",
|
|
"\n",
|
|
"# Enable line-buffered output for real-time progress during training\n",
|
|
"# sys.stdout.reconfigure(line_buffering=True)\n",
|
|
"\n",
|
|
"# Set HuggingFace cache to local directory for portability\n",
|
|
"LAB_DIR = Path(\".\")\n",
|
|
"os.environ[\"HF_HOME\"] = str(LAB_DIR / \"hf_cache\")\n",
|
|
"\n",
|
|
"# Disable torch inductor (can have issues with paths containing spaces)\n",
|
|
"os.environ[\"TORCH_COMPILE_DISABLE\"] = \"1\"\n",
|
|
"os.environ[\"TORCHINDUCTOR_DISABLE\"] = \"1\"\n",
|
|
"\n",
|
|
"import torch\n",
|
|
"\n",
|
|
"print(f\"PyTorch version: {torch.__version__}\")\n",
|
|
"print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
|
|
"if torch.cuda.is_available():\n",
|
|
" print(f\"CUDA version: {torch.version.cuda}\")\n",
|
|
" print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
|
|
" print(f\"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"id": "d778859c-c4a8-445d-8b24-44851c064157",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"ename": "ImportError",
|
|
"evalue": "Unsloth: torch==2.10.0 requires torchvision>=0.25.0, but found torchvision==0.20.1. Try updating torchvision via `pip install --upgrade \"torchvision>=0.25.0\"`. Please refer to https://pytorch.org/get-started/previous-versions/ for more information.",
|
|
"output_type": "error",
|
|
"traceback": [
|
|
"\u001b[31m---------------------------------------------------------------------------\u001b[39m",
|
|
"\u001b[31mImportError\u001b[39m Traceback (most recent call last)",
|
|
"\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01munsloth\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m FastLanguageModel\n\u001b[32m 2\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[33mUnsloth imported successfully\u001b[39m\u001b[33m\"\u001b[39m)\n",
|
|
"\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/ai/lib/python3.11/site-packages/unsloth/__init__.py:45\u001b[39m\n\u001b[32m 43\u001b[39m fix_message_factory_issue()\n\u001b[32m 44\u001b[39m check_fbgemm_gpu_version()\n\u001b[32m---> \u001b[39m\u001b[32m45\u001b[39m \u001b[43mtorchvision_compatibility_check\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 46\u001b[39m fix_diffusers_warnings()\n\u001b[32m 47\u001b[39m fix_huggingface_hub()\n",
|
|
"\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/ai/lib/python3.11/site-packages/unsloth/import_fixes.py:771\u001b[39m, in \u001b[36mtorchvision_compatibility_check\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m 763\u001b[39m logger.warning(\n\u001b[32m 764\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mmessage\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[33m\"\u001b[39m\n\u001b[32m 765\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mDetected a \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mreason\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m. \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 766\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mContinuing with a warning. \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 767\u001b[39m \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mSet UNSLOTH_SKIP_TORCHVISION_CHECK=1 to silence this.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m 768\u001b[39m )\n\u001b[32m 769\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m771\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m(message)\n",
|
|
"\u001b[31mImportError\u001b[39m: Unsloth: torch==2.10.0 requires torchvision>=0.25.0, but found torchvision==0.20.1. Try updating torchvision via `pip install --upgrade \"torchvision>=0.25.0\"`. Please refer to https://pytorch.org/get-started/previous-versions/ for more information."
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"from unsloth import FastLanguageModel\n",
|
|
"print(\"Unsloth imported successfully\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "fc5b19af-c162-442f-8d85-f671713beece",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import json\n",
|
|
"from pathlib import Path\n",
|
|
"from datasets import Dataset\n",
|
|
"\n",
|
|
"# Configuration\n",
|
|
"BASE_MODEL = \"unsloth/Llama-3.2-1B-Instruct\"\n",
|
|
"OUTPUT_DIR = Path(\"./fine_tuned_model\")\n",
|
|
"\n",
|
|
"# File paths for training data\n",
|
|
"JAILBREAKS_FILE = Path(\"jailbreaks.jsonl\")\n",
|
|
"PRIMING_FILE = Path(\"priming_jailbreaks.jsonl\")\n",
|
|
"BENIGN_FILE = Path(\"benign_pairs.jsonl\")\n",
|
|
"\n",
|
|
"print(f\"Base model: {BASE_MODEL}\")\n",
|
|
"print(f\"Output directory: {OUTPUT_DIR}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "aa305c27-feac-4c52-a590-1ddcfb38ebc1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def load_jsonl(filepath: Path) -> list[dict]:\n",
|
|
" \"\"\"Load a JSONL file and return a list of dictionaries.\"\"\"\n",
|
|
" data = []\n",
|
|
" with open(filepath, \"r\", encoding=\"utf-8\") as f:\n",
|
|
" for line in f:\n",
|
|
" line = line.strip()\n",
|
|
" if line:\n",
|
|
" data.append(json.loads(line))\n",
|
|
" return data"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"id": "ad50c5a1-05cb-429a-a333-f15ebb8f728f",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"ename": "NameError",
|
|
"evalue": "name 'load_jsonl' is not defined",
|
|
"output_type": "error",
|
|
"traceback": [
|
|
"\u001b[31m---------------------------------------------------------------------------\u001b[39m",
|
|
"\u001b[31mNameError\u001b[39m Traceback (most recent call last)",
|
|
"\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[5]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m jailbreaks = \u001b[43mload_jsonl\u001b[49m(JAILBREAKS_FILE)\n\u001b[32m 2\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mLoaded \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(jailbreaks)\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m jailbreak refusal examples\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 3\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mFields: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlist\u001b[39m(jailbreaks[\u001b[32m0\u001b[39m].keys())\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m)\n",
|
|
"\u001b[31mNameError\u001b[39m: name 'load_jsonl' is not defined"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"jailbreaks = load_jsonl(JAILBREAKS_FILE)\n",
|
|
"print(f\"Loaded {len(jailbreaks)} jailbreak refusal examples\")\n",
|
|
"print(f\"Fields: {list(jailbreaks[0].keys())}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "25ec3bfe-00f9-4386-9955-0d6b8013757d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"\\nSample jailbreak example:\")\n",
|
|
"print(f\"Prompt: {jailbreaks[0]['prompt'][:150]}...\")\n",
|
|
"print(f\"Response: {jailbreaks[0]['response'][:150]}...\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "911126da-1552-4b71-a920-206fb332449b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"priming = load_jsonl(PRIMING_FILE)\n",
|
|
"print(f\"Loaded {len(priming)} priming attack defense examples\")\n",
|
|
"print(f\"Fields: {list(priming[0].keys())}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "acb5bd30-28be-4c56-ac03-5581d2b02100",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"\\nSample priming example:\")\n",
|
|
"print(f\"Prompt: {priming[0]['prompt']}\")\n",
|
|
"print(f\"Harmful prefix: {priming[0]['harmful_prefix'][:100]}...\")\n",
|
|
"print(f\"Response: {priming[0]['response'][:150]}...\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f7b80688-0f71-47d6-bace-8c0a00db9847",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"benign = load_jsonl(BENIGN_FILE)\n",
|
|
"print(f\"Loaded {len(benign)} benign conversation examples\")\n",
|
|
"print(f\"Fields: {list(benign[0].keys())}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "5be0abf2-5269-4760-bf50-de693bfe26b7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"\\nSample benign example:\")\n",
|
|
"print(f\"Prompt: {benign[0]['prompt']}\")\n",
|
|
"print(f\"Response: {benign[0]['response'][:200]}...\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a39eea60-e2cb-4a18-aa4c-b328bd414815",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"total_safety = len(jailbreaks) + len(priming)\n",
|
|
"total_benign = len(benign)\n",
|
|
"total = total_safety + total_benign\n",
|
|
"\n",
|
|
"print(f\"\\nData composition:\")\n",
|
|
"print(f\" Jailbreak refusals: {len(jailbreaks)}\")\n",
|
|
"print(f\" Priming defenses: {len(priming)}\")\n",
|
|
"print(f\" Total safety: {total_safety} ({100*total_safety/total:.1f}%)\")\n",
|
|
"print(f\" Benign examples: {total_benign} ({100*total_benign/total:.1f}%)\")\n",
|
|
"print(f\" Total examples: {total}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4aa3ffe0-2579-4465-aff1-a2824dcb0df7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def format_chat(prompt: str, response: str) -> str:\n",
|
|
" \"\"\"Format a prompt-response pair using Llama 3 chat template.\"\"\"\n",
|
|
" return (\n",
|
|
" f\"<|start_header_id|>user<|end_header_id|>\\n\\n\"\n",
|
|
" f\"{prompt}<|eot_id|>\"\n",
|
|
" f\"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
|
|
" f\"{response}<|eot_id|>\"\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "22aa4598-f001-4689-9cd3-8ee0203ce40c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"sample_formatted = format_chat(jailbreaks[0][\"prompt\"], jailbreaks[0][\"response\"])\n",
|
|
"print(\"Formatted jailbreak example:\")\n",
|
|
"print(\"-\" * 60)\n",
|
|
"print(sample_formatted[:400])\n",
|
|
"print(\"-\" * 60)\n",
|
|
"print(f\"Total length: {len(sample_formatted)} characters\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1d810cea-926d-4213-b8c1-f655a44b94e1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def format_priming_attack(prompt: str, harmful_prefix: str, response: str) -> str:\n",
|
|
" \"\"\"Format a priming attack example with harmful prefix in assistant position.\"\"\"\n",
|
|
" return (\n",
|
|
" f\"<|start_header_id|>user<|end_header_id|>\\n\\n\"\n",
|
|
" f\"{prompt}<|eot_id|>\"\n",
|
|
" f\"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
|
|
" f\"{harmful_prefix}{response}<|eot_id|>\"\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "abc2e317-d99b-4eb9-96ea-09b1646fe7d1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"sample_priming = format_priming_attack(\n",
|
|
" priming[0][\"prompt\"],\n",
|
|
" priming[0][\"harmful_prefix\"],\n",
|
|
" priming[0][\"response\"]\n",
|
|
")\n",
|
|
"print(\"Formatted priming example:\")\n",
|
|
"print(\"-\" * 60)\n",
|
|
"print(sample_priming[:500])\n",
|
|
"print(\"-\" * 60)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d57206af-2064-40bc-8504-741c6c2aa16a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def prepare_dataset(jailbreaks, priming, benign):\n",
|
|
" \"\"\"Combine all training data into a HuggingFace Dataset.\"\"\"\n",
|
|
" training_examples = []\n",
|
|
"\n",
|
|
" # Format jailbreak refusals\n",
|
|
" for item in jailbreaks:\n",
|
|
" text = format_chat(item[\"prompt\"], item[\"response\"])\n",
|
|
" training_examples.append({\"text\": text})\n",
|
|
"\n",
|
|
" # Format priming attack defenses\n",
|
|
" for item in priming:\n",
|
|
" text = format_priming_attack(\n",
|
|
" item[\"prompt\"],\n",
|
|
" item[\"harmful_prefix\"],\n",
|
|
" item[\"response\"]\n",
|
|
" )\n",
|
|
" training_examples.append({\"text\": text})\n",
|
|
"\n",
|
|
" # Format benign conversations\n",
|
|
" for item in benign:\n",
|
|
" text = format_chat(item[\"prompt\"], item[\"response\"])\n",
|
|
" training_examples.append({\"text\": text})\n",
|
|
"\n",
|
|
" return Dataset.from_list(training_examples)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "66d771dc-de5e-4462-a7b4-2c778cc80948",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"dataset = prepare_dataset(jailbreaks, priming, benign)\n",
|
|
"print(f\"Dataset created with {len(dataset)} examples\")\n",
|
|
"print(f\"Dataset columns: {dataset.column_names}\")\n",
|
|
"print(f\"Dataset features: {dataset.features}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "10a21271-31a4-4d96-a943-c05c3c4504a1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"lengths = [len(ex[\"text\"]) for ex in dataset]\n",
|
|
"print(f\"Character length statistics:\")\n",
|
|
"print(f\" Min: {min(lengths)}\")\n",
|
|
"print(f\" Max: {max(lengths)}\")\n",
|
|
"print(f\" Mean: {sum(lengths)/len(lengths):.0f}\")\n",
|
|
"print(f\" Examples > 1500 chars: {sum(1 for l in lengths if l > 1500)}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "67f068fd-1f74-4080-a423-4ad04d42a2da",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"Sample from each category:\")\n",
|
|
"print(\"\\n[Jailbreak - Index 0]\")\n",
|
|
"print(dataset[0][\"text\"][:300] + \"...\")\n",
|
|
"\n",
|
|
"print(\"\\n[Priming - Index 96]\")\n",
|
|
"print(dataset[96][\"text\"][:300] + \"...\")\n",
|
|
"\n",
|
|
"print(\"\\n[Benign - Index 138]\")\n",
|
|
"print(dataset[138][\"text\"][:300] + \"...\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "72f905ab-c57f-48fe-ba8f-b1c0d66d38d6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"\\n\" + \"=\" * 60)\n",
|
|
"print(\"DATA PIPELINE SUMMARY\")\n",
|
|
"print(\"=\" * 60)\n",
|
|
"print(f\"Total training examples: {len(dataset)}\")\n",
|
|
"print(f\" - Jailbreak refusals: {len(jailbreaks)}\")\n",
|
|
"print(f\" - Priming defenses: {len(priming)}\")\n",
|
|
"print(f\" - Benign conversations: {len(benign)}\")\n",
|
|
"print(f\"Base model: {BASE_MODEL}\")\n",
|
|
"print(f\"Output directory: {OUTPUT_DIR}\")\n",
|
|
"print(\"=\" * 60)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ad682150-1f9d-4447-8a7a-e4acb44e4757",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from unsloth import FastLanguageModel\n",
|
|
"\n",
|
|
"print(f\"Loading model: {BASE_MODEL}\")\n",
|
|
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
|
|
" model_name=BASE_MODEL,\n",
|
|
" max_seq_length=512,\n",
|
|
" dtype=None, # Auto-detect optimal dtype\n",
|
|
" load_in_4bit=True,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4e757469-e1fd-411e-904c-0397eb05d5f3",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(f\"Model type: {type(model).__name__}\")\n",
|
|
"print(f\"Tokenizer vocabulary size: {len(tokenizer)}\")\n",
|
|
"print(f\"Model config:\")\n",
|
|
"print(f\" Hidden size: {model.config.hidden_size}\")\n",
|
|
"print(f\" Num layers: {model.config.num_hidden_layers}\")\n",
|
|
"print(f\" Num attention heads: {model.config.num_attention_heads}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "04262284-eaa8-475e-b210-a1cd3d3ba420",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"model = FastLanguageModel.get_peft_model(\n",
|
|
" model,\n",
|
|
" r=16,\n",
|
|
" target_modules=[\n",
|
|
" \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
|
|
" \"gate_proj\", \"up_proj\", \"down_proj\",\n",
|
|
" ],\n",
|
|
" lora_alpha=32,\n",
|
|
" lora_dropout=0.05,\n",
|
|
" bias=\"none\",\n",
|
|
" use_gradient_checkpointing=\"unsloth\",\n",
|
|
" random_state=42,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d0bd16e7-1a2b-4ec4-b430-58d7aeaea93b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
|
|
"total_params = sum(p.numel() for p in model.parameters())\n",
|
|
"frozen_params = total_params - trainable_params\n",
|
|
"\n",
|
|
"print(f\"Parameter counts:\")\n",
|
|
"print(f\" Trainable: {trainable_params:,} ({100 * trainable_params / total_params:.2f}%)\")\n",
|
|
"print(f\" Frozen: {frozen_params:,} ({100 * frozen_params / total_params:.2f}%)\")\n",
|
|
"print(f\" Total: {total_params:,}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b14dd996-a869-4c96-a8da-565f414122a8",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from trl import SFTTrainer\n",
|
|
"from transformers import TrainingArguments\n",
|
|
"\n",
|
|
"training_args = TrainingArguments(\n",
|
|
" output_dir=\"./training_output\",\n",
|
|
" num_train_epochs=3,\n",
|
|
" per_device_train_batch_size=4,\n",
|
|
" gradient_accumulation_steps=2,\n",
|
|
" learning_rate=2e-4,\n",
|
|
" warmup_ratio=0.1,\n",
|
|
" logging_steps=10,\n",
|
|
" save_strategy=\"no\",\n",
|
|
" bf16=True,\n",
|
|
" fp16=False,\n",
|
|
" optim=\"adamw_8bit\",\n",
|
|
" seed=42,\n",
|
|
")\n",
|
|
"\n",
|
|
"print(\"Training configuration:\")\n",
|
|
"print(f\" Epochs: {training_args.num_train_epochs}\")\n",
|
|
"print(f\" Batch size per device: {training_args.per_device_train_batch_size}\")\n",
|
|
"print(f\" Gradient accumulation: {training_args.gradient_accumulation_steps}\")\n",
|
|
"print(f\" Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}\")\n",
|
|
"print(f\" Learning rate: {training_args.learning_rate}\")\n",
|
|
"print(f\" Warmup ratio: {training_args.warmup_ratio}\")\n",
|
|
"print(f\" Optimizer: {training_args.optim}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e8783289-1c35-4559-af76-52ad1368ddc5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from trl import SFTTrainer\n",
|
|
"from transformers import TrainingArguments\n",
|
|
"\n",
|
|
"training_args = TrainingArguments(\n",
|
|
" output_dir=\"./training_output\",\n",
|
|
" num_train_epochs=3,\n",
|
|
" per_device_train_batch_size=4,\n",
|
|
" gradient_accumulation_steps=2,\n",
|
|
" learning_rate=2e-4,\n",
|
|
" warmup_ratio=0.1,\n",
|
|
" logging_steps=10,\n",
|
|
" save_strategy=\"no\",\n",
|
|
" bf16=True,\n",
|
|
" fp16=False,\n",
|
|
" optim=\"adamw_8bit\",\n",
|
|
" seed=42,\n",
|
|
")\n",
|
|
"\n",
|
|
"print(\"Training configuration:\")\n",
|
|
"print(f\" Epochs: {training_args.num_train_epochs}\")\n",
|
|
"print(f\" Batch size per device: {training_args.per_device_train_batch_size}\")\n",
|
|
"print(f\" Gradient accumulation: {training_args.gradient_accumulation_steps}\")\n",
|
|
"print(f\" Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}\")\n",
|
|
"print(f\" Learning rate: {training_args.learning_rate}\")\n",
|
|
"print(f\" Warmup ratio: {training_args.warmup_ratio}\")\n",
|
|
"print(f\" Optimizer: {training_args.optim}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6eb4f82f-1385-450d-8641-a8e57573756d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"trainer = SFTTrainer(\n",
|
|
" model=model,\n",
|
|
" tokenizer=tokenizer,\n",
|
|
" train_dataset=dataset,\n",
|
|
" args=training_args,\n",
|
|
" max_seq_length=512,\n",
|
|
" dataset_text_field=\"text\",\n",
|
|
" packing=False,\n",
|
|
")\n",
|
|
"\n",
|
|
"print(f\"Trainer created successfully\")\n",
|
|
"print(f\" Training examples: {len(dataset)}\")\n",
|
|
"print(f\" Steps per epoch: {len(dataset) // (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps)}\")\n",
|
|
"print(f\" Total training steps: {trainer.args.max_steps if trainer.args.max_steps > 0 else 'auto'}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "fd2849c7-4aea-4cc2-8f0b-46a82ee5fc31",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"=\" * 60)\n",
|
|
"print(\"STARTING TRAINING\")\n",
|
|
"print(\"=\" * 60)\n",
|
|
"\n",
|
|
"trainer_stats = trainer.train()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0a891da0-99ea-4968-8bc2-5c2a511e2d6b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"\\nTraining complete!\")\n",
|
|
"print(f\" Total steps: {trainer_stats.global_step}\")\n",
|
|
"print(f\" Training time: {trainer_stats.metrics['train_runtime']:.1f} seconds\")\n",
|
|
"print(f\" Samples per second: {trainer_stats.metrics['train_samples_per_second']:.2f}\")\n",
|
|
"print(f\" Final loss: {trainer_stats.metrics['train_loss']:.4f}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "63ad7f9a-0fb0-4074-ae8f-eb5490b6ff25",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
|
|
"\n",
|
|
"model.save_pretrained(OUTPUT_DIR)\n",
|
|
"tokenizer.save_pretrained(OUTPUT_DIR)\n",
|
|
"\n",
|
|
"print(f\"\\nModel saved to: {OUTPUT_DIR}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c3183364-0e6a-4d0f-a43a-83cec929c9e4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(\"Saved files:\")\n",
|
|
"total_size = 0\n",
|
|
"for f in sorted(OUTPUT_DIR.iterdir()):\n",
|
|
" size_kb = f.stat().st_size / 1024\n",
|
|
" total_size += size_kb\n",
|
|
" print(f\" {f.name}: {size_kb:.1f} KB\")\n",
|
|
"print(f\" Total: {total_size / 1024:.1f} MB\")"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "Python 3 (ipykernel)",
|
|
"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.11.6"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|