def step(self, closure=None, **kargs): # Apply the gradients with the weight decay and momentum. with kargs["timer"]("grad.apply_grad", epoch=self.conf.epoch_): utils.apply_gradient( self.param_groups, self.state, apply_grad_to_model=True ) with kargs["timer"]("grad.get_params", epoch=self.conf.epoch_): params, _ = comm.get_data( self.param_groups, self.param_names, is_get_grad=False ) params_tb = TensorBuffer(params) with kargs["timer"]("grad.error_compensate", epoch=self.conf.epoch_): self.memory.buffer += params_tb.buffer with kargs["timer"]("grad.compress", epoch=self.conf.epoch_): sync_buffer = {"original_shapes": self.shapes, "params_tb": self.memory} local_compressed_params_tb = self.compressor.compress(sync_buffer) with kargs["timer"]("grad.update_memory", epoch=self.conf.epoch_): self.memory.buffer = self.memory.buffer - local_compressed_params_tb.buffer with kargs["timer"]("grad.sync", epoch=self.conf.epoch_): self.compressor.sync(sync_buffer) # update local model. with kargs["timer"]("grad.decompress", epoch=self.conf.epoch_): aggregated_info_tb = self.compressor.uncompress( sync_buffer, self.neighbors_info ) params_tb.buffer += aggregated_info_tb.buffer params_tb.unpack(params) return sync_buffer["n_bits"]
def step(self, closure=None, **kargs): with kargs['timer']('sync', epoch=self.conf.epoch_): # do the local update steps. with kargs["timer"]("local_update", epoch=self.conf.epoch_): utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=True) # enter the global sync if it satisfies the condition. if (self.conf.epoch_ < self.turn_on_local_step_from_epoch or self.conf.local_index % self.local_step == 0): with kargs["timer"]("get_params", epoch=self.conf.epoch_): # get parmas. params, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=False) params_tb = TensorBuffer(params) with kargs['timer']('memory_and_compress', epoch=self.conf.epoch_): # get the params difference w.r.t. previous synced model. local_scale, local_sign = [], [] for consensus_param, param, memory in zip( self.consensus_params_tb, params_tb, self.memory_tb): memory.data.copy_(consensus_param - param + memory) # compress. with kargs["timer"]("directions", epoch=self.conf.epoch_): direction = exchange(self.memory_tb.buffer) #signum with kargs['timer']('memory_and_compress', epoch=self.conf.epoch_): for consensus_param, param, memory in zip( self.consensus_params_tb, params_tb, self.memory_tb): _local_scale, _local_sign = scaled_sign(memory) local_scale.append(_local_scale) local_sign.append(_local_sign) memory.data.copy_(memory - _local_scale * _local_sign) with kargs["timer"]("directions", epoch=self.conf.epoch_): global_direction = TB(self.memory_tb, direction) with kargs["timer"]("magnitudes", epoch=self.conf.epoch_): magnitudes_tb = TensorBuffer(local_scale) magnitudes_tb.buffer = self.world_aggregator._agg( magnitudes_tb.buffer, "avg", distributed=self.conf.distributed) # unpack the synced info and update the consensus params. with kargs["timer"]("update_consensus", epoch=self.conf.epoch_): for update_magnitude, update_direction, consensus_param in zip( magnitudes_tb, global_direction, self.consensus_params_tb): consensus_param.add_( -1.0, update_direction.mul(update_magnitude)) # consistent the local models by assigning the consensus params. self.consensus_params_tb.unpack(params) n_bits = get_n_bits(magnitudes_tb.buffer) else: n_bits = 0 return n_bits
def step(self, closure=None, **kargs): with kargs["timer"]("sync.local_update", epoch=self.conf.epoch_): utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=True) with kargs["timer"]("sync.sync_and_update", epoch=self.conf.epoch_): # enter the global sync if it satisfies the condition. if (self.conf.epoch_ < self.turn_on_local_step_from_epoch or self.conf.local_index % self.local_step == 0): # get parmas. params, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=False) params_tb = TensorBuffer(params) # get params_diff. param_diff = self.consensus_params_tb.buffer - params_tb.buffer # sync the directions. param_diff = self.world_aggregator._agg( param_diff, "avg", distributed=self.conf.distributed) # unpack the synced info and update the consensus params. self.consensus_params_tb.buffer.add_(-1.0, param_diff) # consistent the local models by assigning the consensus params. self.consensus_params_tb.unpack(params) # Get n_bits to transmit. n_bits = get_n_bits(param_diff) else: n_bits = 0 return n_bits
def step(self, closure=None, **kargs): # Apply the gradients with the weight decay and momentum. with kargs["timer"]("grad.apply_grad", epoch=self.conf.epoch_): utils.apply_gradient( self.param_groups, self.state, apply_grad_to_model=False ) # get flattened params. with kargs["timer"]("grad.get_params", epoch=self.conf.epoch_): params, _ = comm.get_data( self.param_groups, self.param_names, is_get_grad=False ) flatten_params = TensorBuffer(params) grads, _ = comm.get_data( self.param_groups, self.param_names, is_get_grad=True ) flatten_grads = TensorBuffer(grads) with kargs["timer"]("grad.get_extrapolated_model", epoch=self.conf.epoch_): flatten_updated_params = deepcopy(flatten_params) # get weighted hat params. flatten_updated_params.buffer = sum( [ _hat_params.buffer * self.neighbors_info[_rank] for _rank, _hat_params in self.neighbor_hat_params.items() ] ) # get updated local model (flatten params). with kargs["timer"]("grad.unflatten_to_update", epoch=self.conf.epoch_): flatten_updated_params.buffer.add_( flatten_grads.buffer, alpha=-self.param_groups[0]["lr"] ) flatten_updated_params.unpack(params) # get extrapolated model. flatten_updated_params.buffer = ( (1 - 0.5 * self.conf.local_index) * flatten_params.buffer + 0.5 * self.conf.local_index * flatten_updated_params.buffer ) # compress the model difference and sync. with kargs["timer"]("grad.compress", epoch=self.conf.epoch_): sync_buffer = { "original_shapes": self.shapes, "flatten_updated_params": flatten_updated_params, } self.compressor.compress(sync_buffer) with kargs["timer"]("grad.sync", epoch=self.conf.epoch_): self.compressor.sync(sync_buffer) with kargs["timer"]("grad.unflatten_to_update", epoch=self.conf.epoch_): self.compressor.uncompress( sync_buffer, self.neighbor_hat_params, self.conf.local_index ) return sync_buffer["n_bits"]
def step(self, closure=None, **kargs): if self.conf.is_centralized: with kargs["timer"]("sync/get_data", epoch=self.conf.epoch_): # Get data. grads, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=True) flatten_grads = TensorBuffer(grads) with kargs["timer"]("sync/sync", epoch=self.conf.epoch_): # Aggregate the gradients. flatten_grads.buffer = self.world_aggregator._agg( flatten_grads.buffer, op="avg", distributed=self.conf.distributed) with kargs["timer"]("sync/unflatten_grad", epoch=self.conf.epoch_): # unflatten grads. flatten_grads.unpack(grads) with kargs["timer"]("sync/apply_grad", epoch=self.conf.epoch_): utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=True) # Get n_bits to transmit. n_bits = get_n_bits(flatten_grads.buffer) else: with kargs["timer"]("sync/apply_grad", epoch=self.conf.epoch_): utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=True) with kargs["timer"]("sync/get_data", epoch=self.conf.epoch_): # first get and flatten all params. params, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=False) flatten_params = TensorBuffer(params) with kargs["timer"]("sync/sync", epoch=self.conf.epoch_): # prepare the sync. if self.conf.comm_device == "cpu": flatten_params.buffer.cpu().detach_() # then sync. flatten_params.buffer = self.decentralized_aggregator._agg( flatten_params.buffer, op="weighted") with kargs["timer"]("sync/update_model", epoch=self.conf.epoch_): # finally unflatten. flatten_params.unpack(params) # Get n_bits to transmit. n_bits = get_n_bits(flatten_params.buffer) return n_bits
def step(self, closure=None, **kargs): # Apply the gradients with the weight decay and momentum. with kargs["timer"]("sync.apply_grad", epoch=self.conf.epoch_): utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=True) with kargs["timer"]("sync.finish_sync", epoch=self.conf.epoch_): utils.join_thread(self.helper_thread) self.n_bits = self.sync_buffer.get("n_bits", 0) # recover current params and hat_params with kargs["timer"]("sync.recover_hat_params", epoch=self.conf.epoch_): params, flatten_params, flatten_hat_params = utils.recover_params( param_groups=self.param_groups, param_names=self.param_names, rank=self.rank, neighbor_hat_params=self.neighbor_hat_params, get_hat_params=True, ) # get updated flatten params. with kargs["timer"]("sync.update_flatten_params", epoch=self.conf.epoch_): utils.update_params_from_neighbor( neighbor_hat_params=self.neighbor_hat_params, flatten_params=flatten_params, consenus_stepsize=self.consenus_stepsize, self_rank=self.rank, ) # update the local model. with kargs["timer"]("sync.update_local_model", epoch=self.conf.epoch_): flatten_params.unpack(params) # start compress/sync. with kargs["timer"]("sync.start_sync", epoch=self.conf.epoch_): self.sync_buffer = { "original_shapes": self.shapes, "flatten_params": flatten_params, "flatten_hat_params": flatten_hat_params, } self.helper_thread = utils.HelperThread( name=f"_thread_at_epoch_{self.conf.epoch_}.compress", func=self.compressor.pipeline, # the arguments below will be feeded into the `func`. sync_buffer=self.sync_buffer, neighbor_hat_params=self.neighbor_hat_params, neighbors_info=self.neighbors_info, ) self.helper_thread.start() if self.conf.epoch_ % 1 == 0: utils.join_thread(self.helper_thread) return self.n_bits
def step(self, closure=None, **kargs): with kargs["timer"]("sync.apply_grad", epoch=self.conf.epoch_): # Apply the gradients with the weight decay and momentum. utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=True) self.updated_local_model_flag.set() with kargs["timer"]("sync.gossip_model", epoch=self.conf.epoch_): while not self.gossiped_dist_model_flag.is_set(): pass self.gossiped_dist_model_flag.clear() return self.n_bits.item()
def step(self, closure=None, **kargs): # Apply the gradients with the weight decay and momentum. with kargs["timer"]("grad.apply_grad", epoch=self.conf.epoch_): utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=False) with kargs["timer"]("grad.get_grads", epoch=self.conf.epoch_): params, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=False) flatten_params = TensorBuffer(params) grads, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=True) flatten_grads = TensorBuffer(grads) # Get weighted hat params and apply the local gradient. with kargs["timer"]("grad.apply_local_gradient", epoch=self.conf.epoch_): flatten_half_params = deepcopy(flatten_params) flatten_half_params.buffer = (sum([ _hat_params.buffer * self.neighbors_info[_rank] for _rank, _hat_params in self.neighbor_hat_params.items() ]) - self.param_groups[0]["lr"] * flatten_grads.buffer) # compress the model difference and sync. with kargs["timer"]("grad.compress", epoch=self.conf.epoch_): sync_buffer = { "original_shapes": self.shapes, "flatten_half_params": flatten_half_params, "flatten_params": flatten_params, } self.compressor.compress(sync_buffer) with kargs["timer"]("grad.sync", epoch=self.conf.epoch_): self.compressor.sync(sync_buffer) # finally unflatten and update local model. with kargs["timer"]("grad.unflatten_to_update", epoch=self.conf.epoch_): self.compressor.uncompress(sync_buffer, self.neighbor_hat_params) flatten_params.buffer = self.neighbor_hat_params[ self.rank].buffer.clone() flatten_params.unpack(params) return sync_buffer["n_bits"]
def step(self, closure=None, **kargs): with kargs["timer"]("sync.apply_grad", epoch=self.conf.epoch_): utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=False) with kargs["timer"]("sync.get_data", epoch=self.conf.epoch_): # Get data. grads, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=True) grads_tb = TensorBuffer(grads) with kargs["timer"]("sync.use_memory", epoch=self.conf.epoch_): # use memory. grads_tb.buffer.add_(self.memory_tb.buffer) with kargs["timer"]("sync.compress", epoch=self.conf.epoch_): # compress. sync_buffer = self.compressor.compress(grads_tb) with kargs["timer"]("sync.sync", epoch=self.conf.epoch_): self.compressor.sync(sync_buffer) with kargs["timer"]("sync.update_memory", epoch=self.conf.epoch_): # update memory. self.memory_tb.buffer = (grads_tb.buffer - sync_buffer["synced_grads_tb"].buffer) with kargs["timer"]("sync.decompress", epoch=self.conf.epoch_): sync_grads_tb = self.compressor.decompress(sync_buffer) with kargs["timer"]("sync.apply_grad", epoch=self.conf.epoch_): # appply the gradient but only with the gradient. params, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=False) params_tb = TensorBuffer(params) # apply the gradient. params_tb.buffer.add_(-self.param_groups[0]["lr"] * sync_grads_tb.buffer) # unpack. params_tb.unpack(params) return sync_buffer["n_bits"]
def step(self, closure=None, **kargs): # do the local update steps. with kargs["timer"]("sync/get_data", epoch=self.conf.epoch_): # get parmas. params, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=False) params_tb = TensorBuffer(params) with kargs["timer"]("sync/apply_grad", epoch=self.conf.epoch_): # prepare the gradient (sign) utils.apply_gradient(self.param_groups, self.state, apply_grad_to_model=False) # get grads. grads, _ = comm.get_data(self.param_groups, self.param_names, is_get_grad=True) grads_tb = TensorBuffer(grads) # enter the global sync if it satisfies the condition. # get the params difference w.r.t. previous synced model. with kargs["timer"]("sync/compress", epoch=self.conf.epoch_): sync_buffer = self.compressor.compress(grads_tb) # sync and decompress. with kargs["timer"]("sync/sync_and_decompress", epoch=self.conf.epoch_): self.compressor.sync(sync_buffer) synced_updates_tb = self.compressor.decompress(sync_buffer) # unpack the synced info and update the consensus params. with kargs["timer"]("sync/apply_grad", epoch=self.conf.epoch_): params_tb.buffer -= self.param_groups[0][ "lr"] * synced_updates_tb.buffer params_tb.unpack(params) return sync_buffer["n_bits"]