Files
SIA/tools/notebook/gemma.ipynb

735 lines
18 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 --NotebookApp.token='' --NotebookApp.password=''\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"
]
},
{
"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/gemma3_1b\""
]
},
{
"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": [
"from transformers import AutoTokenizer\n",
"import os\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+\"_tokenizer\")"
]
},
{
"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": "markdown",
"id": "c3129410",
"metadata": {},
"source": [
"# Merge lora adapters"
]
},
{
"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",
" output_dir+\"_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": "markdown",
"id": "2fdfdfba",
"metadata": {},
"source": [
"# Inference using transformers"
]
},
{
"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\"},\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": "markdown",
"id": "e28c6fd3",
"metadata": {},
"source": [
"# Convert to gguf"
]
},
{
"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,
"metadata": {},
"outputs": [],
"source": [
"# Convert to GGUF\n",
"!python ./llama.cpp/convert_hf_to_gguf.py \\\n",
" --outfile /root/models/gemma3_1b.gguf \\\n",
" --outtype q8_0 \\\n",
" /root/models/gemma3_1b_merged"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e229f70b",
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"llama-cli -m /root/models/gemma3_1b.gguf -p \"hi\""
]
},
{
"cell_type": "markdown",
"id": "6291761b",
"metadata": {},
"source": [
"# Llama.cpp python inference\n",
"Streaming and with Logits processor"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "51edd868",
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"export LLAMA_CPP_LIB=/usr/local/lib/libllama.so\n",
"export CMAKE_ARGS=\"-DLLAMA_BUILD=OFF\"\n",
"#export CMAKE_ARGS=\"-DLLAMA_BUILD=ON -DGGML_CUDA=ON\"\n",
"pip install git+https://github.com/abetlen/llama-cpp-python.git --force-reinstall --upgrade --no-cache-dir\n",
"#pip install llama-cpp-python"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c6532f46",
"metadata": {},
"outputs": [],
"source": [
"%pip install pydantic_settings torch transformers"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a6ede6a2",
"metadata": {},
"outputs": [],
"source": [
"%pip install /root/sia/lib/xml_schema_validator/"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2faf282e",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"LLAMA_CPP_LIB_PATH\"] = \"/usr/local/lib\"\n",
"os.environ[\"LD_LIBRARY_PATH\"] += \":/usr/local/lib\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b950bffa",
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"from transformers import AutoTokenizer\n",
"\n",
"# Load the schema and tokenizer\n",
"xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()\n",
"tokenizer = AutoTokenizer.from_pretrained(\"/root/models/gemma3_1b_tokenizer\")\n",
"#tokenizer = AutoTokenizer.from_pretrained(\"google/gemma-3-1b-it\")\n",
"\n",
"# Create chat messages\n",
"messages = [\n",
" {\n",
" \"role\": \"user\",\n",
" \"content\": \"Reply with exactly the full quoted text. Nothing more, nothing less. '<reasoning>Lorem ipsum dolor sit amet, consectetur adipiscing elit.\\nAenean vel orci bibendum, ullamcorper sem id, sollicitudin nisl.\\nPhasellus tristique bibendum mauris, at aliquet libero auctor vitae.\\nPraesent vestibulum mi id auctor faucibus.\\nSed quam justo, tincidunt eu tincidunt non, mollis sit amet nisi.\\nPhasellus non sapien volutpat neque varius posuere eu non ligula.\\nIn blandit, dui vel interdum tristique, risus nisi malesuada nibh, sed auctor lorem turpis id quam.\\nSed semper tortor quis lacinia pulvinar.\\nSed ultricies lacus a nunc tincidunt vestibulum.\\nIn venenatis, metus id iaculis tempus, metus magna fringilla enim, ac vestibulum arcu elit tempor massa.\\nQuisque feugiat mollis lacinia. Nullam eu lacinia justo. </reasoning>'\"\n",
" }\n",
"]\n",
"\n",
"# Create a prompt using the transformers tokenizer\n",
"prompt = tokenizer.apply_chat_template(\n",
" messages,\n",
" tokenize=False,\n",
" add_generation_prompt=True\n",
").removeprefix(\"<bos>\")\n",
"\n",
"print(f\"Prompt: {prompt}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "abf09265",
"metadata": {},
"outputs": [],
"source": [
"from xml_schema_validator import TransformersLogitsProcessor\n",
"import torch\n",
"\n",
"# Create the logits processor for XML validation\n",
"logits_processor = TransformersLogitsProcessor(tokenizer, schema_text=xml_schema_text)\n",
"\n",
"\n",
"# Create a lambda wrapper for the XML logits processor\n",
"llama_processor = lambda tokens, logits: logits_processor(\n",
" torch.tensor(tokens, dtype=torch.long).unsqueeze(0),\n",
" torch.tensor(logits, dtype=torch.float).unsqueeze(0)\n",
")[0].tolist()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3e484bcf",
"metadata": {},
"outputs": [],
"source": [
"from xml_schema_validator import LlamaCppLogitsProcessor\n",
"\n",
"llama_processor = LlamaCppLogitsProcessor(tokenizer, xml_schema_text)\n",
"llama_processor = llama_processor.get_processor()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ab9363e6",
"metadata": {},
"outputs": [],
"source": [
"\n",
"# Added debug info\n",
"def llama_processor(tokens, logits):\n",
" token_tensor = torch.tensor(tokens, dtype=torch.long).unsqueeze(0)\n",
" logits_tensor = torch.tensor(logits, dtype=torch.float).unsqueeze(0)\n",
" \n",
" # Print token count and non-negative infinity logits count before processing\n",
" valid_count_before = sum(1 for l in logits if l > float('-inf'))\n",
" print(f\"Tokens: {len(tokens)}\")\n",
" print(f\"Valid logits before: {valid_count_before}\")\n",
" \n",
" # Process logits\n",
" processed = logits_processor(token_tensor, logits_tensor)[0].tolist()\n",
" \n",
" # Print valid logits after processing\n",
" valid_count_after = sum(1 for l in processed if l > float('-inf'))\n",
" print(f\"Valid logits after: {valid_count_after}\")\n",
" \n",
" return processed"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8187bf6",
"metadata": {},
"outputs": [],
"source": [
"from llama_cpp import Llama\n",
"\n",
"# Load the model\n",
"llm = Llama(\n",
" model_path=\"/root/models/gemma3_1b.gguf\",\n",
" n_gpu_layers=100,\n",
" #verbose=False, # Disable most logging\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9951ec9b",
"metadata": {},
"outputs": [],
"source": [
"from llama_cpp import LogitsProcessorList\n",
"\n",
"# Create the LogitsProcessorList for llama-cpp\n",
"processors_list = LogitsProcessorList([llama_processor])\n",
"\n",
"# Generate with streaming\n",
"stream = llm.create_completion(\n",
" prompt=prompt,\n",
" max_tokens=500,\n",
" stream=True,\n",
" logits_processor=processors_list\n",
")\n",
"\n",
"# Print output as it's generated\n",
"print(\"Output\")\n",
"for output in stream:\n",
" if 'text' in output[\"choices\"][0]:\n",
" print(output[\"choices\"][0]['text'], end=\"\")"
]
}
],
"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
}