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

Nobody has claimed this yet.

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

```text
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](https://github.com/pytorch/pytorch/issues/192812), rather than an independent LitGPT root cause. The proposed upstream fix, [pytorch/pytorch#195077](https://github.com/pytorch/pytorch/pull/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.

```python
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:

```text
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:

```python
torch.set_default_dtype(torch.float32)
```

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

```text
{'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:

```text
{'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:

```text
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](https://github.com/pytorch/pytorch/issues/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:

```python
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:

```text
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.

| seed | float16 zero-probability selections | float32 control |
|---:|---:|---:|
| 0 | 0 | 0 |
| 1 | 2 | 0 |
| 2 | 2 | 0 |
| 3 | 0 | 0 |
| 4 | 1 | 0 |

## Relation to #1715

This has a different local failure mechanism from [#1715](https://github.com/Lightning-AI/litgpt/issues/1715), which was fixed by [#1720](https://github.com/Lightning-AI/litgpt/pull/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

```text
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.

Start at the LLM.generate sampling path described in the report, reproduce the CPU float16 script, and inspect where softmax probabilities are passed to multinomial. Done means generation cannot select a zero-probability token, with a regression check covering the top_k=1 case while preserving the float32 and CUDA controls.

Written by the indexing model from the issue text.