Shared Cuda Tensor Consumes GPU Memory

stevenwjy: <p>I tried to pass a cuda tensor into a multiprocessing spawn. As per my understanding, it will automatically treat the cuda tensor as a shared memory as well (which is supposed to be a no op according to the docs). However, it turns out that such operation makes PyTorch to be unable to reserve quite a significant memory size of my GPUs (2-3 GBs) – which probably is the reserved storage to make the shared memory works. Is this expected?</p>

ptrblck: <p>Could you post a minimal code snippet showing your issue, please?</p>

stevenwjy: <p>Please refer to the code in [1] to reproduce the issue.</p>
<p>Based on my own experiment with 4 GPUs of 16 GB each:</p>
<p>Running: “python3 src/test_shared.py” finishes with:</p>
<pre><code class="lang-auto">[P0] [GPU Mem (MB)] total: 16130.5, reserved = 13826.0, allocated: 13824.0625, free: 1.9375, step: 99
[P1] [GPU Mem (MB)] total: 16130.5, reserved = 13826.0, allocated: 13824.0625, free: 1.9375, step: 99
[P2] [GPU Mem (MB)] total: 16130.5, reserved = 13826.0, allocated: 13824.0625, free: 1.9375, step: 99
[P3] [GPU Mem (MB)] total: 16130.5, reserved = 13826.0, allocated: 13824.0625, free: 1.9375, step: 99
</code></pre>
<p>However, running “python3 src/test_shared.py --shm_cuda”, throws the following error:</p>
<pre><code class="lang-auto">[P1] [GPU Mem (MB)] total: 16130.5, reserved = 9858.0, allocated: 9856.0625, free: 1.9375, step: 68
Traceback (most recent call last):
  File "[...]/test_shared.py", line 84, in &lt;module&gt;
    main()
  File "[...]/test_shared.py", line 80, in main
    mp.spawn(train, nprocs=num_gpus, args=(args, num_gpus, shm_list))
  File "[...]/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "[...]/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
    while not context.join():
  File "[...]/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 150, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException:

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "[...]/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
    fn(i, *args)
  File "[...]/test_shared.py", line 53, in train
    inp = torch.rand((args.batch_size, args.inp_size)).to(rank)
RuntimeError: CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 15.75 GiB total capacity; 9.63 GiB already allocated; 59.88 MiB free; 9.63 GiB reserved in total by PyTorch)
</code></pre>
<p>As could be seen, PyTorch seems to not be able to reserve more memory? The shared tensors are only: (2^10) * (2^14) in size each though… Any thoughts on how it could end up taking &gt;4GB?</p>
<p>[1] Sample code</p>
<pre><code class="lang-python">import os
from argparse import ArgumentParser

import torch
import torch.multiprocessing as mp
import torch.distributed as dist

def load_args():
    parser = ArgumentParser(add_help=False)
    parser.add_argument('--gpu_limit', type=int, default=4, help='Max num of GPUs used')
    parser.add_argument('--num_steps', type=int, default=100, help='Number of training steps')
    parser.add_argument('--batch_size', type=int, default=(2**10), help='Batch size')
    parser.add_argument('--inp_size', type=int, default=(2**14), help='Input size')
    parser.add_argument('--out_size', type=int, default=(2**14), help='Output size')
    parser.add_argument('--shm_cuda', action='store_true')

    args = parser.parse_args()
    return args


def init_dist(rank, world_size, backend="nccl"):
    master_addr = os.environ['MASTER_ADDR']
    master_port = os.environ['MASTER_PORT']

    dist.init_process_group(
        backend,
        init_method="tcp://%s:%s" % (master_addr, master_port),
        rank=rank,
        world_size=world_size)


def log_gpu_mem(local_rank, step, unit="MB"):
    u2d = {"KB": 1024, "MB": 1024 ** 2}
    d = u2d[unit]

    total = torch.cuda.get_device_properties(local_rank).total_memory / d
    reserved = torch.cuda.memory_reserved(local_rank) / d
    allocated = torch.cuda.memory_allocated(local_rank) / d
    free = reserved - allocated

    s = f"[GPU Mem ({unit})] total: {total}, reserved = {reserved}, allocated: {allocated}, free: {free}"
    print(f"[P{local_rank}] " + s + f", step: {step}")


def train(rank, args, world_size, shm_list):
    torch.cuda.set_device(rank)
    init_dist(rank, world_size)

    model = torch.nn.Linear(args.inp_size, args.out_size).to(rank)

    record = []
    for step in range(args.num_steps):
        inp = torch.rand((args.batch_size, args.inp_size)).to(rank)
        out = model(inp)
        record.append(out)

        shm_list[step % len(shm_list)].data.copy_(out.data)

        log_gpu_mem(rank, step)


def main():
    num_gpus = torch.cuda.device_count()
    if num_gpus &lt; 2:
        raise Exception("At least 2 GPUs are required to run this script")

    args = load_args()
    num_gpus = min(num_gpus, args.gpu_limit)

    print(f"Using {num_gpus} GPUs for testing")

    shm_list = []
    for i in range(num_gpus):
        shm = torch.zeros((args.batch_size, args.out_size))
        if args.shm_cuda:
            shm = shm.to(i)

        shm_list.append(shm)

    mp.spawn(train, nprocs=num_gpus, args=(args, num_gpus, shm_list))


if __name__ == '__main__':
    main()
</code></pre>

stevenwjy: <p>Also, maybe does anyone have any reference on what’s happening behind the scene when a Cuda tensor is being converted into shared memory – the docs seem to say that its a no-op; however, is it really that the reference to the GPU memory location is just being passed into multiple processes? Would be great if there’s any recommended reading to better understand about this</p>

ptrblck: <p>Sorry, I don’t quite understand your use case as you are not using <code>tensor.share_memory()</code>, which seems to be the op you are referring to.<br>
In any case, I also see identical memory usage in both use cases (i.e. with the additional argument and without).</p>

stevenwjy: <p>I believe that all tensors passed into multiprocessing spawn will be converted into shared memory?<br>
According to [1], that’s the case when using multiprocessing queue. I actually experimented with this as well before and I believe that the tensors passed as args when spawning sub-processes are indeed converted to shared memory. Otherwise, can also try to add <code>shm.share_memory_()</code> before appending it to the list, and I believe the result will still be the same.</p>
<p>Besides, thanks for taking a look into it! Do you perhaps have any thoughts on what might cause such a different in memory usage?</p>
<p>[1] <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html" class="inline-onebox" rel="noopener nofollow ugc">Multiprocessing best practices — PyTorch 1.9.1 documentation</a></p>

ptrblck: <aside class="quote no-group" data-username="stevenwjy" data-post="6" data-topic="134543">
<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/stevenwjy/48/34212_2.png" class="avatar"> stevenwjy:</div>
<blockquote>
<p>Do you perhaps have any thoughts on what might cause such a different in memory usage?</p>
</blockquote>
</aside>
<p>No, since I’m seeing the same memory usage and thus cannot reproduce the issue.<br>
Could you post your output of both runs for the first ~10 steps?</p>

stevenwjy: <p>Oh, sorry for not communicating this clearly earlier – but I think the issue that I raised in this thread is that PyTorch could not reserve as much memory as what it did for the one without “–shm_cuda” arg.</p>
<p>If you take a look at my sample output before throwing the error:</p>
<pre><code class="lang-auto">[P1] [GPU Mem (MB)] total: 16130.5, reserved = 9858.0, allocated: 9856.0625, free: 1.9375, step: 68
...
RuntimeError: CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 15.75 GiB total capacity; 9.63 GiB already allocated; 59.88 MiB free; 9.63 GiB reserved in total by PyTorch)
</code></pre>
<p>After reserving for ~10GB, it went out of memory even though the total available is ~16GB (there was no other process using the GPU). Meanwhile, in the earlier case you could see that the program managed to finish with reserving ~13GB</p>
<pre><code class="lang-auto">[P0] [GPU Mem (MB)] total: 16130.5, reserved = 13826.0, allocated: 13824.0625, free: 1.9375, step: 99
</code></pre>
<p>The memory usages are indeed the same for both use cases. It’s just that somehow PyTorch couldn’t reserve the same amount of memory.</p>
<p>I purposely made the sample code to get out of memory. Hence if you have GPUs with diff total memory, prob need to reduce/increase the number of steps to be able to reproduce.</p>