Same model and data somehow OOM on better device

JackHu: <p>I’m training a TTS model and the training process is fine on RTX4090, but when moving to a device with 8 *A100-SXM4-40GB, it somehow suddenly get cuda OOM error after thousands of epochs.</p>
<p>Here is the detail situation:<br>
<strong>Both training on RTX4090 or A100 uses exactly same model, config, data, coding, cuda version(12.2)</strong></p>
<ol>
<li>
<p>Epoch 1~4000 training on RTX4090, the GPU memory usage is as picture, everything is fine, no oom, no memory leak, no growing memory usage.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/2/62a9517f3a88580203b596c937ebef60455dfb80.png" alt="4090" data-base62-sha1="e4NypnP9EkHn1JbIbphR1cHpmRG" width="608" height="149"></p>
</li>
<li>
<p>Moving everything to 8 *A100-SXM4-40GB</p>
</li>
<li>
<p>Epoch4001~7000 is still fine on A100 device</p>
</li>
<li>
<p><strong>Epoch 7001~7051 starts to get cuda OOM</strong><br>
GPU memory usage when epoch 7005<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/e/7e3e2b6525d5b60885b07a8593a397cb07efd615.png" data-download-href="https://discuss.pytorch.org/uploads/default/7e3e2b6525d5b60885b07a8593a397cb07efd615" title="f5"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e3e2b6525d5b60885b07a8593a397cb07efd615_2_579x500.png" alt="f5" data-base62-sha1="i0Nmfo6DV4ixMCBdIxmbF1qVKg5" width="579" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e3e2b6525d5b60885b07a8593a397cb07efd615_2_579x500.png, https://discuss.pytorch.org/uploads/default/original/3X/7/e/7e3e2b6525d5b60885b07a8593a397cb07efd615.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/e/7e3e2b6525d5b60885b07a8593a397cb07efd615.png 2x" data-dominant-color="0B0B0B"><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">f5</span><span class="informations">612×528 143 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>
GPU memory usage near epoch 7050<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/3/332d3c4134e52b8f0099776fe9e8110aa8412196.png" data-download-href="https://discuss.pytorch.org/uploads/default/332d3c4134e52b8f0099776fe9e8110aa8412196" title="f50"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/3/332d3c4134e52b8f0099776fe9e8110aa8412196_2_572x500.png" alt="f50" data-base62-sha1="7iJgzBif5xCng8Tku6dZRzZmebQ" width="572" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/3/332d3c4134e52b8f0099776fe9e8110aa8412196_2_572x500.png, https://discuss.pytorch.org/uploads/default/original/3X/3/3/332d3c4134e52b8f0099776fe9e8110aa8412196.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/3/332d3c4134e52b8f0099776fe9e8110aa8412196.png 2x" data-dominant-color="0B0A0B"><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">f50</span><span class="informations">605×528 156 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>
</li>
<li>
<p>When loading data for epoch 7051, cuda OOM will occur<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/8/684575275f760d0ab2889a5d709d5fe0e622d516.png" alt="oom" data-base62-sha1="eSqt5FyihdN6n4kKCq29ZDIuqJo" width="624" height="46"></p>
</li>
</ol>
<p>If there is a memory leak or coding problem, I suppose it should also happen when it is on 4090 or on A100 during epoch 4000 to epoch 7000, but it just suddenly happen after so many stable training progress, so I have no clue how this is happenning.</p>
<p>Any advice or clue is appreciated, thank you!</p>
<p>Here is the source code I used for training.<br>
<a href="https://github.com/p0p4k/vits2_pytorch" rel="noopener nofollow ugc">https://github.com/p0p4k/vits2_pytorch</a></p>

ptrblck: <aside class="quote no-group" data-username="JackHu" data-post="1" data-topic="189469">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/59ef9b/48.png" class="avatar"> JackHu:</div>
<blockquote>
<p>but it just suddenly happen after so many stable training progress</p>
</blockquote>
</aside>
<p>The memory usage doesn’t look stable as in the first screenshot you are already seeing a heavy imbalanced. E.g. could you explain why device 0 uses approx. 9GB of it’s memory while device 3 uses approx. 18GB? Is each device getting different input shapes, which could increase the memory usage?</p>

JackHu: <p>Honestly, I have no idea, it’s my first time training with multi-gpu, the memory usage just keep changing<br>
and each GPU got different memory usage when loading data, I thought it’s normal…<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/e/be19c5b1f5bdfbc40f3aef5f6584db6159316994.gif" alt="gp" data-base62-sha1="r7HW2Dl0CogKKb0eouWYf3rnxQw" width="579" height="500" class="animated"></p>

ptrblck: <p>The memory usage won’t decrease unless you are clearing the cache. So could you explain why you are calling into <code>torch.cuda.empty_cache()</code>?</p>

JackHu: <p>Just for record, I’m not the author of that github repository, and I’m not very experienced with building model from scratch or the torch code.</p>
<p>Here is the main function</p>
<pre><code class="lang-auto">def main():
    """Assume Single Node Multi GPUs Training Only"""
    assert torch.cuda.is_available(), "CPU training is not allowed."

    n_gpus = torch.cuda.device_count()
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "6060"

    hps = utils.get_hparams()
    mp.spawn(
        run,
        nprocs=n_gpus,
        args=(
            n_gpus,
            hps,
        ),
    )
</code></pre>
<p>And part of the run function that initials data loader</p>
<pre><code class="lang-auto">def run(rank, n_gpus, hps):
    global global_step
    if rank == 0:
        logger = utils.get_logger(hps.model_dir)
        logger.info(hps)
        utils.check_git_hash(hps.model_dir)
        writer = SummaryWriter(log_dir=hps.model_dir)
        writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))

    dist.init_process_group(
        backend="nccl", init_method="env://", world_size=n_gpus, rank=rank
    )
    torch.manual_seed(hps.train.seed)
    torch.cuda.set_device(rank)

    train_dataset = TextAudioLoader(hps.data.training_files, hps.data)
    train_sampler = DistributedBucketSampler(
        train_dataset,
        hps.train.batch_size,
        [32, 300, 400, 500, 600, 700, 800, 900, 1000],
        num_replicas=n_gpus,
        rank=rank,
        shuffle=True,
    )

    collate_fn = TextAudioCollate()
    train_loader = DataLoader(
        train_dataset,
        num_workers=8,
        shuffle=False,
        pin_memory=True,
        collate_fn=collate_fn,
        batch_sampler=train_sampler,
    )
    if rank == 0:
        eval_dataset = TextAudioLoader(hps.data.validation_files, hps.data)
        eval_loader = DataLoader(
            eval_dataset,
            num_workers=8,
            shuffle=False,
            batch_size=1,
            pin_memory=True,
            drop_last=False,
            collate_fn=collate_fn,
        )
</code></pre>
<p><code>torch.cuda.empty_cache()</code> is not called in the source code, if decreasing memory is not normal during training, what else could possibly cause it?</p>