Unified memory across NVSwitch'ed GPUs

username: mohotmoz
cooked: <p>Hi everyone,</p>
<p>Fortunately, it is becoming easier and easier to access NVSwitch’ed GPUs in commodity cloud environments, even with some pretty attractive spot pricing. My general understanding is that NVSwitch allows uniform memory access to all GPUs from all GPUs, so you can think of it as “one giant GPU” with a continuous memory space. CUDA unified memory allows for this, using CUDAMallocManaged (though the other “purpose” of CUDA unified memory is to seamlessly integrate with much slower, host memory). My question is this - is there a way to get PyTorch to, when on an NVSwitch system, treat all the memory as one big block? I know that of course its better to do manual data parallelism, and/or shard your model across GPUs if it’s too big, use DeepSpeed, etc. etc. but, for ease of use… wouldn’t it be nice if you could just access more and more memory going across multiple devices? The ease of use would definitely more than make up for a performance penalty (NVSwitch bandwidth still is not the same as on-device bandwidth but it is like relatively close…)… there are a lot of times where it would be nice to port single GPU code over to a bigger system only for the purpose of making the model a little bit bigger than one GPU, making the batch size a little bigger (without dealing with distributed data parallel), etc.</p>
<p>I have tried to recompile pytorch replacing CUDAMalloc with CUDAMallocManaged as in this paper:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://www.mdpi.com/2076-3417/11/21/10377">
  <header class="source">

      <a href="https://www.mdpi.com/2076-3417/11/21/10377" target="_blank" rel="noopener nofollow ugc">MDPI</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:550/204;"><img src="https://www.mdpi.com/applsci/applsci-11-10377/article_deploy/html/images/applsci-11-10377-g001-550.jpg" class="thumbnail" width="550" height="204"></div>

<h3><a href="https://www.mdpi.com/2076-3417/11/21/10377" target="_blank" rel="noopener nofollow ugc">Efficient Use of GPU Memory for Large-Scale Deep Learning Model Training</a></h3>

  <p>To achieve high accuracy when performing deep learning, it is necessary to use a large-scale training model. However, due to the limitations of GPU memory, it is difficult to train large-scale training models within a single GPU. NVIDIA introduced a...</p>


  </article>

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

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

<p>But PyTorch still pins the memory to one device (and if it grows beyond that, it spills into CPU memory not onto the other GPUs).</p>
<p>So anyway figured I would ask folks here that have a more intimate knowledge of PyTorch’s allocation code to see if this was a quick fix?</p>
<p>Thanks everyone for your time!</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="mohotmoz" data-post="1" data-topic="137436">
<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/m/ed655f/48.png" class="avatar"> mohotmoz:</div>
<blockquote>
<p>My general understanding is that NVSwitch allows uniform memory access to all GPUs from all GPUs, so you can think of it as “one giant GPU” with a continuous memory space.</p>
</blockquote>
</aside>
<p>No, <code>NVSwitch</code> uses a more sophisticated connectivity between the devices as described <a href="https://www.nvidia.com/en-us/data-center/nvlink/">here</a>. It doesn’t “automatically” create a single device.</p>
<aside class="quote no-group quote-modified" data-username="mohotmoz" data-post="1" data-topic="137436">
<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/m/ed655f/48.png" class="avatar"> mohotmoz:</div>
<blockquote>
<p>I have tried to recompile pytorch replacing CUDAMalloc with CUDAMallocManaged as in this paper<br>
[…]<br>
But PyTorch still pins the memory to one device (and if it grows beyond that, it spills into CPU memory not onto the other GPUs).</p>
</blockquote>
</aside>
<p>That would be expected, as <a href="https://developer.nvidia.com/blog/unified-memory-cuda-beginners/">Unified memory</a> is used.</p>

username: mohotmoz
cooked: <p>thank you for the fast reply -<br>
reading the link you sent - emphasis added:<br>
It supports full all-to-all communication with direct GPU peer-to-peer memory addressing. These 16 GPUs <strong>can be used as a single high-performance accelerator with unified memory space</strong> and up to 10 petaFLOPS of deep learning compute power.<br>
I guess I’m just wondering if that is possible for DL workloads with PyTorch…? Sounds like no? thank you!</p>

username: ptrblck
cooked: <p>I don’t think the “single accelerator” means it’s usable as a single device, but as an accelerated system due to the better interconnectivity.</p>
<aside class="quote no-group" data-username="mohotmoz" data-post="3" data-topic="137436">
<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/m/ed655f/48.png" class="avatar"> mohotmoz:</div>
<blockquote>
<p>I guess I’m just wondering if that is possible for DL workloads with PyTorch…? Sounds like no?</p>
</blockquote>
</aside>
<p>Yes, you can use NVLink/NVSwitch through <code>DistributedDataParallel</code> or any other parallel approach.</p>