def test_horovod_broadcast_rank_error(self):
"""Test that the broadcast returns an error if different ranks
specify different root rank."""
hvd.init()
rank = hvd.rank()
size = hvd.size()
# This test does not apply if there is only one worker.
if size == 1:
self.skipTest("Only one worker available")
ctx = self._current_context()
shape = (17, 17, 17)
tensor = mx.nd.ones(shape=shape, ctx=ctx)
try:
output = hvd.broadcast(tensor, root_rank=rank)
output.wait_to_read()
assert False, 'hvd.broadcast did not throw rank error'
except (MXNetError, RuntimeError):
pass
def test_horovod_broadcast_error(self):
"""Test that the broadcast returns an error if any dimension besides
the first is different among the tensors being broadcasted."""
hvd.init()
rank = hvd.rank()
size = hvd.size()
# This test does not apply if there is only one worker.
if size == 1:
self.skipTest("Only one worker available")
ctx = self._current_context()
shape = (17, rank + 1)
tensor = mx.nd.ones(shape=shape, ctx=ctx)
try:
output = hvd.broadcast(tensor, 0)
output.wait_to_read()
assert False, 'hvd.broadcast did not throw error'
except (MXNetError, RuntimeError):
pass
def broadcast(self, key, value, out, priority=0):
""" Broadcast the `value` NDArray at rank 0 to all ranks
Parameters
----------
key : str, or int
The key is used to name the tensor for allreduce. Its
usage is different from that of parameter servers.
value : NDArray
The tensor that is to be broadcasted.
out : NDArray, list of NDArray
Output tensor that receives value broadcasted from root process
priority : int, optional
The priority of the operation.
Higher priority operations are likely to be executed before other actions.
Examples
--------
>>> a = mx.nd.ones(shape)
>>> b = mx.nd.zeros(shape)
>>> kv.broadcast('2', value=a, out=b)
>>> print(b.asnumpy)
[[ 1. 1. 1.]
[ 1. 1. 1.]]
"""
import horovod.mxnet as hvd
out = out if isinstance(out, list) else [out]
# TODO (lnyuan): need to copy data to each device memory
for o in out:
o[:] = hvd.broadcast(tensor=value,
root_rank=0,
name=str(key),
priority=priority)
def broadcast_parameters(params):
rank_0_dict = {}
# Run broadcasts.
for key, tensor in params.items():
rank_0_dict[key] = hvd.broadcast(tensor, 0, key)
return rank_0_dict
