Testing Distributed PyTorch code on a single GPU

Alex_McKinney: <p>Greetings,</p>
<p>A rather simple question, that may be complex to answer.</p>
<p>Anyone know how to test distributed PyTorch code on a single GPU? I don’t care how slow it executes, just that it is numerically accurate.</p>
<p>I’m writing a no-nonsense blog on different distributed PyTorch methods, but would rather not have to rent a cluster just to write a blog.</p>
<p>With Jax, I could set XLA environment variable <code>xla_force_host_platform_device_count</code> to force a CPU to appear like many for testing purposes. Anything like this exist in PyTorch?</p>
<p>Thanks!</p>

H-Huang: <p>Hi, thanks for the question. This depends on what aspect of distributed PyTorch you would like to test. For example, if you just need to call collectives (e.g. <code>dist.allreduce</code>, <code>dist.broadcast</code>) or DDP, then you can just use the gloo backend which operates with CPU (<code>dist.init_process_group(backend="gloo")</code>.</p>
<p>If you need something specific to GPU, for example testing NCCL collectives or using FSDP. We don’t currently have a way to virtually split a single GPU to act as multiple.</p>

Alex_McKinney: <blockquote>
<p>This depends on what aspect of distributed PyTorch you would like to test.<br>
Sadly, I want to test… all of them! Lofty goal is 3D parallelism, but would settle for 2D.</p>
</blockquote>
<p>So, that would include things like FSDP, NCCL collectives, as well as pipelining.</p>
<blockquote>
<p>We don’t currently have a way to virtually split a single GPU to act as multiple.<br>
So, the conclusion is I should probably start using my wallet to rent a small cluster <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>
</blockquote>
<p>I saw a few datacentre inference GPUs do have support for splitting GPUs into virtual ones, but nothing on consumer hardware. If PyTorch doesn’t have support, would you know of anything on a consumer GPU that could work?</p>

ptrblck: <aside class="quote no-group" data-username="Alex_McKinney" data-post="3" data-topic="185833">
<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/alex_mckinney/48/61281_2.png" class="avatar"> Alex_McKinney:</div>
<blockquote>
<p>I saw a few datacentre inference GPUs do have support for splitting GPUs into virtual ones</p>
</blockquote>
</aside>
<p>I assume you are referring to MIG on NVIDIA data center GPUs. If so, note that no communication between MIG slices is supported and only a single process per MIG (without IPC) is supported, so MIG won’t help.</p>

Alex_McKinney: <p>MIG was one, yes. I think <code>vgu</code> may also be a candidate (as in, the subcommand <code>nvidia-smi vgpu</code> however this is also not on my hardware. I wasn’t aware that MIG didn’t have support for comms between slices though <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>
<p>Seems this is a dead end for now?</p>