# Copyright 2018 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Utilities for cross_device_ops.""" import copy import threading from typing import Callable, List, Optional, Union from tensorflow.python.distribute import collective_util from tensorflow.python.distribute import values as value_lib from tensorflow.python.eager import backprop_util from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import indexed_slices from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_spec from tensorflow.python.ops import array_ops from tensorflow.python.ops import collective_ops from tensorflow.python.ops import cond from tensorflow.python.ops import math_ops from tensorflow.python.ops import nccl_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.platform import tf_logging as logging from tensorflow.python.types import core INSTANCE_KEY_START_NUMBER = 100 def aggregate_gradients_using_nccl(replica_grads): """Aggregate gradients using nccl allreduce.""" agg_all_g_and_v = [] for single_g_and_v in zip(*replica_grads): single_grads = [g for g, _ in single_g_and_v] agg_grads = nccl_ops.all_sum(single_grads) agg_all_g_and_v.append( [(g, v) for g, (_, v) in zip(agg_grads, single_g_and_v)]) agg_all_g_and_v = list(zip(*agg_all_g_and_v)) return agg_all_g_and_v def aggregate_gradients_using_hierarchical_copy(avail_devices, replica_grads): """Aggregate gradients using hierarchical copies. Args: avail_devices: available GPU devices. replica_grads: List of lists of (gradient, variable) tuples. The outer list is over replicas. The inner list is over individual gradients. Returns: The list of (aggregated_gradient, variable), where the gradient has been summed across all replicas and the variable is chosen from the first replica. """ # This only works for DGX-1 type of machine topology # Device peer to peer matrix # DMA: 0 1 2 3 4 5 6 7 # 0: Y Y Y Y Y N N N # 1: Y Y Y Y N Y N N # 2: Y Y Y Y N N Y N # 3: Y Y Y Y N N N Y # 4: Y N N N Y Y Y Y # 5: N Y N N Y Y Y Y # 6: N N Y N Y Y Y Y # 7: N N N Y Y Y Y Y agg_grads = [] num_devices = len(avail_devices) # In the special case of DGX-1 machine topology, the two groups have equal # size. group_size = num_devices // 2 for i, single_grads in enumerate(zip(*replica_grads)): group_0_main_device = i % num_devices group_1_main_device = (group_0_main_device + group_size) % num_devices if group_0_main_device < group_size: group_0_begin = 0 group_1_begin = group_size else: group_0_begin = group_size group_1_begin = 0 # Aggregate the first group. group_0_device_grads = single_grads[group_0_begin: group_0_begin + group_size] with ops.device(avail_devices[group_0_main_device]): group_0_agg_grads, _ = aggregate_single_gradient_using_copy( group_0_device_grads, False, False) # Aggregate the second group. group_1_device_grads = single_grads[group_1_begin: group_1_begin + group_size] with ops.device(avail_devices[group_1_main_device]): group_1_agg_grads, _ = aggregate_single_gradient_using_copy( group_1_device_grads, False, False) # Aggregate between the groups. with ops.device(avail_devices[group_0_main_device]): (agg_total_grads, _), _ = aggregate_single_gradient_using_copy( [group_0_agg_grads, group_1_agg_grads], False, False) # Broadcast the result back into the root of each group. with ops.device(avail_devices[group_0_main_device]): group_0_agg_grads_bcast = array_ops.identity(agg_total_grads) with ops.device(avail_devices[group_1_main_device]): group_1_agg_grads_bcast = array_ops.identity(agg_total_grads) agg_grads_bcast = [] for j in range(len(single_grads)): with ops.device(avail_devices[j]): # Broadcast the result back to each member in the group from the root. if (group_0_main_device < group_size) == (j < group_size): src_device_grad = group_0_agg_grads_bcast else: src_device_grad = group_1_agg_grads_bcast agg_grads_bcast.append(array_ops.identity(src_device_grad)) agg_grads.append( [(g, v) for g, (_, v) in zip(agg_grads_bcast, single_grads)]) agg_grads = list(zip(*agg_grads)) return agg_grads def aggregate_single_gradient_using_copy(grad_and_vars, use_mean, check_inf_nan): """Calculate the average gradient for a shared variable across all replicas. Note that this function provides a synchronization point across all replicas. Args: grad_and_vars: A list or tuple of (gradient, variable) tuples. Each (gradient, variable) pair within the outer list represents the gradient of the variable calculated for a single replica, and the number of pairs equals the number of replicas. use_mean: if True, mean is taken, else sum of gradients is taken. check_inf_nan: check grads for nans and infs. Returns: The tuple ([(average_gradient, variable),], has_nan_or_inf) where the gradient has been averaged across all replicas. The variable is chosen from the first replica. The has_nan_or_inf indicates the grads has nan or inf. """ grads = [g for g, _ in grad_and_vars] grad = math_ops.add_n(grads) if use_mean and len(grads) > 1: grad = array_ops.multiply(grad, 1.0 / len(grads)) v = grad_and_vars[0][1] if check_inf_nan: has_nan_or_inf = array_ops.logical_not( array_ops.reduce_all(array_ops.is_finite(grads))) return (grad, v), has_nan_or_inf else: return (grad, v), None # TODO(yuefengz): use random key starts to avoid reusing keys? class CollectiveKeys(object): """Class that manages collective keys. We need to manage three different keys for collective: *Group key*: an integer key to identify the set of cooperative devices. Collective ops work under the same set of devices must using the same group key. *Instance key*: an integer key to identify the set of same counterpart of tensors on different devices in a device group that need to be all-reduced. This class is thread safe. """ def __init__(self, group_key_start=1): """Initializes the object. Args: group_key_start: the starting integer of group key. """ self._group_key = group_key_start self._instance_key_table = {} self._lock = threading.Lock() self._known_groups = {} def get_group_key(self, devices): """Returns a group key for the list of local devices. The same group key is returned if the list of local devices is the same. Args: devices: a list of local canonical device strings in a collective group. Returns: a group key. """ with self._lock: devices_key = ','.join(devices) if devices_key not in self._known_groups: self._known_groups[devices_key] = self._get_new_group_key(devices) return self._known_groups[devices_key] def _get_new_group_key(self, devices): """Returns a new group key. The caller should store and reuse the same group key for the same set of devices. Calling this method always returns a new group key. This method is not thread-safe. Args: devices: a list of canonical device strings in a collective group. Returns: a new group key. """ new_key = self._group_key self._group_key += 1 self._instance_key_table[new_key] = {} for device in devices: self._instance_key_table[new_key][device] = INSTANCE_KEY_START_NUMBER return new_key def get_instance_key(self, group_key, device): """Returns a new instance key for use in defining a collective op. You should call this once per each collective op of a collective instance. Args: group_key: the group key returned by get_group_key(). You should not assign the group key yourself. device: a canonical device string. It should be the device this collective op is on. Returns: a new instance key. Raises: ValueError: when the group key is invalid or the device is not in the group. """ with self._lock: group = self._instance_key_table.get(group_key, None) if group is None: raise ValueError(f'Group {group_key} is not found.') if device not in group: raise ValueError(f'Device {device} is not present in group {group_key}') v = group[device] group[device] += 1 return v def __deepcopy__(self, memo): # distribute_coordinator deep-copies the strategy object, so # CollectiveKeys needs to support deep copy as well. copied = CollectiveKeys() copied._group_key = self._group_key copied._instance_key_table = copy.deepcopy(self._instance_key_table, memo) return copied class CollectiveReplicaLauncher(object): """Launch collectives on one replica.""" _prefer_unique_instance_key = True _prefer_ordering_token = True def __init__(self, group_key: int, group_size: int, collective_keys: CollectiveKeys, device: str, options: collective_util.Options): self._group_key = group_key self._group_size = group_size self._collective_keys = collective_keys self._device = device self._options = options if self._use_ordering_token(): with ops.init_scope(), ops.device(device): self._ordering_token = resource_variable_ops.ResourceVariable(0.) else: self._ordering_token = None def _control_input(self, control_input: Union[core.TensorLike, ops.Operation]): if control_input is not None and not self._use_ordering_token(): return ops.control_dependencies([control_input]) return ops.NullContextmanager() def _use_unique_instance_key(self): if not ops.executing_eagerly_outside_functions(): return False return CollectiveReplicaLauncher._prefer_unique_instance_key def _use_ordering_token(self): # We rely on auto control dep to insert control edges between NCCL calls, # but for tf1 graph mode auto control dep is not used. if not ops.executing_eagerly_outside_functions(): return False return CollectiveReplicaLauncher._prefer_ordering_token def _next_instance_key(self): """Returns the next instance key.""" if self._use_unique_instance_key(): # Assigning instance keys at function building time have issues since # different workers may retrace the function at different times. With # collective V2 we can use capture_call_time_value to use a placeholder as # the instance key and feed it at function call time. In this way we also # don't reuse instance keys, which allows for per-instance cancellation. graph = ops.get_default_graph() # Control flow ops don't work with capture_call_time_value, so we put the # capture in the function graph of that control flow op. while getattr(graph, 'is_control_flow_graph', False): graph = graph.outer_graph if not context.executing_eagerly() and graph.building_function: with graph.as_default(): # Capture self._next_instance_key so that when building a function # that calls another tf.function, the instance key assignment is # further delayed until we actually call the function in eager. Note # that capture_call_time_value doesn't automatically propagate the # deferred capture to the outer function. return graph.capture_call_time_value( self._next_instance_key, tensor_spec.TensorSpec([], dtypes.int32)) else: instance_key = self._collective_keys.get_instance_key( self._group_key, self._device) with ops.device('CPU:0'): return ops.convert_to_tensor(instance_key, dtype=dtypes.int32) else: return self._collective_keys.get_instance_key(self._group_key, self._device) def _get_ordering_token(self): if self._use_ordering_token(): return self._ordering_token.handle # pytype: disable=attribute-error def can_order_nccl(self): """Whether this launcher can order NCCL operations.""" return self._use_ordering_token() def all_reduce( self, input_tensor: core.TensorLike, control_input: Optional[Union[core.TensorLike, ops.Operation]] = None, options: Optional[collective_util.Options] = None) -> core.Tensor: """All-reduce a dense tensor. Args: input_tensor: a dense tensor. It must have the same shape on all replicas. control_input: if not None, add control edges between control_input and the all-reduce. options: an optional tf.distribute.experimental.CommunicationOptions. If provided, it overrides the default options. Returns: The reduced tensor. """ instance_key = self._next_instance_key() options = self._options.merge(options) ordering_token = self._get_ordering_token() with ops.device(self._device), \ self._control_input(control_input): return collective_ops.all_reduce_v2( input_tensor, self._group_size, self._group_key, instance_key, communication_hint=options.implementation.value, timeout=options.timeout_seconds, ordering_token=ordering_token) def _all_gather(self, input_tensor: core.TensorLike, options: Optional[collective_util.Options]) -> core.Tensor: """All-gather a dense tensor. Args: input_tensor: a dense tensor. It must have the same shape on all replicas. options: an optional tf.distribute.experimental.CommunicationOptions. If provided, it overrides the default options. Returns: The reduced tensor. """ instance_key = self._next_instance_key() options = self._options.merge(options) ordering_token = self._get_ordering_token() with ops.device(self._device): return collective_ops.all_gather_v2( input_tensor, self._group_size, self._group_key, instance_key, communication_hint=options.implementation.value, timeout=options.timeout_seconds, ordering_token=ordering_token) def batch_all_reduce( self, input_tensor_packs: List[List[core.TensorLike]], options: Optional[collective_util.Options] = None) -> core.Tensor: """Batch all-reduce dense tensors. This takes a list of batches of tensors. Using multiple batches have the benefit that it doesn't need to wait for all inputs to be ready to start the all-reduce. Args: input_tensor_packs: a list of lists of dense tensors. options: an optional tf.distribute.experimental.CommunicationOptions. If provided, it overrides the default options. Returns: A flat list of reduced tensors. """ options = self._options.merge(options) outputs = [] for pack in input_tensor_packs: if context.executing_eagerly(): # We don't batch in eager as it sometimes makes the performance worse # due the concat/split ops. for input_tensor in pack: outputs.append(self.all_reduce(input_tensor, None, options)) else: # TODO(b/169168846): inserts a parallel all_gather to verify packings # are the same on each replica. with ops.device(self._device): flat_tensors = [array_ops.reshape(t, [-1]) for t in pack] shapes = [array_ops.shape(t) for t in pack] if (options.implementation == collective_util.CommunicationImplementation.NCCL and outputs): control_input = outputs[-1] else: control_input = None reduced = self.all_reduce( array_ops.concat(flat_tensors, axis=0), control_input, options) num_elements = [math_ops.reduce_prod(s) for s in shapes] flat_outputs = array_ops.split(reduced, num_elements, axis=0) for shape, flat_output in zip(shapes, flat_outputs): outputs.append(array_ops.reshape(flat_output, shape)) return outputs def all_gather( self, input_tensor: core.TensorLike, axis: core.TensorLike, options: Optional[collective_util.Options] = None) -> core.Tensor: """All-gather a dense tensor. This method must be called inside a tf.function. Args: input_tensor: a dense tensor. It must have the same rank on all replicas, and dimensions other than `axis` need to be the same as well. axis: 0-D int32 Tensor. Dimension along which to gather. Must be in the range [0, rank(value)). options: an optional tf.distribute.experimental.CommunicationOptions. If provided, it overrides the default options. Returns: The gathered Tensor. Raises: RuntimeError: if called in eager mode. """ if context.executing_eagerly(): raise RuntimeError('all_gather is not supported in eager mode.') with ops.device(self._device), \ ops.control_dependencies([array_ops.identity(input_tensor)]): # 1. Transpose # E.g. Given an input_tensor with shape [2,2,5,1] and axis to gather is 3, # we use perm_pre=[3 0 1 2] to reshape it to [1,2,2,5], which # brings the 3rd dim first; afterwards we use perm_after=[1,2,3,0] to # place it back. perm_pre = array_ops.concat( ([axis], math_ops.range(axis), math_ops.range(axis + 1, array_ops.rank(input_tensor))), axis=0) input_tensor_t = array_ops.transpose(input_tensor, perm=perm_pre) # 2. Pad gathered_shape = self._all_gather( array_ops.expand_dims_v2(array_ops.shape_v2(input_tensor_t), axis=0), options) first_dims = gathered_shape[:, 0] full_axis_dim = math_ops.reduce_max(first_dims) padded_input_tensor = _pad_util(input_tensor_t, full_axis_dim) # 3. Gather gather_padded_out_tensor = self._all_gather(padded_input_tensor, options) # 4. Unpad split_tensors = [] for i in range(self._group_size): start_pos = i * full_axis_dim split_tensors.append(gather_padded_out_tensor[start_pos:start_pos + first_dims[i]]) out_tensor_t = array_ops.concat(split_tensors, 0) # 5. Transpose back perm_after = array_ops.concat( (math_ops.range(1, axis + 1), [0], math_ops.range(axis + 1, array_ops.rank(input_tensor_t))), axis=0) return array_ops.transpose(out_tensor_t, perm=perm_after) def all_reduce_indexed_slices( self, input_slices: indexed_slices.IndexedSlices, options: Optional[collective_util.Options] = None ) -> indexed_slices.IndexedSlices: """All-reduce an IndexedSlices. This method can be called outside tf.function. Args: input_slices: an IndexedSlices. options: an optional tf.distribute.experimental.CommunicationOptions. If provided, it overrides the default options. Returns: The reduced IndexedSlices. """ # Current CollectiveAllGather implementations require input IndexedSlices to # have consistent length across the board, we handle the reduction of # IndexedSlices as follows: # 1. Gather the lengths of IndexedSlices from all participants. # 2. If they have consistent length, apply all_gather. # 3. Otherwise pad IndexedSlices to be the same length across all # participants and apply_gather. options = self._options.merge(options) with ops.device(self._device): def all_gather_indexed_slices( all_gather_fn: Callable[ [core.TensorLike, Optional[collective_util.Options]], core.Tensor] ) -> indexed_slices.IndexedSlices: """Use all_gather_fn to aggregate `IndexedSlices`.""" all_values = all_gather_fn(input_slices.values, options) # Add control dependency to order the all-gather. if (options.implementation == collective_util.CommunicationImplementation.NCCL): control = [all_values] else: control = [] with ops.control_dependencies(control): all_indices = all_gather_fn(input_slices.indices, options) return indexed_slices.IndexedSlices( values=all_values, indices=all_indices, dense_shape=input_slices.dense_shape) length = array_ops.shape(input_slices.indices) all_lengths = self._all_gather(length, options) def all_gather_with_padding( input_tensor: core.TensorLike, options: Optional[collective_util.Options]) -> core.Tensor: """all_gather tensors of different sizes using padding.""" max_length = math_ops.reduce_max(all_lengths) padded_tensor = _pad_util(input_tensor, max_length) all_padded_tensors = self._all_gather(padded_tensor, options) split_tensors = [] for i in range(self._group_size): start_pos = i * max_length split_tensors.append(all_padded_tensors[start_pos:start_pos + all_lengths[i]]) return array_ops.concat(split_tensors, 0) return cond.cond( math_ops.equal( math_ops.reduce_max(all_lengths), math_ops.reduce_min(all_lengths)), lambda: all_gather_indexed_slices(self._all_gather), lambda: all_gather_indexed_slices(all_gather_with_padding)) def aggregate_tensors_or_indexed_slices(values, accumulation_fn=math_ops.add_n): """Aggregate tensors using `accumulation_fn` and IndexedSlices via concat.""" if any(isinstance(v, indexed_slices.IndexedSlices) for v in values): return backprop_util.AggregateIndexedSlicesGradients(values) else: return accumulation_fn(values) def divide_by_n_tensors_or_indexed_slices(value, n): if isinstance(value, indexed_slices.IndexedSlices): value = backprop_util.FlattenNestedIndexedSlices(value) return indexed_slices.IndexedSlices(value.values / n, value.indices, value.dense_shape) else: return value / n def copy_tensor_or_indexed_slices_to_device(value, device): """Copies a tensor or IndexedSlices to a device.""" with ops.device(device): if isinstance(value, indexed_slices.IndexedSlices): copied_values = array_ops.identity(value.values) copied_indices = array_ops.identity(value.indices) if value.dense_shape is not None: copied_shape = array_ops.identity(value.dense_shape) else: copied_shape = None result = indexed_slices.IndexedSlices(copied_values, copied_indices, copied_shape) else: result = array_ops.identity(value) return result def is_indexed_slices(value): if isinstance(value, indexed_slices.IndexedSlices): return True if isinstance(value, value_lib.DistributedValues): return all( isinstance(v, indexed_slices.IndexedSlices) for v in value.values) return False def split_by_sparsity(values): """Split values into dense and sparse values. Args: values: a list of tensors or `PerReplica`s. Returns: Four lists: a list of dense values, a list of their indices in `values` and a list of sparse values, a list of their indices in `values`. """ dense_values = [] dense_indices = [] sparse_values = [] sparse_indices = [] for i, v in enumerate(values): if is_indexed_slices(v): sparse_values.append(v) sparse_indices.append(i) else: dense_values.append(v) dense_indices.append(i) return dense_values, dense_indices, sparse_values, sparse_indices def stitch_values(values_and_indices_list): """Stitch values together according to their indices. Args: values_and_indices_list: a list of tuples of values and indices indicating the values and positions in the returned list. Returns: a stitched list of values. """ length = 0 for values_and_indices in values_and_indices_list: length += len(values_and_indices[0]) result = [None] * length for values_and_indices in values_and_indices_list: if values_and_indices and values_and_indices[0]: for v, i in zip(*values_and_indices): assert result[i] is None result[i] = v return result def group_by_size(input_tensors, bytes_per_pack): """Groups `input_tensors` into chunks of `bytes_per_pack`. The method preserves the original order of `input_tensors`. The grouping is best effort, each pack could have more or less bytes than `bytes_per_pack`. It only groups values with known shape. Args: input_tensors: a list of Tensor. bytes_per_pack: an integer. Returns: A list of packs of Tensor. All values are grouped into one pack if `bytes_per_pack` is zero or any of the value has unknown shape. """ if bytes_per_pack == 0: return [input_tensors] packs = [] last_pack_size = 0 for value in input_tensors: num_elements = value.shape.num_elements() if num_elements is None: # Can't pack values with unknown shape. logging.warning( 'not packing values due to the unknown or inconsistent shape of %s', value) return [input_tensors] size = num_elements * value.dtype.size # Try to keep each pack as close to bytes_per_pack as possible, while each # pack is at least bytes_per_pack large. I.E. we err on the side of having # few but large packs. if not packs or last_pack_size > bytes_per_pack: packs.append([]) last_pack_size = 0 packs[-1].append(value) last_pack_size += size return packs def _pad_util(input_tensor, full_axis_dim): """Pad the `input_tensor`'s first dimension to be `full_axis_dim`.""" missing_axis_dim = full_axis_dim - array_ops.shape_v2(input_tensor)[0] tensor_rank = array_ops.rank(input_tensor) paddings_axis = [[0, missing_axis_dim]] paddings = array_ops.concat([ paddings_axis, array_ops.zeros(shape=(tensor_rank - 1, 2), dtype=dtypes.int32) ], axis=0) padded_input_tensor = array_ops.pad(input_tensor, paddings) return padded_input_tensor