def _allreduce_grads(self):
for i, param in enumerate(self._params):
if param.grad_req != 'null':
byteps_push_pull(param.list_grad()[0],
is_average=False,
name="gradient_" + str(i),
priority=-i)
def _do_push_pull(self, index, grad):
if isinstance(index, (tuple, list)):
for i in range(len(index)):
byteps_declare_tensor(grad[i], "gradient_"+str(index[i]))
byteps_push_pull(grad[i], version=0, priority=-index[i], name="gradient_"+str(index[i]), is_average=True)
else:
byteps_declare_tensor(grad, "gradient_"+str(index))
byteps_push_pull(grad, version=0, priority=-index, name="gradient_"+str(index), is_average=True)
def wrapped_init_impl(self, *args, **kwargs):
init_impl(*args, **kwargs)
# Broadcast is implemented as push + pull in BytePS
byteps_push_pull(self.data(),
version=0,
priority=0,
name="parameter_" + str(index),
is_average=False)
self.data().wait_to_read()
def broadcast_parameters(params, root_rank=0):
"""
Broadcasts the parameters from root rank to all other processes.
Typical usage is to broadcast the `Module.get_params()` or the
`Block.collect_params()`.
Arguments:
params: One of the following:
- dict of parameters to broadcast
- ParameterDict to broadcast
root_rank: The rank of the process from which parameters will be
broadcasted to all other processes.
"""
tensors = []
if isinstance(params, dict):
tensors = [p for _, p in sorted(params.items())]
elif isinstance(params, mx.gluon.parameter.ParameterDict):
for _, p in sorted(params.items()):
try:
tensors.append(p.data())
except mx.gluon.parameter.DeferredInitializationError:
# Inject wrapper method with post-initialization broadcast to
# handle parameters with deferred initialization
global parameter_index
byteps_declare_tensor(p.data(),
"parameter_" + str(parameter_index))
new_init = _append_broadcast_init(p, root_rank,
parameter_index)
parameter_index += 1
p._init_impl = types.MethodType(new_init, p)
else:
raise ValueError('invalid params of type: %s' % type(params))
# Run tensor initilization
for i in range(len(tensors)):
byteps_declare_tensor(tensors[i], "parameter_" + str(parameter_index))
# Broadcast is implemented as push + pull in BytePS
# To broadcast: we should zero-out all non-root tensors, and disable push_pull average
if rank() != root_rank:
tensors[i].__imul__(0)
byteps_push_pull(tensors[i],
version=0,
priority=0,
name="parameter_" + str(parameter_index),
is_average=False)
parameter_index += 1
# Make sure tensors pushed to MXNet engine get processed such that all
# workers are synced before starting training.
for tensor in tensors:
tensor.wait_to_read()
def _do_push_pull_param(self, index, delta_weight):
if isinstance(index, (tuple, list)):
for i in range(len(index)):
byteps_declare_tensor("weight_" + str(index[i]))
byteps_push_pull(delta_weight[i],
version=0,
priority=-index[i],
name="weight_" + str(index[i]),
is_average=False)
else:
byteps_declare_tensor("weight_" + str(index))
byteps_push_pull(delta_weight,
version=0,
priority=-index,
name="weight_" + str(index),
is_average=False)
def _init_params(self):
tensors = []
for param in self._params_to_init:
if param._deferred_init:
tensors.append(param)
else:
param_arrays = param._check_and_get(param._data, list)
idx = self._param2idx[param.name]
if rank() != self.root_rank:
param_arrays[0].__imul__(0)
byteps_push_pull(param_arrays[0],
version=0,
priority=0,
name="parameter_" + str(idx),
is_average=False)
self._params_to_init = tensors
def broadcast_parameters(params, root_rank=0):
"""
Broadcasts the parameters from root rank to all other processes.
Typical usage is to broadcast the `Module.get_params()`.
Arguments:
params: dict of parameters to broadcast
root_rank: The rank of the process from which parameters will be
broadcasted to all other processes.
"""
global parameter_index
if isinstance(params, dict):
tensors = [p for _, p in sorted(params.items())]
# Run tensor initilization
for i in range(len(tensors)):
byteps_declare_tensor(tensors[i],
"parameter_" + str(parameter_index))
# Broadcast is implemented as push + pull in BytePS
# To broadcast: we should zero-out all non-root tensors, and disable push_pull average
if rank() != root_rank:
tensors[i].__imul__(0)
byteps_push_pull(tensors[i],
version=0,
priority=0,
name="parameter_" + str(parameter_index),
is_average=False)
parameter_index += 1
# Make sure tensors pushed to MXNet engine get processed such that all
# workers are synced before starting training.
for tensor in tensors:
tensor.wait_to_read()
elif isinstance(params, mx.gluon.parameter.ParameterDict):
raise TypeError("For gluon users, you should not call this function. "
"DistributedTrainer will broadcast all parameters at "
"the first training step.")
else:
raise ValueError('Invalid params of type: %s' % type(params))
def _allreduce_grads(self):
# update lr
if local_rank() == 0:
self._f.seek(0)
ba = struct.pack("d", self.learning_rate)
self._f.write(ba)
self._f.flush()
for i, param in enumerate(self._params):
if param.grad_req != 'null':
# normalized with batch_size and num_workers
nd._internal._mul_scalar(param._grad[0],
1.0 / self._scale / self._bps_size,
out=param._grad[0])
compressed, ctx = self._intra_compressors[param.name].compress(
param._grad[0])
byteps_push_pull(compressed,
is_average=False,
name="gradient_" + str(i),
priority=-i)
param._grad[0][:] = self._intra_compressors[
param.name].decompress(compressed, ctx, x=param._data[0])
