DEV Community

Jaydeep Shah (JD)
Jaydeep Shah (JD)

Posted on

My Fine-Tuned Gemma 4 Loaded Fine, Then Broke on the First Message

#edgeai #android #litertlm #finetuning

I fine-tuned Gemma 4 E2B. The adapter merged cleanly. The export to .litertlm completed without errors. I pushed the model to my phone, initialized the engine, and everything looked green. Then I tried to create a conversation and got this:

Failed to apply template: unknown method: map has no method named get (in template:238)
Enter fullscreen mode Exit fullscreen mode

No model loading failure. No quantization error. The model initialized, the tokenizer loaded, and then the runtime choked on a Jinja template feature it does not support. This failure only surfaces when you actually try to run inference, not when you load the model. If you are demoing at a hackathon, this is the worst possible time to discover a compatibility issue.

I hit this exact bug while building Redacto, a zero-trust PII redaction app that runs Gemma 4 E2B entirely on-device. This post walks through the full fine-tune-to-deploy pipeline: how to QLoRA a model on Colab, export it for LiteRT-LM, and avoid the undocumented template trap that will block your deployment.

The Full Pipeline

Here is what the fine-tune-to-deploy pipeline looks like end to end:

HuggingFace base weights
    -> QLoRA fine-tune (Colab)
    -> Merge adapter into base
    -> Patch chat template   <-- the step nobody tells you about
    -> Quantize + export to .litertlm
    -> Push to device
Enter fullscreen mode Exit fullscreen mode

Each stage has its own failure modes. The template patch step is the one that was undocumented at the time, and it is the one that will cost you hours if you do not know it exists.

A note on framing before we dig in: this was an under-resourced fine-tune. I trained on 3,000 of the 400,000 samples in the ai4privacy/pii-masking-400k dataset for a single epoch, and the label format did not fully match what Redacto expected downstream. The point of this post is not the fine-tune's accuracy - it is the deployment mechanics I had to work through to get any fine-tuned model onto the device at all.

Step 1: QLoRA Fine-Tuning on Colab

QLoRA (Quantized Low-Rank Adaptation) lets you fine-tune a quantized model by training small adapter matrices instead of updating all parameters. Gemma 4 E2B is 5.1B total parameters but only 2.3B effective (the E stands for effective, achieved via Per-Layer Embeddings). That small effective footprint is why a single Colab GPU can handle the fine-tune.

Here is the training configuration that worked for us:

from peft import LoraConfig, get_peft_model, TaskType

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    task_type=TaskType.CAUSAL_LM,
)

model = get_peft_model(base_model, lora_config)
Enter fullscreen mode Exit fullscreen mode

The key parameters:

  • Rank 8 - the rank of the low-rank adapter matrices. Lower rank means fewer trainable parameters and faster training, at the cost of expressiveness. Rank 8 is a common starting point for small models.
  • Alpha 16 - the scaling factor applied to the adapter output. A ratio of alpha/rank = 2 is standard practice. Higher alpha amplifies the adapter's contribution relative to the frozen base weights.
  • Dropout 0.05 - light regularization on the adapter layers. With only 3,000 training examples, some regularization helps prevent overfitting.
  • Target modules - we targeted all the attention projection layers (q/k/v/o_proj) plus the MLP layers (gate/up/down_proj). This gives the adapter influence over both how the model attends to context and how it transforms representations. For Gemma 4's architecture, these module names include a .linear suffix in the full path.

With this configuration, the adapter adds 2,850,816 trainable parameters, 0.06% of the model's total. I trained on 3,000 examples from the ai4privacy/pii-masking-400k dataset (out of 400,000 available) for 1 epoch on a Colab NVIDIA RTX PRO 6000. Total training time: 217 seconds. Final training loss: 4.3064. This was deliberately a small, quick pass, not a rigorous training run - quantity and alignment of the training data both matter, and with the full dataset the outcome could have looked very different.

After training, merge the adapter back into the base model:

from peft import PeftModel

model = PeftModel.from_pretrained(base_model, adapter_path)
merged_model = model.merge_and_unload()
merged_model.save_pretrained("./merged_model")
tokenizer.save_pretrained("./merged_model")
Enter fullscreen mode Exit fullscreen mode

That last line, tokenizer.save_pretrained(), is where the trap is planted. More on that in a moment.

Step 2: Export to .litertlm

Once you have a merged model directory, export it using LiteRT's conversion tool:

python -m litert_torch.generative.export_hf \
  --hf_model_dir ./merged_model \
  --quantize dynamic_wi4_afp32 \
  -p 256 \
  --cache_length 1024 \
  --externalize_embedder \
  --output gemma4_ft.litertlm
Enter fullscreen mode Exit fullscreen mode

The flags:

  • dynamic_wi4_afp32 - INT4 weight quantization with FP32 activations. This is the standard quantization scheme for LiteRT-LM models. It compresses the model aggressively while keeping activation precision high enough for quality inference.
  • -p 256 - prefill chunk size. Controls how many tokens are processed in parallel during the initial prompt encoding.
  • --cache_length 1024 - deployed KV-cache size. This sets the maximum number of in-flight conversation tokens this build can hold, and it is distinct from the model's trained context length. Gemma 4 E2B is trained for a 128K context, but I compiled this build with a 1024-token cache to keep memory usage in check. Do not confuse the deployed cache size with the model's context window.
  • --externalize_embedder - separates the embedding table from the main model bundle for memory efficiency during loading.

This export completes without errors. The .litertlm file is generated. You push it to the device:

adb push gemma4_ft.litertlm \
  /sdcard/Android/data/com.example.starterhack/files/gemma4_ft.litertlm
Enter fullscreen mode Exit fullscreen mode

The engine initializes. No crash. You think you are done. You are not.

One more thing worth flagging about this export: the resulting fine-tuned .litertlm came out to about 4.7 GB and was GPU-only. It could not be compiled down to the NPU path, so this build never got the NPU acceleration that the pre-compiled community models get. That is a real constraint of the fine-tune-and-export route as it stands today, separate from the template trap below.

Step 3: The Chat Template Trap

When you call tokenizer.save_pretrained() during the adapter merge step, it writes out the tokenizer files, including tokenizer_config.json and a standalone chat_template.jinja file. These files contain the chat template: the Jinja2 template that formats user/assistant messages into the token sequence the model expects.

The problem: HuggingFace's Gemma 4 repository (google/gemma-4-E2B-it) ships a chat template that uses map.get(), a Jinja2 method for safely accessing dictionary values with a default fallback. This is perfectly valid Jinja2. Python's Jinja2 library handles it fine. HuggingFace's transformers library handles it fine.

LiteRT-LM's on-device Jinja parser does not support it.

The .litertlm export tool reads the chat template from your model directory and embeds it into the bundle. When the LiteRT-LM runtime on Android tries to apply the template to format a conversation, it fails:

Failed to apply template: unknown method: map has no method named get (in template:238)
Enter fullscreen mode Exit fullscreen mode

This error has three properties that make it particularly nasty:

  1. It only appears when you create a conversation, not when you load the model. Engine initialization succeeds. The model loads into GPU memory. The tokenizer initializes. Everything looks green until you actually try to send a message.

  2. There is no warning during export. The export_hf tool does not validate whether the embedded template is compatible with the on-device parser. It bundles whatever template it finds in the model directory.

  3. It was not documented at the time. During the hackathon, I could not find any published documentation listing which Jinja2 features the LiteRT-LM template parser supports, so I discovered the boundaries by hitting them. I later compared notes with a Google engineer who confirmed the same approach: he independently had the exact same finding written down in his own notes, which was reassuring but also told me this was tribal knowledge rather than something you could look up.

Step 4: The Fix

The fix is to replace the Gemma 4 chat template with an older, compatible one before export. The Gemma 3 family's chat template (google/gemma-3-1b-it) avoids map.get() and uses only Jinja features that LiteRT-LM's parser supports. This makes sense: the litert-community pre-compiled models were built against Gemma 3-era templates.

Here is the re-export workflow:

import json
from huggingface_hub import hf_hub_download
import os

# 1. Download the compatible template from Gemma 3
compatible_config_path = hf_hub_download(
    repo_id="google/gemma-3-1b-it",
    filename="tokenizer_config.json"
)

with open(compatible_config_path, "r") as f:
    compatible_config = json.load(f)

compatible_template = compatible_config["chat_template"]

# 2. Patch the merged model's tokenizer config
merged_config_path = "./merged_model/tokenizer_config.json"
with open(merged_config_path, "r") as f:
    merged_config = json.load(f)

merged_config["chat_template"] = compatible_template

with open(merged_config_path, "w") as f:
    json.dump(merged_config, f, indent=2)

# 3. Remove the standalone template file (export_hf may prefer it over the config)
jinja_path = "./merged_model/chat_template.jinja"
if os.path.exists(jinja_path):
    os.remove(jinja_path)
    print("Removed standalone chat_template.jinja")

print("Template patched. Ready for re-export.")
Enter fullscreen mode Exit fullscreen mode

After patching, re-run the export:

python -m litert_torch.generative.export_hf \
  --hf_model_dir ./merged_model \
  --quantize dynamic_wi4_afp32 \
  -p 256 \
  --cache_length 1024 \
  --externalize_embedder \
  --output gemma4_ft_fixed.litertlm
Enter fullscreen mode Exit fullscreen mode

The resulting .litertlm file now contains a compatible template. Conversations work.

A note on template choice. What I did under hackathon pressure was borrow the Gemma 3 family's template, because I already knew it was compatible with the on-device parser. When I later compared notes with a Google engineer, he pointed at a cleaner source: rather than reaching back to a Gemma 3 template, pull the chat template that ships with the LiteRT-LM build of gemma-4-e2b-it (the litert-community packaged model), which is already known-good for the on-device parser, instead of the raw google/gemma-4-e2b HuggingFace template that caused the failure in the first place. Same idea I had arrived at (use a template the on-device parser is known to accept), but staying within the Gemma 4 family instead of borrowing from Gemma 3.

Trap 2: The Missing Vision Encoder

Once the template issue is fixed, you may hit a second wall:

LiteRtLmJniException: NOT_FOUND: TF_LITE_VISION_ENCODER not found in the model
Enter fullscreen mode Exit fullscreen mode

The fine-tuned .litertlm was compiled text-only: no vision encoder is included. But if your engine configuration specifies vision support (which it does by default for Gemma 4, a multimodal model), the runtime tries to find the vision encoder section in the bundle and fails.

The fix is to configure the engine for text-only mode. In Kotlin:

fun initializeEngine(modelPath: String, textOnly: Boolean = false) {
    val engineConfig = EngineConfig.builder()
        .setModelPath(modelPath)

    if (textOnly) {
        engineConfig
            .setVisionBackend(null)
            .setAudioBackend(null)
            .setMaxNumImages(null)
    } else {
        engineConfig
            .setVisionBackend(Backend.CPU())
            .setMaxNumImages(1)
    }

    engine = LlmEngine.create(engineConfig.build())
}
Enter fullscreen mode Exit fullscreen mode

Note the setMaxNumImages(null), not setMaxNumImages(0). Setting it to 0 triggers a separate validation error: max number of images must be positive or null. The engine treats 0 as an invalid value, not as "no images." You must pass null to indicate the model does not handle images.

Why This Matters Beyond My Project

The chat template trap is not specific to Redacto or to PII redaction. It affects anyone who fine-tunes a Gemma 4 model (or any model whose HuggingFace template uses Jinja features beyond LiteRT-LM's supported subset) and tries to deploy it on-device. Redacto was built for the Qualcomm x Google LiteRT Developer Hackathon 2026, and this trap was one of the last things standing between a green engine init and a working demo.

The root cause is an impedance mismatch between two ecosystems:

  • HuggingFace's transformers library uses Python's full Jinja2 implementation. Template authors can use any Jinja2 feature: filters, methods, macros, the full spec.
  • LiteRT-LM's on-device runtime uses a lightweight Jinja parser written for embedded/mobile use. It supports a subset of Jinja2, but the subset is not documented anywhere.

When you fine-tune through the standard HuggingFace workflow (AutoModelForCausalLM.from_pretrained(), train with peft, merge with merge_and_unload(), save with save_pretrained()) you get whatever chat template the upstream model ships. If that template uses an unsupported feature, your export will silently embed a broken template.

The pre-compiled models from litert-community on HuggingFace do not have this problem because they were compiled with compatible templates. The problem only surfaces when you fine-tune and re-export yourself.

This is the kind of issue that has no answer in forum posts and community discussions. The error message (map has no method named get) does not appear in Google's documentation. The fix (swap the template with an older version) is unintuitive: you are replacing a component of your Gemma 4 model with one from Gemma 3, and it works because the actual template logic is functionally equivalent, just expressed without map.get().

Checklist: Fine-Tune to LiteRT-LM Deployment

For anyone following this path, here is the condensed checklist:

  1. Fine-tune with QLoRA on Colab (rank 8, alpha 16 is a solid starting point for a model with Gemma 4 E2B's effective footprint)
  2. Merge the adapter with merge_and_unload()
  3. Save the merged model with save_pretrained()
  4. Patch the chat template: download tokenizer_config.json from google/gemma-3-1b-it, extract its chat_template, overwrite your merged model's template
  5. Delete chat_template.jinja from the merged model directory if it exists
  6. Export with litert_torch.generative.export_hf using dynamic_wi4_afp32
  7. If the fine-tuned model is text-only, set visionBackend = null, audioBackend = null, maxNumImages = null in the engine config
  8. Test by creating a conversation and sending a message: do not assume init success means inference will work

What I Would Like to See

LiteRT-LM is a powerful runtime for on-device LLMs. The NPU performance on Snapdragon 8 Elite is genuinely impressive. But the fine-tuning deployment story has gaps:

  • Document the supported Jinja2 subset. Even a simple list of supported filters, methods, and control structures would save developers hours. At the time I hit this, the only way to discover the boundaries was to hit them at runtime.
  • Validate templates during export. The export_hf tool could parse the chat template against the same Jinja parser the runtime uses and flag incompatible features before producing the .litertlm file.
  • Ship compatible templates with the export toolchain. If the tool detects that the model's template uses unsupported features, it could offer to substitute a known-compatible template for the same model family.

The pre-compiled models from litert-community work out of the box. The moment you fine-tune, you are on your own. That gap is the engineering story of this blog post.

Related in this series of "Edge AI from the Trenches"

Jaydeep Shah is a developer with roots in embedded systems, Android platform internals, and silicon-level AI optimization. He now explores on-device AI inference - bringing models from the cloud to phones and edge hardware. Along with his team Edge Artists, he builds applications using LiteRT-LM and Gemma models on mobile hardware, and writes about what works, what breaks, and what he learns along the way. This post is part of the Edge AI from the Trenches series.

Sources:

Last updated: July 2026
16th of 23 posts in the "Edge AI from the Trenches" series

Top comments (0)