CUDA Memory Profiling

username: bruno-darochac1
cooked: <p>I’m currently using the <code>torch.profiler.profile</code> to analyze memory peak on my GPUs. I fristly use the argument <code>on_trace_ready</code> to generate a tensorboard and read the information by hand, but now I want to read those information directly in my code.</p>
<p>So I’ve setup my profiler as :</p>
<pre><code class="lang-python">self.prof = torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU
        torch.profiler.ProfilerActivity.CUDA
    ],
    record_shapes=True,
    profile_memory=True
)
</code></pre>
<p>And then I used the functions <code>start()</code>, <code>step()</code> and <code>stop()</code>on my training function and try to read the information using <code>self.prof.key_averages()</code> but I didn’t see any data corresponding to the data read on tensorboard and in addition no dissociation between multiple GPU…</p>
<p>Have you any tips on how I can get those information ?</p>
<p>Thanks in advance</p>

username: bruno-darochac1
cooked: <p>I’ve try a new approach using torch.cuda.max_memory_allocated() for this I’ve created a class as mentionned on this <a href="https://discuss.pytorch.org/t/measuring-peak-memory-usage-tracemalloc-for-pytorch/34067/11">topic</a> and try to call it in my training loop. The fact is that even if I reorganize my layer repartition between gpus the memory allocated on both gpus continue to grow and the peak are always egual to the first run.</p>
<p>As example if I have a model with 14 layers and I split them for the first run with GPU0 11 and GPU1 3 I will have peak at [1371, 527] MB and even if during the same process I swap some layers to have 7 layers on bot I will always have the same peaks value. And if I run it firstly with 7 on both I will have other memory peak value [785, 673] MB</p>
<p>So the first issue that I imagine is that the memory is never released and never reallocated how to do it and none of torch.cuda reset method works…</p>
<p>Here is my profiling class :</p>
<pre><code class="lang-python">import torch
def BToMb(x): return (x // (2*1024))

class TraceMalloc():
    def __init__(self, nb_gpu):
        self.nb_gpu = nb_gpu
        self.begin  = [0] * nb_gpu
        self.end    = [0] * nb_gpu
        self.peak   = [0] * nb_gpu
        self.peaked = [0] * nb_gpu

    def __enter__(self):
        for device in range(self.nb_gpu):
            torch.cuda.reset_accumulated_memory_stats(device)
            self.begin[device] = torch.cuda.memory_allocated(device)


        return self
    
    def __exit__(self, *exc):
        for device in range(self.nb_gpu):
            self.end[device]    = torch.cuda.memory_allocated(device)
            self.peak[device]   = torch.cuda.max_memory_allocated(device)
            self.peaked[device] = BToMb(self.peak[device] - self.begin[device])
            torch.cuda.reset_accumulated_memory_stats(device)

        for device in range(self.nb_gpu):
            print(f"GPU n°{device}")
            print(f"    Memory begin -&gt; {BToMb(self.begin[device])} MB")
            print(f"    Memory end   -&gt; {BToMb(self.end[device])} MB")
            # print(f"Memory used  -&gt; {self.used[device]} MB")
            print(f"    Memory peak  -&gt; {self.peaked[device]} MB")

</code></pre>
<p>Here my training loop where it is call :</p>
<pre><code class="lang-python">def evaluate_mem(self):
        from .dataset import PipelineDataset


        dataset = PipelineDataset(1024, self.input_shape[1:], [1] if len(self.output_shape) == 1 else self.output_shape[1:])
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=self.input_shape[0], shuffle=True)

        criterion = torch.nn.CrossEntropyLoss()

        model = self.get_modules()
        model = torch.distributed.pipeline.sync.Pipe(model, 2)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

        # self.prof.start()
        # with torch.no_grad():
        with self.trace_gpu_alloc:
            for epoch in range(10):
                running_loss = 0.0
                for i, (inputs, labels) in enumerate(dataloader):

                    # optimizer.zero_grad()
                    inputs = inputs.to(0)
                    labels = labels.to(self.nb_gpu - 1)

                    outputs = model(inputs).local_value()

                    # Forward pass
                    loss = criterion(outputs, labels.squeeze())

                    loss.backward()
                    # self.prof.step()
                    optimizer.step()


        torch.cuda.empty_cache()
        print(f"Outer loop --------&gt; {self.trace_gpu_alloc.peaked}")

        return self.trace_gpu_alloc.peaked
</code></pre>
<p>And finally an example of model used :</p>
<pre><code class="lang-python">class x_layer(nn.Module):
    def forward(self, input):
        ret = input
        return input

class conv1_layer(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.fc = nn.Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
    def forward(self, input):
        ret = self.fc(input)
        return ret

class relu_layer(nn.Module):
    def forward(self, input):
        ret = torch.nn.functional.relu(input, inplace=False)
        return ret

class pool_layer(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.fc = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    def forward(self, input):
        ret = self.fc(input)
        return ret

class conv2_layer(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.fc = nn.Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
    def forward(self, input):
        ret = self.fc(input)
        return ret

class relu_1_layer(nn.Module):
    def forward(self, input):
        ret = torch.nn.functional.relu(input, inplace=False)
        return ret

class pool_1_layer(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.fc = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    def forward(self, input):
        ret = self.fc(input)
        return ret

class flatten_layer(nn.Module):
    def forward(self, input):
        ret = torch.flatten(input, 1)
        return ret

class fc1_layer(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.fc = nn.Linear(in_features=400, out_features=120, bias=True)
    def forward(self, input):
        ret = self.fc(input)
        return ret

class relu_2_layer(nn.Module):
    def forward(self, input):
        input = input.clone()
        ret = torch.nn.functional.relu(input, inplace=False)
        return ret

class fc2_layer(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.fc = nn.Linear(in_features=120, out_features=84, bias=True)
    def forward(self, input):
        ret = self.fc(input)
        return ret

class relu_3_layer(nn.Module):
    def forward(self, input):
        ret = torch.nn.functional.relu(input, inplace=False)
        return ret

class fc3_layer(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.fc = nn.Linear(in_features=84, out_features=10, bias=True)
    def forward(self, input):
        ret = self.fc(input)
        return ret

class output_layer(nn.Module):
    def forward(self, input):
        ret = input
        return ret

class PipelinedModel(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.s0 = nn.Sequential(x_layer(), conv1_layer(), relu_layer(), pool_layer(), conv2_layer(), relu_1_layer(), pool_1_layer(), flatten_layer(), fc1_layer(), relu_2_layer()).cuda(0)
        self.s1 = nn.Sequential(fc2_layer(), relu_3_layer(), fc3_layer(), output_layer()).cuda(1)
    def forward(self, input):
        ret = input
        ret = self.s0(ret.to(0))
        ret = self.s1(ret.to(1))
        return ret
    def get_modules(self):
        return  nn.Sequential(*[nn.Sequential(*self.s0),nn.Sequential(*self.s1)])
</code></pre>
<p>Hope I was clear in my explanation :x</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="bruno-darochac1" data-post="2" data-topic="182065">
<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/bruno-darochac1/48/61268_2.png" class="avatar"> bruno-darochac1:</div>
<blockquote>
<p>So the first issue that I imagine is that the memory is never released and never reallocated how to do it and none of torch.cuda reset method works…</p>
</blockquote>
</aside>
<p>I’m not sure I understand the question completely, but PyTorch uses a caching mechanism to avoid expensive and synchronizing <code>cudaMalloc/Free</code> calls.<br>
In case you want to disable this caching allocator and are fine with bad performance you could use <code>PYTORCH_NO_CUDA_MEMORY_CACHING=1</code> (note that we are using this env variable for debugging only due to the expected poor performance).</p>

username: bruno-darochac1
cooked: <p>Hello ! I will try to be more precise. I’m currently trying to split a model into multiple GPU. My objectif is to find automatically what is the best splitting with the less memory peak on both GPU. This is why I want to know after each run of my training what is the memory peak to be able to swap the layers between GPU. The model that I give in the previous message is a typical example of splitting.</p>
<p>What I really want is to be able too see memory peaks and difference between runs if I change the layer disposition.</p>
<p>I’m able to see difference if I change it by hand and rerun my script by hand this is why I think that the memory is never really released…</p>
<p>In addition I’ve tried to used <code>torch.cuda.empty_cache()</code> and <code>gc.collect()</code> but nothing change.</p>
<p>I’ve try to change the env variable you suggested and if I set it to 1 I’m not able to analyze anithing about cuda memory with torch function</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="bruno-darochac1" data-post="4" data-topic="182065">
<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/bruno-darochac1/48/61268_2.png" class="avatar"> bruno-darochac1:</div>
<blockquote>
<p>In addition I’ve tried to used <code>torch.cuda.empty_cache()</code> and <code>gc.collect()</code> but nothing change.</p>
</blockquote>
</aside>
<p>This would indicate that te memory is indeed <em>used</em> and is not in the cache.<br>
I would recommend scaling down the use case and checking the memory stats in a simple example as was discussed in e.g. <a href="https://discuss.pytorch.org/t/why-max-memory-allocated-after-reset-doesnt-return-0/52626/2">this topic</a>.</p>

username: bruno-darochac1
cooked: <p>So I’ve tried what you suggest with this code :</p>
<pre><code class="lang-python">import torch

print(f"Should be empty {torch.cuda.max_memory_allocated(0)}")

print("Create x Tensor")
x = torch.rand(100, 100, 100, device='cuda:0')

print(f"    Memory allocated -&gt; {torch.cuda.memory_allocated(0)}")
print(f"    Memory peak -&gt; {torch.cuda.max_memory_allocated(0)}")

print("Create y Tensor")
y = torch.rand(100, 100, 100, device='cuda:0')

print(f"    Memory allocated -&gt; {torch.cuda.memory_allocated(0)}")
print(f"    Memory peak -&gt; {torch.cuda.max_memory_allocated(0)}")

del y

print("After deleting y tensor")

print(f"    Memory allocated -&gt; {torch.cuda.memory_allocated(0)}")
print(f"    Memory peak -&gt; {torch.cuda.max_memory_allocated(0)}")

print("Reseting memory allocated")
torch.cuda.reset_peak_memory_stats(0)
print(f"    Memory allocated -&gt; {torch.cuda.memory_allocated(0)}")
print(f"    Memory peak -&gt; {torch.cuda.max_memory_allocated(0)}")
</code></pre>
<p>So in my training <a href="https://discuss.pytorch.org/t/cuda-memory-profiling/182065/2">code</a> I should, after all epoch done delete all value that was send to gpus (model, labels, input) and then reset the memory peak and it should be okay ?</p>