BatchNorm layers cause training error when track_running_stats=True with DistributedDataParallel

banton: <p>When using DDP (pytorch 12.1) some of my batch norm layers cause the training to fail due to an inplace operation with the following error:</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: torch.cuda.FloatTensor [65]] is at version 3; expected version 2 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
</code></pre>
<p>The operation that failed was simply:</p>
<pre><code class="lang-auto">class Conv(nn.Module):
    
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2, **bn_params)
        self.act = nn.SiLU(inplace=True) if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)    # &lt;--------------------------------------- this fails
        x = self.act(x)
        return x
</code></pre>
<p>On a whim I tried passing:<br>
<code>self.bn = nn.BatchNorm2d(c2, track_running_stats=False, **bn_params)</code><br>
to all my batch norm layers and the training ran, but of course this is not a viable solution.</p>
<p>For the record, I also tried cloning <code>x</code> and setting <code>nn.SiLU(inplace=False)</code> but got the same error.</p>

ptrblck: <p>I cannot reproduce the issue using your module in:</p>
<pre><code class="lang-python">class Conv(nn.Module):
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, 1, groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU(inplace=True) if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        return x
    
model = Conv(1, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for _ in range(10):
    optimizer.zero_grad()
    x = torch.randn(1, 1, 24, 24)
    out = model(x)
    out.mean().backward()
    optimizer.step()
</code></pre>

banton: <p>Thanks for the response <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<p>While creating a minimal reproducible example, I was able to narrow down the error source: During my training I do two forward passes and a backward pass on the resulting losses. If I do only one, there is no error.</p>
<p>Do I assume correctly that the batch norm layers update the running stats during each forward pass? That could explain the error.<br>
If this is the case, why does it only trigger with DDP?</p>
<p>One more thing: I use Huggingface Accelerate for convenience to wrap DDP, but I don’t think it’s relevant in this case.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from accelerate import Accelerator

class Conv(nn.Module):
    def __init__(self, c1, c2, k=1, s=1, g=1):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, 1, groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        return x

accelerate = Accelerator()
model = Conv(1, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model, optimizer = accelerate.prepare(model, optimizer)

with torch.autograd.set_detect_anomaly(True):
    for _ in range(10):
        optimizer.zero_grad()
        x1 = torch.randn(1, 1, 24, 24)
        x2 = torch.randn(1, 1, 24, 24)
        out1 = model(x1)
        out2 = model(x2)
        out = out1 + out2
        out.mean().backward()
        optimizer.step()
</code></pre>

ptrblck: <p>Your code snippet still works for me and I don’t get the error after removing <code>Accelerator</code> (I haven’t installed the needed dependencies for it). Do you see the error using your code snippet? If so, does dropping <code>Accelerator</code> help?</p>
<aside class="quote no-group" data-username="banton" data-post="3" data-topic="165033">
<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/b/eada6e/48.png" class="avatar"> banton:</div>
<blockquote>
<p>During my training I do two forward passes and a backward pass on the resulting losses. If I do only one, there is no error.</p>
<p>Do I assume correctly that the batch norm layers update the running stats during each forward pass? That could explain the error.</p>
</blockquote>
</aside>
<p>You are correct that the running stats get updated in each forward pass, but this won’t explain the error.<br>
The running stats are buffers (i.e. they won’t be updated by Autograd) and are not even used during training.</p>

banton: <blockquote>
<p>Do you see the error using your code snippet? If so, does dropping <code>Accelerator</code> help?</p>
</blockquote>
<p>If I omit <code>Accelerator</code> the error disappears, but all the package does is wrap DDP around the training.<br>
The following code using pytorch-only modules also produces the error. Tested it on pytorch 1.11.0 and 1.12.1. I set <code>CUDA_VISIBLE_DEVICES=0</code> and still get the error.</p>
<pre><code class="lang-auto">import os
import torch
import torch.nn as nn
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group

class Conv(nn.Module):
    def __init__(self, c1, c2, k=1, s=1, g=1):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, 1, groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        return x

def ddp_setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    init_process_group(backend="nccl", rank=rank, world_size=world_size)

def main(rank: int, world_size: int):
    ddp_setup(rank, world_size)
    model = Conv(1, 1).to(rank)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    model = DDP(model, device_ids=[rank])

    with torch.autograd.set_detect_anomaly(True):
        for _ in range(10):
            optimizer.zero_grad()
            x1 = torch.randn(1, 1, 24, 24).to(rank)
            x2 = torch.randn(1, 1, 24, 24).to(rank)
            out1 = model(x1)
            out2 = model(x2)
            out = out1 + out2
            out.mean().backward()
            optimizer.step()

    destroy_process_group()


if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(main, args=(world_size,), nprocs=world_size)
</code></pre>
<p>This yields the following error message:</p>
<pre><code class="lang-auto">/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/autograd/__init__.py:173: UserWarning: Error detected in CudnnBatchNormBackward0. Traceback of forward call that caused the error:
  File "&lt;string&gt;", line 1, in &lt;module&gt;
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/multiprocessing/spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/multiprocessing/spawn.py", line 129, in _main
    return self._bootstrap(parent_sentinel)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/anba/catkin_ws/src/tas_dev/dev/anba/superpoint/DDP_bug_2.py", line 37, in main
    out1 = model(x1)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1008, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 969, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/anba/catkin_ws/src/tas_dev/dev/anba/superpoint/DDP_bug_2.py", line 17, in forward
    x = self.bn(x)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/nn/modules/batchnorm.py", line 168, in forward
    return F.batch_norm(
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/nn/functional.py", line 2438, in batch_norm
    return torch.batch_norm(
 (Triggered internally at  /home/conda/feedstock_root/build_artifacts/pytorch-recipe_1660083882787/work/torch/csrc/autograd/python_anomaly_mode.cpp:102.)
  Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
Traceback (most recent call last):
  File "DDP_bug_2.py", line 48, in &lt;module&gt;
    mp.spawn(main, args=(world_size,), nprocs=world_size)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 240, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 198, in start_processes
    while not context.join():
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 160, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/anba/catkin_ws/src/tas_dev/dev/anba/superpoint/DDP_bug_2.py", line 40, in main
    out.mean().backward()
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/_tensor.py", line 396, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/anba/anaconda3/envs/superpoint_anba/lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [1]] is at version 3; expected version 2 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
</code></pre>

banton: <p>Were you able to reproduce the error, <a class="mention" href="/u/ptrblck">@ptrblck</a> ?<br>
If not then I give up.</p>

ptrblck: <p>Yes, I’m able to reproduce the issue and it seems it might be raised by usage of different forward passed.<br>
I.e. if I split the forward/backward passes it works for me:</p>
<pre><code class="lang-python">            out1 = model(x1)
            out1.mean().backward()
            out2 = model(x2)
            out2.mean().backward()
            optimizer.step()
</code></pre>
<p>I also don’t know if this is a known limitation of DDP, but I think <a class="mention" href="/u/kwen2501">@kwen2501</a> might know.</p>

fduwjj: <p>Thanks for asking, this is indeed an interesting case. Not sure if bucketizing can affect this. cc: <a class="mention" href="/u/rvarm1">@rvarm1</a></p>

H-Huang: <p>Is this issue related <a href="https://github.com/pytorch/pytorch/issues/66504" class="inline-onebox" rel="noopener nofollow ugc">BatchNorm runtimeError: one of the variables needed for gradient computation has been modified by an inplace operation · Issue #66504 · pytorch/pytorch · GitHub</a>? There are a few workarounds mentioned in the issue</p>

banton: <p>Thanks for link. The issues seem to be related but neither setting <code>broadcast_buffers=False</code> nor using <code>SyncBatchNorm</code> solves my problem.<br>
<code>SyncBatchNorm</code> gives me the same error.<br>
Not broadcasting the buffers lets the training run for a few iterations (the number varies between trainings) but then I get<br>
<code>RuntimeError: Function 'CudnnBatchNormBackward0' returned nan values in its 0th output.</code></p>
<p>The only solution that kind of worked was wrapping the training in Fully Sharded DDP.</p>
<p>Edit: I tried training on my other machine using pytorch 1.12 (was using 1.11 before) and it worked fine with setting <code>broadcast_buffers=False</code>.</p>