Help understanding PyTorch memory model

username: Prashant
cooked: <p>Hi I don’t really understand when pytorch internally allocates memory / does copies.</p>
<p>I’m using the fairseq multihead attention code <a href="https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py" rel="noopener nofollow ugc">https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py</a> to build a transformer model that can do incremental encoding. However, I notice encoding with incremental state is 5x slower than encoding the entire sequence at once.</p>
<p>My current hypothesis is its because of this line<br>
</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py#L266" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py#L266" target="_blank" rel="noopener nofollow ugc">pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py#L266</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="256" style="counter-reset: li-counter 255 ;">
<li>if saved_state is not None:</li><li>    # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)</li><li>    if "prev_key" in saved_state:</li><li>        _prev_key = saved_state["prev_key"]</li><li>        assert _prev_key is not None</li><li>        prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)</li><li>        if static_kv:</li><li>            k = prev_key</li><li>        else:</li><li>            assert k is not None</li><li class="selected">            k = torch.cat([prev_key, k], dim=1)</li><li>    if "prev_value" in saved_state:</li><li>        _prev_value = saved_state["prev_value"]</li><li>        assert _prev_value is not None</li><li>        prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)</li><li>        if static_kv:</li><li>            v = prev_value</li><li>        else:</li><li>            assert v is not None</li><li>            v = torch.cat([prev_value, v], dim=1)</li><li>    prev_key_padding_mask: Optional[Tensor] = None</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
This line cats the new key row to the end of the previously saved key matrix. When this operation is done, does pytorch internally do a malloc, copy the old matrix over, and then copy the new row over? If so, is there a way to pre allocate space for rows like this to avoid a copy when cat is called?

username: ptrblck
cooked: <p><code>k</code> will be a new tensor and thus use new memory for the concatenated tensors.<br>
PyTorch uses a caching allocator internally, which should reuse the GPU memory (assuming you are using the GPU), but this operation would still copy the tensors.<br>
You could append the tensors to a <code>list</code> and call <code>torch.cat</code> or <code>torch.stack</code> at the end, if possible.</p>

username: Prashant
cooked: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>.</p>
<p>If I call cat on a list of tensors repeatedly, won’t each cat call require a new malloc + copy? The list would have a new element for every input token in the sequence.</p>
<p>I was planning on making a much larger (say, 500 column tensor) and fill in columns of that tensor one by one. And resize to a larger tensor when my original one was full.</p>

username: ptrblck
cooked: <p>Yes, calling <code>torch.cat</code> multiple times would create new tensors, so you could append all tensors to a <code>list</code> and create the tensor once it’s ready.<br>
Based on your description, it seems the size of the final tensor is unknown and you would need to reallocate new memory nevertheless (by calling <code>torch.cat</code> multiple times)?</p>

username: Prashant
cooked: <p>Yea my plan was to just preallocate a tensor with say 500 rows and double it in size every time I need to resize so that I don’t end up doing quite as many copies. I was looking for the cheapest way to do that.</p>
<p>Unfortunately at each input step I need to construct the cat’ed tensor ( in order to compute the attention scores at that timestep) so I can’t just use a list and cat it once at the end.</p>