AI-Hypercomputer / AI-Hypercomputer/maxtext

Gemma3 Model Token Logits Mismatch with Different max_prefill_predict_length

未關閉
#2,179 2 則留言 0 個 reaction 已指派 1 人 已被 @gagika 認領 在 GitHub 檢視
bug
主要語言
Python
星號
2.4k
分支
607
平均合併
2 天 19 小時
30 天內合併 PR
158

描述

### 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_

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。