Reproducibility of sampling from a subset of logits

tcm03: <p>Hi community,</p>
<p>I’d like to ask about the reason and a solution (if any) to the problem of different results when sampled from the same logits, even under the same rng state. The toy code for this situation is as follows:</p>
<pre><code class="lang-auto">import torch

def main():
    rng = torch.Generator().manual_seed(42)
    logits = torch.randn(2, 2, 4096, generator=rng)
    logits_gen_state = rng.get_state()
    print(f"After logits, rng state = {logits_gen_state}")
    idx1 = torch.multinomial( 
        logits.softmax(dim=-1).view(-1, 4096), 
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    )
    print(idx1)
    idx1_gen_state = rng.get_state()
    print(f"After idx1, rng state = {idx1_gen_state}")
    print(f"logits gen state != idx1 gen state: {torch.any(logits_gen_state != idx1_gen_state)}")
    rng.set_state(logits_gen_state)
    idx2 = torch.multinomial( 
        logits[0, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096), 
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    ) 
    print(idx2)
    idx2_gen_state = rng.get_state()
    print(f"After idx2, rng state = {idx2_gen_state}")
    print(f"logits gen state != idx2 gen state: {torch.any(logits_gen_state != idx2_gen_state)}")
    print(f"idx1 gen state != idx2 gen state: {torch.any(idx1_gen_state != idx2_gen_state)}")

if __name__ == "__main__":
    main()
</code></pre>
<p>What the code snippet above does is that it first uses the same rng for all sampling, and then creates a logits tensor to sample two other indices tensors. For the sampling of the second indices tensor, I only use the first half of the logits tensor. And below is the result:</p>
<pre><code class="lang-auto">tensor([[1284],
        [1828],
        [3472],
        [2623]])
logits gen state != idx1 gen state: True
tensor([[1284],
        [1828]])
logits gen state != idx2 gen state: True
idx1 gen state != idx2 gen state: False
</code></pre>
<p>So <code>idx2</code> coincides with the first half of <code>idx1</code>, which is correct to my expectation. But if I use the second half of the logits tensor to sample:</p>
<pre><code class="lang-auto">idx2 = torch.multinomial( 
        logits[1, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096), 
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    )
</code></pre>
<p>I don’t get the same indices as the second half of <code>idx1</code></p>
<pre><code class="lang-auto">tensor([[1284],
        [1828],
        [3472],
        [2623]])
logits gen state != idx1 gen state: True
tensor([[1284],
        [2704]])
logits gen state != idx2 gen state: True
idx1 gen state != idx2 gen state: False
</code></pre>
<p>More generally, sampling from a subset of logits does not guarantee reproducibility even under the same rng state. And this is what I desire for my usecase.</p>
<p>Hence, in this post, I’d like to ask about the reason for the discrepancy between the two cases I illustrated above and the solution for reproducibility when sampling from a subset of logits, if any.</p>
<p>Really appreciate your help.</p>
<p>Thanks in advance.</p>

KFrank: <p>Hi Minh!</p>
<aside class="quote no-group" data-username="tcm03" data-post="1" data-topic="223454" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tcm03/48/78360_2.png" class="avatar"> tcm03:</div>
<blockquote>
<p>different results when sampled from the same logits, even under the same rng state.</p>
</blockquote>
</aside>
<p>The problem is that the samples that come out “wrong” weren’t really sampled<br>
using the same rng state.</p>
<p>Roughly speaking, every time you generate a random sample, you advance the<br>
state of the rng by “one unit.”</p>
<p>(Pytorch doesn’t <em>have</em> to do it this way, but it often does.  For example, if<br>
you want, say, seven samples, pytorch could, for efficiency reasons, actually<br>
generate eight samples, discarding the last one.  If it did it this way, it would<br>
advance the rng state by eight units instead of only seven.)</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">    idx1 = torch.multinomial( 
        logits.softmax(dim=-1).view(-1, 4096), 
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    )
</code></pre>
</blockquote>
</aside>
<p>This generates four samples (because you are generating a batch of four).<br>
The first sample is generated using the initial rng state, the second after the<br>
state has been advanced once, the third after the state has been advanced<br>
twice, and the fourth after the state has been advanced three times.</p>
<p>That is the, the last two samples – those corresponding to the “second half”<br>
of your logits – are generated using “initial state + 2.”</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">    rng.set_state(logits_gen_state)
    idx2 = torch.multinomial( 
        logits[0, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096), 
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    ) 
</code></pre>
</blockquote>
</aside>
<p>Here you are generating two samples for the “first half” of the logits starting<br>
with the initial state and you get the right answer.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">idx2 = torch.multinomial( 
        logits[1, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096), 
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    )
</code></pre>
<p>I don’t get the same indices as the second half of <code>idx1</code></p>
</blockquote>
</aside>
<p>Presumably when you run this chunk of code, it is also preceded by:</p>
<pre data-code-wrap="python"><code class="lang-python">    rng.set_state(logits_gen_state)
</code></pre>
<p>When you do this, you are generating two samples for the “second half” of the<br>
logits, but starting with the initial state.  These will not match the third and<br>
fourth samples in <code>idx1</code> because those were generated with “initial state + 2”.</p>
<p>If you advance the rng state properly – for example by generating two samples<br>
that you ignore – before generating <code>idx2</code>, you will get the expected result.</p>
<p>I’ve modified your code modestly to call the generation of <code>idx2</code> twice.  The<br>
first time advances the rng state by “two units” so that the second call will use<br>
the correct state and generate samples that match the second half of <code>idx1</code>.</p>
<p>Here is your modified script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

def main():
    rng = torch.Generator().manual_seed(42)
    logits = torch.randn(2, 2, 4096, generator=rng)
    logits_gen_state = rng.get_state()
    print(f"After logits, rng state = {logits_gen_state}")
    idx1 = torch.multinomial(                       # advances rng state for each of the four samples
        logits.softmax(dim=-1).view(-1, 4096),
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    )
    print(idx1)
    idx1_gen_state = rng.get_state()
    print(f"After idx1, rng state = {idx1_gen_state}")
    print(f"logits gen state != idx1 gen state: {torch.any(logits_gen_state != idx1_gen_state)}")
    rng.set_state(logits_gen_state)
    idx2 = torch.multinomial(                       # generates two samples, advancing rng state
        # logits[0, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096),
        logits[1, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096),
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    ) 
    print(idx2)
    idx2 = torch.multinomial(                       # generates two more samples with the properly advanced rng state
        # logits[0, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096),
        logits[1, ...].unsqueeze(0).softmax(dim=-1).view(-1, 4096),
        num_samples=1, 
        replacement=True, 
        generator=rng, 
    ) 
    print(idx2)
    idx2_gen_state = rng.get_state()
    print(f"After idx2, rng state = {idx2_gen_state}")
    print(f"logits gen state != idx2 gen state: {torch.any(logits_gen_state != idx2_gen_state)}")
    print(f"idx1 gen state != idx2 gen state: {torch.any(idx1_gen_state != idx2_gen_state)}")

if __name__ == "__main__":
    main()
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.8.0+cu129
After logits, rng state = tensor([42,  0,  0,  ...,  0,  0,  0], dtype=torch.uint8)
tensor([[1037],
        [3152],
        [1221],
        [4032]])
After idx1, rng state = tensor([42,  0,  0,  ...,  0,  0,  0], dtype=torch.uint8)
logits gen state != idx1 gen state: True
tensor([[2951],
        [3768]])
tensor([[1221],
        [4032]])
After idx2, rng state = tensor([42,  0,  0,  ...,  0,  0,  0], dtype=torch.uint8)
logits gen state != idx2 gen state: True
idx1 gen state != idx2 gen state: False
</code></pre>
<p>Best.</p>
<p>K. Frank</p>