def init_group(self, world_size, rank, backend=Backend.NCCL, group_name="default"): col.init_collective_group(world_size, rank, backend, group_name) return True
def __init__(self, workers, world_size, rank): self.params = dict() self.optimizer = None self.workers = workers self.world_size = world_size self.rank = rank self.grad_counts = 0 collective.init_collective_group(self.world_size, self.rank, "nccl", "default") for i in range(len(self.workers)): recv = torch.zeros(1, ).cuda() collective.recv(recv, i, "default") for i in range(len(self.workers)): recv = torch.zeros(1, ).cuda() collective.send(recv, i, "default")
def __init__(self, model, batch_size, world_size, rank, num_ps): self.model_type = model print("=> creating model '{}'".format(model)) self.model = torchmodels.__dict__[model]().cuda() self.criterion = nn.CrossEntropyLoss().cuda() self.batch_size = batch_size self.train_loader = self.get_data_loader(self.batch_size) self.world_size = world_size self.rank = rank self.num_ps = num_ps self.num_workers = self.world_size - self.num_ps self.assignments = None # index i of this list stores the names of params in ith server. self.name_list = [[] for i in range(num_ps)] collective.init_collective_group(world_size, rank, "nccl", "default") for i in range(num_ps): send = torch.ones(1, ).cuda() collective.send(send, self.num_workers + i, "default") for i in range(num_ps): send = torch.ones(1, ).cuda() collective.recv(send, self.num_workers + i, "default")
def setup(self, world_size, rank): collective.init_collective_group(world_size, rank, "nccl", "177") return True
def setup(self, world_size, rank): self.rank = rank collective.init_collective_group(world_size, rank, "nccl", "8") return True
def setup(self, world_size, rank): collective.init_collective_group("nccl", world_size, rank, "default") return True
def init_gloo_group(rank: int, world_size: int, group_name: str): col.init_collective_group(world_size, rank, Backend.GLOO, group_name) return True