{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "EAsv3uyVZERg" }, "source": [ "# Quantization Aware Training (QAT)\n", "\n", "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/gemma/blob/main/colabs/quantization_aware_training.ipynb)\n", "\n", "This is an example on how to obtain and run quantized versions of Gemma models. It's best to first read the [finetuning](https://gemma-llm.readthedocs.io/en/latest/finetuning.html) colab to understand this one.\n", "\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "executionInfo": { "elapsed": 3, "status": "ok", "timestamp": 1741556668372, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "7EM9c9_kZERg" }, "outputs": [], "source": [ "!pip install -q gemma" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "executionInfo": { "elapsed": 5222, "status": "ok", "timestamp": 1741556673717, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "j85GsCFMZERg" }, "outputs": [], "source": [ "# Common imports\n", "import os\n", "import treescope\n", "import optax\n", "\n", "# Gemma imports\n", "from kauldron import kd\n", "from gemma import gm\n", "from gemma import peft" ] }, { "cell_type": "markdown", "metadata": { "id": "4Chiw4WkZERg" }, "source": [ "By default, Jax do not utilize the full GPU memory, but this can be overwritten. See [GPU memory allocation](https://docs.jax.dev/en/latest/gpu_memory_allocation.html):" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "executionInfo": { "elapsed": 52, "status": "ok", "timestamp": 1741556673898, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "yh5RMMLy4DR1" }, "outputs": [], "source": [ "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"]=\"1.00\"" ] }, { "cell_type": "markdown", "metadata": { "id": "-apwpcjNZERg" }, "source": [ "## Config updates\n", "\n", "If you're familiar with the [finetuning](https://gemma-llm.readthedocs.io/en/latest/finetuning.html) tutorial, switching to QAT only require 1 change to the trainer.\n", "\n", "This is slightly different to LoRA (we discuss the difference below)" ] }, { "cell_type": "markdown", "metadata": { "id": "Id0PQE16ZERg" }, "source": [ "### 1. Model\n", "\n", "Wrap the model in the `gm.nn.QuantizationAwareTrainingWrapper`. This will apply model surgery to replace all the linear and compatible layers with Simulation for Quantized layers. You can choose among several options for quantization:\n", "\n", "* SFP8: switched floating point in 8 bits (very efficient with gemma.cpp)\n", "* Q4_0: per-block integer quantization (equivalent to 4.5 bits per weights), very popular on llama.cpp\n", "* INT4: per-channel weight quantization (almost exactly 4 bits per weights)" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "executionInfo": { "elapsed": 9380, "status": "ok", "timestamp": 1741556683410, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "rKHZ6d3-ZERg" }, "outputs": [], "source": [ "model = gm.nn.QuantizationAwareWrapper(\n", " method = peft.QuantizationMethod.INT8,\n", " model=gm.nn.Gemma3_4B(tokens=\"batch.input\", text_only=True),\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "9MBCCglxdwht" }, "source": [ "Internally, this uses the [`gemma.peft`](https://github.com/google-deepmind/gemma/blob/main/gemma/peft) mini-library to perform model surgery." ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "executionInfo": { "elapsed": 54, "status": "ok", "timestamp": 1741556683592, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "lS_OkUkWgsvP" }, "outputs": [], "source": [ "init_transform = wrapped=gm.ckpts.LoadCheckpoint(\n", " path=gm.ckpts.CheckpointPath.GEMMA3_4B_IT,\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "O9kiQmjgdwht" }, "source": [ "## Training" ] }, { "cell_type": "markdown", "metadata": { "id": "ybKFzaE_dwht" }, "source": [ "### Data pipeline\n", "\n", "Like for the [finetuning](https://gemma-llm.readthedocs.io/en/latest/finetuning.html) example, we recreate the tokenizer:" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "executionInfo": { "elapsed": 822, "status": "ok", "timestamp": 1741556684537, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "u5MZYuTBdwht", "outputId": "7eaa77fc-179e-4f8a-cfa7-cb5cd63e2597" }, "outputs": [ { "data": { "text/plain": [ "[<_Gemma3SpecialTokens.BOS: 2>, 2094, 563, 614, 2591, 13315]" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tokenizer = gm.text.Gemma3Tokenizer()\n", "\n", "tokenizer.encode('This is an example sentence', add_bos=True)" ] }, { "cell_type": "markdown", "metadata": { "id": "0QNkBrD8dwht" }, "source": [ "And the data pipeline:" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "colab": { "height": 394 }, "executionInfo": { "elapsed": 17257, "status": "ok", "timestamp": 1741556701922, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "ULN7oQ_8dwht", "outputId": "781108d4-eb01-479e-8028-96320966967c" }, "outputs": [ { "data": { "text/html": [ "
" ], "text/plain": [ "{\n", " 'input': # np.ndarray int64(8, 200) [≥0, ≤237_167] zero:1_148 nonzero:452\n", " array([[ 2, 105, 2364, ..., 0, 0, 0],\n", " [ 2, 105, 2364, ..., 0, 0, 0],\n", " [ 2, 105, 2364, ..., 0, 0, 0],\n", " ...,\n", " [ 2, 105, 2364, ..., 0, 0, 0],\n", " [ 2, 105, 2364, ..., 0, 0, 0],\n", " [ 2, 105, 2364, ..., 0, 0, 0]], shape=(8, 200))\n", " ,\n", " 'loss_mask': ,\n", " 'target': ,\n", "}" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "ds = kd.data.py.Tfds(\n", " name='mtnt/en-fr',\n", " split='train',\n", " shuffle=True,\n", " batch_size=8,\n", " transforms=[\n", " # Create the model inputs/targets/loss_mask.\n", " gm.data.Seq2SeqTask(\n", " # Select which field from the dataset to use.\n", " # https://www.tensorflow.org/datasets/catalog/mtnt\n", " in_prompt='src',\n", " in_response='dst',\n", " # Output batch is {'input': ..., 'target': ..., 'loss_mask': ...}\n", " out_input='input',\n", " out_target='target',\n", " out_target_mask='loss_mask',\n", " tokenizer=tokenizer,\n", " # Padding parameters\n", " max_length=200,\n", " truncate=True,\n", " ),\n", " ],\n", ")\n", "\n", "ex = ds[0]\n", "\n", "treescope.show(ex)" ] }, { "cell_type": "markdown", "metadata": { "id": "AH7sBVr8dwht" }, "source": [ "We can decode an example from the batch to inspect the model input and check it is properly formatted:" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "executionInfo": { "elapsed": 52, "status": "ok", "timestamp": 1741556702133, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "4krpPavGdwht", "outputId": "b2da5cd2-4231-4dfe-b591-138ff0d96b84" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "user\n", "Is this a good place to ask about the ethnicity and intelligence debate?\n", "model\n", "Est-ce un bon endroit pour poser des questions sur le débat à propos de l'ethnicité et le renseignement ?\n" ] } ], "source": [ "text = tokenizer.decode(ex['input'][0])\n", "\n", "print(text)" ] }, { "cell_type": "markdown", "metadata": { "id": "V0nJX7uddwht" }, "source": [ "### Trainer" ] }, { "cell_type": "markdown", "metadata": { "id": "P0lMcUindwht" }, "source": [ "We then create the trainer, reusing the `model`, `init_transform` and `optimizer` created above:" ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "executionInfo": { "elapsed": 54, "status": "ok", "timestamp": 1741556702316, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "XGj5dw_Gdwht" }, "outputs": [], "source": [ "trainer = kd.train.Trainer(\n", " seed=42, # The seed of enlightenment\n", " workdir='/tmp/ckpts', # TODO(epot): Make the workdir optional by default\n", " # Dataset\n", " train_ds=ds,\n", " # Model\n", " model=model,\n", " init_transform=init_transform,\n", " # Training parameters\n", " num_train_steps=500,\n", " train_losses={\n", " \"loss\": kd.losses.SoftmaxCrossEntropyWithIntLabels(\n", " logits=\"preds.logits\",\n", " labels=\"batch.target\",\n", " mask=\"batch.loss_mask\",\n", " ),\n", " },\n", " optimizer=optax.adafactor(learning_rate=0.005),\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "FT0plMVmdwht" }, "source": [ "Trainning can be launched with the `.train()` method.\n", "\n", "Note that the trainer like the model are immutables, so it does not store the state nor params. Instead the state containing the trained parameters is returned." ] }, { "cell_type": "code", "execution_count": 10, "metadata": { "colab": { "height": 171, "referenced_widgets": [ "d50e55815c444675a8e669e781035b34", "a8fb8f12675e450cba2901cc78b2981a", "e4593e87157347bb8a9df5936d8b73e5", "6361fa8889444bc8a45a00488493fee1", "1600c7d9b6c3448c925102245c5d4b10", "7894811af08543d29481b3bec06b736b", "9f743749221b4cb281a3d2514e695840", "65f711cd25734c838b1fbc24ef914c1e", "c25865c509ac4517a1efe17e7b2063a4", "18e141276be24e00875769cbee883832", "8c81664c18894eadaf847cd00983330b" ] }, "executionInfo": { "elapsed": 304719, "status": "ok", "timestamp": 1741557007194, "user": { "displayName": "", "userId": "" }, "user_tz": -60 }, "id": "8UmU7u-Hdwht", "outputId": "375d18b6-e07b-4621-920a-7de5efe01488" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Configuring ...\n", "Initializing ...\n", "Disabling pygrain multi-processing (unsupported in colab).\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Starting training loop at step 0\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "d50e55815c444675a8e669e781035b34", "version_major": 2, "version_minor": 0 }, "text/plain": [ "train: 0%| | 0/501 [00:00'" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "sampler = gm.text.Sampler(\n", " model=quantized_model,\n", " params=quantized_params,\n", " tokenizer=tokenizer,\n", ")\n", "\n", "prompt = \"\"\"\\\n", "user\n", "I'm feeling happy!\n", "model\n", "\"\"\"\n", "\n", "sampler.sample(prompt, max_new_tokens=30)" ] }, { "cell_type": "markdown", "metadata": { "id": "8sxaHy2Zdwht" }, "source": [] } ], "metadata": { "colab": { "last_runtime": {}, "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" }, "widgets": { "application/vnd.jupyter.widget-state+json": { "1600c7d9b6c3448c925102245c5d4b10": { "model_module": "@jupyter-widgets/base", "model_module_version": "1.2.0", "model_name": "LayoutModel", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "18e141276be24e00875769cbee883832": { "model_module": "@jupyter-widgets/base", "model_module_version": "1.2.0", "model_name": "LayoutModel", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "6361fa8889444bc8a45a00488493fee1": { "model_module": "@jupyter-widgets/controls", "model_module_version": "1.5.0", "model_name": "HTMLModel", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "HTMLModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "HTMLView", "description": "", "description_tooltip": null, "layout": "IPY_MODEL_18e141276be24e00875769cbee883832", "placeholder": "​", "style": "IPY_MODEL_8c81664c18894eadaf847cd00983330b", "value": " 501/501 [04:27<00:00, 5.36s/it]" } }, "65f711cd25734c838b1fbc24ef914c1e": { "model_module": "@jupyter-widgets/base", "model_module_version": "1.2.0", "model_name": "LayoutModel", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "7894811af08543d29481b3bec06b736b": { "model_module": "@jupyter-widgets/base", "model_module_version": "1.2.0", "model_name": "LayoutModel", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "8c81664c18894eadaf847cd00983330b": { "model_module": "@jupyter-widgets/controls", "model_module_version": "1.5.0", "model_name": "DescriptionStyleModel", "state": { "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "DescriptionStyleModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "StyleView", "description_width": "" } }, "9f743749221b4cb281a3d2514e695840": { "model_module": "@jupyter-widgets/controls", "model_module_version": "1.5.0", "model_name": "DescriptionStyleModel", "state": { "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "DescriptionStyleModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "StyleView", "description_width": "" } }, "a8fb8f12675e450cba2901cc78b2981a": { "model_module": "@jupyter-widgets/controls", "model_module_version": "1.5.0", "model_name": "HTMLModel", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "HTMLModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "HTMLView", "description": "", "description_tooltip": null, "layout": "IPY_MODEL_7894811af08543d29481b3bec06b736b", "placeholder": "​", "style": "IPY_MODEL_9f743749221b4cb281a3d2514e695840", "value": "train: 100%" } }, "c25865c509ac4517a1efe17e7b2063a4": { "model_module": "@jupyter-widgets/controls", "model_module_version": "1.5.0", "model_name": "ProgressStyleModel", "state": { "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "ProgressStyleModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "StyleView", "bar_color": null, "description_width": "" } }, "d50e55815c444675a8e669e781035b34": { "model_module": "@jupyter-widgets/controls", "model_module_version": "1.5.0", "model_name": "HBoxModel", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "HBoxModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "HBoxView", "box_style": "", "children": [ "IPY_MODEL_a8fb8f12675e450cba2901cc78b2981a", "IPY_MODEL_e4593e87157347bb8a9df5936d8b73e5", "IPY_MODEL_6361fa8889444bc8a45a00488493fee1" ], "layout": "IPY_MODEL_1600c7d9b6c3448c925102245c5d4b10" } }, "e4593e87157347bb8a9df5936d8b73e5": { "model_module": "@jupyter-widgets/controls", "model_module_version": "1.5.0", "model_name": "FloatProgressModel", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "FloatProgressModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "ProgressView", "bar_style": "success", "description": "", "description_tooltip": null, "layout": "IPY_MODEL_65f711cd25734c838b1fbc24ef914c1e", "max": 501, "min": 0, "orientation": "horizontal", "style": "IPY_MODEL_c25865c509ac4517a1efe17e7b2063a4", "value": 501 } } } } }, "nbformat": 4, "nbformat_minor": 0 }