AI-Hypercomputer / AI-Hypercomputer/maxtext

Gemma3 Model Token Logits Mismatch with Different max_prefill_predict_length

Open
#2,179 2 comments 0 reactions 1 assignee Claimed by @gagika View on GitHub
bug
Dominant language
Python
Stars
2.4k
Forks
607
Avg merge
2d 19h
Merged PRs (30d)
158

Description

### Bug report

# Gemma3 Model Token Logits Mismatch with Different `max_prefill_predict_length`

## Description
When running **prefill** with the Gemma3 4B model, the output token logits differ depending on the `max_prefill_predict_length` value.
In contrast, other models (e.g., Qwen3, LLaMA 2, LLaMA 3) produce consistent logits regardless of the prefill length.

## Steps to Reproduce
1. Load the Gemma3 4B model.
2. Use the same input prompt: The capital of France is
3. Run inference with:
- **Case 1:** `max_prefill_predict_length = 1024`, `max_target_length = 2048`
- **Case 2:** `max_prefill_predict_length = 4096`, `max_target_length = 8192`

## Observed Behavior (Gemma3 4B)

**Case 1:**
`[null, {"The": -1.68888}, {" capital": -20.57046}, {" of": -0.15805}, {" France": -3.24782}, {" is": -1.67389}] `
`Full output: The capital of France is Paris.\n\nParis is a global center for art `

**Case 2:**
`[null, {"The": -1.68888}, {" capital": -20.30916}, {" of": -0.16284}, {" France": -3.24021}, {" is": -1.64998}] `
`Full output: The capital of France is Paris�Ar¡ 3--- 3 `

Note: The very first token logits matches! Remaining tokens start to divergence. Strongly suspect related to kv cache.

Note: The small cumulative logit differences lead to large divergence in generated output for longer sequences.

## Expected Behavior
Token logits for the same input prompt should remain **identical** across different `max_prefill_predict_length` values, as observed with Qwen3, LLaMA 2, and LLaMA 3 models.

## Additional Information
- **Affected Model:** Gemma3 4B
- **Verified Models (working as expected):** Qwen3, LLaMA 2, LLaMA 3

### Logs/Output

_No response_

### Environment Information

_No response_

### Additional Context

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.