Visualize and understand GPU memory in PyTorch

While it's easy to see that GPU memory is full, understanding why and how to fix it can be more challenging. In this tutorial, we'll go step by step on how to visualize and understand GPU memory usage in PyTorch during training. We’ll also see how to estimate memory requirements and optimize GPU memory usage.

PyTorch provides a handy tool for visualizing GPU memory usage:

Running this code generates a profile.pkl file that contains a history of GPU memory usage during execution. You can visualize this history at: https://pytorch.org/memory_viz .

Model Initialization ( model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda") ): The first step involves loading the model onto the GPU. The model parameters (in blue) occupy memory and remain there until the training ends.

Optimizer Step ( optimizer.step() ): Gradients are used to update the model’s parameters. Initially, the optimizer itself is initialized (green zone). This initialization is only done once. After that, the optimizer uses the gradients to update the model’s parameters. To update the parameters, the optimizer temporarily stores intermediate values (red zone). After the update, both the gradients (yellow) and the intermediate optimizer values (red) are discarded, freeing memory.

For example, a model with 1.5 billion parameters and a precision of 4 bytes requires:

The memory required for the optimizer state depends on the optimizer type and the model parameters. For instance, the AdamW optimizer stores two moments (first and second) per parameter. This makes the optimizer state size:

Input Tensor Creation (1st Loop) : Memory increases by 200 MB matching the input tensor size:

Input Tensor Creation (2nd Loop) : Memory increases by 200 MB for a new input tensor. At this point, you might expect the input tensor from step 2 to be freed. Still, it isn't: the model retains its activation, so even if the tensor is no longer assigned to the variable inputs , it remains referenced by the model's forward pass computation. The model retains its activations because these tensors are required for the backpropagation process in neural networks. Try with torch.no_grad() to see the difference.

Forward Pass (2nd Loop) : Memory increases by 1 GB for the new output tensor, calculated as in step 3.