Large spike in memory usage (torch.cuda.OutOfMemoryError)

Brett_Clark: <p>Hi there,</p>
<p>I’m building a 3D U-Net for medical image segmentation and am seeing an OOM error for large input sizes when hitting one layer of the model in particular. The layer is a <code>Conv3D</code> acting on a large input tensor <code>(1, 64, 370, 370, 250)</code> - this input tensor is a concatenation of the residual feature maps (32 channels) and the upsampled feature maps (32 channels).</p>
<p>I have isolated the layer in question and created a script that reproduces the error. I’m running the script on an NVIDIA A100 GPU with 80GB memory. Note that the script starts a thread to measure GPU usage alongside the layer operation.</p>
<p>Versions:</p>
<pre><code class="lang-auto">cuda: 11.7
python: 3.8.6
torch: 1.13.1
</code></pre>
<p><strong>Script:</strong></p>
<pre><code class="lang-auto">from datetime import datetime
import numpy as np
import os
import pandas as pd
from pynvml.smi import nvidia_smi
from threading import Thread
from time import sleep
import torch
from torch import nn

def record_gpu_usage(
    name: str,
    time: float,
    interval: float) -&gt; None:

    # Create results table.
    nvsmi = nvidia_smi.getInstance()
    n_gpus = len(nvsmi.DeviceQuery('gpu_name')['gpu'])
    cols = {
        'time': str
    }
    for i in range(n_gpus):
        device_name = f'cuda:{i}'
        cols[f'{device_name}-usage'] = float
    df = pd.DataFrame(columns=cols.keys())

    # Add usage.
    n_intervals = int(np.ceil(time / interval))
    start_time = datetime.now()
    for i in range(n_intervals):
        # Record GPU usage.
        data = {
            'time': (datetime.now() - start_time).total_seconds()
        }
        usages_mb = [g['fb_memory_usage']['used'] for g in nvsmi.DeviceQuery('memory.used')['gpu']]
        for j, usage_mb in enumerate(usages_mb):
            device_name = f'cuda:{j}'
            data[f'{device_name}-usage'] = usage_mb
        df = pd.concat((df, pd.DataFrame([data])), axis=0)

        # Wait for time interval to pass.
        time_passed = (datetime.now() - start_time).total_seconds()
        if time_passed &gt; time:
            break
        time_to_wait = ((i + 1) * interval) - time_passed
        if time_to_wait &gt; 0:
            sleep(time_to_wait)
        elif time_to_wait &lt; 0:
            # Makes time problem worse if we log.
            # logging.warning(f"GPU usage recording took longer than allocated interval '{interval}' (seconds).")
            pass

    # Save results.
    filepath = f'{name}.csv'
    df.to_csv(filepath, index=False)

n_channels = 64
name = f'testing-{n_channels}'
input_shape = (1, n_channels, 370, 370, 250)
device = torch.device('cuda:0')
input = torch.rand(input_shape)

# Kick off GPU memory recording.
thread = Thread(target=record_gpu_usage, args=(name, 10, 1e-3))
thread.start()

input = input.half().to(device)
input_GB = input.numel() * input.element_size() / 1e9
print('input GB: ', input_GB)
input.requires_grad = False
layer = nn.Conv3d(in_channels=n_channels, out_channels=32, kernel_size=3, stride=1, padding=1).half().to(device)
for i, param in enumerate(layer.parameters()):
    param_GB = param.numel() * param.element_size() / 1e9
    print(f'param_{i} GB: {param_GB}')
    param.requires_grad = False
output = layer(input)

torch.cuda.empty_cache()
</code></pre>
<p><strong>Plotting script:</strong></p>
<pre><code class="lang-auto">from matplotlib import pyplot as plt
import pandas as pd

def plot_gpu_usage(name: str) -&gt; None:
    filepath = f'{name}.csv'
    df = pd.read_csv(filepath)
    x = df['time']
    n_gpus = len(df.columns) - 1
    for i in range(n_gpus):
        device_name = f'cuda:{i}'
        key = f'{device_name}-usage'
        y = df[key]
        plt.plot(x, y, label=key)
    plt.title(name)
    plt.xlabel('time [seconds]')
    plt.ylabel('GPU usage [MB]')
    plt.show()
</code></pre>
<p><strong>Error:</strong></p>
<pre><code class="lang-auto">torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 110.16 GiB (GPU 0; 79.21 GiB total capacity; 6.12 GiB already allocated; 72.55 GiB free; 6.12 GiB reserved in total by PyTorch) If reserved memory is &gt;&gt; allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
</code></pre>
<p>The amount that PyTorch is trying to allocate (<code>110.16GB</code>) seems very large. I varied the number of input channels (<code>n_channels</code>) from 32 to 64. At <code>n_channels=62</code>, the convolution operation consumes a maximum of ~ 14GB. When <code>n_channels&gt;=63</code> the memory consumption spikes by nearly 10x. Figures show memory usage over time for <code>n_channels=32,40,48,56,62</code>.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/0/a0de763f21b6c6741a25a638b1b97483422f73b1.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/a0de763f21b6c6741a25a638b1b97483422f73b1" title="Screen Shot 2023-02-15 at 12.39.15 pm"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a0de763f21b6c6741a25a638b1b97483422f73b1_2_690x203.jpeg" alt="Screen Shot 2023-02-15 at 12.39.15 pm" data-base62-sha1="mX7161uJ6m630Jdfj6gOxdhOOUV" width="690" height="203" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a0de763f21b6c6741a25a638b1b97483422f73b1_2_690x203.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a0de763f21b6c6741a25a638b1b97483422f73b1_2_1035x304.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a0de763f21b6c6741a25a638b1b97483422f73b1_2_1380x406.jpeg 2x" data-dominant-color="F7F7F8"><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">Screen Shot 2023-02-15 at 12.39.15 pm</span><span class="informations">2562×756 127 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>Any ideas why the memory usage increases by such a large margin for a small increase in <code>n_channels</code>? I’ve tried running with <code>PYTORCH_NO_CUDA_MEMORY_CACHING=1</code> and this didn’t affect the amount of memory requested.</p>
<p>Thanks!<br>
Brett</p>