def init(self, topology_fn: Optional[Callable[[int], networkx.DiGraph]] = None, is_weighted: bool = False): """A function that initializes BlueFog. Args: topology_fn: A callable function that takes size as input and return networkx.DiGraph object to decide the topology. If not provided a default exponential graph (base 2) structure is called. is_weighted: If set to true, the neighbor ops like (win_update, neighbor_allreduce) will execute the weighted average instead, where the weight is the value used in topology matrix (including self). """ self._MPI_LIB_CTYPES.bluefog_init() if topology_fn: topo = topology_fn(self.size()) else: topo = topology_util.ExponentialGraph(self.size()) self.set_topology(topo, is_weighted) atexit.register(self.shutdown)
args = parser.parse_args() args.cuda = not args.no_cuda and torch.cuda.is_available() bf.init() torch.random.manual_seed(args.seed * bf.rank()) if args.cuda: device = bf.local_rank() % torch.cuda.device_count() x = torch.randn(args.data_size, device=device, dtype=torch.double) else: x = torch.randn(args.data_size, dtype=torch.double) if args.virtual_topology == "expo2": pass elif args.virtual_topology == "expo3": bf.set_topology(topology_util.ExponentialGraph(bf.size(), base=3)) elif args.virtual_topology == "expo4": bf.set_topology(topology_util.ExponentialGraph(bf.size(), base=4)) elif args.virtual_topology == "ring": bf.set_topology(topology_util.RingGraph(bf.size(), connect_style=1)) elif args.virtual_topology == "mesh": bf.set_topology(topology_util.RingGraph(bf.size(), connect_style=0), is_weighted=True) elif args.virtual_topology == "star": bf.set_topology(topology_util.StarGraph(bf.size()), is_weighted=True) elif args.virtual_topology == "full": bf.set_topology(topology_util.FullyConnectedGraph(bf.size())) else: raise ValueError("Unknown args.virtual_topology, supporting options are " + "[expo2(Default), ring, mesh, star].")
if bf.rank() == 0: mse.append(torch.norm(x.data - w_opt, p=2)) bf.barrier() w = bf.win_update_then_collect(name="w_buff") x.data = w[:n] / w[-1] return x, mse # ======================= Code starts here ======================= bf.init() if args.topology == 'mesh': bf.set_topology(topology_util.MeshGrid2DGraph(bf.size()), is_weighted=True) elif args.topology == 'expo2': bf.set_topology(topology_util.ExponentialGraph(bf.size())) elif args.topology == 'star': bf.set_topology(topology_util.StarGraph(bf.size()), is_weighted=True) elif args.topology == 'ring': bf.set_topology(topology_util.RingGraph(bf.size())) else: raise NotImplementedError( 'Topology not supported. This example only supports' + ' mesh, star, ring and expo2') # Generate data for logistic regression (synthesized data) torch.random.manual_seed(123417 * bf.rank()) m, n = 20, 5 rho = 1e-2 X, y = generate_data(m, n, task=args.task)
def set_topology(self, topology: Optional[networkx.DiGraph] = None, is_weighted: bool = False) -> bool: """A function that sets the virtual topology MPI used. Args: Topo: A networkx.DiGraph object to decide the topology. If not provided a default exponential graph (base 2) structure is used. is_weighted: If set to true, the win_update and neighbor_allreduce will execute the weighted average instead, where the weights are the value used in topology matrix (including self weight). Note win_get/win_put/win_accumulate do not use this weight since win_update already uses these weights. Returns: A boolean value that whether topology is set correctly or not. Example: >>> import bluefog.torch as bf >>> from bluefog.common import topology_util >>> bf.init() >>> bf.set_topology(topology_util.RingGraph(bf.size())) """ if topology is None: topology = topology_util.ExponentialGraph(size=self.size()) if self.local_rank() == 0: logger.info( "Topology is not specified. Default Exponential Two topology is used." ) if not isinstance(topology, networkx.DiGraph): raise TypeError("topology must be a networkx.DiGraph obejct.") if topology_util.IsTopologyEquivalent(topology, self._topology): if self.local_rank() == 0: logger.debug( "Topology to set is the same as old one. Skip the setting." ) return True # We remove the self-rank for any cases because MPI graph_comm do not include it. destinations = sorted( [r for r in topology.successors(self.rank()) if r != self.rank()]) sources = sorted([ r for r in topology.predecessors(self.rank()) if r != self.rank() ]) indegree = len(sources) outdegree = len(destinations) sources_type = ctypes.c_int * indegree destinations_type = ctypes.c_int * outdegree if not is_weighted: self._MPI_LIB_CTYPES.bluefog_set_topology.argtypes = ([ ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.POINTER(ctypes.c_int) ]) ret = self._MPI_LIB_CTYPES.bluefog_set_topology( indegree, sources_type(*sources), outdegree, destinations_type(*destinations)) else: # Here the source_weights is a vector containing weights from source, i.e., # (in-)neighbors, converted from the neighbor_weights dictionary. self_weight, neighbor_weights = topology_util.GetRecvWeights( topology, self.rank()) source_weights = [ neighbor_weights[r] for r in sorted(neighbor_weights.keys()) ] source_weights_type = ctypes.c_float * indegree self._MPI_LIB_CTYPES.bluefog_set_topology_with_weights.argtypes = ( [ ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.c_float, ctypes.POINTER(ctypes.c_float) ]) ret = self._MPI_LIB_CTYPES.bluefog_set_topology_with_weights( indegree, sources_type(*sources), outdegree, destinations_type(*destinations), self_weight, source_weights_type(*source_weights)) if ret != 1: if self.local_rank() == 0: logger.error( "Cannot set topology correctly. Three common reasons caused this. \n" "1. Has Bluefog been initialized? use bf.init(). \n" "2. The win_create has been called. It is not allowed to change\n" " the topology after that. You can call win_free() to unregister\n" " all window object first, then set the topology. \n" "3. Make sure all previous MPI ops are done. It is not allowed to \n" " change the topology while there is undone MPI ops.") return False self._topology = topology self._is_topo_weighted = is_weighted return True