LLM.generate() can silently select a token excluded by top_k on CPU when sampling probabilities are float16

Bug description

On CPU, LLM.generate() can silently return a token whose probability is exactly zero when the probability tensor passed to sampling is float16.

In the deterministic example below, top_k=1 leaves exactly one legal token. For the prompt "The capital of France is", that token is ID 253 (" the"). With seed 20, however, the normal stochastic generation path returns ID 5833 (" complaint"). There is no warning or exception.

actual:   [5833]  " complaint"
expected: [253]   " the"

The example uses the public LLM.load(), LLM.distribute(), and LLM.generate() APIs. It does not modify LitGPT, the model architecture, or the checkpoint. The required torch.set_default_dtype(torch.float16) call is an explicit public PyTorch runtime setting; it is not LitGPT's default configuration. With the default float32 dtype, the same test passes.

This appears to be a downstream manifestation of PyTorch issue pytorch/pytorch#192812, rather than an independent LitGPT root cause. The proposed upstream fix, pytorch/pytorch#195077, is still open at the time of writing.

Reproduction

Run the following script. LLM.load() downloads and converts the small public EleutherAI/pythia-14m checkpoint when it is not already present.

import torch
from litgpt import LLM

torch.set_num_threads(1)
torch.set_default_dtype(torch.float16)

llm = LLM.load("EleutherAI/pythia-14m", distribute=None)
llm.distribute(accelerator="cpu", devices=1, precision="16-true")

prompt = "The capital of France is"
options = dict(max_new_tokens=1, top_k=1, top_p=1.0, return_as_token_ids=True)

torch.manual_seed(20)
actual = llm.generate(prompt, temperature=1.0, **options)

# With top_k=1, this is the only token retained by the top-k operation.
expected = llm.generate(prompt, temperature=0.0, **options)

print(
    {
        "torch": torch.__version__,
        "default_dtype": str(torch.get_default_dtype()),
        "actual": actual.tolist(),
        "expected": expected.tolist(),
        "actual_text": llm.preprocessor.decode(actual),
        "expected_text": llm.preprocessor.decode(expected),
    }
)

assert torch.equal(actual, expected), (
    f"top_k=1 selected {actual.tolist()} outside its unique candidate {expected.tolist()}"
)

Observed output:

Using 1 device(s)
Precision set
Fabric launched
{'torch': '2.14.0+cpu', 'default_dtype': 'torch.float16', 'actual': [5833], 'expected': [253], 'actual_text': ' complaint', 'expected_text': ' the'}
AssertionError: top_k=1 selected [5833] outside its unique candidate [253]

I reproduced this result in three fresh Python processes. A newly created environment and newly downloaded checkpoint produced the same result.

Float32 control

Change only the default dtype:

torch.set_default_dtype(torch.float32)

The same model, prompt, seed, precision setting, and generation arguments then produce:

{'torch': '2.14.0+cpu', 'default_dtype': 'torch.float32', 'actual': [253], 'expected': [253], 'actual_text': ' the', 'expected_text': ' the'}

LitGPT's completely automatic/default loading path also produced float32 sampling probabilities in this environment and did not violate the invariant. I am therefore not claiming that default LitGPT generation is generally affected.

CUDA control

I repeated the test on an NVIDIA GeForce RTX 3070 with torch 2.14.0+cu130, while keeping LitGPT, the checkpoint, global float16 setting, precision="16-true", prompt, and generation parameters unchanged.

The clean public-API run passed:

{'torch': '2.14.0+cu130', 'cuda_runtime': '13.0', 'gpu': 'NVIDIA GeForce RTX 3070', 'default_dtype': 'torch.float16', 'actual': [253], 'expected': [253], 'actual_text': ' the', 'expected_text': ' the'}

A diagnostic run confirmed that LitGPT passed a finite, normalized torch.float16 probability tensor on cuda:0 to torch.multinomial; its support contained only token 253, and token 253 was selected with probability 1.0.

The less specialized top_k=50, temperature=0.8 CUDA control produced no zero-probability selections in 1,729 sampling calls over seeds 0 through 4. A direct CUDA torch.multinomial control with 151,936 float16 categories produced 0 violations in 20,000 draws.

This CUDA result is specific to the lambda=1 exponential sampling used by torch.multinomial. In a separate direct test of 10 million CUDA float16 exponential_ samples, the exact-zero counts were 0 at lambda=1, 1 at lambda=2, and 319 at lambda=1024. It therefore does not establish that float16 exponential_ is free of underflow on CUDA for other rate parameters.

Evidence that the returned token has zero probability

For a separate diagnostic run, I wrapped torch.multinomial without changing its output and recorded the tensor passed by LitGPT. At the failing call:

device:               cpu
dtype:                torch.float16
shape:                [50304]
all values finite:    True
sum:                  1.0
positive support:     {253}
zero entries:         50303
selected token:       5833
probability[5833]:    0.0

Replaying the operator from the recorded RNG state produced an exact zero at exponential[5833]. Consequently, the score for that category was 0 / 0 = NaN, and argmax returned 5833. Replacing only that zero exponential sample with the smallest positive float16 value made the replay return token 253. The replay consumed the same RNG state as the original call.

This matches the mechanism documented in pytorch/pytorch#192812: on CPU, a small exponential sample underflows to exactly zero when narrowed to float16; torch.multinomial divides the input probability by this value, so a zero-probability category can become NaN and win argmax.

LitGPT's current path computes softmax and forwards the result without dtype promotion:

probs = torch.nn.functional.softmax(logits, dim=-1)
return multinomial_num_samples_1(probs)

Less specialized sampling configuration

The deterministic top_k=1, seed-20 case makes the oracle easy to inspect, but neither condition is required. I also ran generation with:

top_k=50
temperature=0.8
400 generated tokens per seed
seeds=0,1,2,3,4

With float16 sampling probabilities, LitGPT selected a zero-probability token 5 times in 2,000 sampling calls, affecting seeds 1, 2, and 4. The corresponding float32 control had 0 violations in 2,000 calls.

Relation to #1715

This has a different local failure mechanism from #1715, which was fixed by #1720. In #1715, float16 computation produced NaNs during the model forward pass and torch.multinomial raised an error. Here, the sampling distribution is finite and normalized, and generation silently returns a token with probability zero. The change in #1720 makes a safer dtype the default on the affected macOS path, but it does not correct the explicit CPU-float16 sampling operation.

Expected behavior

Sampling must never return a category whose probability is zero. In particular, top_k=1 must always return the sole retained token regardless of the random seed.

I am reporting this here to document the observable LitGPT-level failure and to ask whether a local dtype promotion or regression test is desirable while the upstream PyTorch fix is pending.

Environment

OS: Windows 11 10.0.22631
Python: 3.12.6
LitGPT: 0.5.13
LitGPT commit: 7962f28f2dfb93d029c42450237268d9602e6d7b
PyTorch: 2.14.0+cpu
Lightning: 2.6.5
Hugging Face Hub: 0.36.0 for the clean automatic-download replay
Model: EleutherAI/pythia-14m
Model revision: cf967c0a9a04383db6f7b1108d86b2962634b4ac
Device: CPU

The LitGPT checkout had no tracked source modifications. The converted checkpoint SHA-256 was f83b1a27ed0068442db0e23e71ba390a89a27730c8dc3d3d9d72bb34cb70adcd, and model parameter hashes were unchanged before and after generation.

AI assistance disclosure: This report was prepared with AI assistance. Every numerical result above was produced by executing the described code, and the evidence was checked against the saved structured outputs.