Most teams that ask us to fine-tune a language model do not need a fine-tuned language model. They need a retrieval step, a better prompt, and a scoring harness. Fine-tuning earns its place in a narrower set of cases: when you need a consistent output format, a domain vocabulary the base model fumbles, a smaller model that matches a larger one on one narrow task, or latency and cost that a hosted API cannot hit. That is a real list, and KerasHub makes the work tractable on a single GPU. This post is the path we follow — with the numbers we check before spending the budget.
Price the alternatives first
Fine-tuning is the third thing to try, not the first. Build the evaluation harness before anything else: 150–300 real examples with the outputs you actually want, and a scoring function that is either exact-match, a regex/schema check, or a rubric applied by a stronger model. Without it you cannot tell a good adapter from a confident one.
Then measure, in order: the base model zero-shot; the base model with a tightened prompt and 3–5 few-shot examples; the base model with retrieval if the failures are knowledge failures. Write the three numbers down. If prompting closes 90% of the gap, fine-tuning is buying you the last 10% at the cost of a training pipeline, an artifact to version, and a regression surface. Sometimes that is worth it. You should know which case you are in.
Why LoRA, concretely
Full fine-tuning of a 2B-parameter model in float32 needs weights, gradients, and two Adam moments — roughly 16 bytes per parameter, about 32 GB before activations. That does not fit on a 24 GB card. LoRA (low-rank adaptation) freezes the pretrained weights and trains small rank-r matrices injected into the attention projections. For a 2B model at rank 4, you are training on the order of a few million parameters instead of two billion, so optimizer state becomes negligible and the frozen weights can sit in bfloat16. The same job now fits comfortably, and the artifact you ship is a few megabytes of adapter weights rather than a multi-gigabyte checkpoint.
The trade-off is capacity. LoRA adapts behavior and format well. It is a poor tool for injecting large volumes of new knowledge — that is what retrieval is for.
The setup
Pick the backend before you write code. KerasHub runs on JAX, TensorFlow, or PyTorch, and the choice is set by an environment variable before keras is imported:
import os
os.environ["KERAS_BACKEND"] = "jax"
import keras
import keras_hub
keras.mixed_precision.set_global_policy("bfloat16")
model = keras_hub.models.CausalLM.from_preset("gemma_instruct_2b_en")
model.preprocessor.sequence_length = 256
Two notes on that snippet. First, bfloat16 rather than float16: it has the same exponent range as float32, so loss scaling is unnecessary and the NaN failure mode that bites float16 training largely disappears. Second, sequence_length is the single biggest lever on memory in the whole script. Attention activations grow with sequence length, and the default preset value is often far longer than your data needs. Measure the token-length distribution of your training set and set the cap at roughly the 95th percentile — padding every example to 512 when 190 covers almost all of them can easily double memory for no benefit.
Enabling LoRA is one call on the backbone:
model.backbone.enable_lora(rank=4)
model.summary()
Read the trainable-parameter count in that summary before you start training. If it reads in the billions, LoRA did not take effect and you are about to OOM or, worse, silently train everything.
Hyperparameters that actually matter
Four knobs, in order of impact:
- Rank. Start at 4. Raise to 8 or 16 only if training loss plateaus above where you need it — higher rank means more capacity and more overfitting risk on a few hundred examples. Ranks above 32 rarely pay for themselves on formatting and style tasks.
- Learning rate. LoRA tolerates learning rates one to two orders of magnitude higher than full fine-tuning.
1e-4to5e-4with AdamW is the usual range;1e-5, borrowed from full-model habits, mostly wastes epochs. - Epochs. One to three. With a few hundred examples, loss keeps dropping long after the outputs have started to degrade into memorized copies of the training set. Check generations, not just loss.
- Weight decay exclusions. Exclude biases and normalization scales from decay — a small detail that avoids a class of quiet degradation.
optimizer = keras.optimizers.AdamW(learning_rate=2e-4, weight_decay=0.01)
optimizer.exclude_from_weight_decay(var_names=["bias", "scale"])
model.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=optimizer,
weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
model.fit(train_ds, epochs=2)
Format your training data exactly as you will prompt at inference — the same instruction template, the same delimiters, the same whitespace. A mismatch between training format and serving format is the most common reason a fine-tune that looked fine in a notebook underperforms behind the API. Mask the loss to the response tokens if your harness supports it; training the model to predict its own prompts wastes capacity.
Evaluate against the recorded baselines
Run the harness you built in step one against the adapter, and compare to the three numbers you wrote down. Two checks beyond the headline score:
- Regression on general behavior. Narrow fine-tuning can degrade capabilities you never tested. Keep a small held-out set of off-task prompts and confirm the answers are still sane.
- Format compliance rate. If the goal was valid JSON, measure the percentage that parses. It is usually the number the business actually cares about, and it is often the one that moves most.
If the adapter does not clearly beat the prompt-engineered baseline, ship the prompt. That is a successful outcome — you spent a day instead of a quarter.
Serving and rollback
The deployment story is the reason LoRA is pleasant in production. The adapter weights are small and separable, so you keep one copy of the base preset and version adapters against it:
model.backbone.save_lora_weights("adapters/support-fmt-v3.lora.h5")
Record the preset name, the KerasHub and Keras versions, the rank, the sequence length, and the training data commit alongside the file. Without the preset and rank, an adapter is an unloadable blob six months from now. Rolling back is then swapping a small file, not redeploying a model server.
If you need maximum inference throughput and have fixed on one adapter, merge the LoRA weights into the base weights and export a single artifact — you lose cheap rollback and gain a simpler serving path. Decide that consciously; the default should be to keep them separate until latency measurements say otherwise.
The short version
Build the eval harness. Measure prompting and retrieval first. Then LoRA at rank 4, bfloat16, sequence length set from your data's token distribution, learning rate around 2e-4, one to three epochs, and a trainable-parameter count you verified before training. Ship the adapter with its metadata. The discipline is ordinary; it is also the difference between a fine-tune you can maintain and one nobody dares touch.