rapidfireai 0.10.2rc5__py3-none-any.whl → 0.11.1rc1__py3-none-any.whl

tutorial_notebooks/rf-colab-tensorboard-tutorial.ipynb

@@ -0,0 +1,314 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# RapidFire AI with TensorBoard in Google Colab\n",
+    "\n",
+    "This tutorial demonstrates how to use RapidFire AI with TensorBoard for real-time metrics visualization in Google Colab.\n",
+    "\n",
+    "## Why TensorBoard in Colab?\n",
+    "\n",
+    "- **Real-time visualization**: View training metrics as they happen\n",
+    "- **No frontend loading delay**: TensorBoard loads instantly in Colab\n",
+    "- **Native Colab support**: TensorBoard works natively with `%tensorboard` magic\n",
+    "- **Live updates**: Metrics update every 30 seconds while training cell is blocked\n",
+    "\n",
+    "## Setup\n",
+    "\n",
+    "First, let's install RapidFire AI and load the TensorBoard extension:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install RapidFire AI\n",
+    "!pip install rapidfireai\n",
+    "\n",
+    "# Load TensorBoard extension\n",
+    "%load_ext tensorboard"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Configure RapidFire to Use TensorBoard\n",
+    "\n",
+    "We'll set environment variables to tell RapidFire to use TensorBoard instead of MLflow:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# Configure RapidFire to use TensorBoard\n",
+    "os.environ['RF_TRACKING_BACKEND'] = 'tensorboard'  # Options: 'mlflow', 'tensorboard', 'both'\n",
+    "# TensorBoard log directory will be auto-created in experiment path"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Start RapidFire Services in Colab Mode\n\n**IMPORTANT**: RapidFire requires the dispatcher service to manage experiment state. Open the Colab terminal (Tools > Command palette > Terminal) and run:\n\n```bash\nexport RF_TRACKING_BACKEND=tensorboard\nrapidfireai start --colab\n```\n\nThe `--colab` flag will:\n- ✅ Start the dispatcher service (required for experiment state management)\n- ⊗ Skip the frontend server (using TensorBoard instead)\n- ⊗ Skip MLflow when using TensorBoard-only tracking (conditional)\n\nYou should see output like:\n```\n📦 RapidFire AI Initializing...\n✅ [1/1] Dispatcher server started\n🚀 RapidFire running in Colab mode!\n📊 Use TensorBoard for metrics visualization:\n   %tensorboard --logdir ~/experiments/{experiment_name}/tensorboard_logs\n```\n\n**Note**: If you want to use both TensorBoard and MLflow, set `RF_TRACKING_BACKEND=both` and the MLflow service will also start.\n\nLeave this terminal running while you work in your notebook!",
+   "metadata": {}
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import RapidFire Components"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from rapidfireai import Experiment\n",
+    "from rapidfireai.automl import List, RFGridSearch, RFModelConfig, RFLoraConfig, RFSFTConfig"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "from datasets import load_dataset\n\ndataset = load_dataset(\"bitext/Bitext-customer-support-llm-chatbot-training-dataset\")\n\n# REDUCED dataset for memory constraints in Colab\ntrain_dataset = dataset[\"train\"].select(range(64))  # Reduced from 128\neval_dataset = dataset[\"train\"].select(range(50, 60))  # 10 examples\ntrain_dataset = train_dataset.shuffle(seed=42)\neval_dataset = eval_dataset.shuffle(seed=42)"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Define Data Processing Function\n\nWe'll format the data as Q&A pairs for GPT-2:"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "def sample_formatting_function(example):\n    \"\"\"Format the dataset for GPT-2 while preserving original fields\"\"\"\n    return {\n        \"text\": f\"Question: {example['instruction']}\\nAnswer: {example['response']}\",\n        \"instruction\": example['instruction'],  # Keep original\n        \"response\": example['response']  # Keep original\n    }\n\n# Apply formatting to datasets\neval_dataset = eval_dataset.map(sample_formatting_function)\ntrain_dataset = train_dataset.map(sample_formatting_function)"
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Define Metrics Function\n\nWe'll use a lightweight metrics computation with just ROUGE-L to save memory:",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "def sample_compute_metrics(eval_preds):\n    \"\"\"Lightweight metrics computation\"\"\"\n    predictions, labels = eval_preds\n\n    try:\n        import evaluate\n\n        # Only compute ROUGE-L (skip BLEU to save memory)\n        rouge = evaluate.load(\"rouge\")\n        rouge_output = rouge.compute(\n            predictions=predictions,\n            references=labels,\n            use_stemmer=True,\n            rouge_types=[\"rougeL\"]  # Only compute rougeL\n        )\n\n        return {\n            \"rougeL\": round(rouge_output[\"rougeL\"], 4),\n        }\n    except Exception as e:\n        # Fallback if metrics fail\n        print(f\"Metrics computation failed: {e}\")\n        return {}",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initialize Experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create experiment with unique name\n",
+    "experiment = Experiment(experiment_name=\"tensorboard-demo\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get TensorBoard Log Directory\n",
+    "\n",
+    "The TensorBoard logs are stored in the experiment directory. Let's get the path:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Get experiment path\nfrom rapidfireai.db.rf_db import RfDb\n\ndb = RfDb()\nexperiment_name = \"tensorboard-demo\"\nexperiment_path = db.get_experiments_path(experiment_name)\ntensorboard_log_dir = f\"{experiment_path}/{experiment_name}/tensorboard_logs\"\n\nprint(f\"TensorBoard logs will be saved to: {tensorboard_log_dir}\")"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Start TensorBoard\n",
+    "\n",
+    "**IMPORTANT**: Start TensorBoard BEFORE running training, so you can watch metrics update in real-time!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Define Model Configuration\n\nWe'll use GPT-2 (124M parameters) which is perfect for Colab's memory constraints:"
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "# GPT-2 specific LoRA configs - different module names!\npeft_configs_lite = List([\n    RFLoraConfig(\n        r=8,\n        lora_alpha=16,\n        lora_dropout=0.1,\n        target_modules=[\"c_attn\"],  # GPT-2 combines Q,K,V in c_attn\n        bias=\"none\"\n    ),\n    RFLoraConfig(\n        r=32,\n        lora_alpha=64,\n        lora_dropout=0.1,\n        target_modules=[\"c_attn\", \"c_proj\"],  # c_attn (QKV) + c_proj (output)\n        bias=\"none\"\n    )\n])\n\n# 2 configs with GPT-2 (124M params)\nconfig_set_lite = List([\n    RFModelConfig(\n        model_name=\"gpt2\",  # Only 124M params\n        peft_config=peft_configs_lite,\n        training_args=RFSFTConfig(\n            learning_rate=5e-4,  # Lower learning rate for GPT-2 stability\n            lr_scheduler_type=\"linear\",\n            per_device_train_batch_size=2,  # Reduced for memory\n            per_device_eval_batch_size=2,\n            max_steps=128,\n            gradient_accumulation_steps=2,  # Effective batch size = 4\n            logging_steps=2,\n            eval_strategy=\"steps\",\n            eval_steps=4,\n            fp16=True,\n            gradient_checkpointing=True,  # Save memory\n            report_to=\"none\",  # Disables wandb\n        ),\n        model_type=\"causal_lm\",\n        model_kwargs={\n            \"device_map\": \"auto\",\n            \"torch_dtype\": \"float16\",  # Explicit fp16\n            \"use_cache\": False\n        },\n        formatting_func=sample_formatting_function,\n        compute_metrics=sample_compute_metrics,\n        generation_config={\n            \"max_new_tokens\": 128,  # Reduced from 256\n            \"temperature\": 0.7,\n            \"top_p\": 0.9,\n            \"top_k\": 40,\n            \"repetition_penalty\": 1.1,\n            \"pad_token_id\": 50256,  # GPT-2's EOS token\n        }\n    ),\n    RFModelConfig(\n        model_name=\"gpt2\",\n        peft_config=peft_configs_lite,\n        training_args=RFSFTConfig(\n            learning_rate=2e-4,  # Even more conservative\n            lr_scheduler_type=\"cosine\",  # Try cosine schedule\n            per_device_train_batch_size=2,\n            per_device_eval_batch_size=2,\n            max_steps=128,\n            gradient_accumulation_steps=2,\n            logging_steps=2,\n            eval_strategy=\"steps\",\n            eval_steps=4,\n            fp16=True,\n            gradient_checkpointing=True,\n            report_to=\"none\",  # Disables wandb\n            warmup_steps=10,  # Add warmup for stability\n        ),\n        model_type=\"causal_lm\",\n        model_kwargs={\n            \"device_map\": \"auto\",\n            \"torch_dtype\": \"float16\",\n            \"use_cache\": False\n        },\n        formatting_func=sample_formatting_function,\n        compute_metrics=sample_compute_metrics,\n        generation_config={\n            \"max_new_tokens\": 128,\n            \"temperature\": 0.7,\n            \"top_p\": 0.9,\n            \"top_k\": 40,\n            \"repetition_penalty\": 1.1,\n            \"pad_token_id\": 50256,\n        }\n    )\n])",
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "def sample_create_model(model_config):\n    \"\"\"Function to create model object with GPT-2 adjustments\"\"\"\n    from transformers import AutoModelForCausalLM, AutoTokenizer\n\n    model_name = model_config[\"model_name\"]\n    model_type = model_config[\"model_type\"]\n    model_kwargs = model_config[\"model_kwargs\"]\n\n    if model_type == \"causal_lm\":\n        model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)\n    else:\n        # Default to causal LM\n        model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)\n\n    tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n    # GPT-2 specific: Set pad token (GPT-2 doesn't have one by default)\n    if \"gpt2\" in model_name.lower():\n        tokenizer.pad_token = tokenizer.eos_token\n        tokenizer.padding_side = \"left\"  # GPT-2 works better with left padding\n        model.config.pad_token_id = model.config.eos_token_id\n\n    return (model, tokenizer)",
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": "# Simple grid search across all config combinations = 4 total (2 LoRA configs × 2 training configs)\nconfig_group = RFGridSearch(\n    configs=config_set_lite,\n    trainer_type=\"SFT\"\n)",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Interactive Run Controller\n\nUse the Interactive Controller to monitor and manage training runs in real-time:",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Simple grid search\n",
+    "config_group = RFGridSearch(\n",
+    "    configs=config_set,\n",
+    "    trainer_type=\"SFT\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Run Training\n\nNow let's start training! The metrics will appear in TensorBoard above in real-time:",
+   "metadata": {}
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Interactive Run Controller\n\nRapidFire provides an Interactive Controller that lets you manage training runs in real-time from your notebook:\n\n- ▶️ **Resume**: Resume a stopped run\n- ⏹️ **Stop**: Gracefully stop a running experiment\n- 🗑️ **Delete**: Remove a run from the database\n- 📋 **Clone**: Create a new run with modified hyperparameters (with optional warm start)\n- 🔄 **Refresh**: Update run status and metrics\n\nThe controller uses ipywidgets and is compatible with both Colab (ipywidgets 7.x) and Jupyter (ipywidgets 8.x).",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Launch training\nexperiment.run_fit(\n    config_group, \n    sample_create_model, \n    train_dataset, \n    eval_dataset, \n    num_chunks=4,  # 4 chunks for parallel execution\n    seed=42\n)"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## End Experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# View final logs\n",
+    "%tensorboard --logdir {tensorboard_log_dir}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## View TensorBoard Logs\n\nAfter training completes, you can view the full logs in TensorBoard:",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "# View final logs\n%tensorboard --logdir {tensorboard_log_dir}",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Using Both MLflow and TensorBoard\n",
+    "\n",
+    "You can also log to both backends simultaneously by setting:\n",
+    "\n",
+    "```python\n",
+    "os.environ['RF_TRACKING_BACKEND'] = 'both'\n",
+    "```\n",
+    "\n",
+    "This gives you:\n",
+    "- **TensorBoard**: Real-time visualization during training\n",
+    "- **MLflow**: Experiment comparison and model registry\n",
+    "\n",
+    "## Tips for Colab + TensorBoard\n",
+    "\n",
+    "1. **Start TensorBoard first**: Always start TensorBoard before training\n",
+    "2. **Frequent logging**: Set `logging_steps` to a small value (e.g., 2-5) for responsive updates\n",
+    "3. **Refresh rate**: TensorBoard polls logs every 30 seconds in Colab\n",
+    "4. **Multiple experiments**: Use different experiment names for different runs\n",
+    "5. **Clean logs**: Delete old logs with `!rm -rf {tensorboard_log_dir}` to start fresh\n",
+    "\n",
+    "## Comparison: TensorBoard vs MLflow in Colab\n",
+    "\n",
+    "| Feature | TensorBoard | MLflow |\n",
+    "|---------|-------------|--------|\n",
+    "| Real-time updates | ✅ Yes (30s polling) | ❌ No (frontend load time) |\n",
+    "| Colab native | ✅ %tensorboard magic | ❌ Requires tunneling |\n",
+    "| Load time | ✅ Instant | ❌ 3-5 minutes via tunnel |\n",
+    "| Model registry | ❌ No | ✅ Yes |\n",
+    "| Experiment comparison | ✅ Basic | ✅ Advanced |\n",
+    "\n",
+    "**Recommendation**: Use `'both'` backend to get the best of both worlds!\n",
+    "\n",
+    "## Next Steps\n",
+    "\n",
+    "- Try different model configs and compare in TensorBoard\n",
+    "- Experiment with `'both'` backend for comprehensive tracking\n",
+    "- Check out other RapidFire tutorials for DPO and GRPO training\n",
+    "\n",
+    "Happy training! 🚀"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12"
+  },
+  "colab": {
+   "provenance": [],
+   "gpuType": "T4"
+  },
+  "accelerator": "GPU"
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}