Is emptying the cache a good practice when dealing with huge data from the Dataloader?

RKnowledge: <p>Hi!</p>
<p>I wanted to ask if it’s a good practice or if there are better ways to free memory when using a Dataloader with a huge dataset.</p>
<p>For the same model, the same batch size I’m able to train on a smaller dataset so I’d like to know if there are any ways to leverage the bigger dataset.</p>
<p>Thank you!</p>

ptrblck: <p>Which cache exactly would you like to free and how would you do it?<br>
Huge datasets are usually loaded lazily so I’m unsure which cache you are using.</p>

RKnowledge: <p>Thank you for your reply!</p>
<p>I am talking about the CUDA cache because I run of memory when using the bigger dataset, and it doesn’t happen with the smaller one.</p>
<p>Are huge datasets loaded lazily by default? Otherwise I’m not sure I have enable such behavior. For the moment I have the classic Dataset, Dataloader configuration and in the training loop I do all the loading and processing on the CPU then I send to the GPU using <code>.to(device)</code>. The batches themselves are small in size, tensors of size (batch size, 700) but it is the Dataset that is big. (Huge length)</p>
<p>Maybe I’m doing it wrong and I should proceed otherwise for this case?</p>

ptrblck: <aside class="quote no-group" data-username="RKnowledge" data-post="3" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>I am talking about the CUDA cache because I run of memory when using the bigger dataset, and it doesn’t happen with the smaller one.</p>
</blockquote>
</aside>
<p>Unsure if I understand this claim correctly, but are you running our of GPU memory if you are loading more samples from your dataset into the host RAM or are you trying to load the data directly to the GPU? In the latter case, why don’t you move the batches to the device inside the <code>DataLoader</code> loop as is the common approach?</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="3" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>Are huge datasets loaded lazily by default?</p>
</blockquote>
</aside>
<p>No, as the data loading logic is defined in your <code>Dataset</code> so you would need to check it and see how exactly the data is loaded and processed.</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="3" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>For the moment I have the classic Dataset</p>
</blockquote>
</aside>
<p>There is no “classic Dataset”, so you are either implementing a custom <code>torch.utils.data.Dataset</code> or use a built-in class such as <code>ImageFolder</code>.</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="3" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>and in the training loop I do all the loading and processing on the CPU then I send to the GPU using <code>.to(device)</code></p>
</blockquote>
</aside>
<p>This is the common approach so I’m unsure how the GPU memory should be changed by the size of the dataset.</p>
<p>To your original question: no, clearing the CUDA cache won’t help as the GPU memory usage is unrelated to the size of the dataset (it’s rather related to the batch size, the model, etc.).</p>

RKnowledge: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="184875">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Unsure if I understand this claim correctly, but are you running our of GPU memory if you are loading more samples from your dataset into the host RAM or are you trying to load the data directly to the GPU? In the latter case, why don’t you move the batches to the device inside the <code>DataLoader</code> loop as is the common approach?</p>
</blockquote>
</aside>
<p>It is the first case, my training loop is basically this</p>
<pre><code class="lang-auto">for batch in data_loader:
   examples, labels = batch
   examples = examples.to(device)
   examples = torch.squeeze(examples)
   labels = labels.to(device)
</code></pre>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="184875">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>No, as the data loading logic is defined in your <code>Dataset</code> so you would need to check it and see how exactly the data is loaded and processed.</p>
</blockquote>
</aside>
<p>Thank you, I’ll check how to change that then, but since it seems that the issue doesn’t come from here I’ll focus on other areas first.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="184875">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>There is no “classic Dataset”, so you are either implementing a custom <code>torch.utils.data.Dataset</code> or use a built-in class such as <code>ImageFolder</code>.</p>
</blockquote>
</aside>
<p>Yeah you’re right, my bad sorry. I have a custom Dataset, and it is defined like this:</p>
<pre><code class="lang-auto">class SentenceCLS(Dataset):
  def __init__(self, data, labels, tokenizer):
    self.data = data
    self.labels = labels
    self.tokenizer = tokenizer

  def __len__(self):
    return len(self.data)

  def __getitem__(self, idx):
    if isinstance(idx, int):
      idx = [idx]
    if torch.is_tensor(idx):
      idx = idx.tolist()

    examples = [self.data[i] for i in idx]
    labels = [self.labels[i] for i in idx]

    examples = self.tokenizer(examples, padding=True, return_tensors="pt")['input_ids']
    return examples, labels
</code></pre>
<p>I also use a <code>Sampler</code> but I don’t know if it’s relevant as it only outputs lists of indices. I can put the code in here too but I didn’t want to surcharge the post.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="184875">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>This is the common approach so I’m unsure how the GPU memory should be changed by the size of the dataset.</p>
</blockquote>
</aside>
<p>Won’t PyTorch cache in the memory all the samples that are sent to the GPU for easier allocation later on? And if the dataset is bigger then it’ll cache in more samples.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="184875">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>To your original question: no, clearing the CUDA cache won’t help as the GPU memory usage is unrelated to the size of the dataset (it’s rather related to the batch size, the model, etc.).</p>
</blockquote>
</aside>
<p>Can I ask you please if it can be related to intermediate tensors in the model? The model is too big to put it is composed of many components, and some components contain residual connections so there are many forward loops of this type:</p>
<pre><code class="lang-auto">def forward(self, x):
   # x goes through a bunch of layers
   x = some_computations(x)

   # we use another variable for next layers while storing x for the residual connection
   y = some_computations2(x)

   # residual layer
   x = layer_norm_concat(x, y)
</code></pre>
<p>And it goes on. So I was wondering if these intermediate <code>y</code> tensors could be the issue.</p>

ptrblck: <aside class="quote no-group" data-username="RKnowledge" data-post="5" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>It is the first case, my training loop is basically this</p>
</blockquote>
</aside>
<p>The loop looks fine and matches the common approach.</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="5" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>I have a custom Dataset, and it is defined like this:</p>
</blockquote>
</aside>
<p>Also the custom <code>Dataset</code> looks alright assuming the <code>self.tokenizer</code> processes samples on the CPU.</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="5" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>I also use a <code>Sampler</code> but I don’t know if it’s relevant as it only outputs lists of indices.</p>
</blockquote>
</aside>
<p>Sounds also correct as I again assume the indices are purely on the CPU.</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="5" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>Won’t PyTorch cache in the memory all the samples that are sent to the GPU for easier allocation later on? And if the dataset is bigger then it’ll cache in more samples.</p>
</blockquote>
</aside>
<p>PyTorch will push device memory which is free to the cache to be able to reuse it as it avoids calling the expensive <code>cudaMalloc/Free</code> calls. The host memory is irrelevant for it and in your example only the batch size would matter, since only a single batch is moved to the device. Once the input data is no longer needed, the memory will be moved the the cache and reused for the next allocation.<br>
The size of the dataset won’t matter as you are just iterating more batches. The GPU memory usage would only be defined by the single batch, the model, intermediate activations, the optimizer etc., but is unrelated to the number of samples in the dataset.</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="5" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>Can I ask you please if it can be related to intermediate tensors in the model?</p>
</blockquote>
</aside>
<p>Intermediates are stored if they are needed for the gradient computation. However, the size of these intermediate tensors (and thus also their GPU memory usage) is defined by the batch size and the model (or rather the actual operation performed in the layer).</p>

RKnowledge: <p>Thank you for your answers. Do you think the issue may come from having batches holding tensors of different sizes? For the same batch size, some batches have the size <code>(batch_size, max_sequence_length1)</code>, others <code>(batch_size, max_sequence_length2)</code> …</p>
<p>It’s because in the huge dataset the variance in the distribution of sequence lengths is so huge I have decided to bin the sequences into categories of lengths so that the padding is more consistent. That’s why I’m using a <code>Sampler</code>.</p>
<p>I have understood from the following two quotes from your reply that the GPU memory usage would be only defined by the first single batch, but then during the loading when it meets batches of different shapes, is it safe to assume that PyTorch will case the first batch of the new sizes? If it is the case I think what explodes my GPU RAM is this, plus the intermediate the activations that will have different sizes now depending on the batch too (since after the embedding layer I’m using 3D tensors but only operating on the last dimension, the embedding dimension, so the middle dimension corresponding to sequence length and since this changes then the intermediate activations’ sizes change too).</p>
<aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="184875">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>PyTorch will push device memory which is free to the cache to be able to reuse it as it avoids calling the expensive <code>cudaMalloc/Free</code> calls. The host memory is irrelevant for it and in your example only the batch size would matter, since only a single batch is moved to the device. Once the input data is no longer needed, the memory will be moved the the cache and reused for the next allocation.<br>
The size of the dataset won’t matter as you are just iterating more batches. The GPU memory usage would only be defined by the single batch, the model, intermediate activations, the optimizer etc., but is unrelated to the number of samples in the dataset.</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="184875">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Intermediates are stored if they are needed for the gradient computation. However, the size of these intermediate tensors (and thus also their GPU memory usage) is defined by the batch size and the model (or rather the actual operation performed in the layer).</p>
</blockquote>
</aside>

ptrblck: <p>Yes, your explanation makes sense and a variable input shape will also cause different memory requirements. A large sequence length in one batch could increase the memory usage significantly and could yield an out of memory error. You could try to check the min/max sequence lengths and maybe clip them to a shape which fits your setup and especially the GPU memory.</p>

RKnowledge: <p>Thank you a lot for your answers. I am able to fit for high batch sizes (1028) my maximum length sequences (around 800). For now I’ll stick to a small batch size so that I avoid this effect of having huge different intermediate activations cached and different batches sizes too.</p>
<p>Can I ask you one last question please, is there any way to the inner workings of the memory management please? I couldn’t understand well just from PyTorch’s documentation.<br>
I have run the following loop on my smaller dataset, with a variable max sequence length, just to iterate fast and see how that will affect the allocated and cached memory:</p>
<pre><code class="lang-auto">print(torch.cuda.memory_allocated()/1024**2)
print(torch.cuda.memory_cached()/1024**2)
print()

for batch in data_loader:
  examples, labels = batch
  examples = torch.squeeze(examples)

  print(examples.size())

  examples = examples.to(device)

  print(torch.cuda.memory_allocated()/1024**2)
  print(torch.cuda.memory_cached()/1024**2)
  print("-----")
</code></pre>
<p>And I got the following</p>
<pre><code class="lang-auto">0.0
0.0
/usr/local/lib/python3.10/dist-packages/torch/cuda/memory.py:416: FutureWarning: torch.cuda.memory_cached has been renamed to torch.cuda.memory_reserved
  warnings.warn(

torch.Size([4092, 21])
0.65576171875
2.0
-----
torch.Size([4092, 18])
1.2177734375
2.0
-----
torch.Size([4092, 19])
1.1552734375
2.0
-----
torch.Size([4092, 21])
1.2490234375
2.0
-----
torch.Size([4092, 23])
1.3740234375
2.0
-----
torch.Size([4092, 30])
1.6552734375
2.0
-----
torch.Size([4092, 35])
2.02978515625
22.0
-----
torch.Size([357, 48])
1.06787109375
22.0
-----
</code></pre>
<p>I would love tu understand why the allocated memory decreased going from <code>torch.Size([4092, 35])</code> to <code>torch.Size([357, 48])</code> and I’d love to be able to compute by myself when the cached memory will increase (like why at <code>torch.Size([4092, 35])</code> and not <code>torch.Size([4092, 30])</code>). (All batches have the same data type <code>torch.int64</code>)</p>
<p>Thank you a lot again, it feels satisfying to be able to pinpoint the issue ^^</p>

ptrblck: <aside class="quote no-group" data-username="RKnowledge" data-post="9" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>I would love tu understand why the allocated memory decreased going from <code>torch.Size([4092, 35])</code> to <code>torch.Size([357, 48])</code></p>
</blockquote>
</aside>
<p>The batch size decreased by a large amount while the sequence length (or feature dimension) increased by a bit. Depending on the model architecture the decrease in memory would be expected and you would need to check what layers are used and how their memory usage is defined by the a) batch size and b) feature dimension input.</p>
<aside class="quote no-group" data-username="RKnowledge" data-post="9" data-topic="184875">
<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/rknowledge/48/62329_2.png" class="avatar"> RKnowledge:</div>
<blockquote>
<p>and I’d love to be able to compute by myself when the cached memory will increase (like why at <code>torch.Size([4092, 35])</code> and not <code>torch.Size([4092, 30])</code>)</p>
</blockquote>
</aside>
<p>Some caching settings can be adjusted via <code>PYTORCH_CUDA_ALLOC_CONF</code> as described in the docs. Based on your use case the second allocation needs <code>&gt;=2MB</code> and is thus allocating from the <code>large pool</code> as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">print(torch.cuda.memory_allocated()/1024**2)
# 0.0
print(torch.cuda.memory_reserved()/1024**2)
# 0.0

# 1MB
x = torch.randn(1024*1024//4 * 1, device="cuda")
print(torch.cuda.memory_summary())
# ...
# |---------------------------------------------------------------------------|
# | Allocations           |       1    |       1    |       1    |       0    |
# |       from large pool |       0    |       0    |       0    |       0    |
# |       from small pool |       1    |       1    |       1    |       0    |
# |---------------------------------------------------------------------------|
# ...
print(torch.cuda.memory_allocated()/1024**2)
# 1.0
print(torch.cuda.memory_reserved()/1024**2)
# 2.0
del x

# 2MB
x = torch.randn(1024*1024//4 * 2, device="cuda")
print(torch.cuda.memory_summary())
# ...
# |---------------------------------------------------------------------------|
# | Allocations           |       1    |       1    |       2    |       1    |
# |       from large pool |       1    |       1    |       1    |       0    |
# |       from small pool |       0    |       1    |       1    |       1    |
# |---------------------------------------------------------------------------|
# ...
print(torch.cuda.memory_allocated()/1024**2)
# 2.0
print(torch.cuda.memory_reserved()/1024**2)
# 22.0
</code></pre>

RKnowledge: <p>I see. Thank you for your answers and your help. Learned a lot from them and clarified my understanding about the PyTorch/CUDA interaction.</p>