DeepAdvisors | Fine-Tuning Gemma with KerasHub on TensorFlow

Open-weight language models changed what a small team can build. You no longer need a frontier-scale budget to get a model that speaks your domain: you take a released checkpoint, fine-tune it on a few thousand examples, and serve it yourself. On the TensorFlow and Keras stack, the tool for this is KerasHub, the successor to KerasNLP and KerasCV, and the model family we reach for first is Google's Gemma. This tutorial walks through a LoRA fine-tune end to end.

Setup

pip install -U keras keras-hub
# plus a backend: jax[cuda12] for training speed, or tensorflow

Gemma checkpoints are gated: accept the license on Kaggle and set your Kaggle credentials (KAGGLE_USERNAME / KAGGLE_KEY) so KerasHub can download the preset. The KerasHub Gemma docs list the available presets and sizes; Google's own LoRA tuning guide is the reference this tutorial follows.

We train on JAX and export for TensorFlow serving - the pattern from our Keras 3 multi-backend tutorial:

import os
os.environ["KERAS_BACKEND"] = "jax"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.9"   # let JAX use the GPU

import keras
import keras_hub

Load the model

gemma = keras_hub.models.GemmaCausalLM.from_preset("gemma2_2b_en")
gemma.summary()

GemmaCausalLM bundles the tokenizer, the preprocessor, and the backbone, so you can prompt it immediately:

print(gemma.generate("Explain what a LoRA adapter is in two sentences.", max_length=128))

Pick the smallest preset that can plausibly do the job. A 2B-class model fine-tunes on a single 24 GB GPU with LoRA; larger presets need model parallelism (see keras.distribution) or a TPU.

Prepare the dataset

Fine-tuning data is prompt/response pairs. Keep the format consistent, because the model will learn the format as much as the content:

import json

def load_examples(path):
    with open(path) as f:
        for line in f:
            ex = json.loads(line)
            yield (
                f"Instruction:\n{ex['instruction']}\n\n"
                f"Response:\n{ex['response']}"
            )

train_texts = list(load_examples("train.jsonl"))
val_texts = list(load_examples("val.jsonl"))

A few hundred high-quality examples beat tens of thousands of noisy ones. Hold out a validation set and a separate evaluation set you never train or tune on.

Set the sequence length on the preprocessor to the longest example you actually need; shorter is dramatically cheaper:

gemma.preprocessor.sequence_length = 512

Enable LoRA

LoRA freezes the original weights and trains small low-rank adapter matrices inside the attention projections. KerasHub turns it on with one call:

gemma.backbone.enable_lora(rank=8)
gemma.summary()   # trainable params drop from billions to a few million

Rank 4-16 covers most domain-adaptation tasks. Higher rank means more capacity and more risk of overfitting a small dataset; start at 8.

Train

gemma.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.AdamW(learning_rate=5e-5, weight_decay=0.01),
    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)

gemma.fit(
    train_texts,
    validation_data=val_texts,
    batch_size=2,
    epochs=2,
    callbacks=[keras.callbacks.EarlyStopping(patience=1, restore_best_weights=True)],
)

Practical notes from doing this for clients:

Learning rate for LoRA is higher than full fine-tuning; 1e-5 to 1e-4 is the range.
Mixed precision. keras.mixed_precision.set_global_policy("mixed_bfloat16") before loading halves memory on hardware that supports bfloat16.
Gradient accumulation if batch size 2 is all that fits: AdamW(..., gradient_accumulation_steps=8).
Watch validation loss, not train loss. Language models memorize small datasets within a couple of epochs.

Evaluate like you mean it

Loss is not quality. Before you call a fine-tune done, score it on the held-out evaluation set with a task-appropriate metric (exact match, F1, or a rubric scored by a human sample), and compare against the base model with the same prompt format. If the fine-tune does not beat the base model by a margin you can defend, you have a data problem, not a training problem.

for prompt in eval_prompts[:5]:
    print(gemma.generate(prompt, max_length=256))
    print("-" * 80)

Save and export

Save the fine-tuned model (LoRA weights merged into the preset) in the backend-neutral format:

gemma.save_to_preset("./gemma2_2b_domain")

To serve with TF Serving, reload on the TensorFlow backend and export a SavedModel:

# KERAS_BACKEND=tensorflow
import keras_hub
gemma = keras_hub.models.GemmaCausalLM.from_preset("./gemma2_2b_domain")
gemma.export("./serving/gemma_domain/1")

Alternatively, serve the preset directly from a Python service (FastAPI plus generate) on the JAX backend; for low-traffic internal tools that is simpler and just as good. For on-device use, the Google AI Edge stack has its own LLM inference path for Gemma-class models; that is a separate conversion pipeline and out of scope here.

Cost expectations

Rough numbers for a 2B-class model with LoRA, a few thousand examples at 512 tokens, two epochs:

Single 24 GB GPU (e.g. an L4 or similar): an hour or two of training; single-digit dollars of cloud time.
Single high-end 80 GB GPU: minutes to tens of minutes; you can raise batch size and sequence length.
TPU v5e pod slice with keras.distribution: the way to go for 9B-class and larger presets, or for frequent retraining.

The expensive part is never the compute. It is building the evaluation set and the data-cleaning pipeline, and that is where we tell clients to spend their budget first.

Keeping it reproducible

Treat the fine-tune as a pipeline, not a notebook. Pin keras, keras-hub, and the backend in a lockfile; record the preset name, LoRA rank, learning rate, sequence length, and dataset hash alongside the saved preset; and keep the evaluation script in version control so the next retrain produces a comparable score. Most of the fine-tunes we are asked to rescue went wrong not in training but in nobody being able to say which data and settings produced the model that is in production.

When not to do this

If the task needs frontier-model capability (multi-step reasoning across long documents, broad world knowledge), a 2B fine-tune will disappoint, and we will say so in the first call. If your team is PyTorch-native, Keras 3 runs on PyTorch too, but you may be happier in that ecosystem's native tooling. See our Generative AI on the TensorFlow Stack page for how we scope these projects - and contact us if you want a feasibility sprint on your data.