Pytorch cuda oom, memory usage keep increasing

bach_vu_viet: <p>I am training a model on a few shot problem. However my gpu consumption keep increasing after every iteration. At each iteration, I use only 1 few shot task. I have read other posts on this gpu mem increase issue and implement the suggestions including</p>
<ol>
<li>use total_loss += lose.item() instead of total_loss += loss.</li>
<li>delete variable loss</li>
<li>use torch.cuda.empty_cache()</li>
</ol>
<p>However, it still doesn’t work. I would appreciate your help.<br>
Below is my train loop</p>
<pre data-code-wrap="python"><code class="lang-python">def compute_loss(cfg, input, output):
    loss = F.cross_entropy(output['logits'], input['labels'])
    return loss

def train_epoch(
    cfg,
    train_loader,
    val_loader,
    model,
    optimizer,
    scheduler,
):
    """
    Perform the video training for one epoch.
    Args:
        train_loader (loader): video training loader.
        model (model): the video model to train.
        model_ema (model): the ema model to update.
        optimizer (optim): the optimizer to perform optimization on the model's
            parameters.
        train_meter (TrainMeter): training meters to log the training performance.
        cur_epoch (int): current epoch of training.
        cfg (Config): The global config object.
    """
    # Enable train mode.
    model.to(DEVICE)
    model.train()

    if cfg.TRAIN.type == 'iter':
        num_iter = cfg.TRAIN.NUM_TRAIN_ITER    
    elif cfg.TRAIN.type == 'meta':
        num_iter = cfg.TRAIN.NUM_TRAIN_TASKS
    loader_iter = iter(itertools.cycle(train_loader))


    running_acc = 0
    running_loss = 0

    log_interval = cfg.STEP_PER_LOG
    val_interval = cfg.TRAIN.VAL_FRE_ITER
    assert val_interval % log_interval == 0
    gradient_accumulation_steps = cfg.GRADIENT_ACCUMULATION_STEPS
    
    last_checkpoint_path = None
    keep_num_best = cfg.num_best_checkpoint if hasattr(cfg, "num_best_checkpoint") else 3
    ckpt_saver = CheckpointSaver(keep_num_best=keep_num_best)
    
    # warm up
    input = next(loader_iter)
    input = squeeze(input, DEVICE)
    
    input["split"] = "train"

    output = model(input)
    # debug
    with open("debug.txt", 'w') as f:
        pass
    
    for cur_iter in tqdm(range(num_iter)):
        wandb_log = {}

        input = next(loader_iter)
        input = squeeze(input, DEVICE)
        
        input["split"] = "train"
        
        with profiler.profile(with_stack=True, profile_memory=True) as prof:
            output = model(input)
        
        with open("debug.txt", 'a') as f:    
            f.write(prof.key_averages().table(sort_by="cpu_time_total"))
            f.write("\n")
            
        loss = compute_loss(cfg, input, output)
        loss.backward()

        # acc = compute_accuracy(cfg, input, output)
        
        # running_acc += acc.item()
        running_loss += loss.item()
        
        del loss

        if (cur_iter + 1) % gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()
        
        # evaluate
        if (cur_iter + 1) % val_interval == 0:
            with torch.no_grad():
                res = inductive_evaluate(
                    cfg, model, val_loader, cfg.TRAIN.NUM_VAL_TASKS, "val"
                )

                assert isinstance(res, dict)
                wandb_log.update({"val_acc": res["acc"], "val_loss": res["loss"]})
                stat = {"iter": cur_iter}

                ckpt_saver(
                    res["acc"],
                    cfg.OUTPUT_DIR + f"/best_{cur_iter+1:05}_acc{res['acc']:.3f}.pt",
                    partial(
                        save_checkpoint,
                        model=model,
                        optimizer=optimizer,
                        scheduler=scheduler,
                        stat=stat,
                    ),
                )

                remove_file(last_checkpoint_path)
                last_checkpoint_path = (
                    cfg.OUTPUT_DIR + f"/last_{cur_iter+1:05}_acc{res['acc']:.3f}.pt"
                )
                save_checkpoint(
                    model, optimizer, scheduler, stat=stat, save_path=last_checkpoint_path
                )
            model.train()

        if (cur_iter + 1) % log_interval == 0:
            lr = optimizer.param_groups[0]["lr"]
            wandb_log["lr"] = lr
            wandb_log.update(
                {
                    f"train_loss_{log_interval}": running_loss / log_interval,
                    f"train_acc_{log_interval}": running_acc / log_interval,
                }
            )
            running_loss = 0
            running_acc = 0
            wandb.log(wandb_log)
        
        # del loss, acc, output, input
        # torch.cuda.empty_cache()


    return {"best_checkpoint": ckpt_saver.get_best_checkpoint_path()}
</code></pre>

ptrblck: <p>The first point is valid, the second and third points should not be needed.<br>
Could you post a minimal and executable code snippet showing the memory increase even after the loss is detached?</p>

bach_vu_viet: <p>Thanks for your reply. Here is my minimal code to reproduce the issue.</p><aside class="onebox githubgist" data-onebox-src="https://gist.github.com/vuvietbach/ca103b38b5d2af092d4bfd8f19526bd7">
  <header class="source">

      <a href="https://gist.github.com/vuvietbach/ca103b38b5d2af092d4bfd8f19526bd7" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/vuvietbach/ca103b38b5d2af092d4bfd8f19526bd7" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/vuvietbach/ca103b38b5d2af092d4bfd8f19526bd7</a></h4>

  <h5>env.yml</h5>
  <pre><code class="YAML">name: test_env
channels:
  - pytorch
  - nvidia
  - anaconda
  - conda-forge
  - defaults
dependencies:
  - pytorch=2.0.1=py3.11_cuda11.8_cudnn8.7.0_0
  - torchvision=0.15.2=py311_cu118</code></pre>
    This file has been truncated. <a href="https://gist.github.com/vuvietbach/ca103b38b5d2af092d4bfd8f19526bd7" target="_blank" rel="noopener nofollow ugc">show original</a>
  <h5>oom.py</h5>
  <pre><code class="Python">import torch
import torch.nn.functional as F
import itertools
from torch.utils.data import Dataset, DataLoader
from easydict import EasyDict as edict
import torch.nn as nn
from typing import Tuple, Union
from torch.utils.checkpoint import checkpoint as grad_ckpt
from collections import OrderedDict
from einops import rearrange</code></pre>
    This file has been truncated. <a href="https://gist.github.com/vuvietbach/ca103b38b5d2af092d4bfd8f19526bd7" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


bach_vu_viet: <p>After following this post <a href="https://discuss.pytorch.org/t/how-to-debug-causes-of-gpu-memory-leaks/6741/9" class="inline-onebox">How to debug causes of GPU memory leaks? - #9 by Carson_McNeil</a>, I was able to identify the cause of the memory leak. It seems that the input to my model was being accumulated, resulting in an out-of-memory issue. Here is my log output after printing out allocated tensor.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/1/b11b6c7259215df421a2ce8b260500f49e0c32f5.png" data-download-href="https://discuss.pytorch.org/uploads/default/b11b6c7259215df421a2ce8b260500f49e0c32f5" title="Screenshot 2024-01-10 at 17.49.46"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/1/b11b6c7259215df421a2ce8b260500f49e0c32f5.png" alt="Screenshot 2024-01-10 at 17.49.46" data-base62-sha1="pgLgyKlRS35jAqBroJ2xoI1bFaJ" width="690" height="170" data-dominant-color="2F2F2F"><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">Screenshot 2024-01-10 at 17.49.46</span><span class="informations">906×224 5.71 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><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/1/1139610259648647fc0cb42414187ae6d7331050.png" data-download-href="https://discuss.pytorch.org/uploads/default/1139610259648647fc0cb42414187ae6d7331050" title="Screenshot 2024-01-10 at 17.50.10"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/1/1139610259648647fc0cb42414187ae6d7331050.png" alt="Screenshot 2024-01-10 at 17.50.10" data-base62-sha1="2sn39kqm0EWEGN1KCIRq9jikmn6" width="690" height="316" data-dominant-color="2F2F2F"><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">Screenshot 2024-01-10 at 17.50.10</span><span class="informations">906×416 6.96 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><br>
Still, I have not been able to resolve it yet. Any suggestions would be appreciated.</p>

bach_vu_viet: <p>I believe the error is caused by wrapping the dataloader in <code>itertools.cycle</code> . After removing <code>itertools.cycle</code> , the functionality works correctly.</p>

ptrblck: <p>Great debugging and thanks for sharing the workaround. I had it in my backlog but didn’t have a chance to run it yet.</p>