{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "qKlB5QTDIV6S"
},
"source": [
"# Quantization (Sampling)\n",
"\n",
"[](https://colab.research.google.com/github/google-deepmind/gemma/blob/main/colabs/quantization_sampling.ipynb)\n",
"\n",
"Example on using quantization with Gemma (for inference). For an example of quantization aware training (QAT), see [QAT finetuning](https://gemma-llm.readthedocs.io/en/latest/quantization_aware_training.html) example."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TR-L25KVKT_F"
},
"outputs": [],
"source": [
"!pip install -q gemma"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "I6fEKB1tISVW"
},
"outputs": [],
"source": [
"# Common imports\n",
"import os\n",
"import jax\n",
"import jax.numpy as jnp\n",
"import treescope\n",
"\n",
"# Gemma imports\n",
"from gemma import gm\n",
"from gemma import peft # Parameter fine-tuning module"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cxGT2XeU4L47"
},
"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": null,
"metadata": {
"id": "o4MidM--4L47"
},
"outputs": [],
"source": [
"os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"]=\"1.00\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-kdAZkvOIryQ"
},
"source": [
"## Initializing the model\n",
"\n",
"To use Gemma with quantization, simply wrap any Gemma model in `gm.nn.IntWrapper` (in the example for int8 weight inference):"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "x-BbrzCVIupV"
},
"outputs": [],
"source": [
"model = gm.nn.IntWrapper(model=gm.nn.Gemma3_4B(text_only=True), dtype=jnp.int8)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hI3Lg07SJff4"
},
"source": [
"Initialize the weights:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1shC1DpiJfsw"
},
"outputs": [],
"source": [
"token_ids = jnp.zeros((1, 256,), dtype=jnp.int32) # Create the (batch_size, seq_length)\n",
"\n",
"params = model.init(jax.random.key(0), token_ids)\n",
"\n",
"params = params['params']"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bGJl5YpKKOf-"
},
"source": [
"\n",
"\n",
"```\n",
"# This is formatted as code\n",
"```\n",
"\n",
"Restore the pre-trained params. We use `peft.quantize` to quantize the checkpoint.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "AcO6oBuLKNjb"
},
"outputs": [],
"source": [
"del params\n",
"# Load the params from the checkpoint\n",
"original = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA3_4B_IT)\n",
"# edit params`\n",
"params = peft.quantize(original, method='INT8', checkpoint_kernel_key='w')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6jYRSpv2IwIY"
},
"source": [
"## Fine-tuning\n",
"\n",
"See our [finetuning guide](https://gemma-llm.readthedocs.io/en/latest/lora_finetuning.html) for more info."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MvsQbQM4I4Cs"
},
"source": [
"## Inference\n",
"\n",
"Here's an example of running a single model call:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"height": 542
},
"executionInfo": {
"elapsed": 11145,
"status": "ok",
"timestamp": 1741513754866,
"user": {
"displayName": "",
"userId": ""
},
"user_tz": -60
},
"id": "eqU7a4eCI5Wr",
"outputId": "e21147d0-a206-493e-fff2-0590d5c7053e"
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
" \n",
"\n",
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"tokenizer = gm.text.Gemma3Tokenizer()\n",
"\n",
"prompt = tokenizer.encode('The capital of France is')\n",
"prompt = jnp.asarray([tokenizer.special_tokens.BOS] + prompt)\n",
"\n",
"\n",
"# Run the model\n",
"out = model.apply(\n",
" {'params': params},\n",
" tokens=prompt,\n",
" return_last_only=True, # Only predict the last token\n",
")\n",
"\n",
"\n",
"# Show the token distribution\n",
"tokenizer.plot_logits(out.logits)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6dOSL9MHuMUa"
},
"source": [
"To sample an entire sentence:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"height": 35
},
"executionInfo": {
"elapsed": 10172,
"status": "ok",
"timestamp": 1741513765357,
"user": {
"displayName": "",
"userId": ""
},
"user_tz": -60
},
"id": "_ckwREdyqown",
"outputId": "4f46f48b-4ad5-4a0b-d332-f47690a6f135"
},
"outputs": [
{
"data": {
"application/vnd.google.colaboratory.intrinsic+json": {
"type": "string"
},
"text/plain": [
"' Paris.\\n\\nParis is a global center for art, fashion, gastronomy, and culture. It is known for its iconic landmarks such as the Eiffel Tower'"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sampler = gm.text.Sampler(\n",
" model=model,\n",
" params=params,\n",
" tokenizer=tokenizer,\n",
")\n",
"\n",
"sampler.sample('The capital of France is', max_new_tokens=30)"
]
}
],
"metadata": {
"colab": {
"last_runtime": {},
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}