{ "cells": [ { "cell_type": "markdown", "id": "62a637c4", "metadata": {}, "source": [ "# Supervised fine-tuning for LFM2.5-VL-3B with Halo\n", "\n", "This notebook fine-tunes [LFM2.5-VL-3B](https://huggingface.co/LiquidAI/LFM2.5-VL-3B) with a LoRA adapter using [Halo](https://whitecircle.com/research/halo), which is a distributed-training layer for Hugging Face models. The [Halo repository](https://github.com/whitecircle/halo) contains the source code.\n", "\n", "It uses [CPPE-5](https://huggingface.co/datasets/rishitdagli/cppe-5), which contains five classes of personal protective equipment. The notebook makes a fixed split with 850 training images and 150 test images. It evaluates the first 30 test images.\n", "\n", "The evaluation matches boxes by class at an intersection over union (IoU) threshold of 0.5.\n" ] }, { "cell_type": "markdown", "id": "cd93e55c", "metadata": {}, "source": [ "## Prerequisites\n", "\n", "- A GPU that matches the published Halo images, which are built for Hopper (H100, H200) and Blackwell (B200, B300).\n", "- An NVIDIA driver that provides CUDA 13.0 or higher\n", "- Docker with the NVIDIA Container Toolkit\n", "- A large writable volume for model, dataset, adapter, and evaluation files\n", "\n", "This notebook does not run in Google Colab. Use a compatible GPU machine from your cloud provider or your own cluster. \n", "The validated B300 run completed 150 optimizer steps in 482 seconds. Other GPUs can take longer.\n", "\n", "## Recipe at a Glance\n", "\n", "This recipe uses a LoRA adapter on a single GPU.\n", "\n", "| Setting | Value |\n", "| --- | --- |\n", "| GPUs | 1 |\n", "| Method | LoRA, rank 16 |\n", "| Dataset | CPPE-5, 850 train / 150 test |\n", "| Sequence length | 4,096 |\n", "| Effective batch size | 8 |\n", "| Training steps | 150 |\n", "| Output | LoRA adapter |" ] }, { "cell_type": "markdown", "id": "330f8e8f", "metadata": {}, "source": [ "## Installation & Setup\n", "\n", "This notebook uses the public Halo image that matches the detected GPU. Define these host paths before you start the container:\n", "\n", "- `HALO_DIR`: an optional Halo clone for source reference\n", "- `COOKBOOK_DIR`: a clone of this cookbook repository\n", "- `DATA_DIR`: a large persistent volume for downloads and run files\n", "\n", "~~~bash\n", "export HALO_DIR=/path/to/halo\n", "export COOKBOOK_DIR=/path/to/cookbook\n", "export DATA_DIR=/path/to/large/volume\n", "\n", "git clone --recurse-submodules https://github.com/whitecircle/halo.git \"$HALO_DIR\"\n", "git clone https://github.com/Liquid4All/cookbook.git \"$COOKBOOK_DIR\"\n", "~~~\n", "\n", "Pull the Docker image for the detected GPUs.\n", "~~~bash\n", "GPU_NAME=\"$(nvidia-smi --query-gpu=name --format=csv,noheader | head -n 1)\"\n", "case \"$GPU_NAME\" in\n", " *H100*|*H200*) HALO_IMAGE=public.ecr.aws/whitecircle/halo:hopper ;;\n", " *B200*|*B300*|*GB200*|*GB300*) HALO_IMAGE=public.ecr.aws/whitecircle/halo:blackwell ;;\n", " *) echo \"Unsupported GPU: $GPU_NAME\"; exit 1 ;;\n", "esac\n", "export HALO_IMAGE\n", "\n", "docker pull \"$HALO_IMAGE\"\n", "~~~\n", "\n", "Launch the Docker image:\n", "~~~bash\n", "docker run --rm -it --gpus '\"device=0\"' \\\n", " --ipc=host --shm-size=128g \\\n", " --ulimit memlock=-1 --ulimit stack=67108864 \\\n", " -p 8888:8888 \\\n", " -e HF_HOME=/mnt/hf \\\n", " -e HF_DATASETS_CACHE=/mnt/hf/datasets \\\n", " -e TMPDIR=/mnt/tmp \\\n", " -e HALO_DATA_ROOT=/mnt \\\n", " -v \"$COOKBOOK_DIR\":/cookbook \\\n", " -v \"$DATA_DIR\":/mnt \\\n", " -w /workspace \\\n", " \"$HALO_IMAGE\" \\\n", " jupyter lab --ip=0.0.0.0 --port=8888 --no-browser \\\n", " --allow-root --notebook-dir=/cookbook/finetuning/notebooks\n", "~~~\n", "\n", "The container runs the Halo version packaged in the image. \n", "The `HALO_DIR` clone is for source reference and is not mounted over that package. \n", "`-p 8888:8888` publishes only the Jupyter port. \n", "This single-GPU run does not use DeepEP, so it does not need `CUDA_DEVICE_MAX_CONNECTIONS=1`.\n", "\n", "Open the URL printed by Jupyter. Then open this notebook.\n" ] }, { "cell_type": "markdown", "id": "b3abf961", "metadata": {}, "source": [ "## Check the Runtime\n", "\n", "This cell sets the run paths, confirms that CUDA is available, and creates the data directory. It stops if no GPU is visible.\n" ] }, { "cell_type": "code", "execution_count": 1, "id": "fe1f0aad", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "NVIDIA B300 SXM6 AC\n" ] } ], "source": [ "from pathlib import Path\n", "\n", "import torch\n", "\n", "DATA_ROOT = Path(\"/mnt/lfm25-cppe5\")\n", "DATASET_PATH = DATA_ROOT / \"cppe5-halo\"\n", "CONFIG_PATH = DATA_ROOT / \"lfm25-vl-3b-cppe5.yaml\"\n", "ADAPTER_PATH = DATA_ROOT / \"lfm25-vl-3b-cppe5-halo\"\n", "BASELINE_PATH = DATA_ROOT / \"baseline.json\"\n", "TUNED_PATH = DATA_ROOT / \"tuned.json\"\n", "\n", "assert torch.cuda.is_available(), \"CUDA is not available\"\n", "DATA_ROOT.mkdir(parents=True, exist_ok=True)\n", "\n", "print(torch.cuda.get_device_name())\n" ] }, { "cell_type": "markdown", "id": "0994b20a", "metadata": {}, "source": [ "## Data Prep\n", "\n", "The [CPPE-5](https://huggingface.co/datasets/rishitdagli/cppe-5) source has one training split with images and bounding-box annotations. This section creates a repeatable 850-image training split and 150-image test split. It converts each training row into a user prompt, an image, and the expected JSON response. The model returns each class and its normalized bounding box.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5db99a53", "metadata": {}, "outputs": [], "source": [ "import json\n", "\n", "from datasets import DatasetDict, Image, load_dataset\n", "\n", "DATASET_ID = \"rishitdagli/cppe-5\"\n", "CLASS_NAMES = [\"coverall\", \"face shield\", \"gloves\", \"goggles\", \"mask\"]\n", "SEED = 1337\n", "\n", "\n", "def detection_prompt():\n", " labels = \", \".join(CLASS_NAMES)\n", " return (\n", " f\"Detect all instances of: {labels}. Response must be a JSON array: \"\n", " '[{\"label\": \"...\", \"bbox\": [x1, y1, x2, y2]}, ...]. '\n", " \"Coordinates are normalized to [0,1].\"\n", " )\n", "\n", "\n", "def training_messages(row):\n", " width, height = row[\"width\"], row[\"height\"]\n", " answer = []\n", " for category, box in zip(\n", " row[\"objects\"][\"category\"], row[\"objects\"][\"bbox\"], strict=True\n", " ):\n", " x, y, box_width, box_height = box\n", " answer.append(\n", " {\n", " \"label\": CLASS_NAMES[category],\n", " \"bbox\": [\n", " round(max(0.0, x / width), 4),\n", " round(max(0.0, y / height), 4),\n", " round(min(1.0, (x + box_width) / width), 4),\n", " round(min(1.0, (y + box_height) / height), 4),\n", " ],\n", " }\n", " )\n", " return [\n", " {\"role\": \"user\", \"content\": detection_prompt()},\n", " {\n", " \"role\": \"assistant\",\n", " \"content\": json.dumps(answer, separators=(\",\", \":\")),\n", " },\n", " ]\n", "\n", "\n", "def prepare_split(dataset):\n", " messages = [training_messages(row) for row in dataset]\n", " return (\n", " dataset.add_column(\"messages\", messages)\n", " .select_columns([\"image\", \"messages\"])\n", " .cast_column(\"image\", Image(mode=\"RGB\"))\n", " )\n", "\n", "\n", "source = load_dataset(DATASET_ID)\n", "split = source[\"train\"].train_test_split(test_size=0.15, seed=SEED)\n", "formatted = DatasetDict(\n", " {name: prepare_split(dataset) for name, dataset in split.items()}\n", ")\n", "\n", "assert len(formatted[\"train\"]) == 850\n", "assert len(formatted[\"test\"]) == 150\n", "if DATASET_PATH.exists():\n", " raise FileExistsError(\n", " f\"{DATASET_PATH} already exists. Remove it before you rebuild the split.\"\n", " )\n", "formatted.save_to_disk(DATASET_PATH)\n", "print({\"train\": 850, \"test\": 150, \"seed\": SEED})\n" ] }, { "cell_type": "markdown", "id": "71336d23", "metadata": {}, "source": [ "## Define the Evaluation\n", "\n", "The evaluation uses greedy generation on 30 test images. It matches predicted boxes to reference boxes from the same class.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "4e138c41", "metadata": {}, "outputs": [], "source": [ "import gc\n", "from collections import defaultdict\n", "\n", "from datasets import load_dataset\n", "from peft import PeftModel\n", "from transformers import AutoModelForImageTextToText, AutoProcessor\n", "\n", "MODEL_ID = \"LiquidAI/LFM2.5-VL-3B\"\n", "EVALUATION_IMAGES = 30\n", "IOU_THRESHOLD = 0.5\n", "\n", "\n", "def empty_boxes():\n", " return {label: [] for label in CLASS_NAMES}\n", "\n", "\n", "def parse_boxes(text):\n", " start = text.find(\"[\")\n", " end = text.rfind(\"]\")\n", " if start < 0 or end < start:\n", " return empty_boxes()\n", " try:\n", " items = json.loads(text[start : end + 1])\n", " except json.JSONDecodeError:\n", " return empty_boxes()\n", " if not isinstance(items, list):\n", " return empty_boxes()\n", "\n", " boxes = empty_boxes()\n", " for item in items:\n", " if not isinstance(item, dict) or not isinstance(item.get(\"bbox\"), list):\n", " continue\n", " box = item[\"bbox\"]\n", " if len(box) != 4 or not all(\n", " isinstance(value, (int, float)) for value in box\n", " ):\n", " continue\n", " label = str(item.get(\"label\", \"\")).strip().casefold().replace(\"_\", \" \")\n", " if label not in boxes:\n", " continue\n", " x1, y1, x2, y2 = (float(value) for value in box)\n", " if x2 <= x1 or y2 <= y1:\n", " continue\n", " boxes[label].append(\n", " [\n", " max(0.0, min(1.0, x1)),\n", " max(0.0, min(1.0, y1)),\n", " max(0.0, min(1.0, x2)),\n", " max(0.0, min(1.0, y2)),\n", " ]\n", " )\n", " return boxes\n", "\n", "\n", "def ground_truth(row):\n", " width, height = row[\"width\"], row[\"height\"]\n", " boxes = empty_boxes()\n", " for category, box in zip(\n", " row[\"objects\"][\"category\"], row[\"objects\"][\"bbox\"], strict=True\n", " ):\n", " x, y, box_width, box_height = box\n", " boxes[CLASS_NAMES[category]].append(\n", " [\n", " max(0.0, x / width),\n", " max(0.0, y / height),\n", " min(1.0, (x + box_width) / width),\n", " min(1.0, (y + box_height) / height),\n", " ]\n", " )\n", " return boxes\n", "\n", "\n", "def box_iou(first, second):\n", " left = max(first[0], second[0])\n", " top = max(first[1], second[1])\n", " right = min(first[2], second[2])\n", " bottom = min(first[3], second[3])\n", " intersection = max(0.0, right - left) * max(0.0, bottom - top)\n", " first_area = (first[2] - first[0]) * (first[3] - first[1])\n", " second_area = (second[2] - second[0]) * (second[3] - second[1])\n", " union = first_area + second_area - intersection\n", " return intersection / union if union > 0 else 0.0\n", "\n", "\n", "def match_boxes(predicted, expected):\n", " unmatched = set(range(len(expected)))\n", " true_positives = 0\n", " for prediction in predicted:\n", " if not unmatched:\n", " break\n", " best = max(unmatched, key=lambda index: box_iou(prediction, expected[index]))\n", " if box_iou(prediction, expected[best]) >= IOU_THRESHOLD:\n", " unmatched.remove(best)\n", " true_positives += 1\n", " return true_positives, len(predicted) - true_positives, len(unmatched)\n", "\n", "\n", "def resolve_adapter(path):\n", " if (path / \"adapter_config.json\").is_file():\n", " return path\n", " checkpoints = list(path.glob(\"checkpoint-*/adapter_config.json\"))\n", " if not checkpoints:\n", " raise FileNotFoundError(f\"No adapter found under {path}\")\n", " return max(\n", " checkpoints,\n", " key=lambda item: int(item.parent.name.removeprefix(\"checkpoint-\")),\n", " ).parent\n", "\n", "\n", "def evaluate(output_path, adapter_path=None):\n", " dataset = load_dataset(DATASET_ID, split=\"train\")\n", " dataset = dataset.train_test_split(test_size=0.15, seed=SEED)[\"test\"]\n", " dataset = dataset.select(range(EVALUATION_IMAGES))\n", " processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)\n", " model = AutoModelForImageTextToText.from_pretrained(\n", " MODEL_ID,\n", " dtype=torch.bfloat16,\n", " attn_implementation=\"sdpa\",\n", " trust_remote_code=True,\n", " ).cuda()\n", "\n", " adapter = resolve_adapter(adapter_path) if adapter_path else None\n", " if adapter:\n", " model = PeftModel.from_pretrained(model, adapter)\n", " model.eval()\n", "\n", " records = []\n", " totals = defaultdict(int)\n", " for image_index, row in enumerate(dataset):\n", " message = [\n", " {\n", " \"role\": \"user\",\n", " \"content\": [\n", " {\"type\": \"image\"},\n", " {\"type\": \"text\", \"text\": detection_prompt()},\n", " ],\n", " }\n", " ]\n", " text = processor.apply_chat_template(\n", " message, tokenize=False, add_generation_prompt=True\n", " )\n", " inputs = processor(\n", " text=[text],\n", " images=[[row[\"image\"].convert(\"RGB\")]],\n", " return_tensors=\"pt\",\n", " add_special_tokens=False,\n", " ).to(\"cuda\")\n", " with torch.inference_mode():\n", " generated = model.generate(\n", " **inputs, max_new_tokens=512, do_sample=False, use_cache=True\n", " )\n", " new_tokens = generated[:, inputs[\"input_ids\"].shape[1] :]\n", " raw_prediction = processor.batch_decode(\n", " new_tokens, skip_special_tokens=True\n", " )[0]\n", " predicted = parse_boxes(raw_prediction)\n", " expected = ground_truth(row)\n", "\n", " image_totals = defaultdict(int)\n", " for label in CLASS_NAMES:\n", " tp, fp, fn = match_boxes(predicted[label], expected[label])\n", " totals[\"tp\"] += tp\n", " totals[\"fp\"] += fp\n", " totals[\"fn\"] += fn\n", " image_totals[\"tp\"] += tp\n", " image_totals[\"fp\"] += fp\n", " image_totals[\"fn\"] += fn\n", " records.append(\n", " {\n", " \"image\": image_index,\n", " \"expected\": expected,\n", " \"predicted\": predicted,\n", " \"tp\": image_totals[\"tp\"],\n", " \"fp\": image_totals[\"fp\"],\n", " \"fn\": image_totals[\"fn\"],\n", " \"raw_prediction\": raw_prediction,\n", " }\n", " )\n", " print(f\"evaluated {image_index + 1}/{len(dataset)}\")\n", "\n", " precision = totals[\"tp\"] / max(1, totals[\"tp\"] + totals[\"fp\"])\n", " recall = totals[\"tp\"] / max(1, totals[\"tp\"] + totals[\"fn\"])\n", " f1 = 2 * precision * recall / max(1e-12, precision + recall)\n", " result = {\n", " \"model\": MODEL_ID,\n", " \"adapter\": str(adapter) if adapter else None,\n", " \"images\": len(dataset),\n", " \"true_positives\": totals[\"tp\"],\n", " \"false_positives\": totals[\"fp\"],\n", " \"false_negatives\": totals[\"fn\"],\n", " \"precision_at_iou_0_5\": precision,\n", " \"recall_at_iou_0_5\": recall,\n", " \"f1_at_iou_0_5\": f1,\n", " \"predictions\": records,\n", " }\n", " output_path.write_text(json.dumps(result, indent=2))\n", "\n", " del model\n", " gc.collect()\n", " torch.cuda.empty_cache()\n", " return result\n" ] }, { "cell_type": "markdown", "id": "7705d562", "metadata": {}, "source": [ "## Measure the Base Model\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ba60c6fe", "metadata": {}, "outputs": [], "source": [ "baseline = evaluate(BASELINE_PATH)\n", "print(\n", " {\n", " \"precision\": baseline[\"precision_at_iou_0_5\"],\n", " \"recall\": baseline[\"recall_at_iou_0_5\"],\n", " \"f1\": baseline[\"f1_at_iou_0_5\"],\n", " }\n", ")\n" ] }, { "cell_type": "markdown", "id": "b977e4ee", "metadata": {}, "source": [ "## Create the Training Configuration\n", "\n", "The LoRA target list covers attention, feed-forward, vision-tower, and projector linear layers. Halo updates only the adapter weights and trains only assistant response tokens. The local dataset already contains the fixed train and test splits, so `test_size` remains null.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "59aea5d2", "metadata": {}, "outputs": [], "source": [ "CONFIG_PATH.write_text(\n", " r\"\"\"\n", "# --- Model ---\n", "model_name_or_path: LiquidAI/LFM2.5-VL-3B\n", "trust_remote_code: true\n", "attn_implementation: sdpa\n", "\n", "# --- Dataset ---\n", "dataset:\n", "- /mnt/lfm25-cppe5/cppe5-halo\n", "conversation_field: messages\n", "images_field: image\n", "test_size: null\n", "train_on_completions_only: true\n", "assistant_message_template: \"<|im_start|>assistant\\n\"\n", "\n", "# --- LoRA ---\n", "use_peft: true\n", "lora_r: 16\n", "lora_alpha: 32\n", "lora_dropout: 0.05\n", "lora_target_modules:\n", "- q_proj\n", "- k_proj\n", "- v_proj\n", "- out_proj\n", "- in_proj\n", "- w1\n", "- w2\n", "- w3\n", "- linear_1\n", "- linear_2\n", "- fc1\n", "- fc2\n", "lora_task_type: CAUSAL_LM\n", "\n", "# --- Sequence & packing ---\n", "max_length: 4096\n", "packing: false\n", "padding_free: false\n", "bf16: true\n", "\n", "# --- Training ---\n", "per_device_train_batch_size: 1\n", "per_device_eval_batch_size: 1\n", "gradient_accumulation_steps: 8\n", "max_steps: 150\n", "gradient_checkpointing: true\n", "gradient_checkpointing_kwargs:\n", " use_reentrant: false\n", "seed: 42\n", "\n", "# --- Optimizer & LR schedule ---\n", "optim: adamw_torch_fused\n", "learning_rate: 2.0e-05\n", "weight_decay: 0.01\n", "lr_scheduler_type: cosine\n", "warmup_steps: 10\n", "max_grad_norm: 1.0\n", "\n", "# --- Output ---\n", "output_dir: /mnt/lfm25-cppe5/lfm25-vl-3b-cppe5-halo\n", "save_strategy: steps\n", "save_steps: 150\n", "eval_strategy: \"no\"\n", "save_total_limit: 1\n", "save_only_model: true\n", "\n", "# --- Logging ---\n", "logging_steps: 1\n", "logging_first_step: true\n", "report_to: none\n", "\n", "# --- Data ---\n", "dataloader_num_workers: 0\n", "dataset_num_proc: 1\n", "remove_unused_columns: false\n", "\"\"\".lstrip()\n", ")\n", "print(CONFIG_PATH.read_text())" ] }, { "cell_type": "markdown", "id": "80ba0e5b", "metadata": {}, "source": [ "## Check the Launch Command\n", "\n", "The dry run prints the one-process command without loading the model.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "695b7c0c", "metadata": {}, "outputs": [], "source": [ "!halo launch sft {CONFIG_PATH} --dry-run\n" ] }, { "cell_type": "markdown", "id": "4d97fc0d", "metadata": {}, "source": [ "## Launch Training\n", "\n", "Halo uses one process for this LoRA run.\n", "\n", "Any configuration field can be overridden on the command line, and arguments after `--` take precedence over the file. For a shorter check before the full run:\n", "\n", "~~~bash\n", "halo launch sft /mnt/lfm25-cppe5/lfm25-vl-3b-cppe5.yaml -- --max_steps=20 --save_steps=20\n", "~~~\n" ] }, { "cell_type": "code", "execution_count": null, "id": "b79ce888", "metadata": {}, "outputs": [], "source": [ "!halo launch sft {CONFIG_PATH}\n" ] }, { "cell_type": "markdown", "id": "ad38825c", "metadata": {}, "source": [ "## Measure the Fine-Tuned Model\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e2b1eca3", "metadata": {}, "outputs": [], "source": [ "tuned = evaluate(TUNED_PATH, ADAPTER_PATH)\n", "gain = tuned[\"f1_at_iou_0_5\"] - baseline[\"f1_at_iou_0_5\"]\n", "print(\n", " {\n", " \"base_precision\": baseline[\"precision_at_iou_0_5\"],\n", " \"tuned_precision\": tuned[\"precision_at_iou_0_5\"],\n", " \"base_recall\": baseline[\"recall_at_iou_0_5\"],\n", " \"tuned_recall\": tuned[\"recall_at_iou_0_5\"],\n", " \"base_f1\": baseline[\"f1_at_iou_0_5\"],\n", " \"tuned_f1\": tuned[\"f1_at_iou_0_5\"],\n", " \"f1_gain_percentage_points\": 100 * gain,\n", " }\n", ")\n" ] }, { "cell_type": "markdown", "id": "e2a5a00c", "metadata": {}, "source": [ "## Inference\n", "\n", "Show one evaluated test image and the fine-tuned model's response. The `tuned` result comes from the same greedy-generation path used by the metric.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e3374425", "metadata": {}, "outputs": [], "source": [ "from IPython.display import display\n", "\n", "preview_dataset = load_dataset(DATASET_ID, split=\"train\")\n", "preview_dataset = preview_dataset.train_test_split(\n", " test_size=0.15, seed=SEED\n", ")[\"test\"]\n", "preview = tuned[\"predictions\"][0]\n", "\n", "display(preview_dataset[preview[\"image\"]][\"image\"])\n", "print(preview[\"raw_prediction\"])\n" ] }, { "cell_type": "markdown", "id": "822de4e2", "metadata": {}, "source": [ "## Validated Result\n", "\n", "The run trained 33.3 million parameters, or 1.06% of the model. It used one NVIDIA B300 GPU.\n", "\n", "The adapter reached 45.58% F1 on the 30-image evaluation set. Precision was 52.76%, and recall was 40.12%.\n" ] }, { "cell_type": "markdown", "id": "81d650df", "metadata": {}, "source": [ "## Cleanup\n", "\n", "Shut down Jupyter, then press `Ctrl+C` in the terminal that runs the container. The `--rm` option removes the stopped container. The files under `DATA_DIR` remain available. Stop or terminate the cloud GPU machine when you finish to avoid further charges.\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }