def MPIBroadcast(inputs, root, mpi_ranks=None, **kwargs): """Broadcast a tensor to all nodes in the ``MPIGroup``. Parameters ---------- inputs : Tensor The tensor to broadcast. root : int The world rank of root node. mpi_ranks: sequence of int, optional The world rank of nodes in group. Default is ``None`` (Use All). Returns ------- Tensor The broadcast output. Notes ----- For root, the output **shares** the input. For others, the input is **inaccessible**. """ arguments = ParseArgs(locals()) if mpi_ranks is None: num_nodes = mpi.Size() mpi_ranks = [i for i in range(0, num_nodes)] if not isinstance(mpi_ranks, list): mpi_rank = [mpi_ranks] comm, group = mpi.CreateGroup(root, incl=mpi_ranks) arguments = {'inputs': arguments['inputs'], 'comm': comm, 'group': group} return Tensor.CreateOperator('MPIBroadcast', **arguments)
def _inject_update_ops(graph_def, updater): """Inject the update ops GraphDef. The ``updater`` should generate update targets before. Parameters ---------- graph_def : GraphDef The definition of graph. updater : BaseUpdater The updater. Returns ------- None """ if updater is None: return updater.register_in_workspace() grads, update_ops = [], [] extra_arguments = updater._extra_kwargs extra_arguments['slot'] = updater._slot # Build update ops according to the updater for e in updater._param_group: (param, grad), arguments = e if _workspace.HasTensor(grad): grads.append(grad) arguments = dict(arguments, **extra_arguments) update_ops.append( _proto_utils.MakeOperatorDef( op_type=updater.type(), inputs=[grad], outputs=[param], name=_helper.OperatorHelper.get_name(), **arguments)) else: _logging.info('Skip to update Tensor({}).'.format(param)) # Check data parallel if necessary if _mpi.Is_Init(): (rank, group), arguments = _mpi.AllowParallel(), {} if rank != -1: arguments['mode'] = '%s_ALLREDUCE' % _mpi.GetParallelMode() arguments['root'], (arguments['comm'], arguments['group']) \ = group[0], _mpi.CreateGroup(root=group[0], incl=group) update_ops.insert( 0, _proto_utils.MakeOperatorDef( op_type='CollectiveUpdate', inputs=grads, outputs=grads, name=_helper.OperatorHelper.get_name(), **arguments)) graph_def.op.extend(update_ops)
def GraphDef_Update(meta_graph, updater): """Inject the update targets into GraphDef. The ``updater`` should generate update targets before. Parameters ---------- meta_graph : dragon_pb2.GraphDef The definition of meta graph. updater : BaseUpdater The updater. Returns ------- None """ if updater is None: return updater._prefix = meta_graph.name + '_' extra_arguments = updater._extra_kwargs extra_arguments['domain'] = updater._prefix parallel_arguments = {} # wrap hyper-parameters as Tensor for CC for k, v in updater._hyper_params.items(): ws.FeedTensor(updater._prefix + k, np.array([v], dtype=np.float32)) # check data parallel if necessary if mpi.Is_Init(): idx, group = mpi.AllowParallel() if idx != -1: parallel_arguments['parallel_mode'] = mpi.GetParallelMode() parallel_arguments['comm'], parallel_arguments['group'] \ = mpi.CreateGroup(root=group[0], incl=group) parallel_arguments['root'] = group[0] for k, v in parallel_arguments.items(): meta_graph.arg.add().CopyFrom(MakeArgument(k, v)) for tuple in updater._tuples: tensors = tuple[0] arguments = tuple[1] kwargs = dict(arguments, **extra_arguments) u_target = pb.UpdateTarget() u_target.type = updater._type _, u_target.name = GetOperatorName() for tensor in tensors: u_target.tensor.append(tensor) for k, v in kwargs.items(): u_target.arg.add().CopyFrom(MakeArgument(k, v)) meta_graph.u_target.extend([u_target])
def GraphDef_Update(meta_graph, updater): """Inject the update targets into GraphDef. The ``updater`` should generate update targets before. Parameters ---------- meta_graph : dragon_pb2.GraphDef The definition of meta graph. updater : BaseUpdater The updater. Returns ------- None """ if updater is None: return # use graph name if missing slot if updater._slot is None: updater._slot = meta_graph.name extra_arguments = updater._extra_kwargs extra_arguments['slot'] = updater._slot parallel_arguments = {} updater.register_in_workspace() # check data parallel if necessary if mpi.Is_Init(): idx, group = mpi.AllowParallel() if idx != -1: parallel_arguments['parallel_mode'] = mpi.GetParallelMode() parallel_arguments['comm'], parallel_arguments['group'] \ = mpi.CreateGroup(root=group[0], incl=group) parallel_arguments['root'] = group[0] for k, v in parallel_arguments.items(): meta_graph.arg.add().CopyFrom(MakeArgument(k, v)) for e in updater._param_group: pair, arguments = e kwargs = dict(arguments, **extra_arguments) u_target = pb.UpdateTarget() u_target.type = updater.type() _, u_target.name = GetOperatorName() for t in pair: u_target.tensor.append(t) for k, v in kwargs.items(): u_target.arg.add().CopyFrom(MakeArgument(k, v)) meta_graph.u_target.extend([u_target])
def register_op(self): idx, group = mpi.AllowParallel() if idx == -1: raise RuntimeError('The mpi node({}) dost not in ' 'parallel groups. \nSet it using mpi.Parallel([..]).'.format(mpi.Rank())) mpi_comm, mpi_group = mpi.CreateGroup(root=group[0], incl=group) self.op_meta = { 'op_type': 'CollectiveUpdate', 'arguments': { 'mode': self.mode, 'comm': mpi_comm, 'group': mpi_group, 'root': group[0], # Assume the 1st node of group as root }, }
def MPIGather(inputs, root, mpi_ranks=None, **kwargs): """Gather a tensor from all nodes to root in the ``MPIGroup``. Parameters ---------- inputs : Tensor The tensor to gather. root : int The world rank of root node. mpi_ranks: int, list of int or None The world rank of nodes in group. Default is ``None`` (Use All). Returns ------- Tensor or list of Tensor The gathered outputs. Notes ----- The number of outputs is decided on the number of ``mpi_ranks``. The outputs are **accessible** only for root and vice versa. """ CheckInputs(inputs, 1) arguments = ParseArguments(locals()) if mpi_ranks is None: num_nodes = mpi.Size() mpi_ranks = [i for i in range(0, num_nodes)] if not isinstance(mpi_ranks, list): mpi_ranks = [mpi_ranks] comm, group = mpi.CreateGroup(root, incl=mpi_ranks) arguments = {'inputs': arguments['inputs'], 'comm': comm, 'group': group} outputs = Tensor.CreateOperator(nout=len(mpi_ranks), op_type='MPIGather', **arguments) if inputs.shape is not None: if isinstance(outputs, list): for output in outputs: output.shape = inputs.shape[:] else: outputs.shape = inputs.shape[:] return outputs
def GraphDef_Update(graph_def, updater): """Inject the update targets into GraphDef. The ``updater`` should generate update targets before. Parameters ---------- graph_def : GraphDef The definition of graph. updater : BaseUpdater The updater. Returns ------- None """ if updater is None: return extra_arguments = updater._extra_kwargs extra_arguments['slot'] = updater._slot parallel_arguments = {} updater.register_in_workspace() # Check data parallel if necessary if mpi.Is_Init(): idx, group = mpi.AllowParallel() if idx != -1: parallel_arguments['parallel_mode'] = mpi.GetParallelMode() parallel_arguments['comm'], parallel_arguments['group'] \ = mpi.CreateGroup(root=group[0], incl=group) parallel_arguments['root'] = group[0] for k, v in parallel_arguments.items(): graph_def.arg.add().CopyFrom(MakeArgument(k, v)) for e in updater._param_group: pair, arguments = e kwargs = dict(arguments, **extra_arguments) u_target = pb.UpdaterProto() u_target.type = updater.type() u_target.name = OperatorHelper.get_name() u_target.tensor.extend(pair) for k, v in kwargs.items(): u_target.arg.add().CopyFrom(MakeArgument(k, v)) graph_def.updater.extend([u_target])
def MPIGather(inputs, root, mpi_ranks=None, **kwargs): """Gather a tensor from all nodes to root in the ``MPIGroup``. Parameters ---------- inputs : Tensor The tensor to gather. root : int The world rank of root node. mpi_ranks: sequence of int, optional The world rank of nodes in group. Default is ``None`` (Use All). Returns ------- sequence of Tensor The gathered outputs. Notes ----- The number of outputs is decided on the number of ``mpi_ranks``. The outputs are **accessible** only for root and vice versa. """ arguments = ParseArgs(locals()) if mpi_ranks is None: num_nodes = mpi.Size() mpi_ranks = [i for i in range(0, num_nodes)] if not isinstance(mpi_ranks, list): mpi_ranks = [mpi_ranks] comm, group = mpi.CreateGroup(root, incl=mpi_ranks) arguments = { 'inputs': arguments['inputs'], 'comm': comm, 'group': group, 'num_outputs': len(mpi_ranks) } return Tensor.CreateOperator('MPIGather', **arguments)