Surviving OOM events in distributed training

username: yoelshoshan
cooked: <p>Dealing with varying input size, I catch OOM exceptions during training (in my setting roughly 1 in few hundred minibatches). Due to domain specific reasons, I prefer not to crop/resize inputs to a constant size. Also, there is not a clear way to know in advance which input sizes will cause an OOM.<br>
This is generally fine by me as long as I can recover from the OOM events and continue training.</p>
<p>If I detect an OOM event, I’m “cleaning up” using <code>torch.cuda.empty_cache()</code>, zero gradients, and then continue training as usual. This works great in a non-distributed setup, but creates problems in a distributed setting.</p>
<p>note - I am following the suggested way to deal with OOM as mentioned here:</p><aside class="quote" data-post="1" data-topic="12754">
  <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/cab0a1/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/is-it-safe-to-recover-from-cuda-oom/12754">Is it safe to recover from CUDA OOM?</a> 
  </div>
  <blockquote>
    I’m looking at trying to improve the robustness of our trainer code.  If I select batch parameters that are not tight enough, I may run out of memory.  Is it safe to catch the OOM, reduce the batch size, and try again? 
Thanks 
Jerry
  </blockquote>
</aside>

<p>To deal with OOM in a distributed setting, I do something like this:</p>
<pre><code class="lang-auto">            if problems_occured:
                success_tens = torch.ones(0)
            else:
                success_tens = torch.ones(1)

            dist.all_reduce(success_tens, op=dist.reduce_op.SUM)  ###error happens here
</code></pre>
<p>and then, only if success_tens reached the size of the world_size, I do an additional all_reduce over the gradients to sum them.<br>
This is to make sure that all workers succeeded in calculating their own gradient before combining the gradient.</p>
<p>However, after I catch the OOM event in the worker that caught this OOM event I get the following error:</p>
<pre><code class="lang-auto">miniconda3/envs/py36torch/lib/python3.6/site-packages/torch/distributed/distributed_c10d.py", line 838, in all_reduce
    work.wait()
RuntimeError: [enforce fail at /opt/conda/conda-bld/pytorch_1544174967633/work/third_party/gloo/gloo/allreduce.cc:29] opts.elements &gt; 0
</code></pre>
<p>note: as can be seen in the error message - I’m currently using gloo as the distributed backend.</p>
<p>Any suggestions on how to solve this are very welcome <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Note - while I’m currently using a simple syncronized distributed gradients calculation, I’m open to any suggestions, as long as they help survive occasional (relatively rare) OOM events.</p>

username: yoelshoshan
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> any idea? the tl;dr is that I can survive OOM events in a none distributed setting (just discarding this training minibatch, cleaning up memory and continuing), but using distributed setting I can’t.<br>
This is important for me as I’m in medical imaging setting in which resolution is important, and cropping is too destructive.<br>
Any suggestions and/or pointing me to relevant people if necessary is very welcome <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: ptrblck
cooked: <p>Sorry, I’m not really familiar with distributed training and gloo, so I can’t give any useful input to fix this issue. <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=6" title=":confused:" class="emoji" alt=":confused:"></p>
<p>However, have you thought about other approaches to avoid the OOM issues?<br>
<a href="https://pytorch.org/docs/stable/checkpoint.html" rel="nofollow noopener">torch.utils.checkpoint</a> might be worth a try (although I’m not sure how it behaves in a distributed setup) or <a href="https://github.com/NVIDIA/apex" rel="nofollow noopener">NVIDIA’s apex - mixed precision training</a>.</p>
<p>Let me know, this would be an option.</p>

username: yoelshoshan
cooked: <p>Thanks! I’m already using both checkpointing and mixed precision, which helped to make the OOM events pretty rare, but they still exist here and there. Perhaps it’s reasonable to consider this “a bug” or feature request and just report it on the github channel.</p>

username: pietern
cooked: <p>Hi <a class="mention" href="/u/yoelshoshan">@yoelshoshan</a>! The error message indicates that the tensor that you’re passing to allreduce is empty. Is it possible that the “success_tens” itself is somehow empty? Not sure this is possible, but since you’re already dealing with an OOM…</p>

username: yoelshoshan
cooked: <p>Hi! first of all thanks for trying to assist <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
success_tens is not empty.</p>
<p>I managed to build NCCL and when using it as the backend for the distributed functions, this issue does not happen, so I believe that this is gloo specific problem.</p>

username: pietern
cooked: <p>OK, I just realized what’s going on here. I misunderstood the code snippet you list in the original post. If you see an OOM, you create an <strong>empty tensor</strong>. This is why the error triggers. Instead of <code>torch.ones(0)</code>, you’ll want to use <code>torch.zeros(1)</code>.</p>

username: yoelshoshan
cooked: <p>Like you suggested, using torch.tensor(0.0) or torch.tensor(1.0) does not trigger that issue.</p>
<p>Thanks for the help! &lt;3</p>