Gemma 3 1B: Model is corrupted after converting to .tflite

[colfrog]

Description of the bug:

I have a fine-tuned version of the Gemma 3 1B model in safetensors format., fine-tuned with the transformers library. The fine-tuned model shows good results when executed with PyTorch.

However, after converting to .tflite and testing the model via the AI Edge Gallery application, the outputs are distressing. I'm not one to think that LLMs are alive, but if this one were it would be suffering.

Actual vs expected behavior:

The expected behaviour is a model that performs similarly between torch format and tflite format.

My fine-tuned model is a poetry assistant. This is the output of the model in torch format when prompted "Write a poem."

I’m not sure what I’m seeing when I see her—
the one that I thought I’d only ever come to know.
I’ve been standing here a little while, watching
her out of the corner of my eye. Like all other people,
she’s a person, she’s human, she’s a girl.
She’s wearing an awkward dress, awkward shoes—
and in the end, in my mind’s eye, she’s always going to be awkward,
a girl that nobody ever wants to like,
a girl always going to be awkward.
All of this, it seems, is wrong, and the thing is the girl
is what it’s like. It’s something she’s always been.
A girl that nobody ever wants to like.
A girl that nobody ever wants to write poems about.
It’s something she’s always been.
She’s human.
No, I’m telling myself something
that you need to remember that it’s all fine.
But what if I do want to?
I don’t see any difference between me
and her now. And my friends say, the girls,
that even if you think you know me
because I’m wearing the clothes I’m wearing,
I’m just as awkward as you,
and I just as awkward as you.
And that’s the thing about people,
in any generation,
that’s what it means,
what it means,
and it means nothing.
If I’ve only ever seen her, I will always see her.<end_of_turn>

And this is example output of the same prompt with the converted model in AI Edge Gallery

Any other information you'd like to share?

This is the script I'm converting the model with:

export TF_ENABLE_ONEDNN_OPTS=0 # To avoid numerical imprecisions, tested with and without
export CUDA_VISIBLE_DEVICES="-1" # Not enough VRAM to use the GPU, tested with and without
python -m ai_edge_torch.generative.examples.gemma3.convert_gemma3_to_tflite --checkpoint_path="gemma3_poetry_merged" --output_path="tflite-model-v7" --kv_cache_max_len=2048 # Also tested with --prefill_seq_lens=2048 and --kv_cache_max_len=4096. I tried multiple quantization methods.

[colfrog]

I found the issue. I was quantizing the model twice, once during training, saved quantized, then converted to tflite with quantization. Removing the quantization from the training code made the model perform as expected.