Tensor permutation inplace to save memory

dsesto: <p>We are using <a href="https://pytorch.org/memory_viz" rel="noopener nofollow ugc">MemoryViz</a> tool to profile the GPU utilization of our application, following the “<a href="https://pytorch.org/blog/understanding-gpu-memory-1/" rel="noopener nofollow ugc">Understanding GPU Memory</a>” blogpost.<br>
This tool has helped us identify this pattern at the very beginning of an epoch:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/e/9e8e70d9deb1cd93574e3aff7c7ed34f1c19c331.png" data-download-href="https://discuss.pytorch.org/uploads/default/9e8e70d9deb1cd93574e3aff7c7ed34f1c19c331" title="Screenshot from 2024-02-07 12-05-08"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/e/9e8e70d9deb1cd93574e3aff7c7ed34f1c19c331_2_690x280.png" alt="Screenshot from 2024-02-07 12-05-08" data-base62-sha1="mCECl5BRODSanE6kruR9KGj9jdn" width="690" height="280" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/e/9e8e70d9deb1cd93574e3aff7c7ed34f1c19c331_2_690x280.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/e/9e8e70d9deb1cd93574e3aff7c7ed34f1c19c331_2_1035x420.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/e/9e8e70d9deb1cd93574e3aff7c7ed34f1c19c331.png 2x" data-dominant-color="E0D7BE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-02-07 12-05-08</span><span class="informations">1167×474 16.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
All the big chunks of memory are input tensors. The short-lived ones (on the left) are the tensors loaded in GPU pre-shuffling. Then we shuffle those tensors on every epoch using something like below, and new sorted vectors are created, which persist through the epoch.</p>
<pre><code class="lang-auto">class CustomDataloader:
    ...
    if self.shuffle:
        permutation = torch.randperm(self.data_len)
            self.data = [tensor[permutation] for tensor in self.data]
</code></pre>
<p>Our understanding of the previous plot is that some tensors (e.g. light blue, dark green) are cleaned by the Garbage Collector after they are replaced with the sorted versions (e.g. yellow, pink).</p>
<p>If we disable shuffling on every epoch (which happens on the <code>__iter__()</code> call of our DataLoader), the memory utilization does not have this spike:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/3/e33e619a8d308f3328dff58fef82f29dbb102351.png" data-download-href="https://discuss.pytorch.org/uploads/default/e33e619a8d308f3328dff58fef82f29dbb102351" title="Screenshot from 2024-02-07 12-09-43"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/3/e33e619a8d308f3328dff58fef82f29dbb102351_2_690x280.png" alt="Screenshot from 2024-02-07 12-09-43" data-base62-sha1="wqi2imqgzAjWjXlneJ6uLgL8n29" width="690" height="280" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/3/e33e619a8d308f3328dff58fef82f29dbb102351_2_690x280.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/3/e33e619a8d308f3328dff58fef82f29dbb102351_2_1035x420.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/3/e33e619a8d308f3328dff58fef82f29dbb102351.png 2x" data-dominant-color="A1C4AB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-02-07 12-09-43</span><span class="informations">1167×474 9.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Would there be any way to do the tensor shuffling in-place, so that there is no duplication of the tensors, generating the undesired and temporary peak in memory utilization?</p>
<p>Thanks in advance!</p>
<p><em><strong>Further context:</strong></em><br>
We are training a model using tabular data stored in several GBs of parquet files.<br>
After initially using PyTorch’s DataLoader with a custom Dataset implementation, we ended up implementing our own DataLoader too, because the random access to tensor positions with the <code>__getitem__()</code> call was too slow for us.<br>
Our custom implementation loads the whole Parquet dataset into GPU tensors at once, and then the DataLoader iterator <code>__next__()</code> call returns a batch of elements by slicing the tensors through the appropriate indices.</p>

ccaven: <p>Hi, <a class="mention" href="/u/dsesto">@dsesto</a>! You could try implementing an in place shuffling algorithm, such as the <a href="https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle" rel="noopener nofollow ugc">Fisher-Yates Shuffle</a>. For example, something like this:</p>
<pre><code class="lang-auto">from random import randint
for i in range(0, x.size(dim) - 1):
    j = randint(i, x.size(dim) - 1)
    tmp = x[j].clone()
    x[j] = x[i]
    x[i] = tmp
</code></pre>
<p>It’s likely to be a bit slower, but should use less memory.</p>

KFrank: <p>Hi David!</p>
<aside class="quote no-group" data-username="dsesto" data-post="1" data-topic="196758" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dsesto/48/65875_2.png" class="avatar"> dsesto:</div>
<blockquote>
<pre><code class="lang-auto">            self.data = [tensor[permutation] for tensor in self.data]
</code></pre>
</blockquote>
</aside>
<p>As I understand it, <code>self.data</code> is a <code>list</code> of (references to) tensors.  That<br>
list holds on to those references until such time as the list itself is released.<br>
This doesn’t happen until the full list comprehension has completed, at<br>
which time the reassignment of <code>self.data</code> permits the original list and the<br>
tensors it references to be freed.  So you basically have two copies of<br>
each tensor in memory until the full list comprehension is finished.</p>
<p>Try something like:</p>
<pre data-code-wrap="python"><code class="lang-python">for  i, t in enumerate (self.data):
    self.data[i] = t[permutation]
</code></pre>
<p>As soon as <code>self.data[i]</code> is reassigned (within the loop), the tensor to which<br>
<code>self.data[i]</code> refers can be freed from memory (without waiting for the full<br>
loop to finish).</p>
<p>Best.</p>
<p>K. Frank</p>

dsesto: <p>Thank you both!<br>
It is not exactly what I had in mind, but they are both definitely good tips! I was expecting PyTorch to have a functionality that allowed sorted views of an existing vector, or something on that line, but probably nothing like that exists and/or is feasible.</p>
<p><a class="mention" href="/u/kfrank">@KFrank</a> I’ve tested your suggestion with good results. Now the maximum GPU utilization is bounded by the largest vector to sort (left is the list-comprehension implementation, right is the for-loop one).<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/f/bfa2865d42c9be0cd12d7ed97ede6497c9b6d5dc.png" data-download-href="https://discuss.pytorch.org/uploads/default/bfa2865d42c9be0cd12d7ed97ede6497c9b6d5dc" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bfa2865d42c9be0cd12d7ed97ede6497c9b6d5dc_2_690x154.png" alt="image" data-base62-sha1="rlhp5UuBgPNbja05N88eIK8zL64" width="690" height="154" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bfa2865d42c9be0cd12d7ed97ede6497c9b6d5dc_2_690x154.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bfa2865d42c9be0cd12d7ed97ede6497c9b6d5dc_2_1035x231.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bfa2865d42c9be0cd12d7ed97ede6497c9b6d5dc_2_1380x308.png 2x" data-dominant-color="C7C8BF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2069×463 31.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><a class="mention" href="/u/ccaven">@ccaven</a> I have not yet tried to implement the in-place shuffling myself, but I may give it a try in the near future.</p>

dsesto: <p>Just in case this ends up being useful to somebody else, we iterated on top of this to do per-dimension shuffling instead of per-tensor shuffling.</p>
<p>Our <code>data</code> are 2D tensors, where each row represents an instance in the dataset, and each column is a different feature. By shuffling each feature individually instead of the whole tensor together, we reduced our peak memory utilization even further (left is the per-tensor for-loop, right is the per-feature for-loop).<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/4/44b965516ca327e21573719ade7671c056cab78d.png" data-download-href="https://discuss.pytorch.org/uploads/default/44b965516ca327e21573719ade7671c056cab78d" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44b965516ca327e21573719ade7671c056cab78d_2_690x133.png" alt="image" data-base62-sha1="9NXFWRCpg7LvIEfXpkEUapiwFPf" width="690" height="133" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44b965516ca327e21573719ade7671c056cab78d_2_690x133.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44b965516ca327e21573719ade7671c056cab78d_2_1035x199.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/44b965516ca327e21573719ade7671c056cab78d_2_1380x266.png 2x" data-dominant-color="B7BDB4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2863×553 37 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The general idea is we ended up having something like this:</p>
<pre><code class="lang-auto">for tensor_idx in range(len(self.data)):
    for feature_idx in range(self.data[tensor_idx].shape[1]):
        tensor = self.data[tensor_idx][:, feature_idx]
        self.data[tensor_idx][:, feature_idx] = tensor[permutation]
</code></pre>
<p>Performance-wise, execution times of all implementations are identical.</p>