Fine-Tuning Gemma 3 270M with Keras, JAX, and LoRA on Medical Q&A

A practical guide to parameter-efficient fine-tuning with LoRA on GCP using modern Keras and JAX tooling

Fine-tuning Gemma 3 with Keras and JAX for a Medical Q&A task — image obtained with Gemini by the author

Introduction

I wanted to see how far I could push a relatively small model in a specialized domain without turning the whole thing into an infrastructure project. The goal was simple: take Gemma 3 (270M), fine-tune it on a medical Q&A dataset, and make it produce answers that feel at least somewhat grounded in that domain.

You are not a Member of Medium? Here is a Friendly Link for you.

The constraint was just as important as the goal. I didn’t want to rely on huge GPUs, distributed training, or a stack of frameworks glued together. The idea was to keep it practical: one GPU, a managed environment, and tooling that doesn’t fight you too much. That’s how I ended up with Keras 3, a JAX backend, and LoRA for the actual fine-tuning.

The whole thing runs on a Google Cloud Agent Platform (former Vertex AI) Workbench instance, mostly because it removes a lot of the friction around environment setup and GPU access. You spin up an instance, pick an instance with around 32GB RAM and one L4 and you are ready to go.

Note: The article and the GitHub repo associated were created as part of a project in the frame of TPU Sprint.

Resources and prerequisites

The project will fine-tune the Gemma 4 270M language model using LoRA, Keras, and a JAX backend. The notebook is available in a GitHub project:

[embed]GitHub - gabrielpreda/fine-tune-gemma3-keras-jax-medical-q-a: Fine tune Gemma 3 270M with Keras &… Fine tune Gemma 3 270M with Keras & JAX for Medical Q&A - gabrielpreda/fine-tune-gemma3-keras-jax-medical-q-agithub.com

After spining out a Agent Platform Workbench instance (I used one with 8 vCPUs, 32GB RAM, 1 NVIDIA L4 GPU), start a Terminal and clone the GitHub project:

 git clone https://github.com/gabrielpreda/fine-tune-gemma3-keras-jax-medical-q-a.git

Next, you will need to download the dataset. It is a Kaggle dataset and, to avoid downloading and uploading, you can use Kaggle CLI to do it programatically. If it is the first time you will use it, let’s see what are the steps to install and use it.

Step 1 — Generate a Kaggle API Token

1. Log in to kaggle.com.
2. Click on your profile avatar (top-right) → Settings.
3. Scroll to the API section and click Create New Token.
4. A file named kaggle.json will be downloaded to your computer. It contains your credentials in this format:

{
  "username": "your_kaggle_username",
  "key": "your_kaggle_api_key"
}

Step 2 — Copy the Token to the GCP Instance

You need to place kaggle.json in the ~/.kaggle directory on the Workbench instance.

Option A — Copy via JupyterLab Upload + Terminal

1. In JupyterLab, use the file browser to upload kaggle.json to your home directory.
2. Open a Terminal tab and run:

mkdir -p ~/.kaggle
mv ~/kaggle.json ~/.kaggle/kaggle.json
chmod 600 ~/.kaggle/kaggle.json

Option B — Paste content directly in the terminal

1. Open a Terminal tab in JupyterLab.
2. Create the .kaggle directory and the file:

mkdir -p ~/.kaggle
cat > ~/.kaggle/kaggle.json << 'EOF'
{"username":"your_kaggle_username","key":"your_kaggle_api_key"}
EOF
chmod 600 ~/.kaggle/kaggle.json

Replace your_kaggle_username and your_kaggle_api_key with the actual values from your downloaded kaggle.json.

Security note: The chmod 600 command restricts the file so only your user can read it, which is required by the Kaggle CLI.

Installing the Kaggle CLI

In a JupyterLab terminal on the Workbench instance, install the Kaggle CLI using pip:

pip install -upgrade kaggle

Verify the installation:

kaggle -version

Downloading the Dataset

The notebook uses the MedQuAD dataset hosted on Kaggle at:

[embed]MedQuAD Medical Questions and Answerswww.kaggle.com

To download it locally on the GCP instance, open a Terminal and run:

# Navigate to the project directory (or wherever you want the data)
cd ~/fine-tune-gemma3-keras-jax-medical-q-a
# Download and unzip the dataset
kaggle datasets download -d gpreda/medquad - unzip

This will place medquad.csv in the current directory. The notebook expects this file to be present in the same directory from which the notebook is opened (i.e., the working directory at runtime).

You are now ready to run the Notebook and start finetuning the model.

Install the updated libraries

Before starting to run the Notebook cells, we will need to also install updated libraries.

1. In JupyterLab open finetune-gemma-3–270m-using-lora-for-medical-q-a-gcp-run.ipynb.

2. Run the first two code cells to install/update the required packages: keras-nlp, keras, jax[cuda12], and kagglehub.

!pip install -q -U keras-nlp keras
!pip install -q -U "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

!pip install -q -U kagglehub

3. Restart the kernel when prompted (there is a prominent reminder cell in the notebook).

Load the model

With Keras 3, we can run workflows on one of three backends: TensorFlow, JAX, and PyTorch.

For this Notebook, we will configure the backend for JAX. We are also setting the parameter XLA_PYTHON_CLIENT_MEM_FRACTION to avoid memory fragmentation on JAX backend.

import os
os.environ["KERAS_BACKEND"] = "jax" # you can also use tensorflow or torch
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.90" # avoid memory fragmentation on JAX backend.
os.environ["JAX_PLATFORMS"] = ""

Once installed the additional packages and configured JAX backend, we can import the packages we updated.

import keras
import keras_nlp
from keras_nlp.samplers import TopKSampler
from time import time
import csv

print("KerasNLP version: ", keras_nlp.__version__)
print("Keras version: ", keras.__version__)

We load the model using Gemma3CausalLM from keras_nlp. In order to succesfully load from preset the model, kaglehub should be installed.

gemma_lm = keras_nlp.models.Gemma3CausalLM.from_preset("gemma3_270m")

Note: the model is available with different frameworks and on various platforms, including Kaggle, HuggingFace, Unsloth, Ollama, and Agent Platform (former Vertex AI). Here is the path to Kaggle Models Gemma 3 model card (for Keras framework):

[embed]Keras | Gemma 3 | Kaggle Keras implementation of the Gemma3 model. Can run on JAX, TensorFlow and PyTorch.www.kaggle.com

Running model.summary() we can inspect the preprocessor and model parameters.

Model summary shows 270M parameters, all trainable — image by author

Prepare the data

We already downloaded locally the dataset for finetuning using Kaggle CLI. After reading the csv, we select a subset for training and validation and one for testing.

data = []

# The CSV file contains two columns 'question' and 'answer'
with open("medquad.csv", mode='r', encoding='utf-8') as file:
    reader = csv.DictReader(file)
    for row in reader:
        # we replace with 'prompts' and 'responses'
        data.append({"prompts": row['question'], 'responses': row['answer']})

Using tensorflow.data.Dataset, we prepare the datasets for training and testing. From the training dataset we keep a subset only for validation.

import tensorflow as tf

dataset = tf.data.Dataset.from_generator(
    lambda: (item for item in train_data),
    output_signature={
        "prompts": tf.TensorSpec(shape=(), dtype=tf.string),
        "responses": tf.TensorSpec(shape=(), dtype=tf.string),
    }
)

test_dataset = tf.data.Dataset.from_generator(
    lambda: (item for item in test_data),
    output_signature={
        "prompts": tf.TensorSpec(shape=(), dtype=tf.string),
        "responses": tf.TensorSpec(shape=(), dtype=tf.string),
    }
)

We prepare the data batches for training and validation:

train_dataset = dataset.skip(valid_size)
val_dataset = dataset.take(valid_size)

train_dataset = (
    train_dataset
    .shuffle(train_size)
    .batch(cfg.batch_size)
    .repeat()
)
val_dataset = (
    val_dataset
    .batch(cfg.batch_size)
    .repeat()
)

Fine-tuning using LoRA

With Keras, fine-tuning with LoRA is made very simple. First we need just enable LoRA:

gemma_lm.backbone.enable_lora(rank=cfg.lora_rank)

If we run again model.summary(), we can see that now a number of parameters are trainable (less than 0.2%), while the majority of parameters will be frozen. If we substract the number of trainable parameters from the total number of parameters, we see that the trainable parameters were added to the initial, now frozen, paramters.

Model summary after enabling LoRA shows 534K trainable parameters (less than 0.2% from total), and 270M non-trainable parameters— image by author

After we set the sequence length and the optimizer, we compile the model. We are using SparseCategoricalCrossentropy loss and set asSparseCategoricalAccuracy as metric.

# Limit the input sequence length to 512 (to control memory usage).
gemma_lm.preprocessor.sequence_length = cfg.max_length
# Use AdamW (a common optimizer for transformer models).
optimizer = keras.optimizers.AdamW(
    learning_rate=5e-5,
    weight_decay=0.01,
)
# Exclude layernorm and bias terms from decay.
optimizer.exclude_from_weight_decay(var_names=["bias", "scale"])

gemma_lm.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizer,
    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)

Then, we start the training.

history = gemma_lm.fit(train_dataset,
                       validation_data=val_dataset,
                       epochs=cfg.epochs,
                       steps_per_epoch=steps_per_epoch,
                       validation_steps=validation_steps,
                      )

In the Notebook we set to run for 10 epochs. You can rerun with different number of epochs, sequence length, learning rate, and weight decay. As well, you can experiment with various data dimmensions.

The next two figures capture the train and validation loss, as well as the train and validation loss.

Train and validation loss (10 epochs) — image by author

Train and validation accuracy — image by author

Model evaluation

To evaluate the model, the metric used for training and validation (SparseCategoricalAccuracy) is irelevant. Since we generate text for answers, and we want to check if its meaning is similar with the expected value (available in the dataset).

We first prepared the test predictions and references:

from tqdm import tqdm

predictions = []
references = []

for batch in tqdm(test_dataset):
    outputs = gemma_lm.generate(batch["prompts"])

    predictions.extend(outputs)
    references.extend(batch["responses"])

In the Notebook we used two metrics for evaluating the predictions against references:

• F1 scores
• Rouge scores (prediction, recall, fmeasure)

Here is the code for computation of f1 score and Rouge score, respectively:

from collections import Counter

def f1_score(pred, ref):
    pred_tokens = pred.split()
    ref_tokens = ref.split()

    common = Counter(pred_tokens) & Counter(ref_tokens)
    num_same = sum(common.values())

    if num_same == 0:
        return 0.0

    precision = num_same / len(pred_tokens)
    recall = num_same / len(ref_tokens)

    return 2 * precision * recall / (precision + recall)

from rouge_score import rouge_scorer

scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=True)

def rouge_scores(pred, ref):
    scores = scorer.score(ref, pred)
    return scores

The averaged scores are obtained with:

f1s = [f1_score(str(p), str(r)) for p, r in zip(predictions, references)]
r_s = [rouge_scores(str(p), str(r)) for p, r in zip(predictions, references)]

print("F1:", sum(f1s)/len(f1s))

precision_I = [r_s_item.get("rouge1").precision for r_s_item in r_s]
recall_I = [r_s_item.get("rouge1").recall for r_s_item in r_s]
f1_score_I = [r_s_item.get("rouge1").fmeasure for r_s_item in r_s]

precisionL = [r_s_item.get("rougeL").precision for r_s_item in r_s]
recallL = [r_s_item.get("rougeL").recall for r_s_item in r_s]
f1_scoreL = [r_s_item.get("rougeL").fmeasure for r_s_item in r_s]

print("RougeI precision: ", sum(precision_I)/len(precision_I))
print("RougeI recall: ", sum(recall_I)/len(recall_I))
print("RougeI f1_score: ", sum(f1_score_I)/len(f1_score_I))

print("RougeL precision: ", sum(precisionL)/len(precisionL))
print("RougeL recall: ", sum(recallL)/len(recallL))
print("RougeL f1_score: ", sum(f1_scoreL)/len(f1_scoreL))

Let’s look also to one example of model response, after fine-tuning.

Question: What are the complications of Paget’s Disease of Bone ?

Answer: Paget’s disease of bone can cause bone pain, bone loss, and bone deformity. It can also cause bone fractures. Paget’s disease of bone can be caused by a number of things. It can be caused by a disease, such as a virus, cancer, or infection. It can also be caused by a reaction to certain medicines. Paget’s disease of bone can be difficult to diagnose. It is usually diagnosed by a bone scan. Treatment for Paget’s disease of bone depends on the cause. It may include medicines, surgery, or bone marrow transplant. NIH: National Institute of Arthritis and Musculoskeletal and Skin Diseases

Final remarks

What I found most interesting about this project wasn’t the model itself, but how many small details matter in practice. The environment setup, the dataset pipeline, when to repeat, how to shuffle, when to restart the kernel — none of these are complicated individually, but together they determine whether the whole thing works smoothly or becomes frustrating.

Keras with JAX is surprisingly capable, but it’s not as forgiving as some other stacks. You have to be explicit about things like backends and execution flow. On the other hand, once it works, it’s clean and fast.

If I were to extend this, I’d probably focus on evaluation first. Adding a stronger semantic metric like BERTScore or even using another model as a judge would give a better sense of actual quality. I’d also try a slightly larger Gemma variant just to see how much the gains justify the extra cost.

But even at 270M parameters, you can already get something useful out of this setup. That’s probably the most important takeaway. You don’t need a massive model or a complicated pipeline to adapt an LLM to a specific domain — you just need to be careful with the details.