Use torch.nn.SyncBatchNorm only for running mean variance update without DDP backward

username: xuef
cooked: <p>I have a need to update only the batch norm running mean and variance during the training process across GPUs. So want to leverage SyncBatchNorm for that. I don’t need the gradient for any of the parameters in the model (its parameters are updated from the online model). Is there a way to use SyncBatchNorm setting for that purpose only?</p>
<p>I have attempted to use DDP without backward, but that create sync issues. One thing I’m thinking is to create a fake 0 loss, and perform the backward as a work around. I wonder if there are better ideas.</p>

username: xuef
cooked: <p>After some attempts, dist._all_gather_base in SyncBatchNorm kept being timed out after some training time for the exponential moving average model that wrapped in DDP.</p>
<p>Anyone have any suggestions on what to look into?</p>

username: kumpera
cooked: <p>Try setting the <code>TORCH_DISTRIBUTED_DEBUG</code> env var to <code>DETAIL</code> and look at the log files for clues.</p>
<p>More detail here: <a href="https://pytorch.org/docs/master/distributed.html#debugging-torch-distributed-applications" class="inline-onebox" rel="noopener nofollow ugc">Distributed communication package - torch.distributed — PyTorch master documentation</a></p>
<p>Timeouts usually happens when collectives are issues out of order or at least one rank is not participating.</p>
<p>Have you tried to call SyncBatchNorm under the <code>torch.no_grad()</code> context manager? IE:</p>
<p>You can pick the code from BatchNorm and extract the sync_batch_norm path for your own scenario: <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/batchnorm.py#L684" class="inline-onebox" rel="noopener nofollow ugc">pytorch/batchnorm.py at master · pytorch/pytorch · GitHub</a></p>

username: xuef
cooked: <p>Thank you <a class="mention" href="/u/kumpera">@kumpera</a> for your reply. I’m calling the model with SyncBatchNorm under the <code>torch.no_grad()</code> context manager. The timeout happens after the training works fine for a while (not at a particular epoch/step). I will try to use the <code>TORCH_DISTRIBUTED_DEBUG</code> env and look into the details.</p>
<aside class="quote no-group" data-username="kumpera" data-post="3" data-topic="172209">
<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/kumpera/48/48311_2.png" class="avatar"> kumpera:</div>
<blockquote>
<p>You can pick the code from BatchNorm and extract the sync_batch_norm path for your own scenario</p>
</blockquote>
</aside>
<p>Do you mean I change that part of the code to fit to my needs?</p>

username: kumpera
cooked: <aside class="quote no-group" data-username="xuef" data-post="4" data-topic="172209">
<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/x/ecd19e/48.png" class="avatar"> xuef:</div>
<blockquote>
<p>Do you mean I change that part of the code to fit to my needs?</p>
</blockquote>
</aside>
<p>Yes, that code imposes some restrictions that might not work for you.</p>

username: xuef
cooked: <p>In my case that sync is timed out. SyncBatchNorm works as expected in the forward pass. The issue in my case is that I have multiple data loaders and each with multiple workers, they turn out can create a deadlock situation. Once I reduce the num_workers to 1, issue is resolved.</p>