Using apex AMP (Automatic Mixed Precision) with model parallelism

username: Caesar
cooked: <p>My model has a few LSTMs which run out of Cuda memory when run on large sequences with one GPU. So I shifted a few components of the model to another GPU. I tried 2 things with Apex AMP:</p>
<ol>
<li>Move the model components to another GPU before invoking  <code>amp.initialize</code> . In this case, I get NaNs soon after first backpropagation.</li>
<li>First invoke  <code>amp.initialize</code> , and then move the model components to another GPU. In this case, its like the model backpropagation runs on a single GPU. It runs out of Cuda memory.</li>
</ol>
<p>The model training runs fine without Apex, so I suppose I am missing some step where the loss is backpropagated on both GPUs. I looked through the <a href="https://nvidia.github.io/apex/parallel.html" rel="nofollow noopener">documentations</a> of Apex, however, it only talks about DataParallelism, and not ModelParallelism.</p>
<p>Any ideas?</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/caesar">@Caesar</a>, have you tried the native AMP in PyTorch? I haven’t tried that with model parallel yet, but if it does not work, that will be a bug that we need to fix.</p>
<p><a href="https://pytorch.org/docs/stable/amp.html" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/docs/stable/amp.html</a><br>
<a href="https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples</a></p>

username: Caesar
cooked: <p>Thanks for your response. I am constrained to use an older version of PyTorch which does not support AMP natively. So I am usingNVIDIA apex. <a href="https://github.com/NVIDIA/apex" rel="nofollow noopener">https://github.com/NVIDIA/apex</a></p>

username: mrshenli
cooked: <p>I see.</p>
<p>cc AMP author <a class="mention" href="/u/mcarilli">@mcarilli</a></p>

username: mcarilli
cooked: <p>I don’t think the apex amp API supports this without complex/undocumented hacks.  Apex amp is in maintenance mode now, no new features will be added.</p>
<p>However, <code>torch.cuda.amp</code> is designed to support model parallelism (ie different layers on different device) out of the box.  Please consider upgrading if at all possible.</p>

username: Caesar
cooked: <p>Thanks for your response.</p>

username: mcarilli
cooked: <p>There’s a <a href="https://github.com/pytorch/pytorch/blob/1a21c92364ea0527d7e8d3b0c6728e8e768e727e/test/test_cuda.py#L2127" rel="nofollow noopener">multi-GPU <code>torch.cuda.amp.GradScaler</code> test case</a> that ensures ordinary GradScaler usage supports networks with layers on different devices.</p>
<p><code>torch.cuda.amp.autocast</code> locally enables/disables autocast for all devices used by the invoking thread.  (However, the autocast state is thread local, so if you spawn a thread to control each device, you must re-invoke autocast in the side thread(s).  This affects usage with <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#working-with-multiple-gpus" rel="nofollow noopener"><code>torch.nn.DataParallel</code> and <code>torch.nn.parallel.DistributedDataParallel</code> with multiple GPUs per process</a>.)</p>