Repository: horovod/horovod

Issue:
mxnet+hvd with different random seed won't work with Gluon API
**Environment:**
1. Framework: (MXNet)
2. Framework version: 1.4
3. Horovod version: v0.16.0 
4. MPI version: 
5. CUDA version:
6. NCCL version:
7. Python version: 2.7/3.6
8. OS and version: 

**Checklist:**
1. Did you search issues to find if somebody asked this question before?
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.md)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.md)?

**Bug report:**
Please describe errorneous behavior you're observing and steps to reproduce it.

If the shape of a parameter has to be inferred after a batch of data is seen, the parameter will not be broadcast. Therefore, the correctness of the program depends on whether all workers are initialized with the same random seed. So the following program will not work:

```
random.seed(hvd.local_rank())
data = ..
model = mx.gluon.nn.Dense(10)
model.initialize()

// no parameters are broadcast, because initialization is deferred.
hvd.broadcast_parameters(model.collect_parameters())

for batch in data:
    // params are initialized after shape is known
    pred = model(batch.data[0])
    ...
...
```



Discussion:
@eric-haibin-lin append "using Gluon API" in the title to avoid confusion?
Added. Thanks

Resolving patch:
diff --git a/examples/mxnet_mnist.py b/examples/mxnet_mnist.py
--- a/examples/mxnet_mnist.py
+++ b/examples/mxnet_mnist.py
@@ -112,13 +112,12 @@ def evaluate(model, data_iter, context):
 model.cast(args.dtype)
 model.hybridize()
 
-# Define hyper parameters
+# Create optimizer
 optimizer_params = {'momentum': args.momentum,
                     'learning_rate': args.lr * hvd.size(),
                     'rescale_grad': 1.0 / args.batch_size}
-
-# Add Horovod Distributed Optimizer
 opt = mx.optimizer.create('sgd', **optimizer_params)
+# Horovod: wrap optimizer with DistributedOptimizer
 opt = hvd.DistributedOptimizer(opt)
 
 # Initialize parameters
diff --git a/horovod/mxnet/__init__.py b/horovod/mxnet/__init__.py
--- a/horovod/mxnet/__init__.py
+++ b/horovod/mxnet/__init__.py
@@ -30,6 +30,7 @@
 from horovod.mxnet.mpi_ops import mpi_threads_supported
 
 import mxnet as mx
+import types
 
 
 # This is where Horovod's DistributedOptimizer wrapper for MXNet goes
@@ -68,6 +69,16 @@ def set_wd_mult(self, args_wd_mult):
         self._optimizer.set_wd_mult(args_wd_mult)
 
 
+# Wrapper to inject Horovod broadcast after parameter initialization
+def _append_broadcast_init(param, root_rank):
+    init_impl = getattr(param, '_init_impl')
+    def wrapped_init_impl(self, *args, **kwargs):
+        init_impl(*args, **kwargs)
+        broadcast_(self.data(), root_rank=root_rank)
+        self.data().wait_to_read()
+    return wrapped_init_impl
+
+
 def broadcast_parameters(params, root_rank=0):
     """
     Broadcasts the parameters from root rank to all other processes.
@@ -89,8 +100,10 @@ def broadcast_parameters(params, root_rank=0):
             try:
                 tensors.append(p.data())
             except mx.gluon.parameter.DeferredInitializationError:
-                # skip broadcasting deferred init param
-                pass
+                # Inject wrapper method with post-initialization broadcast to
+                # handle parameters with deferred initialization
+                new_init = _append_broadcast_init(p, root_rank)
+                p._init_impl = types.MethodType(new_init, p)
     else:
         raise ValueError('invalid params of type: %s' % type(params))
 


Test patch:
diff --git a/test/test_mxnet.py b/test/test_mxnet.py
--- a/test/test_mxnet.py
+++ b/test/test_mxnet.py
@@ -355,7 +355,6 @@ def test_horovod_broadcast_grad(self):
         shapes = [(), (17), (17, 17), (17, 17, 17)]
         root_rank = 1
         tensor_dict = {}
-        broadcast_dict = {}
         root_dict = {}
         for dtype, dim, in itertools.product(dtypes, dims):
             tensor_dict[count] = mx.nd.ones(shapes[dim], ctx=ctx) * rank
@@ -445,5 +444,32 @@ def test_horovod_broadcast_rank_error(self):
         except (MXNetError, RuntimeError):
             pass
 
+    def test_horovod_broadcast_deferred_init_parameters(self):
+        """Test that the deferred initialized parameters are broadcasted."""
+        hvd.init()
+        root_rank = 0
+        rank = hvd.rank()
+
+        # This test does not apply if there is only one worker.
+        if hvd.size() == 1:
+            return
+
+        mx.random.seed(rank)
+        layer = mx.gluon.nn.Conv2D(10, 2)
+        layer.initialize()
+        hvd.broadcast_parameters(layer.collect_params(), root_rank=root_rank)
+
+        x = mx.nd.ones((5, 4, 10, 10))
+        layer(x)
+        tensors = [p.data() for _, p in sorted(layer.collect_params().items())]
+        root_tensors = []
+        for tensor in tensors:
+            root_tensors.append(hvd.broadcast(tensor, root_rank=root_rank))
+
+        for tensor, root_tensor in zip(tensors, root_tensors):
+            assert same(tensor.asnumpy(), root_tensor.asnumpy()), \
+                'horovod did not broadcast deferred initialized parameter correctly'
+
+
 if __name__ == '__main__':
     unittest.main()