Repository: horovod/horovod

Issue:
Sync Batch Norm for Pytorch
There are sync batch norm implementations for PyTorch distributed and Apex but can't seem to find one for Horovod. Is there a technical reason why there isn't one? 




Discussion:
@msalvaris sync batch norm is not implemented for horovod yet but it is in our todo list coming out soon.
@abditag2 That is great to hear. Would the Horovod implementation of sync batch norm simply be like the PyTorch one except replacing the distributed methods such as all_reduce with Horovod ones or is there some further technical issues that need to be overcome?
Most likely. It will also be implemented for keras and Tensorflow.

On Fri, Sep 6, 2019 at 5:59 AM Mat <notifications@github.com> wrote:

> @abditag2
> <https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_abditag2&d=DwMCaQ&c=r2dcLCtU9q6n0vrtnDw9vg&r=Jcn0WXylthWxLPEr38PPjw&m=Zv2I8tZBRh0dpnA6y6x9FvkyMxPOF_CD1eqhPyCzZjk&s=3ULBmBWSR2yUnmO8TiVuYTVIsAMshaxRxLg6yoLVO8g&e=>
> That is great to hear. Would the Horovod implementation of sync batch norm
> simply be like the PyTorch one except replacing the distributed methods
> such as all_reduce with Horovod ones or is there some further technical
> issues that need to be overcome?
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_horovod_horovod_issues_1384-3Femail-5Fsource-3Dnotifications-26email-5Ftoken-3DAAW4JGRGNDPRUMTUT35HON3QIJH2LA5CNFSM4IUCHM72YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOD6CYODA-23issuecomment-2D528844556&d=DwMCaQ&c=r2dcLCtU9q6n0vrtnDw9vg&r=Jcn0WXylthWxLPEr38PPjw&m=Zv2I8tZBRh0dpnA6y6x9FvkyMxPOF_CD1eqhPyCzZjk&s=fk9wfmCisBIHcSSe3EVy3mvmZ14OWqbmI5zhGvsbDuQ&e=>,
> or mute the thread
> <https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_notifications_unsubscribe-2Dauth_AAW4JGXRW5WE7BZW5DRPJRTQIJH2LANCNFSM4IUCHM7Q&d=DwMCaQ&c=r2dcLCtU9q6n0vrtnDw9vg&r=Jcn0WXylthWxLPEr38PPjw&m=Zv2I8tZBRh0dpnA6y6x9FvkyMxPOF_CD1eqhPyCzZjk&s=eD9kEM147p0laLeo4th7npJQk5iov4S3hvdXW6qgBWk&e=>
> .
>

Hello, I have implemented a sync-bn based on pytorch and horovod.
See: https://github.com/atranitell/Synchronized-BatchNorm-PyTorch-Horovod/blob/master/sync_bn.py

```bash
[1,0]<stdout>:affine: True track_running_stats False
[1,0]<stdout>:[INFERENCE PHASE-1]
[1,0]<stdout>:[NAME]: model.outputs.0-108.2299  [SHAPE]: [8, 1000]       [SUM]: 190.1940002, [MEAN]: 0.0237742, [MAX]: 2.6216030, [MIN]: -2.8100750, [VAR]: 0.5494183
[1,0]<stdout>:[NAME]: model.outputs.1-56.0818   [SHAPE]: [8, 1000]       [SUM]: 165.8392944, [MEAN]: 0.0207299, [MAX]: 2.7186546, [MIN]: -3.5223341, [VAR]: 0.5672572
[1,0]<stdout>:[NAME]: model.outputs.2-62.6330   [SHAPE]: [8, 1000]       [SUM]: 181.9670258, [MEAN]: 0.0227459, [MAX]: 2.6616516, [MIN]: -2.7214265, [VAR]: 0.5471250
[1,0]<stdout>:[NAME]: sync_model.outputs.0-78.4819 [SHAPE]: [8, 1000]       [SUM]: 190.1936188, [MEAN]: 0.0237742, [MAX]: 2.6216090, [MIN]: -2.8100815, [VAR]: 0.5494181
[1,0]<stdout>:[NAME]: sync_model.outputs.1-56.4089 [SHAPE]: [8, 1000]       [SUM]: 165.8392487, [MEAN]: 0.0207299, [MAX]: 2.7186563, [MIN]: -3.5223362, [VAR]: 0.5672570
[1,0]<stdout>:[NAME]: sync_model.outputs.2-51.5120 [SHAPE]: [8, 1000]       [SUM]: 181.9673615, [MEAN]: 0.0227459, [MAX]: 2.6616564, [MIN]: -2.7214241, [VAR]: 0.5471250
[1,0]<stdout>:[TRAINING PHASE]
[1,0]<stdout>:[NAME]: model.outputs.0-55.6757   [SHAPE]: [8, 1000]       [SUM]: 190.1940002, [MEAN]: 0.0237742, [MAX]: 2.6216030, [MIN]: -2.8100750, [VAR]: 0.5494183
[1,0]<stdout>:[NAME]: model.outputs.1-68.0540   [SHAPE]: [8, 1000]       [SUM]: -117.8934021, [MEAN]: -0.0147367, [MAX]: 2.8053782, [MIN]: -3.3515174, [VAR]: 0.5563220
[1,0]<stdout>:[NAME]: model.outputs.2-122.4730  [SHAPE]: [8, 1000]       [SUM]: -394.6956787, [MEAN]: -0.0493370, [MAX]: 2.5790358, [MIN]: -2.8156824, [VAR]: 0.5413557
[1,0]<stdout>:[NAME]: sync_model.outputs.0-166.4059 [SHAPE]: [8, 1000]       [SUM]: 190.1939545, [MEAN]: 0.0237742, [MAX]: 2.6216037, [MIN]: -2.8100758, [VAR]: 0.5494182
[1,0]<stdout>:[NAME]: sync_model.outputs.1-240.0510 [SHAPE]: [8, 1000]       [SUM]: -117.8935699, [MEAN]: -0.0147367, [MAX]: 2.8053713, [MIN]: -3.3515224, [VAR]: 0.5563222
[1,0]<stdout>:[NAME]: sync_model.outputs.2-247.0720 [SHAPE]: [8, 1000]       [SUM]: -394.6961365, [MEAN]: -0.0493370, [MAX]: 2.5790341, [MIN]: -2.8156407, [VAR]: 0.5413551
[1,0]<stdout>:[INFERENCE PHASE-2]
[1,0]<stdout>:[NAME]: model.outputs.0-92.9980   [SHAPE]: [8, 1000]       [SUM]: -838.2194214, [MEAN]: -0.1047774, [MAX]: 2.7204847, [MIN]: -2.9403963, [VAR]: 0.5573729
[1,0]<stdout>:[NAME]: model.outputs.1-74.8370   [SHAPE]: [8, 1000]       [SUM]: -830.7557373, [MEAN]: -0.1038445, [MAX]: 2.6922336, [MIN]: -3.0616376, [VAR]: 0.5771688
[1,0]<stdout>:[NAME]: model.outputs.2-79.0393   [SHAPE]: [8, 1000]       [SUM]: -798.6572266, [MEAN]: -0.0998322, [MAX]: 3.2130013, [MIN]: -2.9527130, [VAR]: 0.5481955
[1,0]<stdout>:[NAME]: sync_model.outputs.0-59.4788 [SHAPE]: [8, 1000]       [SUM]: -838.2209473, [MEAN]: -0.1047776, [MAX]: 2.7205055, [MIN]: -2.9404104, [VAR]: 0.5573736
[1,0]<stdout>:[NAME]: sync_model.outputs.1-58.3360 [SHAPE]: [8, 1000]       [SUM]: -830.7567749, [MEAN]: -0.1038446, [MAX]: 2.6922214, [MIN]: -3.0616069, [VAR]: 0.5771680
[1,0]<stdout>:[NAME]: sync_model.outputs.2-84.6570 [SHAPE]: [8, 1000]       [SUM]: -798.6591797, [MEAN]: -0.0998324, [MAX]: 3.2129552, [MIN]: -2.9527016, [VAR]: 0.5481985
[1,0]<stdout>:
[1,0]<stdout>:
[1,0]<stdout>:affine: True track_running_stats True
[1,0]<stdout>:[INFERENCE PHASE-1]
[1,0]<stdout>:[NAME]: model.outputs.0-43.8871   [SHAPE]: [8, 1000]       [SUM]: 73.2726212, [MEAN]: 0.0091591, [MAX]: 0.5324938, [MIN]: -0.5004993, [VAR]: 0.0269264
[1,0]<stdout>:[NAME]: model.outputs.1-29.7771   [SHAPE]: [8, 1000]       [SUM]: 79.4624710, [MEAN]: 0.0099328, [MAX]: 0.5239193, [MIN]: -0.5046086, [VAR]: 0.0266207
[1,0]<stdout>:[NAME]: model.outputs.2-35.8410   [SHAPE]: [8, 1000]       [SUM]: 77.2687073, [MEAN]: 0.0096586, [MAX]: 0.5608573, [MIN]: -0.4926656, [VAR]: 0.0273252
[1,0]<stdout>:[NAME]: sync_model.outputs.0-38.7599 [SHAPE]: [8, 1000]       [SUM]: 73.2726212, [MEAN]: 0.0091591, [MAX]: 0.5324938, [MIN]: -0.5004993, [VAR]: 0.0269264
[1,0]<stdout>:[NAME]: sync_model.outputs.1-42.7032 [SHAPE]: [8, 1000]       [SUM]: 79.4624710, [MEAN]: 0.0099328, [MAX]: 0.5239193, [MIN]: -0.5046086, [VAR]: 0.0266207
[1,0]<stdout>:[NAME]: sync_model.outputs.2-48.8539 [SHAPE]: [8, 1000]       [SUM]: 77.2687073, [MEAN]: 0.0096586, [MAX]: 0.5608573, [MIN]: -0.4926656, [VAR]: 0.0273252
[1,0]<stdout>:[TRAINING PHASE]
[1,0]<stdout>:[NAME]: model.outputs.0-57.8671   [SHAPE]: [8, 1000]       [SUM]: 190.1940002, [MEAN]: 0.0237742, [MAX]: 2.6216030, [MIN]: -2.8100750, [VAR]: 0.5494183
[1,0]<stdout>:[NAME]: model.outputs.1-112.4310  [SHAPE]: [8, 1000]       [SUM]: -117.8934021, [MEAN]: -0.0147367, [MAX]: 2.8053782, [MIN]: -3.3515174, [VAR]: 0.5563220
[1,0]<stdout>:[NAME]: model.outputs.2-97.1479   [SHAPE]: [8, 1000]       [SUM]: -394.6956787, [MEAN]: -0.0493370, [MAX]: 2.5790358, [MIN]: -2.8156824, [VAR]: 0.5413557
[1,0]<stdout>:[NAME]: sync_model.outputs.0-203.5007 [SHAPE]: [8, 1000]       [SUM]: 190.1939545, [MEAN]: 0.0237742, [MAX]: 2.6216037, [MIN]: -2.8100758, [VAR]: 0.5494182
[1,0]<stdout>:[NAME]: sync_model.outputs.1-256.1831 [SHAPE]: [8, 1000]       [SUM]: -117.8931427, [MEAN]: -0.0147366, [MAX]: 2.8053741, [MIN]: -3.3515232, [VAR]: 0.5563222
[1,0]<stdout>:[NAME]: sync_model.outputs.2-277.2660 [SHAPE]: [8, 1000]       [SUM]: -394.6949768, [MEAN]: -0.0493369, [MAX]: 2.5790620, [MIN]: -2.8156748, [VAR]: 0.5413552
[1,0]<stdout>:[INFERENCE PHASE-2]
[1,0]<stdout>:[NAME]: model.outputs.0-50.8077   [SHAPE]: [8, 1000]       [SUM]: -82.1022491, [MEAN]: -0.0102628, [MAX]: 0.2733500, [MIN]: -0.3195015, [VAR]: 0.0071264
[1,0]<stdout>:[NAME]: model.outputs.1-66.3068   [SHAPE]: [8, 1000]       [SUM]: -79.3727341, [MEAN]: -0.0099216, [MAX]: 0.2692815, [MIN]: -0.2908769, [VAR]: 0.0072672
[1,0]<stdout>:[NAME]: model.outputs.2-48.0020   [SHAPE]: [8, 1000]       [SUM]: -81.5532227, [MEAN]: -0.0101942, [MAX]: 0.2574194, [MIN]: -0.3002077, [VAR]: 0.0071647
[1,0]<stdout>:[NAME]: sync_model.outputs.0-67.1532 [SHAPE]: [8, 1000]       [SUM]: -81.9084091, [MEAN]: -0.0102386, [MAX]: 0.2692602, [MIN]: -0.3135209, [VAR]: 0.0068271
[1,0]<stdout>:[NAME]: sync_model.outputs.1-46.2432 [SHAPE]: [8, 1000]       [SUM]: -79.1543961, [MEAN]: -0.0098943, [MAX]: 0.2645385, [MIN]: -0.2849320, [VAR]: 0.0069433
[1,0]<stdout>:[NAME]: sync_model.outputs.2-35.9240 [SHAPE]: [8, 1000]       [SUM]: -81.3666458, [MEAN]: -0.0101708, [MAX]: 0.2528106, [MIN]: -0.2988353, [VAR]: 0.0068605
[1,0]<stdout>:
[1,0]<stdout>:
```

There is a slight error (about 1e-6), and I'm not clear currently.

KJ,
Best
@atranitell Do you want to make a PR? 
@abditag2 Emmm.. Surely, I am glad to do that. But, should I place the "sync_bn" codes at "/examples"?
In addition, I just implement a pytorch version.

Best,
KJ
@atranitell Nice work! Any updates on this feature? Could you please update the it with latest version?
@weberxie probably, I will update the latest version and test it at this weekend, compared with the torch.dist.

Resolving patch:
diff --git a/docs/mocks.py b/docs/mocks.py
--- a/docs/mocks.py
+++ b/docs/mocks.py
@@ -65,6 +65,9 @@ def _dummy():
     'keras.backend',
 
     'torch',
+    'torch.autograd.function',
+    'torch.nn.functional',
+    'torch.nn.modules.batchnorm',
     'torch.utils',
     'torch.utils.data',
     'torch.utils.tensorboard',
diff --git a/horovod/torch/__init__.py b/horovod/torch/__init__.py
--- a/horovod/torch/__init__.py
+++ b/horovod/torch/__init__.py
@@ -52,6 +52,7 @@
 from horovod.torch.mpi_ops import gloo_enabled, gloo_built
 from horovod.torch.mpi_ops import nccl_built, ddl_built, ccl_built
 from horovod.torch.mpi_ops import Average, Sum, Adasum
+from horovod.torch.sync_batch_norm import SyncBatchNorm
 
 import torch
 import collections
diff --git a/horovod/torch/sync_batch_norm.py b/horovod/torch/sync_batch_norm.py
new file mode 100644
--- /dev/null
+++ b/horovod/torch/sync_batch_norm.py
@@ -0,0 +1,194 @@
+# Based on https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/_functions.py
+# Modifications copyright 2020 Maka Autonomous Robotic Systems
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from horovod.torch.mpi_ops import allgather_async, allreduce_async, Sum, size, synchronize
+
+from distutils.version import LooseVersion
+
+import torch
+from torch.autograd.function import Function
+import torch.nn.functional as F
+from torch.nn.modules.batchnorm import _BatchNorm
+
+
+# Backward compat for old PyTorch
+if not hasattr(torch.jit, 'unused'):
+    torch.jit.unused = lambda x: x
+
+
+_SYNC_BN_V2 = LooseVersion(torch.__version__) >= LooseVersion('1.6.0')
+
+
+class SyncBatchNorm(_BatchNorm):
+    """
+    Applies synchronous version of N-dimensional BatchNorm.  In this version, normalization
+    parameters are synchronized across workers during forward pass.  This is very useful in
+    situations where each GPU can fit a very small number of examples.
+
+    See https://pytorch.org/docs/stable/nn.html#batchnorm2d for more details about BatchNorm.
+
+    Arguments:
+        num_features: number of channels `C` from the shape `(N, C, ...)`
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to `None` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to `True`, this module has
+            learnable affine parameters. Default: `True`
+        track_running_stats: a boolean value that when set to `True`, this
+            module tracks the running mean and variance, and when set to `False`,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: `True`
+    
+    NOTE: only GPU input tensors are supported in the training mode.
+    """
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True):
+        super().__init__(num_features, eps, momentum, affine, track_running_stats)
+
+    def _check_input_dim(self, input):
+        if input.dim() < 2:
+            raise ValueError('expected at least 2D input (got {}D input)'.format(input.dim()))
+
+    def _run_bn(self, input):
+        return F.batch_norm(
+            input, self.running_mean, self.running_var, self.weight, self.bias,
+            self.training or not self.track_running_stats, self.momentum, self.eps)
+
+    @torch.jit.unused
+    def _maybe_run_sync_bn(self, input):
+        if size() == 1:
+            return self._run_bn(input)
+        return _SyncBatchNorm.apply(
+            input, self.weight, self.bias, self.running_mean, self.running_var,
+            self.eps, self.momentum)
+
+    def forward(self, input):
+        # currently only GPU input is supported by underlying kernel from PyTorch
+        if not input.is_cuda:
+            raise ValueError('SyncBatchNorm expected input tensor to be on GPU')
+
+        self._check_input_dim(input)
+
+        if self.training and self.track_running_stats:
+            self.num_batches_tracked = self.num_batches_tracked + 1
+
+        if not self.training and self.track_running_stats:
+            return self._run_bn(input)
+        else:
+            return self._maybe_run_sync_bn(input)
+
+
+class _SyncBatchNorm(Function):
+    @staticmethod
+    def forward(self, input, weight, bias, running_mean, running_var, eps, momentum):
+        input = input.contiguous()
+
+        size = input.numel() // input.size(1)
+        count = torch.tensor([size])
+
+        # calculate mean/invstd for input.
+        mean, invstd = torch.batch_norm_stats(input, eps)
+
+        count_handle = allgather_async(count.unsqueeze(0), name='sync_batch_norm.count')
+        mean_handle = allgather_async(mean.unsqueeze(0), name='sync_batch_norm.mean')
+        invstd_handle = allgather_async(invstd.unsqueeze(0), name='sync_batch_norm.invstd')
+
+        # wait on the async communication to finish
+        count_all = synchronize(count_handle)
+        mean_all = synchronize(mean_handle)
+        invstd_all = synchronize(invstd_handle)
+
+        if _SYNC_BN_V2:
+            counts_for_bngswc = count_all.view(-1).float().to(input.device)
+        else:
+            # backwards compatibility
+            counts_for_bngswc = count_all.view(-1).tolist()
+
+        # calculate global mean & invstd
+        mean, invstd = torch.batch_norm_gather_stats_with_counts(
+            input,
+            mean_all,
+            invstd_all,
+            running_mean,
+            running_var,
+            momentum,
+            eps,
+            counts_for_bngswc
+        )
+
+        self.save_for_backward(input, weight, mean, invstd, count_all)
+
+        # apply element-wise normalization
+        return torch.batch_norm_elemt(input, weight, bias, mean, invstd, eps)
+
+    @staticmethod
+    def backward(self, grad_output):
+        grad_output = grad_output.contiguous()
+        saved_input, weight, mean, invstd, count_all = self.saved_tensors
+        need_input_grad, need_weight_grad, need_bias_grad = self.needs_input_grad[0:3]
+
+        # calculate local stats as well as grad_weight / grad_bias
+        sum_dy, sum_dy_xmu, grad_weight, grad_bias = torch.batch_norm_backward_reduce(
+            grad_output,
+            saved_input,
+            mean,
+            invstd,
+            weight,
+            need_input_grad,
+            need_weight_grad,
+            need_bias_grad
+        )
+
+        if need_input_grad:
+            # synchronizing stats used to calculate input gradient.
+            sum_dy_handle = allreduce_async(sum_dy, op=Sum, name='sync_batch_norm.sum_dy')
+            sum_dy_xmu_handle = allreduce_async(sum_dy_xmu, op=Sum, name='sync_batch_norm.sum_dy_xmu')
+
+            # wait on the async communication to finish
+            sum_dy = synchronize(sum_dy_handle)
+            sum_dy_xmu = synchronize(sum_dy_xmu_handle)
+
+            if _SYNC_BN_V2:
+                mean_dy = sum_dy / count_all.sum()
+                mean_dy_xmu = sum_dy_xmu / count_all.sum()
+            else:
+                # before 1.6.0, sum_dy was sum of means from every worker, so we just 
+                # need to divide it by number of workers
+                mean_dy = sum_dy / size()
+                mean_dy_xmu = sum_dy_xmu / size()
+
+            # backward pass for gradient calculation
+            grad_input = torch.batch_norm_backward_elemt(
+                grad_output,
+                saved_input,
+                mean,
+                invstd,
+                weight,
+                mean_dy,
+                mean_dy_xmu
+            )
+        else:
+            grad_input = None
+
+        # synchronizing of grad_weight / grad_bias is not needed as distributed
+        # training would handle all reduce.
+        if weight is None or not need_weight_grad:
+            grad_weight = None
+
+        if weight is None or not need_bias_grad:
+            grad_bias = None
+
+        return grad_input, grad_weight, grad_bias, None, None, None, None, None, None


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -24,6 +24,7 @@
 import itertools
 import os
 import pytest
+import sys
 import unittest
 import warnings
 
@@ -1667,6 +1668,63 @@ def test_horovod_join_broadcast(self):
                 ret = hvd.join(hvd.local_rank())
             else:
                 ret = hvd.join()
+    
+    def test_horovod_sync_batch_norm(self):
+        """Tests Horovod version of SyncBatchNorm."""
+        if not torch.cuda.is_available():
+            self.skipTest("No GPUs available")
+
+        if sys.version_info < (3,):
+            # TODO: remove this check after Py2 deprecation
+            self.skipTest("Python 3 only feature")
+
+        hvd.init()
+
+        ts_list = [
+            torch.stack([
+                torch.tensor([
+                    [r, r + 1],
+                    [r * 2, r * 2 + 1],
+                    [r * 3, r * 3 + 1],
+                    [r * 4, r * 4 + 1]
+                ])
+                for r in range(hvd.size())
+            ]),
+            torch.stack([
+                torch.tensor([
+                    [r + 1],
+                    [r * 2 + 1],
+                    [r * 3 + 1],
+                    [r * 4 + 1]
+                ])
+                for r in range(hvd.size())
+            ]),
+        ]
+
+        for ts in ts_list:
+            sync_bn = hvd.SyncBatchNorm(num_features=4)
+            sync_bn.cuda(hvd.local_rank())
+
+            bn = torch.nn.BatchNorm1d(num_features=4)
+            bn.cuda(hvd.local_rank())
+
+            ts = ts.cuda(hvd.local_rank()).float()
+            ts1 = ts.clone().requires_grad_()
+            ts2 = ts.clone().requires_grad_()
+
+            # Training
+            sync_bn_out = sync_bn(ts1[hvd.rank()].unsqueeze(0))
+            bn_out = bn(ts2)
+            assert (sync_bn_out - bn_out[hvd.rank()].unsqueeze(0)).abs().sum() < 1e-6
+            assert (sync_bn.running_mean - bn.running_mean).abs().sum() < 1e-6
+            assert (sync_bn.running_var - bn.running_var).abs().sum() < 1e-6
+
+            # Gradients
+            sync_bn_out.sum().backward()
+            bn_out.mean(dim=0).sum().backward()
+            assert (hvd.allreduce(sync_bn.weight.grad, name='sync_bn.weight.grad') - bn.weight.grad).abs().sum() < 1e-6
+            assert (hvd.allreduce(sync_bn.bias.grad, name='sync_bn.bias.grad') - bn.bias.grad).abs().sum() < 1e-6
+            assert (hvd.allreduce(ts1.grad, name='ts1.grad') - ts2.grad).abs().sum() < 1e-6
 
 if __name__ == "__main__":
-   unittest.main()
\ No newline at end of file
+   unittest.main()