pytorch / pytorch/pytorch

torch.multinomial selects zero-probability indices on CPU for float16, so exponential_ underflows to exactly 0 in the double to half cast

Open
#192,812 3 comments 0 reactions 1 assignee Claimed by @malfet View on GitHub
bot-triaged module: correctness (silent) module: cpu module: half module: random needs research triaged
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

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.