def _update_gradient_consumers(new_graph_item, consumer_ops, control_consumer_ops, old_tensor_name, new_tensor): """Make gradient's consumers consume the aggregated gradient instead of the original one of replica_0.""" # Get the original tensor (the one from replica 0) to replace old_op_name = strip_replica_prefix(get_op_name(old_tensor_name)) replica_0_op_name = ops.prepend_name_scope(old_op_name, replica_prefix(0)) replica_0_op = new_graph_item.graph.get_operation_by_name( replica_0_op_name) output_idx = get_index_from_tensor_name(old_tensor_name) replica_0_tensor = replica_0_op.outputs[output_idx] update_consumers(consumer_ops, replica_0_tensor, new_tensor) update_control_consumers(control_consumer_ops, replica_0_tensor.op, new_tensor.op)
def _collect_sparse_gradients(self, graph_item, var_op_name): """Append collective ops after the gradient is calculated.""" if self.num_workers > 1 and not ENV.AUTODIST_INTERNAL_TF.value: raise NotImplementedError( 'Currently the collective NCCL AllGather is not supported in TensorFlow release.' 'Please choose another strategy.') conf = {} if self._spec: conf = {'communication_hint': self._spec} if self._compressor_type: logging.warning( 'AllGather currently does not support AutoDist compressor so it skips.' ) if self.num_replicas * self.num_workers <= 1: raise ValueError( 'CollectiveOps requires collective group size > 1') for i in range(0, self.num_replicas): op_name = ops.prepend_name_scope(var_op_name, replica_prefix(i)) graph_item.updated = True grad, _, _ = graph_item.var_op_name_to_grad_info_v2[op_name] # TODO (Tairui): (3) Merge of reduction for performance indices_c_ops = grad.indices.consumers() indices_cc_ops = get_control_consumers(grad.indices.op) values_c_ops = grad.values.consumers() values_cc_ops = get_control_consumers(grad.values.op) with ops.name_scope(replica_prefix(i)): with ops.colocate_with(grad.indices.op): new_indices = collective_ops.all_gather( grad.indices, self.num_replicas * self.num_workers, get_collective_keys().get_group_key( self.all_canonical_replica_devices), get_collective_keys().get_instance_key(var_op_name + '-indices'), **conf) with ops.colocate_with(grad.values.op): new_values = collective_ops.all_gather( grad.values, self.num_replicas * self.num_workers, get_collective_keys().get_group_key( self.all_canonical_replica_devices), get_collective_keys().get_instance_key(var_op_name + '-values'), **conf) update_consumers(indices_c_ops, grad.indices, new_indices) update_control_consumers(indices_cc_ops, grad.indices.op, new_indices.op) update_consumers(values_c_ops, grad.values, new_values) update_control_consumers(values_cc_ops, grad.values.op, new_values)
def _get_accum_apply_and_agg_grad(var_op, grad, indices, dense_shape): if indices is None: tensor = variable_utils.get_read_var_tensor(var_op) grad_accum = data_flow_ops.ConditionalAccumulator( grad.dtype, shape=tensor.get_shape(), shared_name=var_op.name + "/grad_accum") # Get a copy of consumers list before creating accum_apply_op grad_consumers = list(grad.consumers()) accum_apply_op = grad_accum.apply_grad(grad, local_step=MAX_INT64, name=grad.op.name + '_accum_apply_grad') agg_grad = grad_accum.take_grad(num_accum_required, name=var_op.name + '_take_grad') update_consumers(grad_consumers, grad, agg_grad) update_control_consumers(get_control_consumers(grad.op), grad.op, agg_grad.op) else: grad_indexed_slices = ops.IndexedSlices( values=grad, indices=indices, dense_shape=dense_shape) grad_accum = data_flow_ops.SparseConditionalAccumulator( grad.dtype, shape=grad.shape, shared_name=var_op.name + "/grad_accum") # Get a copy of consumers list before creating accum_apply_op indices_consumers = list(indices.consumers()) grad_consumers = list(grad.consumers()) accum_apply_op = grad_accum.apply_indexed_slices_grad( grad_indexed_slices, local_step=MAX_INT64, name=grad.op.name + '_accum_apply_grad') agg_grad = grad_accum.take_indexed_slices_grad( num_accum_required, name=var_op.name + '_take_grad') agg_indices = agg_grad.indices if indices.dtype != agg_grad.indices.dtype: agg_indices = math_ops.cast(agg_grad.indices, indices.dtype) agg_grad = ops.IndexedSlices(values=agg_grad.values, indices=agg_indices, dense_shape=agg_grad.dense_shape) assert isinstance(agg_grad, ops.IndexedSlices) update_consumers(indices_consumers, indices, agg_grad.indices) update_consumers(grad_consumers, grad, agg_grad.values) update_control_consumers(get_control_consumers(indices.op), indices.op, agg_grad.indices.op) update_control_consumers(get_control_consumers(grad.op), grad.op, agg_grad.values.op) return accum_apply_op, agg_grad
def add_sync_op(self, graph_item, var_update_op, variable_replicator=None): """ Adds additional ops needed for synchronous distributed training into current graph. Main purpose of additional ops are: 1. Initialization 2. Synchronization 3. Gradient aggregation Args: graph_item (graph_item.GraphItem): the graph var_update_op: The op variable_replicator: The dictionary of master variable op name -> list of replicated variables, could be None Returns: None """ this_worker_cpu = device_spec.DeviceSpecV2.from_string( self.worker_device) this_worker_cpu = this_worker_cpu.replace(device_type='CPU', device_index=0) var_op = var_update_op.inputs[UPDATE_OP_VAR_POS].op is_trainable = var_op in graph_item.trainable_var_op_to_var source_op = self._get_optimizer_source_op(var_update_op) cc = get_control_consumers(source_op) with ops.device(var_op.device): if self._staleness == 0: queue_ops = self._get_queue_ops(var_update_op, source_op, self.is_chief, is_trainable) elif self._staleness > 0: queue_ops = self._get_queue_ops_stale(var_update_op, source_op, self.is_chief, is_trainable) else: raise ValueError( "staleness should be greater than or equal to 0.") # Only dense trainable variables are replicated locally if variable_replicator: mirror_variable_update_ops = variable_replicator.get_all_update_ops( queue_ops, worker_device=this_worker_cpu) with ops.device(this_worker_cpu): finish_op = control_flow_ops.group( *mirror_variable_update_ops) else: finish_op = control_flow_ops.group(*queue_ops) # Place computation ops of aggregated gradients on PS # Note that even though this is doing a graph traversal, it is called in such a way that it # only traverses from a gradient aggregator op to a gradient application op (or vice versa) -- # these corresponding ops should always be adjacent in the graph. self._place_post_grad_agg_ops( device_spec.DeviceSpecV2.from_string(self.target_device), self._var_op_to_agg_grad, {var_op: var_update_op} if is_trainable else {}) # Replace the control input of train_op to be finish_op # Note(Hao): this cc is stale, i.e. cc \subset get_control_consumers(source_op) update_control_consumers(cc, source_op, finish_op)