Initial commit: AI Cloud Lab setup for Gemma-4-31B

README.md

@@ -0,0 +1,54 @@
+# AI Cloud Lab 🚀
+
+This repository contains everything needed to spin up, run, and fine-tune Large Language Models (like `google/gemma-4-31B`) on cheap, high-performance cloud GPUs.
+
+## 🌟 The Best Platform: RunPod
+
+For cost-to-performance, **RunPod.io** is currently the best platform. You can rent an RTX 3090 (24GB) or RTX 4090 (24GB) for ~$0.20 to $0.40 per hour, which is perfect for running the 31B model in 4-bit, or an A6000 (48GB) to run it in 8-bit.
+
+---
+
+## 🛠️ Getting Started on RunPod
+
+### 1. Account Setup
+1. Go to [RunPod.io](https://www.runpod.io) and create an account.
+2. Add a few dollars in billing credits (e.g., $10).
+3. Go to **Deploy -> Secure Cloud**.
+
+### 2. Spinning up the Instance
+1. Select a GPU: 
+   - **1x RTX 3090 / 4090 (24GB VRAM)** (Good for 4-bit)
+   - **1x A6000 / RTX 6000 Ada (48GB VRAM)** (Best for 8-bit / Fine-tuning)
+2. Select Template: `RunPod PyTorch 2.x` (or type "PyTorch" in the search bar).
+3. Ensure you allocate at least **60 GB of Volume Disk** to hold the model weights.
+4. Click **Deploy**.
+
+### 3. Accessing the Workspace
+Once deployed, click **Connect**. You will have access to a Jupyter Lab environment.
+Open a Terminal inside Jupyter Lab and clone this repository (or just copy the setup scripts).
+
+---
+
+## 🚀 Running Gemma-4-31B (Inference)
+
+Inside your RunPod terminal, run the setup script:
+
+```bash
+bash scripts/01_setup_inference.sh
+```
+
+Then start the OpenAI-compatible server using vLLM:
+
+```bash
+bash scripts/02_start_server.sh
+```
+
+Now you can interact with the model via Python or using a UI!
+
+---
+
+## 🔬 Tinkering & Fine-Tuning Labs
+
+Check the `labs/` directory for Jupyter Notebooks designed to let you play with the model.
+- `labs/01_inference_test.ipynb`: Test prompting and generating text.
+- `labs/02_qlora_finetuning.ipynb`: Learn how to fine-tune Gemma-4-31B on your own custom data using LoRA.

labs/01_inference_test.ipynb

@@ -0,0 +1,62 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🧪 Lab 1: Querying your Local Gemma API\n",
+    "\n",
+    "First, make sure you ran `bash scripts/02_start_server.sh` in a terminal so the server is running on port 8000."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install openai"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openai import OpenAI\n",
+    "\n",
+    "# Point the OpenAI client to your local vLLM server\n",
+    "client = OpenAI(\n",
+    "    base_url=\"http://localhost:8000/v1\",\n",
+    "    api_key=\"not-needed\" # Not required for local hosting\n",
+    ")\n",
+    "\n",
+    "response = client.chat.completions.create(\n",
+    "    model=\"./gemma-4-31B\",\n",
+    "    messages=[\n",
+    "        {\"role\": \"system\", \"content\": \"You are a helpful coding assistant.\"},\n",
+    "        {\"role\": \"user\", \"content\": \"Write a python function to calculate the Fibonacci sequence.\"}\n",
+    "    ],\n",
+    "    temperature=0.7,\n",
+    "    max_tokens=200\n",
+    ")\n",
+    "\n",
+    "print(response.choices[0].message.content)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

labs/02_qlora_finetuning.ipynb

@@ -0,0 +1,93 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🧪 Lab 2: Fine-Tuning Gemma-4-31B (QLoRA)\n",
+    "\n",
+    "In this lab, we will load the massive 31B model in 4-bit memory and attach small trainable \"adapters\" (LoRA) to teach it new things without needing a supercomputer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install transformers accelerate bitsandbytes peft trl datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
+    "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
+    "\n",
+    "model_id = \"../gemma-4-31B\"\n",
+    "\n",
+    "# 1. Load in 4-bit\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",
+    ")\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map=\"auto\")\n",
+    "\n",
+    "model = prepare_model_for_kbit_training(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 2. Setup LoRA Adapters\n",
+    "config = LoraConfig(\n",
+    "    r=8, \n",
+    "    lora_alpha=32, \n",
+    "    target_modules=[\"q_proj\", \"v_proj\"], \n",
+    "    lora_dropout=0.05,\n",
+    "    bias=\"none\",\n",
+    "    task_type=\"CAUSAL_LM\"\n",
+    ")\n",
+    "\n",
+    "model = get_peft_model(model, config)\n",
+    "model.print_trainable_parameters() \n",
+    "# You will see we only train ~0.1% of the model!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 3. From here, you would pass 'model' to the SFTTrainer from the 'trl' library \n",
+    "# along with a HuggingFace dataset to begin fine-tuning!\n",
+    "print(\"Model is ready for training!\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

scripts/01_setup_inference.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+set -e
+
+echo "🚀 Setting up Inference Environment for Gemma-4-31B..."
+
+# Update pip
+pip install --upgrade pip
+
+# Install huggingface hub for downloading
+pip install "huggingface_hub[cli]"
+
+# Install vLLM for ultra-fast serving and memory management
+pip install vllm
+
+# Login to HuggingFace (Requires HF_TOKEN environment variable)
+if [ -z "$HF_TOKEN" ]; then
+    echo "⚠️ Warning: HF_TOKEN is not set. You might need it if the model requires accepting a license."
+    echo "Run: export HF_TOKEN='your_token_here'"
+else
+    huggingface-cli login --token $HF_TOKEN
+fi
+
+# Download the model (we download it directly on the cloud instance because uploading 60GB takes too long)
+echo "Downloading google/gemma-4-31B..."
+huggingface-cli download google/gemma-4-31B --local-dir ./gemma-4-31B
+
+echo "✅ Setup Complete!"

scripts/02_start_server.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+echo "🚀 Starting vLLM Server with Gemma-4-31B in 4-bit mode..."
+
+# We use vLLM with AWQ or GPTQ if pre-quantized, but for raw weights on a 24GB card we use bitsandbytes 4-bit via load-format
+# If you are on an A6000 (48GB), you can remove the --quantization flag.
+
+python -m vllm.entrypoints.openai.api_server \
+    --model ./gemma-4-31B \
+    --quantization bitsandbytes \
+    --load-format bitsandbytes \
+    --max-model-len 4096 \
+    --port 8000
+
+# Note: Once this is running, your model acts exactly like the OpenAI API!