SyncBatchNorm and torch.save behavior

Fst_Msn: <p>I am training a custom implementation of NVAE (<a href="https://github.com/NVlabs/NVAE" class="inline-onebox" rel="noopener nofollow ugc">GitHub - NVlabs/NVAE: The Official PyTorch Implementation of "NVAE: A Deep Hierarchical Variational Autoencoder" (NeurIPS 2020 spotlight paper)</a>) in pytorch.</p>
<p>I am running a distributed environment with 4 gpus and DDP, using SyncBatchNorm as normalization.</p>
<p>Example of SyncBatchNorm initialization in a residual block of my autoencoder:</p>
<pre><code class="lang-auto">class ResidualCellEncoder(nn.Module):

    def __init__(self, in_channels: int, out_channels: int, downsampling: bool, use_SE: bool):
        """
        FIG 4.B of the paper
        :param in_channels:
        :param out_channels:
        :param downsampling:
        :param use_SE:
        """

        super().__init__()

        # skip connection has convs if downscaling
        if downsampling:
            stride = 2
            self.skip_connection = SkipDown(in_channels, out_channels, stride)
        else:
            stride = 1
            self.skip_connection = nn.Identity()

        # (BN - SWISH) + conv 3x3 + (BN - SWISH) + conv 3x3 + SE
        # downsampling in the first conv, depending on stride
        self.residual = nn.Sequential(
            SyncBatchNorm(in_channels, eps=1e-5, momentum=0.05),
            SiLU(),
            Conv2D(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True, weight_norm=True),
            SyncBatchNorm(out_channels, eps=1e-5, momentum=0.05),
            SiLU(),
            Conv2D(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True, weight_norm=True)
        )

        if use_SE:
            self.residual.append(SE(out_channels, out_channels))

    def forward(self, x: torch.Tensor):

        residual = 0.1 * self.residual(x)
        x = self.skip_connection(x)

        return x + residual
</code></pre>
<p>I have other SyncBatchNorm in the Autoencoder, all initialized like above (with custom eps and momentum).</p>
<p>Here Below, the code that I use to save the model</p>
<pre><code class="lang-auto"># Save checkpoint (after validation)
if WORLD_RANK == 0:
    print(f'[INFO] Saving Checkpoint')

    checkpoint_dict = {
        'epoch': epoch + 1,
        'global_step': global_step + 1,
        'configuration': config,
        'state_dict': ddp_model.module.state_dict(),
        'optimizer': optimizer.state_dict(),
        'grad_scaler': grad_scalar.state_dict()
    }
    ckpt_file = f"{args.checkpoint_base_path}/{args.run_name}/epoch={epoch:02d}.pt"
    torch.save(checkpoint_dict, ckpt_file)
</code></pre>
<p><strong>The issue is the following:</strong> the <code>running_mean</code> and <code>running_var</code> parameters are not saved correcly in the state dict.</p>
<p>I first noticed something wrong when trying to load a pre-trained model to continue training. The loss value of the first batch after loading is much higher than the last loss value registered before training stopped.</p>
<p>To confirm that the problem is SyncBatchNorm, I tried to manually save each <code>running_mean</code> and <code>running_var</code> <strong>immediately after</strong> saving the state dict:</p>
<pre><code class="lang-auto">import numpy as np
for i, (n, layer) in enumerate(ddp_model.module.named_modules()):

    if isinstance(layer, torch.nn.SyncBatchNorm):
        file_name = f"{args.checkpoint_base_path}/{args.run_name}/rank={WORLD_RANK}_I={i}.npy"
        np.save(file_name, torch.stack([layer.running_mean, layer.running_var]).cpu().numpy())
</code></pre>
<p>Then, in a new file, I load all the numpy files and the state dict, and run tests to assert two different things:</p>
<ol>
<li>that each rank saves the same numpy arrays (means/vars)</li>
<li>that the means and vars saved in the <code>rank_0</code> numpy arrays are the same as the ones in state dict.</li>
</ol>
<p>Here is the code:</p>
<pre><code class="lang-auto">input_files = sorted(os.listdir('./cifar10_4x3_test'))

ckpt_file = None
rank_0, rank_1, rank_2, rank_3 = [], [], [], []
for file in input_files:
    if file.endswith('.npy'):
        ar = torch.tensor(np.load('./cifar10_4x3_test/' + file))
        if 'rank=0' in file:
            rank_0.append(ar)
        elif 'rank=1' in file:
            rank_1.append(ar)
        elif 'rank=2' in file:
            rank_2.append(ar)
        else:
            rank_3.append(ar)

    else:
        ckpt_file = './cifar10_4x3_test/' + file

# assert sync batch norm works (all ranks have the same stuff)
for r0, r1, r2, r3 in zip(rank_0, rank_1, rank_2, rank_3):
    assert r0[0].sum() == r1[0].sum() == r2[0].sum() == r3[0].sum()
    assert r0[1].sum() == r1[1].sum() == r2[1].sum() == r3[1].sum()

# assert what you load from checkpoint is correct
checkpoint = torch.load(ckpt_file, map_location='cpu')
config = checkpoint['configuration']
model = AutoEncoder(config['autoencoder'], config['resolution'])
model.load_state_dict(checkpoint['state_dict'])

count_n = 0
for _, layer in model.named_modules():

    if isinstance(layer, torch.nn.SyncBatchNorm):
        mean = layer.running_mean
        var = layer.running_var

        assert rank_0[count_n][0] == mean
        assert rank_0[count_n][1] == var
        count_n += 1
</code></pre>
<p><strong>The first assert passes, the second doesn’t!</strong></p>
<p>This means that the running means and vars saved as numpy tensors and the ones saved by torch save are different. I guess that something is wrong in my saving/loading procedure, can someone help ?</p>
<p>Thank you</p>