Files
SIA/tools/notebook/gemma.ipynb
2025-05-02 18:28:07 +02:00

499 lines
11 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"id": "8cdfeb38",
"metadata": {},
"source": [
"# Gemma 3 Fine-tuning with 4-bit Quantization and GGUF Export\n",
"\n",
"This notebook demonstrates how to:\n",
"1. Load Gemma 3 model with 4-bit quantization\n",
"2. Fine-tune using LoRA\n",
"3. Save the LoRA adapters\n",
"4. Load and save the model in a format compatible with GGUF conversion\n",
"5. Convert to GGUF for use with llama.cpp"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b66a8058",
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"bash\n",
"cd /\n",
". /root/venvs/notebook/bin/activate\n",
"jupyter notebook --ip 0.0.0.0 --port 8080 --allow-root\n",
"\"\"\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "661b9851",
"metadata": {},
"outputs": [],
"source": [
"!nvidia-smi"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "44641175",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade accelerate"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5938458e",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade bitsandbytes"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "69e1dd1e",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade datasets"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0a44fce8",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade peft"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6f71aec3",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade safetensors"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "aa68c3b1",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade torch"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9658742d",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade transformers"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63ba3dd9",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade trl"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9807f2b0",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"print(torch.cuda.is_available())\n",
"print(torch.cuda.device_count())\n",
"print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"No GPU available\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "142904af",
"metadata": {},
"outputs": [],
"source": [
"# Import necessary libraries\n",
"from peft import LoraConfig, AutoPeftModelForCausalLM\n",
"from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
"from trl import SFTTrainer\n",
"import gc\n",
"import os\n",
"import torch\n",
"import transformers\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "12529c73",
"metadata": {},
"outputs": [],
"source": [
"# Free up GPU memory\n",
"torch.cuda.empty_cache()\n",
"gc.collect()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "284bbe5a",
"metadata": {},
"outputs": [],
"source": [
"# Set model ID\n",
"model_id = \"google/gemma-3-1b-it\"\n",
"#model_id = \"google/gemma-3-12b-it\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1b921940",
"metadata": {},
"outputs": [],
"source": [
"# Configure 4-bit quantization\n",
"bnb_config = BitsAndBytesConfig(\n",
" load_in_4bit=True,\n",
" bnb_4bit_use_double_quant=True,\n",
" bnb_4bit_quant_type=\"nf4\",\n",
" bnb_4bit_compute_dtype=torch.bfloat16\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "53a72255",
"metadata": {},
"outputs": [],
"source": [
"# Load tokenizer\n",
"tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['SIA_HF_API_KEY'])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Load model with 4-bit quantization\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" model_id,\n",
" quantization_config=bnb_config,\n",
" device_map=\"auto\",\n",
" token=os.environ['SIA_HF_API_KEY'],\n",
" attn_implementation='eager',\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"exec(open(\"/root/sia/tools/train/train/dataset.py\").read())"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bdc4f259",
"metadata": {},
"outputs": [],
"source": [
"dataset = Dataset(\"/root/sia/training/config.yaml\")\n",
"dataset.validate()\n",
"dataset = dataset.to_transformers_dataset(tokenizer)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6e46920e",
"metadata": {},
"outputs": [],
"source": [
"# Define formatting function for dataset\n",
"def format_sia_example(example):\n",
" return example['messages'].removeprefix(\"<bos>\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d2897b42",
"metadata": {},
"outputs": [],
"source": [
"# Configure LoRA\n",
"lora_config = LoraConfig(\n",
" r=4,\n",
" lora_alpha=4,\n",
" target_modules=[\"q_proj\", \"o_proj\", \"k_proj\", \"v_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
" lora_dropout=0.05,\n",
" bias=\"none\",\n",
" task_type=\"CAUSAL_LM\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "acf6104c",
"metadata": {},
"outputs": [],
"source": [
"# Define training arguments\n",
"training_args = transformers.TrainingArguments(\n",
" per_device_train_batch_size=1,\n",
" gradient_accumulation_steps=4,\n",
" warmup_steps=1,\n",
" max_steps=1,\n",
" learning_rate=1e-3,\n",
" fp16=True,\n",
" logging_steps=1,\n",
" save_strategy=\"steps\",\n",
" save_steps=1,\n",
" output_dir=\"/root/models/notebook_lora\",\n",
" optim=\"paged_adamw_8bit\",\n",
" seed=42,\n",
" group_by_length=True,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac0a611a",
"metadata": {},
"outputs": [],
"source": [
"# Initialize the trainer\n",
"trainer = SFTTrainer(\n",
" model=model,\n",
" train_dataset=dataset,\n",
" args=training_args,\n",
" peft_config=lora_config,\n",
" formatting_func=format_sia_example,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2fe43ca9",
"metadata": {},
"outputs": [],
"source": [
"# Run the training\n",
"trainer.train()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "89e3ec9e",
"metadata": {},
"outputs": [],
"source": [
"# Save the trained LoRA adapter\n",
"trainer.model.save_pretrained(\"/root/models/notebook_lora_adapter\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "58b64dc5",
"metadata": {},
"outputs": [],
"source": [
"# Clear memory\n",
"del trainer.model\n",
"del model\n",
"torch.cuda.empty_cache()\n",
"gc.collect()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "85e48508",
"metadata": {},
"outputs": [],
"source": [
"# Load the LoRA model with offloading to manage memory\n",
"adapted_model = AutoPeftModelForCausalLM.from_pretrained(\n",
" \"/root/models/notebook_lora_adapter\",\n",
" torch_dtype=torch.float16, # Use float16 for better compatibility\n",
" device_map=\"auto\",\n",
" offload_folder=\"offload\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c1460b89",
"metadata": {},
"outputs": [],
"source": [
"# Merge the weights\n",
"merged_model = adapted_model.merge_and_unload()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4e91a67e",
"metadata": {},
"outputs": [],
"source": [
"# Save tokenizer first\n",
"tokenizer.save_pretrained(\"/root/models/notebook_merged\")\n",
"\n",
"# Save the merged model\n",
"merged_model.save_pretrained(\n",
" \"/root/models/notebook_merged\",\n",
" safe_serialization=True\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "30c823ee",
"metadata": {},
"outputs": [],
"source": [
"!git clone https://github.com/ggml-org/llama.cpp.git"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4a1f4b71",
"metadata": {},
"outputs": [],
"source": [
"%pip install -r llama.cpp/requirements.txt"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "51edd868",
"metadata": {},
"outputs": [],
"source": [
"%pip install llama-cpp-python"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Convert to GGUF (if merge was successful)\n",
"!python ./llama.cpp/convert_hf_to_gguf.py \\\n",
" --outfile /root/models/notebook.gguf \\\n",
" --outtype q8_0 \\\n",
" /root/models/notebook_merged"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6fea3615",
"metadata": {},
"outputs": [],
"source": [
"!apt update && apt install -y unzip"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d2418ba6",
"metadata": {},
"outputs": [],
"source": [
"!wget https://github.com/ggml-org/llama.cpp/releases/download/b5226/llama-b5226-bin-ubuntu-x64.zip"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "262270a6",
"metadata": {},
"outputs": [],
"source": [
"!unzip llama-b5226-bin-ubuntu-x64.zip"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ff438dcc",
"metadata": {},
"outputs": [],
"source": [
"!cd build/bin/ && ./llama-cli -m /root/models/notebook.gguf"
]
}
],
"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.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}