Some Common Solutions for PyTorch Running Out of GPU Memory

torch.cuda.memory_allocated(device)

By logging the GPU memory usage at each step during forward propagation, you can identify which step consumes the most memory.

for i, batch in enumerate(train_loader):
    print("1:", torch.cuda.memory_allocated(0))

    outputs = model(**batch)
    print("2:", torch.cuda.memory_allocated(0))

    loss = outputs.loss
    print("3:", torch.cuda.memory_allocated(0))

    loss.backward()
    print("4:", torch.cuda.memory_allocated(0))

    optimizer.step()
    optimizer.zero_grad()
    print("5:", torch.cuda.memory_allocated(0))

Accumulated Computation Graphs

Many variables in PyTorch are backed by computation graphs, which can easily lead to accidental accumulation of computation graphs.

For example, appending losses in a for loop, or accumulating losses, will cause the computation graph to accumulate. All previous variables will be retained, eventually causing OOM.

with torch.no_grad() can prevent gradient computation, but it does not release accumulated tensors. Only detach() or .item() can do that.

import swanlab
from swanlab.integration.transformers import SwanLabCallback
from transformers import Trainer, TrainingArguments, DataCollatorWithPadding
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Trainer
from torch.utils.data import DataLoader, Sampler
import torch
import math
import random


swanlab_callback = SwanLabCallback(
    project="map", 
    experiment_name=config['name'],
)

callbacks = []

if ENV == "SERVER":
    callbacks.append(swanlab_callback)



class SupConTrainer(Trainer):
    def __init__(self, contrastive_weight=0.1, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.contrastive_weight = contrastive_weight

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):

        # ⚡ Training phase requires hidden states
        # print("1:", torch.cuda.memory_allocated(0))
        # Standard forward pass
        outputs = model(**inputs, output_hidden_states=True)
        # print("outputs:", outputs)
        # print("2:", torch.cuda.memory_allocated(0))
        logits = outputs.logits
        labels = inputs["labels"]

        # -------- 1. Cross-entropy loss --------
        ce_loss = F.cross_entropy(logits, labels)
        # print("3:", torch.cuda.memory_allocated(0))
        # -------- 2. Contrastive learning loss --------
        # Take the last layer's hidden states (batch_size, hidden_dim)
        hidden_states = outputs.hidden_states[-1][:, 0, :]  # [CLS] vector

        # Normalize
        hidden_states = F.normalize(hidden_states, dim=-1)
        # print("4:", torch.cuda.memory_allocated(0))
        # Similarity matrix (batch, batch)
        similarity_matrix = torch.matmul(hidden_states, hidden_states.T)
        # print("5:", torch.cuda.memory_allocated(0))
        # Only use samples of the same class as positive pairs
        mask = labels.unsqueeze(0) == labels.unsqueeze(1)  # (batch, batch)

        # InfoNCE / SupCon loss
        logits_contrastive = similarity_matrix / 0.1  # temperature parameter 0.1
        contrastive_loss = -torch.log_softmax(logits_contrastive, dim=1)[mask].mean()
        # print("6:", torch.cuda.memory_allocated(0))
        # -------- 3. Total loss --------
        loss = ce_loss + self.contrastive_weight * contrastive_loss
        # print("7:", torch.cuda.memory_allocated(0))
        # else:
        #     # ⚡ Eval phase only needs logits, avoid returning hidden_states
        #     outputs = model(**inputs, output_hidden_states=False)
        #     logits = outputs.logits
        #     labels = inputs["labels"]
        #     loss = self.ce_loss(logits, labels)

        outputs = {
            "loss": outputs['loss'],
            "logits": logits,
        }

        return (loss, outputs) if return_outputs else loss
        

trainer = SupConTrainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    tokenizer=tokenizer,
    compute_metrics=compute_map3,
    # data_collator=data_collator,
    callbacks=callbacks,
    contrastive_weight= config['contrastive_weight'],  # Control contrastive loss weight
)


trainer.train()

The above is a modified trainer function. If the outputs are not modified, the hidden_states will be retained and passed to the outside, causing the hidden states to accumulate continuously and eventually leading to OOM.