Torch.multinomial with generator is not determinstic under certain conditions

username: szvsw
cooked: <p>Hello,</p>
<p>I think there is a bug in <code>torch.multinomial</code> which sporadically appears when the number of input probabilities is large (e.g. greater than 10k, but gets worse as number gets larger) and <code>replacement=True</code>.  I think it is also worse when the dynamic range of the probabilities is very  high - i.e. a lot very very close to zero.  I would love if this is user error, but I’m pretty sure it is not.</p>
<p>My guess is that it is some sort of numerical / floating point error?</p>
<p>What happens is that sometimes, the nth row of the returned indices is off by 1.  e.g. in one trial, the 57th row out of 1000 might be index <code>57291</code> but in another trial it might end up as <code>57292</code> or <code>57290</code> (but never something far away like <code>1272</code>).</p>
<p>Here is a reproducible - nb: we are on <code>torch==2.5.1+cu118</code>.  If this is fixed in another version, apologies.</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F


def run_test(
        device: torch.device,
        n_candidates: int,
        n_draws: int,
        n_trials: int,
        replacement: bool,
        band_width: float,
        band_offset: float,
        generator: torch.Generator,

):
    all_results = []
    for _ in range(n_trials):
        generator.manual_seed(42)

        # Stand-in for some process which randomly generates logits.
        # e.g. randomly sampling designs and running them through a classifier.
        logits = (
            torch.rand(n_candidates, device=device, generator=generator) * band_width
            - band_offset
        )
        # convert logits to probabilities
        probs = F.softmax(logits, dim=-1)

        # draw samples
        results = torch.multinomial(
            probs,
            n_draws,
            replacement=replacement,
            generator=generator,
        )
        all_results.append(results.unsqueeze(1))

    # combine the samples, shape is (n_draws, n_trials)
    results = torch.cat(all_results, dim=1)
    errors = []
    for nth_draw in range(n_draws):
        # get the nth draw across all trials
        nth_draw_across_trials = results[nth_draw]
        unique, counts = torch.unique(nth_draw_across_trials, return_counts=True)

        # We expect each draw to be identical across all trials.
        # which would mean a single unique index was selected across all trials.
        if len(unique) &gt; 1:
            errors.append((nth_draw, unique))

    message = (
        f"There are {len(errors)} row(s) in the output that are not identical across all {n_trials} trials."
        " The following samples are misaligned:"
    )
    if len(errors) &gt; 0:
        print(message)

        for nth_draw, unique in errors:
            print(f"Draw {nth_draw} was not identical across all trials.  {len(unique)} different indices were sampled across {n_trials} trials: {unique.cpu().numpy().tolist()}")
    else:
        print("All samples are identical across all trials.")


device = torch.device("cuda")
generator = torch.Generator(device)
band_width = 5
band_offset = -5
n_candidates = 1_000_000
n_draws = 1000
n_trials = 1000
replacement = True

run_test(
    device=device,
    n_candidates=n_candidates,
    n_draws=n_draws,
    n_trials=n_trials,
    replacement=replacement,
    band_width=band_width,
    band_offset=band_offset,
    generator=generator,

)
</code></pre>
<p>Example output (will change across runs):</p>
<pre><code class="lang-auto">There are 11 row(s) in the output that are not identical across all 1000 trials. The following samples are misaligned:
Draw 65 was not identical across all trials.  2 different indices were sampled across 1000 trials: [197680, 197681]
Draw 111 was not identical across all trials.  2 different indices were sampled across 1000 trials: [128305, 128306]
Draw 115 was not identical across all trials.  2 different indices were sampled across 1000 trials: [181477, 181478]
Draw 306 was not identical across all trials.  2 different indices were sampled across 1000 trials: [66297, 66298]
Draw 465 was not identical across all trials.  2 different indices were sampled across 1000 trials: [131393, 131394]
Draw 550 was not identical across all trials.  2 different indices were sampled across 1000 trials: [90228, 90229]
Draw 585 was not identical across all trials.  2 different indices were sampled across 1000 trials: [129377, 129378]
Draw 802 was not identical across all trials.  2 different indices were sampled across 1000 trials: [154692, 154693]
Draw 846 was not identical across all trials.  2 different indices were sampled across 1000 trials: [175321, 175322]
Draw 856 was not identical across all trials.  2 different indices were sampled across 1000 trials: [135070, 135071]
Draw 954 was not identical across all trials.  2 different indices were sampled across 1000 trials: [184852, 184853]
</code></pre>

username: ptrblck
cooked: <p>This seems to be related to <a href="https://discuss.pytorch.org/t/deterministic-problem-for-torch-multinomial-on-different-device/219069/2">this post</a> so make sure you are enabling deterministic algorithms and are using a new PyTorch version.</p>