Mixed precision training with nn.SyncBatchNorm return NaN for running_var

username: yft123
cooked: <p>I am training a model with mixed precision (amp_level=‘O1’) with a stack of 3D ResNet blocks.</p>
<pre><code class="lang-auto">class BasicBlock3d(nn.Module):
    """ 3x3x3 Resnet Basic Block"""
    expansion = 1
    __constants__ = ['downsample']

    def __init__(self, inplanes, planes, stride=1, downsample=None,dilation=1):
        super(BasicBlock3d, self).__init__()

        self.conv1 = conv3x3x3(inplanes, planes, stride, 1, dilation)
        self.bn1 = nn.SyncBatchNorm(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3x3(planes, planes, 1, 1, dilation)
        self.bn2 = nn.SyncBatchNorm( planes)

        if downsample is not None:
            self.downsample = downsample
        elif stride != 1 or inplanes != planes:
            self.downsample = nn.Sequential(conv1x1x1(inplanes, planes, stride),  get_norm_3d(norm, planes))
        else:
            self.downsample = None

        self.stride = stride

    def forward(self, x):
     
        identity = x

        out1 = self.conv1(x)
        out2 = self.bn1(out1)
        out_ = self.relu(out2)

        out3 = self.conv2(out_)
        out = self.bn2(out3) 

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)


        return out
</code></pre>
<p>After training more than tens of epochs (70+), the <code>self.bn1.running_var</code> turns to be NaN in some middle blocks and crash the training pipeline (loss=NaN). I have found the direct reason is the input feature (<code>out1</code>) of <code>self.bn1</code> got a sum larger than 65,504 in certain channel, which leads its mean becomes <code>inf</code>. One example is shown below</p>
<pre><code class="lang-auto">In[32]:out1[0,15,0,1,:]
Out[32]: 
tensor([14720., 21936., 21776., 21584., 21376., 21120., 20848., 20544., 20272.,
        19984., 19760., 19568., 19376., 19168., 18928., 18656., 18416., 18192.,
        18032., 17904., 17776., 17632., 17504., 17392., 17296., 17264., 17216.,
        17168., 17120., 17088., 17072., 17056., 17040., 17072., 17136., 17216.,
        17264., 17408., 17584., 17824., 17968., 18080., 18176., 18240., 18304.,
        18352., 18400., 18448., 18528., 18624., 18784., 18992., 19232., 19520.,
        19840., 20192., 20560., 20976., 21408., 21888., 22336., 22720., 22960.,
        15320.], dtype=torch.float16, grad_fn=&lt;SliceBackward&gt;)

In[33]:torch.mean(out1[0,15,0,1,:])
Out[33]: tensor(inf, dtype=torch.float16, grad_fn=&lt;MeanBackward0&gt;)

In[35]:torch.sum(out1[0,15,0,1,:])
Out[35]: tensor(inf, dtype=torch.float16, grad_fn=&lt;SumBackward0&gt;)

In[36]:torch.sum(out1[0,15,0,1,:].float())
Out[36]: tensor(1206136., grad_fn=&lt;SumBackward0&gt;)

In[38]:out2[0,15].isnan().all()
Out[38]: tensor(True)
</code></pre>
<p>Is there any way we can prevent the feature growing this large during the training process?</p>

username: fangwei123456
cooked: <p>I meet the same problem.</p>

username: ptrblck
cooked: <p>The activation values are not clipped or manipulated in any way, so that the FP16 bounds can yield these NaN values. I’m not familiar with your workflow and would recommend to normalize the inputs etc. as is done in the majority of the models. This would keep the parameters as well as the activations in a reasonable range and would avoid running into the overflows.</p>
<p>CC <a class="mention" href="/u/yft123">@yft123</a></p>