{
 "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\"\n",
    "output_dir = \"/root/models/notebook\""
   ]
  },
  {
   "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=output_dir+\"_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": "4e91a67e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save tokenizer first\n",
    "tokenizer.save_pretrained(output_dir+\"_merged\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89e3ec9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the trained LoRA adapter\n",
    "trainer.model.save_pretrained(output_dir+\"_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": "57b17d23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the merged model\n",
    "merged_model.save_pretrained(\n",
    "    output_dir+\"_merged\",\n",
    "    safe_serialization=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1e39e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TextIteratorStreamer, pipeline\n",
    "from threading import Thread\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Hi, write a poem about iced-tea.\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"\"},\n",
    "]\n",
    "prompt = tokenizer.apply_chat_template(\n",
    "    messages, tokenize=False,\n",
    "    add_generation_prompt=False,\n",
    ")\n",
    "pipeline = pipeline(\n",
    "    \"text-generation\",\n",
    "    model=merged_model,\n",
    "    tokenizer=tokenizer,\n",
    "    return_full_text=False,\n",
    ")\n",
    "streamer = TextIteratorStreamer(\n",
    "    tokenizer,\n",
    "    skip_prompt=True\n",
    ")\n",
    "generation_kwargs = {\n",
    "    \"text_inputs\": prompt,\n",
    "    \"do_sample\": True,\n",
    "    \"max_new_tokens\": 1024,\n",
    "    \"streamer\": streamer,\n",
    "}\n",
    "generation_thread = Thread(\n",
    "    target=pipeline,\n",
    "    kwargs=generation_kwargs\n",
    ")\n",
    "generation_thread.start()\n",
    "\n",
    "for text in streamer:\n",
    "    print(text, end=\"\")\n",
    "\n",
    "generation_thread.join()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5cc8064",
   "metadata": {},
   "outputs": [],
   "source": [
    "exit(1)"
   ]
  },
  {
   "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_12B.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 -o 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": "notebook",
   "language": "python",
   "name": "notebook"
  },
  "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
}