google-deepmind / google-deepmind/gemma

Bug: Division by zero in sampling methods when temperature is 0.0

Open
#562 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
5.7k
Forks
1k
Avg merge
10h 33m
Merged PRs (30d)
2

Description

Description
While analyzing the sampling implementations in gemma/gm/text/_sampling.py, I identified a potential division by zero issue when temperature is set to 0.0.

The RandomSampling, TopkSampling, and TopPSampling classes all perform division by self.temperature without checking if it is zero. This causes inf or NaN values in the logits, leading to undefined behavior or crashes.

Affected Components
RandomSampling.get_next_tokens (line 61)
TopkSampling.get_next_tokens (line 78)
TopPSampling.get_next_tokens (line 95)
Expected Behavior
When temperature is 0.0 (or effectively zero), the sampler should theoretically behave deterministically, equivalent to Greedy sampling.

Proposed Solution
Add a guard clause to these methods to check if self.temperature is below a small epsilon (e.g., 1e-6). If so, delegate to Greedy().get_next_tokens(...) to ensure stable and deterministic output.

Reproduction
python :
import jax
from gemma.gm.text import _sampling
logits = jax.numpy.array([[1.0, 2.0]])
rng = jax.random.PRNGKey(0)
sampler = _sampling.RandomSampling(temperature=0.0)
# This results in division by zero
sampler.get_next_tokens(logits, rng)

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.