torch.multinomial selects zero-probability indices on CPU for float16, so exponential_ underflows to exactly 0 in the double to half cast
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
`torch.multinomial` on CPU returns indices whose probability is exactly `0` when the input is `float16`. float32, bfloat16 and float64 are unaffected.
```python
import torch
torch.manual_seed(0)
p = torch.zeros(151936, dtype=torch.float16)
p[0] = 1.0 # only index 0 has nonzero probability
for i in range(20000):
idx = torch.multinomial(p, 1)
if p[idx] == 0:
print(f"iter {i}: multinomial returned index {int(idx)} with p = 0")
break
```
```
iter 136: multinomial returned index 15078 with p = 0
```
Expected: only index 0 can ever be returned. Over 20000 draws this picks a zero-probability index **85 times** for float16, and **0 times** for float32.
The cause is in `exponential_`, not in `multinomial`. `multinomial` uses the Gumbel trick (`Distributions.cpp:614` draws `q = at::empty_like(self).exponential_(1)` and takes `argmax(p / q)`). The CPU `exponential_kernel` samples in `double` and then narrows (`cpu/DistributionTemplates.h:341`):
```cpp
at::exponential_distribution exponential(lambda);
cpu_serial_kernel(iter, [&]() -> scalar_t {
return static_cast(exponential(generator));
});
```
float16's smallest subnormal is `2^-24 = 5.96e-08`, so any double sample below `2.98e-08` becomes exactly 0 in the cast. Then `p/q` is `0/0 = NaN`, and `argmax` selects that NaN's index.
The rate follows from the dtype, so it is predictable rather than incidental: `P(Exp(1) < 2^-25) ≈ 2.98e-08`, times a vocabulary of 151936, is `4.53e-03` per draw -- `90.6` expected in 20000 against **85 observed**. It scales with the number of categories, so it shows up in LLM sampling and stays invisible in small tests.
Measured dtype scope (exact zeros produced by `exponential_(1)` in 400 draws, alongside the smallest representable positive value):
| dtype | exact zeros | smallest nonzero |
|---|---|---|
| float16 | 3 | 5.96e-08 |
| bfloat16 | 0 | 9.18e-41 |
| float32 | 0 | 1.40e-45 |
| float64 | 0 | 4.94e-324 |
CUDA is fine, it floors the log at `-epsilon/2`. The `#else` CPU branch is `-1/lambda * log1p(-val)` with no lower bound at all, the divergence the `TODO: must be investigated and unified!!!` above it (#38662) already flags.
Related: #192577 and #192621 are the same failure shape on MPS (`exponential_` emitting `-0.0`). This one is CPU and pre-existing. It surfaced as a CI failure on #192621 (`test_multinomial_zero_probability_regression_cpu_float16`, `AssertionError: 0.0 not greater than 0.0`, failed consistently on rerun); that test has since been removed from the PR, so nothing upstream currently guards it.
I can send a PR addressing this issue if you are interested, but there are two possible fixes:
- clamp inside the CPU branch of `transformation::exponential` (matching CUDA's `eps/2`, which for float16 is also exactly its smallest subnormal)
- floor after the narrowing cast in `exponential_kernel`
The first changes float32/float64 output too; the second touches only the narrow dtypes.
---
*AI usage disclosure: this issue was investigated and drafted with Claude Code. Every measurement in it was produced by running the code shown; I have reviewed the report and am responsible for its contents.*
### Versions
PyTorch version: 2.9.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A
OS: macOS 26.6 (arm64)
GCC version: Could not collect
Clang version: 21.0.0 (clang-2100.1.1.101)
CMake version: version 4.4.2
Libc version: N/A
Python version: 3.12.10 (main, May 22 2025, 01:38:44) [Clang 20.1.4 ] (64-bit runtime)
Python platform: macOS-26.6-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Is XPU available: False
CPU:
Apple M3 Pro
Versions of relevant libraries:
[pip3] numpy==2.2.6
[pip3] torch==2.9.1
[conda] Could not collect
cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @aditew01 @pbelevich
Contributor guide
Assessment
This issue has not been assessed yet.