Larger GPU memory usage during the first batch

username: haowxu
cooked: <p>Hi, I have a training code that runs on a single GPU like this:</p>
<pre><code class="lang-auto">def train():
    model.train()
    torch.cuda.reset_max_memory_allocated()

    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs = inputs.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)

        outputs = model(inputs)
        loss = criterion(outputs, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(f'Max memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024}')
        torch.cuda.reset_max_memory_allocated()


def main():
    model = Net().cuda()
    cudnn.benchmark = True

    train_dataset = MyDataset()
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, pin_memory=True)

    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    print(f'Max memory before training: {torch.cuda.max_memory_allocated() // 1024 // 1024}')
    train(train_loader, model, optimizer, criterion)
</code></pre>
<p>Then the output I got in this way is:</p>
<pre><code class="lang-auto">Max memory before training: 487
Max memory allocated: 15402
Max memory allocated: 13615
Max memory allocated: 13591
Max memory allocated: 13591
Max memory allocated: 13591
Max memory allocated: 13591
...
</code></pre>
<p>What could be the reason of that the GPU memory usage of the first batch is larger than the following ones? This would lead to OOM if I use slightly larger model. Is there any ways to reduce the extra usage of memory?</p>

username: ptrblck
cooked: <p>cudnn benchmarking will try different algorithms for your input shape and model.<br>
Some algorithms trade memory for speed, so you might see a higher memory usage and a slowdown (due to multiple profiling runs) for the first iteration.</p>
<p>This shouldn’t yield to an OOM error, as cudnn shouldn’t use an algo which uses too much memory.</p>

username: haowxu
cooked: <p>Thanks for your reply! The memory reduces especially for the first batch after the cudnn benchmarking is commented. But the running time takes 30% longer.</p>
<p>From the output I posted above, I assume that the memory used by benchmarking algorithms can exceed 1.5GB. For example, if 15 GB out of 16GB is occupied by the model training itself, will this get an OOM error? I just wonder is there a way to so something like making cudnn benchmarking try these algorithms before the training to reduce the peak memory allocated?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="haowxu" data-post="3" data-topic="58768">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/e56c9b/48.png" class="avatar"> haowxu:</div>
<blockquote>
<p>The memory reduces especially for the first batch after the cudnn benchmarking is commented. But the running time takes 30% longer.</p>
</blockquote>
</aside>
<p>That’s expected, since cudnn benchmarking will try to find the fastest algorithm for your workload.</p>
<aside class="quote no-group" data-username="haowxu" data-post="3" data-topic="58768">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/e56c9b/48.png" class="avatar"> haowxu:</div>
<blockquote>
<p>For example, if 15 GB out of 16GB is occupied by the model training itself, will this get an OOM error?</p>
</blockquote>
</aside>
<p>This shouldn’t be the case and cudnn should discard all algorithms which could cause an OOM error. Have you seen an OOM using <code>cudnn.benchmark</code>?</p>
<aside class="quote no-group" data-username="haowxu" data-post="3" data-topic="58768">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/e56c9b/48.png" class="avatar"> haowxu:</div>
<blockquote>
<p>I just wonder is there a way to so something like making cudnn benchmarking try these algorithms before the training to reduce the peak memory allocated?</p>
</blockquote>
</aside>
<p>Benchmarking will run for the first iteration and cache the fastest algorithms for the current setup.<br>
If you don’t want to trigger it in the first training iteration, you could pass a dummy tensor with the same shape as your training samples before entering the training loop.<br>
Note that each different input shape will trigger the benchmarking again, so if you are dealing with a lot of different shapes, your code might run faster, if you disable it.</p>

username: haowxu
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="58768">
<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 shouldn’t be the case and cudnn should discard all algorithms which could cause an OOM error. Have you seen an OOM using <code>cudnn.benchmark</code> ?</p>
</blockquote>
</aside>
<p>Actually my real training code is a complicated version of the posted one. It’s not just a cross entropy loss but a distributed manner in which every process computes hand-crafted loss and gradients. Several tensor are used as intermediate variables. Will cudnn take these into consideration when choosing algorithms?</p>