NCCL timeout when reducing batch size

username: addisonklinke
cooked: <p>Hi all, looking for some debugging insight. I’m successfully using using Lightning’s <a href="https://lightning.ai/docs/pytorch/stable/common/trainer.html?utm_term=&amp;utm_campaign=PMax+Q1+FY26&amp;utm_source=adwords&amp;utm_medium=ppc&amp;hsa_acc=1332059986&amp;hsa_cam=22691136685&amp;hsa_grp=&amp;hsa_ad=&amp;hsa_src=x&amp;hsa_tgt=&amp;hsa_kw=&amp;hsa_mt=&amp;hsa_net=adwords&amp;hsa_ver=3&amp;gad_source=1&amp;gad_campaignid=22681463925&amp;gbraid=0AAAAA9Wu0105E9ASdP_QVjuKCFp8EEwXJ&amp;gclid=CjwKCAjw1dLDBhBoEiwAQNRiQQoshso4tUSx6jcHmcRtmVlbQuy3Ke9yDP-aXzgvx36gO6wyYtj8dRoCvm4QAvD_BwE" rel="noopener nofollow ugc"><code>Trainer(devices=1+)</code></a> to run DDP trainings for a recsys model. However, if I lower the batch size below a critical limit (~1/3rd the normal setting), then I start getting consistent NCCL timeouts. So far, the most direct evidence I’ve found is the below traceback</p>
<pre><code class="lang-auto">[rank0]:   File "/opt/env/lib/python3.11/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 319, in _training_step
[rank0]:     training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
[rank0]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/env/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py", line 328, in _call_strategy_hook
[rank0]:     output = fn(*args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/env/lib/python3.11/site-packages/pytorch_lightning/strategies/strategy.py", line 390, in training_step
[rank0]:     return self._forward_redirection(self.model, self.lightning_module, "training_step", *args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/env/lib/python3.11/site-packages/pytorch_lightning/strategies/strategy.py", line 641, in __call__
[rank0]:     wrapper_output = wrapper_module(*args, **kwargs)
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 1632, in forward
[rank0]:     inputs, kwargs = self._pre_forward(*inputs, **kwargs)
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 1530, in _pre_forward
[rank0]:     self._sync_buffers()
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 2165, in _sync_buffers
[rank0]:     self._sync_module_buffers(authoritative_rank)
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 2169, in _sync_module_buffers
[rank0]:     self._default_broadcast_coalesced(authoritative_rank=authoritative_rank)
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 2191, in _default_broadcast_coalesced
[rank0]:     self._distributed_broadcast_coalesced(bufs, bucket_size, authoritative_rank)
[rank0]:   File "/opt/env/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 2106, in _distributed_broadcast_coalesced
[rank0]:     dist._broadcast_coalesced(
...
[rank0]: RuntimeError: [Rank 0]: Ranks 1, 2, 3, 4, 5, 6, 7 failed to pass monitoredBarrier in 1800000 ms
</code></pre>
<p>It appears <code>_sync_buffers()</code> is hanging on rank 0 causing ranks 1+ to hit the default NCCL timeout at 30 minutes. I’m curious whether anyone has seen similar behavior before? Particularly correlated to changes in batch size. I’ll be digging into the source code more to understand the point of the buffer sync, but would appreciate any pointers. Thanks!</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="addisonklinke" data-post="1" data-topic="221522">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/addisonklinke/48/20034_2.png" class="avatar"> addisonklinke:</div>
<blockquote>
<p>I’m curious whether anyone has seen similar behavior before? Particularly correlated to changes in batch size.</p>
</blockquote>
</aside>
<p>We’ve seen similar issues if the number of batches on all ranks is unequal, which could happen if the batch size is changed and no <code>DistributedSampler</code> is used. You could double check if all ranks yield the same number of batches as the code could otherwise hang.</p>