wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/distribute/v1/BUILD
+98
View File
@@ -0,0 +1,98 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
cuda_py_strict_test(
name = "cross_device_ops_test",
srcs = ["cross_device_ops_test.py"],
tags = [
# "multi_and_single_gpu", # TODO(b/287692888): re-enable once the 2gpu test passes.
"no_windows_gpu", # b/216367668
"no_windows", # b/326464742
],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/distribute:collective_all_reduce_strategy",
"//tensorflow/python/distribute:collective_util",
"//tensorflow/python/distribute:combinations",
"//tensorflow/python/distribute:cross_device_ops",
"//tensorflow/python/distribute:cross_device_utils",
"//tensorflow/python/distribute:device_util",
"//tensorflow/python/distribute:distribute_utils",
"//tensorflow/python/distribute:multi_worker_test_base",
"//tensorflow/python/distribute:multi_worker_util",
"//tensorflow/python/distribute:reduce_util",
"//tensorflow/python/distribute:strategy_combinations",
"//tensorflow/python/distribute:values",
"//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:kernels",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:collective_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "all_reduce",
srcs = [
"all_reduce.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nccl_ops",
],
)
tf_py_strict_test(
name = "all_reduce_test",
srcs = ["all_reduce_test.py"],
deps = [
":all_reduce",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:tf_logging",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "input_lib",
srcs = ["input_lib.py"],
deps = [
"//tensorflow/python/data/experimental/ops:cardinality",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:multi_device_iterator_ops",
"//tensorflow/python/data/ops:optional_ops",
"//tensorflow/python/distribute:input_lib",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/types:data",
"//tensorflow/python/util:deprecation",
],
)
tensorflow/python/distribute/v1/all_reduce.py
+862
View File
@@ -0,0 +1,862 @@
# Copyright 2017 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 to construct a TF subgraph implementing distributed All-Reduce."""
import collections
import math
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nccl_ops
def _flatten_tensors(tensors):
"""Check tensors for isomorphism and flatten.
Args:
tensors: list of `tf.Tensor` which must all have the same shape.
Returns:
tensors: a list of `tf.Tensor` which are flattened (1D) views of tensors
shape: the original shape of each element of input tensors
Raises:
ValueError: tensors are empty or non-isomorphic or have unknown shape.
"""
if not tensors:
raise ValueError("tensors cannot be empty")
shape = tensors[0].shape
for tensor in tensors:
shape = shape.merge_with(tensor.shape)
if not shape.is_fully_defined():
raise ValueError("Tensors must have statically known shape.")
if len(shape) != 1:
reshaped = []
for t in tensors:
with ops.colocate_with(t):
reshaped.append(array_ops.reshape(t, [-1]))
tensors = reshaped
return tensors, shape
def _reshape_tensors(tensors, shape):
"""Reshape tensors flattened by _flatten_tensors.
Args:
tensors: list of `tf.Tensor` of identical length 1D tensors.
shape: list of integers describing the desired shape. Product of
the elements must equal the length of each tensor.
Returns:
list of `tf.Tensor` which are the reshaped inputs.
"""
reshaped = []
for t in tensors:
with ops.colocate_with(t):
reshaped.append(array_ops.reshape(t, shape))
return reshaped
def _padded_split(tensor, pieces):
"""Like split for 1D tensors but pads-out case where len % pieces != 0.
Args:
tensor: `tf.Tensor` that must be 1D.
pieces: a positive integer specifying the number of pieces into which
tensor should be split.
Returns:
list of `tf.Tensor` of length pieces, which hold the values of
thin input tensor, in order. The final tensor may
be zero-padded on the end to make its size equal to those of all
of the other tensors.
Raises:
ValueError: The input tensor is not 1D.
"""
shape = tensor.shape
if 1 != len(shape):
raise ValueError("input tensor must be 1D")
tensor_len = shape.dims[0].value
with ops.colocate_with(tensor):
if tensor_len % pieces != 0:
# pad to an even length
chunk_size = 1 + tensor_len // pieces
if pieces > tensor_len:
# This is an edge case that should not come up in practice,
# i.e. a different reduction algorithm would be better,
# but we'll make it work just for completeness.
pad_len = pieces - tensor_len
extended_whole = array_ops.concat(
[tensor, array_ops.zeros([pad_len], dtype=tensor.dtype)], 0)
parts = array_ops.split(extended_whole, pieces)
return parts, pad_len
elif (pieces - 1) * chunk_size >= tensor_len:
# Another edge case of limited real interest.
pad_len = (pieces * chunk_size) % tensor_len
extended_whole = array_ops.concat(
[tensor, array_ops.zeros([pad_len], dtype=tensor.dtype)], 0)
parts = array_ops.split(extended_whole, pieces)
return parts, pad_len
else:
last_chunk_size = tensor_len - (pieces - 1) * chunk_size
pad_len = chunk_size - last_chunk_size
piece_lens = [chunk_size for _ in range(pieces - 1)] + [last_chunk_size]
parts = array_ops.split(tensor, piece_lens)
parts[-1] = array_ops.concat(
[parts[-1], array_ops.zeros([pad_len], dtype=tensor.dtype)], 0)
return parts, pad_len
else:
return array_ops.split(tensor, pieces), 0
def _strip_padding(tensors, pad_len):
"""Strip the suffix padding added by _padded_split.
Args:
tensors: list of `tf.Tensor` of identical length 1D tensors.
pad_len: number of elements to be stripped from the end of each tensor.
Returns:
list of `tf.Tensor` which are the stripped inputs.
Raises:
ValueError: tensors must be a non-empty list of 1D tensors, and
each must be longer than pad_len.
"""
if not tensors:
raise ValueError("tensors cannot be empty")
shape = tensors[0].shape
if len(shape) > 1:
raise ValueError("tensors must be 1D")
prefix_len = int(shape[0] - pad_len)
if prefix_len < 0:
raise ValueError("pad_len longer than tensor")
stripped = []
for t in tensors:
with ops.colocate_with(t):
stripped.append(array_ops.slice(t, [0], [prefix_len]))
return stripped
def _ragged_split(tensor, pieces):
"""Like split for 1D tensors but allows case where len % pieces != 0.
Args:
tensor: `tf.Tensor` that must be 1D.
pieces: a positive integer specifying the number of pieces into which
tensor should be split.
Returns:
list of `tf.Tensor` of length pieces, which hold the values of
the input tensor, in order. The final tensor may be shorter
than the others, which will all be of equal length.
Raises:
ValueError: input tensor must be 1D.
"""
shape = tensor.shape
if 1 != len(shape):
raise ValueError("input tensor must be 1D")
tensor_len = shape.dims[0].value
chunk_size = tensor_len // pieces
with ops.colocate_with(tensor):
if tensor_len != (pieces * chunk_size):
# last piece will be short
assert pieces > 1
last_chunk_size = tensor_len - ((pieces - 1) * chunk_size)
assert last_chunk_size > 0
piece_lens = [chunk_size for _ in range(pieces - 1)] + [last_chunk_size]
return array_ops.split(tensor, piece_lens)
else:
return array_ops.split(tensor, pieces)
def _ring_permutations(num_workers, num_subchunks, gpu_perm):
""""Generate an array of device index arrays, one for each subchunk.
In the basic ring reduction algorithm there are size(T)/num_devices
data chunks and each device process one chunk per tick, i.e. sending
one chunk and receiving one chunk. The idea of subchunking is that
each device processes num_subchunks smaller data regions per tick,
and the ring rank permutation is different for each subchunk index
so that a device is potentially sending to and receiving from
num_subchunks different other devices at each tick. Where multiple
independent data channels exist between devices, this strategy
supplies a method of using them in parallel.
Args:
num_workers: number of worker tasks
num_subchunks: number of subchunks into which to divide each per-GPU chunk.
gpu_perm: an array of integers in [0, num_gpus-1] giving the default
ring order of GPUs at each worker. Other permutations will be generated
by rotating this array and splicing together per-worker instances.
Raises:
ValueError: the number of subchunks may not exceed the number of GPUs.
Returns:
pred_by_s_d: list of lists that maps (by index) from (subchunk, dev) to
preceding device in the permutation for that subchunk. The
device index of GPU i at worker j is i + (j * num_gpus).
rank_by_s_d: list of lists that maps (by index) from (subchunk, dev) to
local rank of device d in the permutation for that subchunk.
"""
num_gpus = len(gpu_perm)
devices = num_workers * num_gpus
if devices == 0:
return [], []
if num_subchunks > num_gpus:
raise ValueError(
"num_subchunks %d must be <= num_gpus %d" % (num_subchunks, num_gpus))
rotation_interval = max(1, int(num_gpus / num_subchunks))
perms_by_s = []
for s in range(0, num_subchunks):
full_order = []
offset = s * rotation_interval
for w in range(0, num_workers):
default_order = [(w * num_gpus) + i for i in gpu_perm]
dev_order = default_order[offset:] + default_order[:offset]
full_order += dev_order
perms_by_s.append(full_order)
pred_by_s_d = [[-1 for d in range(0, devices)]
for s in range(0, num_subchunks)]
rank_by_s_d = [[-1 for d in range(0, devices)]
for s in range(0, num_subchunks)]
for s in range(0, num_subchunks):
for d in range(0, devices):
for t in range(0, devices):
if d == perms_by_s[s][t]:
rank_by_s_d[s][d] = t
pred_by_s_d[s][d] = perms_by_s[s][(t + devices - 1) % devices]
break
return (pred_by_s_d, rank_by_s_d)
def build_ring_all_reduce(input_tensors, num_workers, num_subchunks,
gpu_perm, red_op, un_op=None):
"""Construct a subgraph performing a ring-style all-reduce of input_tensors.
Args:
input_tensors: a list of `tf.Tensor` objects, which must all
have the same shape and type.
num_workers: number of worker tasks spanned by input_tensors.
num_subchunks: number of subchunks each device should process in one tick.
gpu_perm: a list of ints giving a ring-wise rank ordering of GPUs at
each worker. All workers must have the same number of
GPUs with the same rank ordering. If NVLINK is available, this should
be a ring order supported by NVLINK edges.
red_op: a binary operator for elementwise reduction.
un_op: an optional unary operator to apply to fully reduced values.
Raises:
ValueError: empty input_tensors or they don't all have same
size.
Returns:
a list of `tf.Tensor` identical sum-reductions of input_tensors.
"""
if len(input_tensors) < 2:
raise ValueError("input_tensors must be length 2 or longer")
input_tensors, shape = _flatten_tensors(input_tensors)
devices = [t.device for t in input_tensors]
(pred_by_s_d, rank_by_s_d) = _ring_permutations(
num_workers, num_subchunks, gpu_perm)
chunks_by_dev, pad_len = _build_ring_gather(
input_tensors, devices,
num_subchunks, pred_by_s_d, rank_by_s_d, red_op)
if un_op:
chunks_by_dev = _apply_unary_to_chunks(un_op, chunks_by_dev)
output_tensors = _build_ring_scatter(pred_by_s_d, rank_by_s_d,
chunks_by_dev)
if pad_len > 0:
output_tensors = _strip_padding(output_tensors, pad_len)
if len(shape) != 1:
output_tensors = _reshape_tensors(output_tensors, shape)
return output_tensors
def _build_ring_gather(input_tensors, devices, num_subchunks,
pred_by_s_d, rank_by_s_d, red_op):
"""Construct a subgraph for the first (reduction) pass of ring all-reduce.
Args:
input_tensors: a list of `tf.Tensor` 1D input tensors of same
shape and type.
devices: array of device name strings
num_subchunks: number of subchunks each device should process in one tick.
pred_by_s_d: as produced by _ring_permutations
rank_by_s_d: as produced by _ring_permutations
red_op: a binary operator for elementwise reduction
Raises:
ValueError: tensors must all be one dimensional.
Returns:
list of list of `tf.Tensor` of (partially) reduced values where
exactly num_subchunks chunks at each device are fully reduced.
"""
num_devices = len(input_tensors)
if num_devices == 0:
return []
if num_devices == 1:
return input_tensors
shape = input_tensors[0].shape
if 1 != len(shape):
raise ValueError("input tensors must be 1D")
num_chunks = num_devices * num_subchunks
num_ticks = num_devices - 1
# Initialize chunks_by_dev with splits of the input tensors.
chunks_by_dev = []
split_pad_len = 0
for d in range(0, num_devices):
with ops.device(devices[d]):
splits, split_pad_len = _padded_split(input_tensors[d], num_chunks)
chunks_by_dev.append(splits)
# Reduction phase
for tick in range(0, num_ticks):
# One new partial reduction for every chunk
new_partial_reductions = [None for _ in range(0, num_chunks)]
# Compute reductions with respect to last tick's values
for d in range(0, num_devices):
with ops.device(devices[d]):
for s in range(0, num_subchunks):
rank = rank_by_s_d[s][d]
seg_index = (rank + num_devices - (2 + tick)) % num_devices
pred_dev = pred_by_s_d[s][d]
chunk_index = (seg_index * num_subchunks) + s
new_partial_reductions[chunk_index] = red_op(
chunks_by_dev[pred_dev][chunk_index],
chunks_by_dev[d][chunk_index])
# Update chunks_by_dev with the new values at the end of the tick.
for d in range(0, num_devices):
for s in range(0, num_subchunks):
rank = rank_by_s_d[s][d]
seg_index = (rank + num_devices - (2 + tick)) % num_devices
chunk_index = (seg_index * num_subchunks) + s
chunks_by_dev[d][chunk_index] = new_partial_reductions[chunk_index]
return chunks_by_dev, split_pad_len
def _apply_unary_to_chunks(f, chunks_by_dev):
"""Apply a unary op to each tensor in chunks_by_dev, on same device.
Args:
f: a unary function over `tf.Tensor`.
chunks_by_dev: list of lists of `tf.Tensor`.
Returns:
new list of lists of `tf.Tensor` with the same structure as
chunks_by_dev containing the derived tensors.
"""
output = []
for x in chunks_by_dev:
with ops.colocate_with(x[0]):
output.append([f(t) for t in x])
return output
def _build_ring_scatter(pred_by_s_d, rank_by_s_d,
chunks_by_dev):
"""Construct subgraph for second (scatter) pass of ring all-reduce.
Args:
pred_by_s_d: as produced by _ring_permutations
rank_by_s_d: as produced by _ring_permutations
chunks_by_dev: list of list of `tf.Tensor` indexed by ints
(device, chunk)
Raises:
ValueError: chunks_by_dev is not well-formed
Returns:
list of `tf.Tensor` which are the fully reduced tensors, one
at each device corresponding to the outer dimension of chunks_by_dev.
"""
num_devices = len(chunks_by_dev)
num_chunks = len(chunks_by_dev[0])
if 0 != num_chunks % num_devices:
raise ValueError(
"Expect number of chunks per device to be divisible by num_devices")
num_subchunks = int(num_chunks / num_devices)
num_ticks = num_devices - 1
for tick in range(0, num_ticks):
passed_values = [None for _ in range(0, num_chunks)]
for d in range(0, num_devices):
with ops.colocate_with(chunks_by_dev[d][0]):
for s in range(0, num_subchunks):
rank = rank_by_s_d[s][d]
seg_index = (rank + num_devices - (1 + tick)) % num_devices
pred_dev = pred_by_s_d[s][d]
chunk_index = (seg_index * num_subchunks) + s
passed_values[chunk_index] = array_ops.identity(
chunks_by_dev[pred_dev][chunk_index])
for d in range(0, num_devices):
for s in range(0, num_subchunks):
rank = rank_by_s_d[s][d]
seg_index = (rank + num_devices - (1 + tick)) % num_devices
chunk_index = (seg_index * num_subchunks) + s
chunks_by_dev[d][chunk_index] = passed_values[chunk_index]
# Join chunks at each device.
output = []
for x in chunks_by_dev:
with ops.colocate_with(x[0]):
output.append(array_ops.concat(x, 0))
return output
def build_recursive_hd_all_reduce(input_tensors, red_op, un_op=None):
"""Construct a subgraph for recursive halving-doubling all-reduce.
The recursive halving-doubling algorithm is described in
(Thakur et al., 2015).
The concept is to arrange the participating n devices in
a linear sequence where devices exchange data pairwise
with one other device in each round. During the gather
phase there are lg(n) rounds where devices exchange
increasingly smaller sub-tensors with another device
at increasingly greater distances, until at the top
each device has 1/n of the fully reduced values. During the
scatter phase each device exchanges its fully reduced
sub-tensor (which doubles in length at each round)
with one other device at increasingly smaller distances
until each device has all of the fully reduced values.
Note: this preliminary version requires that len(input_tensors) be a
power of 2. TODO(tucker): relax this restriction. Also, the
number of elements in each tensor must be divisible by 2^h where h
is the number of hops in each phase. This will also be relaxed in
the future with edge-case specific logic.
Args:
input_tensors: list of `tf.Tensor` to be elementwise reduced.
red_op: a binary elementwise reduction Op.
un_op: an optional unary elementwise Op to apply to reduced values.
Returns:
list of `tf.Tensor` which are the fully reduced tensors, one
at each device of input_tensors.
Raises:
ValueError: num_devices not a power of 2, or tensor len not divisible
by 2 the proper number of times.
References:
Optimization of Collective Communication Operations in MPICH:
[Thakur et al., 2005]
(https://journals.sagepub.com/doi/abs/10.1177/1094342005051521)
([pdf](http://wwwi10.lrr.in.tum.de/~gerndt/home/Teaching/HPCSeminar/mpich_multi_coll.pdf))
"""
devices = [t.device for t in input_tensors]
input_tensors, shape = _flatten_tensors(input_tensors)
reduced_shards = _build_recursive_hd_gather(input_tensors, devices, red_op)
if un_op:
reduced_shards = [un_op(t) for t in reduced_shards]
output_tensors = _build_recursive_hd_scatter(reduced_shards, devices)
if len(shape) != 1:
output_tensors = _reshape_tensors(output_tensors, shape)
return output_tensors
def _build_recursive_hd_gather(input_tensors, devices, red_op):
"""Construct the gather phase of recursive halving-doubling all-reduce.
Args:
input_tensors: list of `tf.Tensor` to be elementwise reduced.
devices: a list of strings naming the devices hosting input_tensors,
which will also be used to host the (partial) reduction values.
red_op: a binary elementwise reduction Op.
Returns:
list of `tf.Tensor` which are the fully reduced tensor shards.
Raises:
ValueError: num_devices not a power of 2, or tensor len not divisible
by 2 the proper number of times.
"""
num_devices = len(devices)
num_hops = int(math.log(num_devices, 2))
if num_devices != (2 ** num_hops):
raise ValueError("num_devices must be a power of 2")
chunks = input_tensors
for h in range(0, num_hops):
span = 2 ** h
group_size = span * 2
new_chunks = [[] for _ in devices]
for d in range(0, num_devices):
if (d % group_size) >= (group_size / 2):
# skip right half of a pair
continue
left_dev = devices[d]
right_dev = devices[d + span]
left_split = array_ops.split(chunks[d], 2)
right_split = array_ops.split(chunks[d+span], 2)
with ops.device(left_dev):
new_chunks[d] = red_op(left_split[0], right_split[0])
with ops.device(right_dev):
new_chunks[d + span] = red_op(left_split[1], right_split[1])
chunks = new_chunks
return chunks
def _build_recursive_hd_scatter(input_tensors, devices):
"""Construct the scatter phase of recursive halving-doubling all-reduce.
Args:
input_tensors: list of `tf.Tensor` that are fully-reduced shards.
devices: a list of strings naming the devices on which the reconstituted
full tensors should be placed.
Returns:
list of `tf.Tensor` which are the fully reduced tensors.
"""
num_devices = len(devices)
num_hops = int(math.log(num_devices, 2))
assert num_devices == (2 ** num_hops), "num_devices must be a power of 2"
chunks = input_tensors
for h in reversed(range(0, num_hops)):
span = 2 ** h
group_size = span * 2
new_chunks = [[] for _ in devices]
for d in range(0, num_devices):
if (d % group_size) >= (group_size / 2):
# skip right half of a pair
continue
left_idx = d
right_idx = d + span
left_dev = devices[left_idx]
right_dev = devices[right_idx]
with ops.device(left_dev):
new_chunks[left_idx] = array_ops.concat([chunks[left_idx],
chunks[right_idx]], 0)
with ops.device(right_dev):
new_chunks[right_idx] = array_ops.concat([chunks[left_idx],
chunks[right_idx]], 0)
chunks = new_chunks
return chunks
def build_shuffle_all_reduce(input_tensors, gather_devices, red_op, un_op=None):
"""Construct a subgraph for shuffle all-reduce.
Shuffle reduce is essentially the algorithm implemented when using
parameter servers. Suppose tensor length is n, there are d devices
and g gather shards. Each device sends a n/g length sub-tensor to
each gather shard. The gather shards perform a reduction across d
fragments, then broadcast the result back to each device. The
devices then join the g fully reduced fragments they receive from
the shards. The gather shards could perform d-1 pairwise
reductions, or one d-way reduction. The first is better where
reduction Op time is low compared to transmission time, the second
better in the other case.
Args:
input_tensors: list of `tf.Tensor` values to be reduced.
gather_devices: list of names of devices on which reduction shards
should be placed.
red_op: an n-array elementwise reduction Op
un_op: optional elementwise unary Op to be applied to fully-reduced values.
Returns:
list of `tf.Tensor` which are the fully reduced tensors.
"""
input_tensors, shape = _flatten_tensors(input_tensors)
dst_devices = [t.device for t in input_tensors]
reduced_shards = _build_shuffle_gather(input_tensors, gather_devices,
red_op, un_op)
output_tensors = _build_shuffle_scatter(reduced_shards, dst_devices)
if len(shape) != 1:
output_tensors = _reshape_tensors(output_tensors, shape)
return output_tensors
def _build_shuffle_gather(input_tensors, gather_devices, red_op, un_op=None):
"""Construct the gather (concentrate and reduce) phase of shuffle all-reduce.
Args:
input_tensors: list of `tf.Tensor` values to be reduced.
gather_devices: list of names of devices on which reduction shards
should be placed.
red_op: the binary reduction Op
un_op: optional elementwise unary Op to be applied to fully-reduced values.
Returns:
list of `tf.Tensor` which are the fully reduced shards.
Raises:
ValueError: inputs not well-formed.
"""
num_source_devices = len(input_tensors)
num_gather_devices = len(gather_devices)
shape = input_tensors[0].shape
if len(shape) != 1:
raise ValueError("input_tensors must be 1D")
shards_by_source = []
for d in range(0, num_source_devices):
with ops.colocate_with(input_tensors[d]):
shards_by_source.append(
_ragged_split(input_tensors[d], num_gather_devices))
reduced_shards = []
for d in range(0, num_gather_devices):
with ops.device(gather_devices[d]):
values = [s[d] for s in shards_by_source]
red_shard = red_op(values)
if un_op:
red_shard = un_op(red_shard)
reduced_shards.append(red_shard)
return reduced_shards
def _build_shuffle_scatter(reduced_shards, dst_devices):
"""Build the scatter phase of shuffle all-reduce.
Args:
reduced_shards: list of `tf.Tensor` fully reduced shards
dst_devices: list of names of devices at which the fully-reduced value
should be reconstituted.
Returns:
list of `tf.Tensor` scattered tensors.
"""
num_devices = len(dst_devices)
out_tensors = []
for d in range(0, num_devices):
with ops.device(dst_devices[d]):
out_tensors.append(array_ops.concat(reduced_shards, 0))
return out_tensors
def _split_by_task(devices, values):
"""Partition devices and values by common task.
Args:
devices: list of device name strings
values: list of `tf.Tensor` of same length as devices.
Returns:
(per_task_devices, per_task_values) where both values are
lists of lists with isomorphic structure: the outer list is
indexed by task, and the inner list has length of the number
of values belonging to that task. per_task_devices contains
the specific devices to which the values are local, and
per_task_values contains the corresponding values.
Raises:
ValueError: devices must be same length as values.
"""
num_devices = len(devices)
if num_devices != len(values):
raise ValueError("len(devices) must equal len(values)")
per_task_devices = collections.OrderedDict()
per_task_values = collections.OrderedDict()
for d in range(num_devices):
d_spec = device_lib.DeviceSpec.from_string(devices[d])
if not hasattr(d_spec, "task") or d_spec.task is None:
assert False, "failed to parse device %s" % devices[d]
index = (d_spec.job or "localhost", d_spec.replica or 0, d_spec.task)
if index not in per_task_devices:
per_task_devices[index] = []
per_task_values[index] = []
per_task_devices[index].append(devices[d])
per_task_values[index].append(values[d])
return (list(per_task_devices.values()), list(per_task_values.values()))
def build_nccl_all_reduce(input_tensors, red_op, un_op=None):
"""Build a subgraph that does one full all-reduce, using NCCL.
Args:
input_tensors: list of `tf.Tensor` of same-shape and type values to
be reduced.
red_op: binary elementwise reduction operator. Must be one of
{tf.add}
un_op: optional unary elementwise Op to apply to fully-reduce values.
Returns:
list of `tf.Tensor` of reduced values.
Raises:
ValueError: red_op not supported.
"""
if red_op == math_ops.add:
output_tensors = nccl_ops.all_sum(input_tensors)
else:
raise ValueError("red_op not supported by NCCL all-reduce: ", red_op)
if un_op:
un_op_wrapped = []
for t in output_tensors:
with ops.colocate_with(t):
un_op_wrapped.append(un_op(t))
output_tensors = un_op_wrapped
return output_tensors
def _build_nccl_hybrid(input_tensors, red_op, upper_level_f):
"""Construct a subgraph for NCCL hybrid all-reduce.
Args:
input_tensors: list of `tf.Tensor` of same-shape and type values to
be reduced.
red_op: binary elementwise reduction operator.
upper_level_f: function for reducing one value per worker, across
workers.
Returns:
list of `tf.Tensor` of reduced values.
Raises:
ValueError: inputs not well-formed.
"""
input_tensors, shape = _flatten_tensors(input_tensors)
devices = [t.device for t in input_tensors]
per_worker_devices, per_worker_values = _split_by_task(devices, input_tensors)
num_workers = len(per_worker_devices)
up_values = [None for w in range(0, num_workers)]
up_devices = up_values[:]
down_values = up_values[:]
# First stage: reduce within each worker using NCCL
for w in range(0, num_workers):
worker_values = build_nccl_all_reduce(per_worker_values[w], red_op)
# NOTE: these reductions will not run to completion unless
# every output value is used. Since we only need one, we
# need to put control dependencies on the rest.
with ops.control_dependencies(worker_values):
with ops.device(worker_values[0].device):
up_values[w] = array_ops.identity(worker_values[0])
up_devices[w] = per_worker_devices[w][0]
# Second stage: Apply upper_level_f to reduce across first device at
# each worker
level_2_output = upper_level_f(up_values)
# Third stage: propagate within each worker using NCCL Broadcast
for w in range(0, num_workers):
dst_tensors = []
with ops.device(per_worker_devices[w][0]):
broadcast_src = nccl_ops.broadcast(array_ops.identity(level_2_output[w]))
for d in per_worker_devices[w]:
with ops.device(d):
dst_tensors.append(array_ops.identity(broadcast_src))
down_values[w] = dst_tensors
output_tensors = [v for sublist in down_values for v in sublist]
if len(shape) != 1:
output_tensors = _reshape_tensors(output_tensors, shape)
return output_tensors
def _reduce_non_singleton(input_tensors, red_f, un_op):
"""If len(input_tensors) > 1, apply red_f, else apply un_op."""
if len(input_tensors) > 1:
return red_f(input_tensors)
else:
if not un_op:
return input_tensors
output_tensors = []
for t in input_tensors:
with ops.colocate_with(t):
output_tensors.append(un_op(t))
return output_tensors
def build_nccl_then_ring(input_tensors, subdiv, red_op, un_op=None):
"""Construct hybrid of NCCL within workers, Ring across workers."""
def upper_builder(y):
return build_ring_all_reduce(y, len(y), subdiv, [0], red_op, un_op)
def upper_level_f(x):
return _reduce_non_singleton(x, upper_builder, un_op)
return _build_nccl_hybrid(input_tensors, red_op, upper_level_f)
def build_nccl_then_recursive_hd(input_tensors, red_op, un_op=None):
"""Construct hybrid of NCCL within workers, Recursive-HD across workers."""
upper_level_f = lambda x: build_recursive_hd_all_reduce(x, red_op, un_op)
return _build_nccl_hybrid(input_tensors, red_op, upper_level_f)
def build_nccl_then_shuffle(input_tensors, gather_devices, nccl_red_op,
shuffle_red_op, un_op=None):
"""Construct hybrid of NCCL within workers, Shuffle across workers."""
def upper_level_f(x):
return build_shuffle_all_reduce(x, gather_devices, shuffle_red_op, un_op)
return _build_nccl_hybrid(input_tensors, nccl_red_op, upper_level_f)
def _build_shuffle_hybrid(input_tensors, gather_devices, red_op, upper_level_f):
"""Construct a subgraph for Shuffle hybrid all-reduce.
Args:
input_tensors: list of `tf.Tensor` of same-shape and type values to
be reduced.
gather_devices: list of device names on which to host gather shards.
red_op: binary elementwise reduction operator.
upper_level_f: function for reducing one value per worker, across
workers.
Returns:
list of `tf.Tensor` of reduced values.
Raises:
ValueError: inputs not well-formed.
"""
input_tensors, shape = _flatten_tensors(input_tensors)
# First stage, reduce across each worker using gather_devices.
devices = [t.device for t in input_tensors]
per_worker_devices, per_worker_values = _split_by_task(devices, input_tensors)
num_workers = len(per_worker_devices)
up_values = []
if len(gather_devices) != num_workers:
raise ValueError("For shuffle hybrid, gather_devices must contain one "
"device per worker. ")
for w in range(0, num_workers):
reduced_shards = _build_shuffle_gather(
per_worker_values[w], [gather_devices[w]], red_op)
up_values.append(reduced_shards[0])
# Second stage, apply upper_level_f.
level_2_output = upper_level_f(up_values)
# Third stage, apply shuffle scatter at each worker.
output_tensors = []
for w in range(0, num_workers):
output_tensors += _build_shuffle_scatter(
[level_2_output[w]], per_worker_devices[w])
if len(shape) != 1:
output_tensors = _reshape_tensors(output_tensors, shape)
return output_tensors
def build_shuffle_then_ring(input_tensors, gather_devices, subdiv,
red_n_op, red_op, un_op=None):
"""Construct hybrid of Shuffle within workers, Ring across workers."""
def upper_builder(tensors):
return build_ring_all_reduce(tensors, len(tensors), subdiv, [0],
red_op, un_op)
def upper_level_f(tensors):
return _reduce_non_singleton(tensors, upper_builder, un_op)
return _build_shuffle_hybrid(
input_tensors, gather_devices, red_n_op, upper_level_f)
def build_shuffle_then_shuffle(input_tensors, first_gather_devices,
second_gather_devices, red_op, un_op=None):
"""Construct hybrid of Shuffle within workers, Shuffle across workers."""
def upper_builder(tensors):
return build_shuffle_all_reduce(tensors, second_gather_devices,
red_op, un_op)
def upper_level_f(tensors):
return _reduce_non_singleton(tensors, upper_builder, un_op)
return _build_shuffle_hybrid(
input_tensors, first_gather_devices, red_op, upper_level_f)
tensorflow/python/distribute/v1/all_reduce_test.py
+238
View File
@@ -0,0 +1,238 @@
# Copyright 2017 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.
# ==============================================================================
"""Tests for all_reduce."""
import time
import numpy as np
from tensorflow.core.framework import types_pb2
from tensorflow.python.distribute.v1 import all_reduce as ar
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging
class AllReduceTest(test_util.TensorFlowTestCase):
@test_util.run_deprecated_v1
def testFlattenTensorsShapesDefined(self):
x = array_ops.placeholder(types_pb2.DT_FLOAT, [None])
with self.assertRaisesRegex(ValueError, "must have statically known shape"):
ar._flatten_tensors([x, x])
def testRingPermutations(self):
# 0 devices
pred_by_c_d, rank_by_c_d = ar._ring_permutations(1, 0, [])
self.assertEqual(pred_by_c_d, [])
self.assertEqual(rank_by_c_d, [])
# 1 worker, 1 subchunk cases
pred_by_c_d, rank_by_c_d = ar._ring_permutations(1, 1, [0])
self.assertEqual(pred_by_c_d, [[0]])
self.assertEqual(rank_by_c_d, [[0]])
pred_by_c_d, rank_by_c_d = ar._ring_permutations(1, 1, [0, 1, 2])
self.assertEqual(pred_by_c_d, [[2, 0, 1]])
self.assertEqual(rank_by_c_d, [[0, 1, 2]])
# multiple workers, 1 subchunk cases
pred_by_c_d, rank_by_c_d = ar._ring_permutations(2, 1, [0, 1, 2])
self.assertEqual(pred_by_c_d, [[5, 0, 1, 2, 3, 4]])
self.assertEqual(rank_by_c_d, [[0, 1, 2, 3, 4, 5]])
pred_by_c_d, rank_by_c_d = ar._ring_permutations(3, 1, [0, 1, 2])
self.assertEqual(pred_by_c_d, [[8, 0, 1, 2, 3, 4, 5, 6, 7]])
self.assertEqual(rank_by_c_d, [[0, 1, 2, 3, 4, 5, 6, 7, 8]])
pred_by_c_d, rank_by_c_d = ar._ring_permutations(2, 1, [2, 1, 0])
self.assertEqual(pred_by_c_d, [[1, 2, 3, 4, 5, 0]])
self.assertEqual(rank_by_c_d, [[2, 1, 0, 5, 4, 3]])
# 1 worker, multiple subchunk cases
pred_by_c_d, rank_by_c_d = ar._ring_permutations(1, 2, [0, 1, 2, 3])
self.assertEqual(pred_by_c_d, [[3, 0, 1, 2], [3, 0, 1, 2]])
self.assertEqual(rank_by_c_d, [[0, 1, 2, 3], [2, 3, 0, 1]])
pred_by_c_d, rank_by_c_d = ar._ring_permutations(1, 4, [0, 1, 2, 3])
self.assertEqual(pred_by_c_d, [[3, 0, 1, 2], [3, 0, 1, 2],
[3, 0, 1, 2], [3, 0, 1, 2]])
self.assertEqual(rank_by_c_d, [[0, 1, 2, 3], [3, 0, 1, 2],
[2, 3, 0, 1], [1, 2, 3, 0]])
# multiple worker, multiple subchunk cases
pred_by_c_d, rank_by_c_d = ar._ring_permutations(2, 2, [0, 1, 2, 3])
self.assertEqual(pred_by_c_d, [[7, 0, 1, 2, 3, 4, 5, 6],
[3, 0, 5, 2, 7, 4, 1, 6]])
self.assertEqual(rank_by_c_d, [[0, 1, 2, 3, 4, 5, 6, 7],
[2, 3, 0, 1, 6, 7, 4, 5]])
pred_by_c_d, rank_by_c_d = ar._ring_permutations(2, 2, [0, 3, 2, 1])
self.assertEqual(pred_by_c_d, [[5, 2, 3, 0, 1, 6, 7, 4],
[1, 2, 7, 0, 5, 6, 3, 4]])
self.assertEqual(rank_by_c_d, [[0, 3, 2, 1, 4, 7, 6, 5],
[2, 1, 0, 3, 6, 5, 4, 7]])
def _buildInput(self, num_workers, num_gpus):
t8 = constant_op.constant(
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
types_pb2.DT_FLOAT)
input_tensors = []
device_names = []
for w in range(0, num_workers):
for d in range(0, num_gpus):
dn = "/replica:0/task:%d/device:GPU:%d" % (w, d % num_gpus)
device_names.append(dn)
with ops.device(dn):
input_tensors.append(array_ops.identity(t8))
return input_tensors, device_names
@test_util.run_deprecated_v1
def testBuildRingGatherPassStructure(self):
# 1 worker, 1 device
input_tensors, device_names = self._buildInput(1, 1)
pred_by_c_d, rank_by_c_d = ar._ring_permutations(1, 1, [0])
output_tensors = ar._build_ring_gather(input_tensors, device_names, 1,
pred_by_c_d, rank_by_c_d,
math_ops.add)
self.assertEqual(output_tensors, input_tensors)
# 1 worker, 4 devices, 2 subchunks
input_tensors, device_names = self._buildInput(1, 4)
pred_by_c_d, rank_by_c_d = ar._ring_permutations(1, 2, [0, 1, 2, 3])
output_tensors, pad_len = ar._build_ring_gather(
input_tensors, device_names, 2, pred_by_c_d, rank_by_c_d, math_ops.add)
self.assertEqual(0, pad_len)
# same number outputs as inputs
self.assertEqual(len(output_tensors), len(input_tensors))
num_chunks = 2 * len(input_tensors)
tlen = tensor_shape.dimension_value(input_tensors[0].shape[0])
for otl in output_tensors:
self.assertEqual(len(otl), num_chunks)
for ot in otl:
self.assertEqual(ot.shape, [tlen//num_chunks])
def _buildInitialVars(self, shape, dev_list):
values = []
num_devices = len(dev_list)
dim = np.prod(shape, dtype=int) if shape else 1
for d in range(0, num_devices):
with ops.device(dev_list[d]):
npt = np.zeros(shape).astype(np.float32)
alias = np.frombuffer(npt.data, dtype=np.float32)
for i in range(0, dim):
alias[i] = i + 0.01 * d
var = state_ops.variable_op(shape, types_pb2.DT_FLOAT)
state_ops.init_variable(var, npt).op.run()
values.append(var)
return values
# pylint: disable=g-long-lambda
def _buildRing(self, num_workers, num_gpus, subdiv):
gpu_perm = range(0, num_gpus)
return lambda x, un_op: ar.build_ring_all_reduce(
x, num_workers, subdiv, gpu_perm, math_ops.add, un_op)
def _testAllReduce(self, num_workers, num_gpus, shape, build_f):
# Use local CPU as device for all inputs.
num_devices = num_workers * num_gpus
dev_list = ["/replica:0/task:0/device:CPU:0"
for _ in range(num_devices)]
with self.cached_session():
input_tensors = self._buildInitialVars(shape, dev_list)
un_op = lambda x: math_ops.div(
x, constant_op.constant(num_devices, dtype=types_pb2.DT_FLOAT))
simple_sum = math_ops.add_n(input_tensors)
simple_sum.op.run()
output_tensors = build_f(input_tensors, un_op)
sum_reduced = math_ops.add_n(output_tensors)
sum_reduced.op.run()
self.assertAllClose(sum_reduced, self.evaluate(simple_sum))
def _testRingAllReduce(self, num_workers, num_gpus, shape, subdiv):
start_time = time.time()
build_f = self._buildRing(num_workers, num_gpus, subdiv)
self._testAllReduce(num_workers, num_gpus, shape, build_f)
elapsed = time.time() - start_time
tf_logging.info("RingAllReduce num_workers=%d num_gpus=%d shape=%s "
"subdiv=%d elapsed=%f" %
(num_workers, num_gpus, shape, subdiv, elapsed))
@test_util.run_deprecated_v1
def testRingAllReduce(self):
self._testRingAllReduce(1, 2, [], 1)
self._testRingAllReduce(1, 2, [8], 1)
self._testRingAllReduce(1, 2, [4, 4], 1)
self._testRingAllReduce(6, 1, [8], 1)
self._testRingAllReduce(1, 8, [32], 1)
self._testRingAllReduce(1, 8, [120], 1)
self._testRingAllReduce(2, 8, [7, 13], 1)
self._testRingAllReduce(2, 8, [8, 8], 2)
self._testRingAllReduce(2, 8, [8, 8], 4)
# TODO(tucker): The following test is surprisingly slow.
# Diagnose and fix before re-enabling.
# self._testRingAllReduce(4, 8, [8, 8, 2], 4)
def _buildShuffle(self, num_workers, num_gpus, num_shards):
# Use local CPU for all shuffle shards
gather_devices = ["/replica:0/task:0/device:CPU:0"
for _ in range(num_shards)]
return lambda x, un_op: ar.build_shuffle_all_reduce(
x, gather_devices, math_ops.add_n, un_op)
def _testShuffleAllReduce(self, num_workers, num_gpus, shape, num_shards):
start_time = time.time()
build_f = self._buildShuffle(num_workers, num_gpus, num_shards)
self._testAllReduce(num_workers, num_gpus, shape, build_f)
elapsed = time.time() - start_time
tf_logging.info("ShuffleAllReduce num_workers=%d num_gpus=%d shape=%s "
"elapsed=%f" % (num_workers, num_gpus, shape, elapsed))
@test_util.run_deprecated_v1
def testShuffleAllReduce(self):
self._testShuffleAllReduce(1, 2, [], 1)
self._testShuffleAllReduce(1, 2, [8], 1)
self._testShuffleAllReduce(1, 2, [4, 4], 1)
self._testShuffleAllReduce(1, 8, [32], 1)
self._testShuffleAllReduce(1, 8, [120], 1)
self._testShuffleAllReduce(2, 8, [7, 13], 3)
self._testShuffleAllReduce(2, 8, [8, 8], 2)
self._testShuffleAllReduce(2, 8, [8, 8], 4)
self._testShuffleAllReduce(4, 8, [8, 8, 2], 4)
def _buildRecursiveHD(self, num_workers, num_gpus):
return lambda x, un_op: ar.build_recursive_hd_all_reduce(
x, math_ops.add, un_op)
# pylint: enable=g-long-lambda
def _testRecursiveHDAllReduce(self, num_workers, num_gpus, shape):
start_time = time.time()
build_f = self._buildRecursiveHD(num_workers, num_gpus)
self._testAllReduce(num_workers, num_gpus, shape, build_f)
elapsed = time.time() - start_time
tf_logging.info("RecursiveHDAllReduce num_workers=%d num_gpus=%d "
"shape=%s elapsed=%f" %
(num_workers, num_gpus, shape, elapsed))
@test_util.run_deprecated_v1
def testRecursiveHDAllReduce(self):
self._testRecursiveHDAllReduce(1, 2, [8])
self._testRecursiveHDAllReduce(1, 2, [4, 4])
self._testRecursiveHDAllReduce(1, 8, [32])
self._testRecursiveHDAllReduce(1, 8, [120])
self._testRecursiveHDAllReduce(2, 8, [8, 8])
self._testRecursiveHDAllReduce(4, 8, [8, 8, 2])
if __name__ == "__main__":
test.main()
tensorflow/python/distribute/v1/cross_device_ops_test.py
+857
View File
@@ -0,0 +1,857 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for CrossDeviceOps in v1 graph mode."""
import itertools
import os
import threading
import time
from absl.testing import parameterized
import numpy as np
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.distribute import cluster_resolver
from tensorflow.python.distribute import collective_all_reduce_strategy as mwms_lib
from tensorflow.python.distribute import collective_util
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib
from tensorflow.python.distribute import cross_device_utils
from tensorflow.python.distribute import device_util
from tensorflow.python.distribute import distribute_utils
from tensorflow.python.distribute import multi_worker_test_base
from tensorflow.python.distribute import multi_worker_util
from tensorflow.python.distribute import reduce_util
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.distribute import values as value_lib
from tensorflow.python.eager import context
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import indexed_slices as indexed_slices_lib
from tensorflow.python.framework import kernels
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import collective_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
def _get_devices(devices):
if isinstance(devices, (tuple, list)):
return tuple(device_util.resolve(d) for d in devices)
elif isinstance(devices, value_lib.DistributedValues):
return devices._devices
elif isinstance(devices, tensor_lib.Tensor):
return (device_util.resolve(devices.device),)
return (device_util.resolve(devices),)
def _make_per_replica(values, devices, regroup=False):
devices = _get_devices(devices)
assert len(values) == len(devices)
# We simulate the result of regroup called on PerReplica which strips the
# PerReplica wrapper if it has only one value.
if len(values) == 1 and regroup:
with ops.device(devices[0]):
placed_v = array_ops.identity(values[0])
return placed_v
index = []
for d, v in zip(devices, values):
with ops.device(d):
placed_v = array_ops.identity(v)
index.append(placed_v)
return distribute_utils.regroup(index)
# pylint: disable=g-doc-args,g-doc-return-or-yield
def _fake_mirrored(value, devices):
"""Create a faked Mirrored object for testing.
All components of the returned Mirrored have the same objects, which is not
true in reality.
"""
devices = _get_devices(devices)
values = []
for d in devices:
with ops.device(d):
values.append(array_ops.identity(value))
return distribute_utils.regroup(
values,
wrap_class=value_lib.Mirrored)
def _make_indexed_slices(values, indices, dense_shape, device):
with ops.device(device):
tensor = indexed_slices_lib.IndexedSlices(
values=constant_op.constant(values),
indices=constant_op.constant(indices),
dense_shape=constant_op.constant(dense_shape))
return tensor
def _make_mirrored_indexed_slices(devices, values, indices, dense_shape):
values = [_make_indexed_slices(values, indices, dense_shape, d)
for d in devices]
return distribute_utils.regroup(
values,
wrap_class=value_lib.Mirrored)
_cpu_device = "/device:CPU:0"
class CrossDeviceOpsTestBase(test.TestCase, parameterized.TestCase):
def _assert_indexed_slices_equal(self, left, right):
self.assertIsInstance(left, indexed_slices_lib.IndexedSlices)
self.assertIsInstance(right, indexed_slices_lib.IndexedSlices)
self.assertEqual(
device_util.resolve(left.device), device_util.resolve(right.device))
self.assertAllEqual(
self.evaluate(ops.convert_to_tensor(left)),
self.evaluate(ops.convert_to_tensor(right)))
def _assert_mirrored_equal(self,
left_list,
right_list,
sess=None,
run_options=None):
if not isinstance(left_list, list):
left_list, right_list = [left_list], [right_list]
for left, right in zip(left_list, right_list):
self.assertEqual(type(left), type(right))
# Convert Mirrored to a list since sess.run(Mirrored) only returns one
# value.
if isinstance(left, value_lib.Mirrored):
left, right = left.values, right.values
else:
# When there's only one replica Mirrored is automatically unwrapped.
left, right = [left], [right]
for left_value, right_value in zip(left, right):
self.assertEqual(
device_util.resolve(left_value.device),
device_util.resolve(right_value.device))
# Densify IndexedSlices.
left = [ops.convert_to_tensor(v) for v in left]
right = [ops.convert_to_tensor(v) for v in right]
if not context.executing_eagerly():
left, right = sess.run((left, right), options=run_options)
for left_value, right_value in zip(left, right):
self.assertAllEqual(left_value, right_value)
def _testReductionAndBroadcast(self, cross_device_ops, devices):
if context.num_gpus() < sum(1 for d in devices if "GPU" in d.upper()):
self.skipTest("Not enough GPUs")
with self.cached_session() as sess:
values = [constant_op.constant(float(d)) for d in range(len(devices))]
per_replica = _make_per_replica(values, devices)
mean = (len(devices) - 1.) / 2.
values_2 = [constant_op.constant(d + 1.0) for d in range(len(devices))]
per_replica_2 = _make_per_replica(values_2, devices)
mean_2 = mean + 1.
destination_mirrored = _fake_mirrored(1., devices)
destination_different = _fake_mirrored(1.,
device_util.resolve(_cpu_device))
destination_str = device_util.resolve(_cpu_device)
all_destinations = [
destination_mirrored,
destination_different,
destination_str,
]
# test reduce()
for destinations in all_destinations:
self._assert_mirrored_equal(
cross_device_ops.reduce(
reduce_util.ReduceOp.MEAN,
per_replica,
destinations=destinations), _fake_mirrored(mean, destinations),
sess)
self._assert_mirrored_equal(
cross_device_ops.reduce(
reduce_util.ReduceOp.MEAN,
per_replica_2,
destinations=destinations),
_fake_mirrored(mean_2, destinations), sess)
self._assert_mirrored_equal(
cross_device_ops.reduce(
reduce_util.ReduceOp.SUM,
per_replica,
destinations=destinations),
_fake_mirrored(mean * len(devices), destinations), sess)
self._assert_mirrored_equal(
cross_device_ops.reduce(
reduce_util.ReduceOp.SUM,
per_replica_2,
destinations=destinations),
_fake_mirrored(mean_2 * len(devices), destinations), sess)
# test batch_reduce()
for d1, d2 in itertools.product(all_destinations, all_destinations):
self._assert_mirrored_equal(
cross_device_ops.batch_reduce(reduce_util.ReduceOp.MEAN,
[(per_replica, d1),
(per_replica_2, d2)]),
[_fake_mirrored(mean, d1),
_fake_mirrored(mean_2, d2)], sess)
self._assert_mirrored_equal(
cross_device_ops.batch_reduce(reduce_util.ReduceOp.SUM,
[(per_replica, d1),
(per_replica_2, d2)]),
[
_fake_mirrored(mean * len(devices), d1),
_fake_mirrored(mean_2 * len(devices), d2)
], sess)
# test broadcast()
for destinations in all_destinations:
self._assert_mirrored_equal(
cross_device_ops.broadcast(constant_op.constant(1.), destinations),
_fake_mirrored(1., destinations), sess)
def _testIndexedSlicesAllReduce(self, devices, cross_device_ops_instance,
reduce_op, batch_reduce):
with self.cached_session() as sess:
dense_shape = [5, 2]
t0 = _make_indexed_slices([[1., 2.]], [1], dense_shape, devices[0])
t1 = _make_indexed_slices([[3., 4.], [5., 6.]], [1, 3], dense_shape,
devices[1])
per_replica = value_lib.PerReplica((t0, t1))
if batch_reduce:
result = cross_device_ops_instance.batch_reduce(
reduce_op, [(per_replica, per_replica)])
else:
result = cross_device_ops_instance.reduce(reduce_op, per_replica,
per_replica)
total_indices_with_dups = [1, 1, 3]
total_indices_without_dups = [1, 3]
if reduce_op == reduce_util.ReduceOp.SUM:
total_values_with_dups = [[1., 2.], [3., 4.], [5., 6.]]
total_values_without_dups = [[4., 6.], [5., 6.]]
else:
assert reduce_op == reduce_util.ReduceOp.MEAN
total_values_with_dups = [[0.5, 1.], [1.5, 2.], [2.5, 3.]]
total_values_without_dups = [[2., 3.], [2.5, 3.]]
total_mirrored_with_dups = _make_mirrored_indexed_slices(
devices, total_values_with_dups, total_indices_with_dups, dense_shape)
total_mirrored_without_dups = _make_mirrored_indexed_slices(
devices, total_values_without_dups, total_indices_without_dups,
dense_shape)
# Test that the result is semantically equal to both the concatenated
# IndexedSlices, as well as when the duplicate indices are summed up.
if batch_reduce:
total_mirrored_with_dups = [total_mirrored_with_dups]
total_mirrored_without_dups = [total_mirrored_without_dups]
self._assert_mirrored_equal(total_mirrored_with_dups, result, sess)
self._assert_mirrored_equal(total_mirrored_without_dups, result, sess)
class SingleWorkerCrossDeviceOpsTest(CrossDeviceOpsTestBase):
reduction_to_one_combinations = combinations.combine(
cross_device_ops=[
combinations.NamedObject("DefaultReductionToOneDevice",
cross_device_ops_lib.ReductionToOneDevice()),
combinations.NamedObject(
"ReductionToCPUDeviceCrossDeviceOps",
cross_device_ops_lib.ReductionToOneDevice(
reduce_to_device=_cpu_device)),
combinations.NamedObject(
"AccumulateNCrossDeviceOp",
cross_device_ops_lib.ReductionToOneDevice(
accumulation_fn=math_ops.add_n)),
],
devices=[
["/cpu:0"],
["/cpu:0", "/gpu:0"],
["/gpu:0", "/gpu:1"],
],
mode=["graph", "eager"])
allreduce_combinations = combinations.combine(
cross_device_ops=[
combinations.NamedObject(
"AllReduce",
cross_device_ops_lib.AllReduceCrossDeviceOps("nccl", 1)),
combinations.NamedObject(
"AllReduceNoGradientRepacking",
cross_device_ops_lib.AllReduceCrossDeviceOps("nccl", 0)),
combinations.NamedObject("NcclAllReduce",
cross_device_ops_lib.NcclAllReduce()),
combinations.NamedObject(
"HierarchicalCopy",
cross_device_ops_lib.HierarchicalCopyAllReduce(8)),
],
devices=[
["/gpu:0", "/gpu:1"],
],
mode=["graph", "eager"])
@combinations.generate(reduction_to_one_combinations + allreduce_combinations)
def testReductionAndBroadcast(self, cross_device_ops, devices):
if isinstance(
cross_device_ops._obj, # pylint: disable=protected-access
cross_device_ops_lib.AllReduceCrossDeviceOps
) and context.executing_eagerly():
self.skipTest("b/149881884")
self._testReductionAndBroadcast(cross_device_ops, devices)
def testChooseAlgorithm(self):
# Not use nccl if there is any cpu device.
self.assertIsInstance(
cross_device_ops_lib.select_cross_device_ops(["/cpu:0"]),
cross_device_ops_lib.ReductionToOneDevice)
# Not use nccl if requested device is not visible to TensorFlow.
# TODO(yuefengz): make `select_cross_device_ops` work with device strings
# self.assertIsInstance(
# cross_device_ops_lib.select_cross_device_ops(["/gpu:100"]),
# cross_device_ops_lib.ReductionToOneDevice)
if context.num_gpus() < 1:
return
devices = ["/gpu:0"]
def mock_get_registered_kernels_for_op(op):
if op == "NcclAllReduce":
return [object]
else:
return []
# Use nccl if nccl kernel is found.
with test.mock.patch.object(kernels, "get_registered_kernels_for_op",
mock_get_registered_kernels_for_op):
self.assertIsInstance(
cross_device_ops_lib.select_cross_device_ops(devices),
cross_device_ops_lib.NcclAllReduce)
# Not use nccl if nccl kernel is not found.
with test.mock.patch.object(kernels,
"get_registered_kernels_for_op", lambda _: []):
self.assertIsInstance(
cross_device_ops_lib.select_cross_device_ops(devices),
cross_device_ops_lib.ReductionToOneDevice)
@combinations.generate(combinations.combine(
mode=["graph", "eager"],
required_gpus=1))
def testSimpleReduceWithIndexedSlices(self):
devices = ["/cpu:0", "/gpu:0"]
t0 = _make_indexed_slices([[1., 2.]], [1], [5, 2], devices[0])
t1 = _make_indexed_slices([[3., 4.], [5., 6.]], [1, 3], [5, 2], devices[1])
per_replica = value_lib.PerReplica((t0, t1))
result = cross_device_ops_lib._simple_reduce(
per_replica, devices[0], math_ops.add_n, reduce_util.ReduceOp.SUM)
# Test that the result is semantically equal to both the concatenated
# IndexedSlices with and without duplicate indices.
total_with_dups = _make_indexed_slices(
[[1., 2.], [3., 4.], [5., 6.]], [1, 1, 3], [5, 2], devices[0])
total_without_dups = _make_indexed_slices(
[[4., 6.], [5., 6.]], [1, 3], [5, 2], devices[0])
self._assert_indexed_slices_equal(total_with_dups, result)
self._assert_indexed_slices_equal(total_without_dups, result)
@combinations.generate(
combinations.combine(
cross_device_ops_instance=[
combinations.NamedObject(
"ReductionToOneDevice",
cross_device_ops_lib.ReductionToOneDevice()),
combinations.NamedObject(
"AllReduceCrossDeviceOps",
cross_device_ops_lib.AllReduceCrossDeviceOps())
],
reduce_op=[reduce_util.ReduceOp.SUM, reduce_util.ReduceOp.MEAN],
batch_reduce=[True, False],
mode=["graph", "eager"],
required_gpus=1))
def testIndexedSlicesAllReduce(self, cross_device_ops_instance, reduce_op,
batch_reduce):
devices = ["/cpu:0", "/gpu:0"]
self._testIndexedSlicesAllReduce(devices, cross_device_ops_instance,
reduce_op, batch_reduce)
@combinations.generate(
combinations.combine(
distribution=strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
cross_device_ops_instance=[
combinations.NamedObject(
"ReductionToOneDevice",
cross_device_ops_lib.ReductionToOneDevice()),
combinations.NamedObject(
"AllReduceCrossDeviceOps",
cross_device_ops_lib.AllReduceCrossDeviceOps("ring"))
],
batch_reduce=[True, False],
mode=["graph", "eager"]))
def testReduceDistributedVariable(self, distribution,
cross_device_ops_instance, batch_reduce):
with distribution.scope():
v = variables.Variable(1.)
if batch_reduce:
result = cross_device_ops_instance.batch_reduce(reduce_util.ReduceOp.MEAN,
[(v, v)])[0]
else:
result = cross_device_ops_instance.reduce(reduce_util.ReduceOp.MEAN, v, v)
for v in result.values:
self.assertIsInstance(v, tensor_lib.Tensor)
self.evaluate(variables.global_variables_initializer())
self.assertAllEqual(self.evaluate(result.values), [1.0, 1.0])
NUM_WORKERS = 3
CollectiveCommunication = collective_util.CollectiveCommunication
class CollectiveAllReduceTest(multi_worker_test_base.MultiWorkerTestBase,
CrossDeviceOpsTestBase):
collective_key_base = 100000
@classmethod
def setUpClass(cls):
"""Create a local cluster with 3 workers."""
cls._cluster_spec = multi_worker_test_base.create_in_process_cluster(
num_workers=NUM_WORKERS, num_ps=0)
def setUp(self):
super(CollectiveAllReduceTest, self).setUp()
# Reusing keys is not supported well. So we have to give a different
# collective key base for different tests.
CollectiveAllReduceTest.collective_key_base += 100000
mwms_lib.CollectiveAllReduceStrategy._collective_key_base = (
CollectiveAllReduceTest.collective_key_base)
def _get_test_objects(self,
task_type,
task_id,
num_gpus=0,
communication=CollectiveCommunication.AUTO,
use_strategy_object=False,
local_mode=False):
collective_keys = cross_device_utils.CollectiveKeys(
group_key_start=10 + CollectiveAllReduceTest.collective_key_base)
if local_mode:
if num_gpus:
devices = ["/device:GPU:%d" % i for i in range(num_gpus)]
else:
devices = ["/device:CPU:0"]
comm_options = collective_util.Options(implementation=communication)
if use_strategy_object:
strategy = (mwms_lib.CollectiveAllReduceStrategy
._from_local_devices(devices, comm_options)) # pylint: disable=protected-access
return strategy, devices, ""
else:
collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce(
devices=devices,
group_size=len(devices),
options=comm_options,
collective_keys=collective_keys)
return collective_all_reduce_ops, devices, ""
else:
# NCCL requires physical GPUs for every replica, which we can't do with
# simulated multi host set up now.
assert communication != CollectiveCommunication.NCCL
if num_gpus:
devices = [
"/job:%s/task:%d/replica:0/device:GPU:%d" % (task_type, task_id, i)
for i in range(num_gpus)
]
else:
devices = [
"/job:%s/task:%d/replica:0/device:CPU:0" % (task_type, task_id)
]
comm_options = collective_util.Options(implementation=communication)
if use_strategy_object:
resolver = cluster_resolver.SimpleClusterResolver(
cluster_spec=multi_worker_util.normalize_cluster_spec(
self._cluster_spec),
task_type=task_type,
task_id=task_id,
num_accelerators={"GPU": num_gpus})
strategy = mwms_lib.CollectiveAllReduceStrategy(
communication_options=comm_options, cluster_resolver=resolver)
return (strategy, devices,
"grpc://" + self._cluster_spec[task_type][task_id])
else:
collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce(
devices=devices,
group_size=len(devices) * NUM_WORKERS,
options=comm_options,
collective_keys=collective_keys)
return (collective_all_reduce_ops, devices,
"grpc://" + self._cluster_spec[task_type][task_id])
def _assert_mirrored_equal(self, left_list, right_list, sess=None):
if context.executing_eagerly():
run_options = None
else:
# TODO(b/151025792): figure out why missing run options would make the
# test flaky and whether this is a problem in TF 2.
run_options = config_pb2.RunOptions()
run_options.experimental.collective_graph_key = 5
super(CollectiveAllReduceTest, self)._assert_mirrored_equal(
left_list, right_list, sess, run_options=run_options)
def _test_reduction(self,
task_type,
task_id,
num_gpus,
communication,
use_strategy_object=False,
local_mode=False,
hints=None):
collective_all_reduce, devices, master_target = self._get_test_objects(
task_type,
task_id,
num_gpus,
communication=communication,
use_strategy_object=use_strategy_object,
local_mode=local_mode)
if local_mode:
num_workers = 1
worker_device = None
else:
num_workers = len(self._cluster_spec.get("chief", [])) + len(
self._cluster_spec.get("worker", []))
worker_device = "/job:%s/task:%d" % (task_type, task_id)
def _reduce(test_object, reduce_op, per_replica, destinations):
if use_strategy_object:
with test_object.scope():
return test_object.extended.reduce_to(reduce_op, per_replica,
destinations, hints)
else:
return test_object.reduce(reduce_op, per_replica, destinations, hints)
def _batch_reduce(test_object, reduce_op, value_destination_pairs):
if use_strategy_object:
with test_object.scope():
return test_object.extended.batch_reduce_to(reduce_op,
value_destination_pairs,
hints)
else:
return test_object.batch_reduce(reduce_op, value_destination_pairs,
hints)
with ops.Graph().as_default(), \
ops.device(worker_device), \
self.cached_session(target=master_target) as sess:
# Collective ops doesn't support scalar tensors, so we have to construct
# 1-d tensors.
values = [constant_op.constant([float(d)]) for d in range(len(devices))]
per_replica = _make_per_replica(values, devices)
mean = np.array([(len(devices) - 1.) / 2.])
values_2 = [constant_op.constant([d + 1.0]) for d in range(len(devices))]
per_replica_2 = _make_per_replica(values_2, devices)
mean_2 = np.array([mean[0] + 1.])
destination_mirrored = _fake_mirrored(1., devices)
destination_different = _fake_mirrored(1., _cpu_device)
destination_str = _cpu_device
all_destinations = [
destination_different, destination_mirrored, destination_str
]
# test reduce()
for destinations in all_destinations:
self._assert_mirrored_equal(
_reduce(
collective_all_reduce,
reduce_util.ReduceOp.MEAN,
per_replica,
destinations=destinations), _fake_mirrored(mean, destinations),
sess)
self._assert_mirrored_equal(
_reduce(
collective_all_reduce,
reduce_util.ReduceOp.MEAN,
per_replica_2,
destinations=destinations),
_fake_mirrored(mean_2, destinations), sess)
self._assert_mirrored_equal(
_reduce(
collective_all_reduce,
reduce_util.ReduceOp.SUM,
per_replica,
destinations=destinations),
_fake_mirrored(mean * len(devices) * num_workers, destinations),
sess)
self._assert_mirrored_equal(
_reduce(
collective_all_reduce,
reduce_util.ReduceOp.SUM,
per_replica_2,
destinations=destinations),
_fake_mirrored(mean_2 * len(devices) * num_workers, destinations),
sess)
# test batch_reduce()
for d1, d2 in itertools.product(all_destinations, all_destinations):
self._assert_mirrored_equal(
_batch_reduce(collective_all_reduce, reduce_util.ReduceOp.MEAN,
[(per_replica, d1), (per_replica_2, d2)]),
[_fake_mirrored(mean, d1),
_fake_mirrored(mean_2, d2)], sess)
self._assert_mirrored_equal(
_batch_reduce(collective_all_reduce, reduce_util.ReduceOp.SUM,
[(per_replica, d1), (per_replica_2, d2)]),
[
_fake_mirrored(mean * len(devices) * num_workers, d1),
_fake_mirrored(mean_2 * len(devices) * num_workers, d2)
], sess)
def _get_indexed_slices(self,
devices,
start_i,
variable_length,
as_per_replica=True):
dense_shape = [10, 2]
values = ([[1., 2.]], [[3., 4.]], [[2., 1.]], [[0., 0.]], [[3., 1.]],
[[2., 1.]])
indices = ([1], [2], [3], [4], [5], [6])
# values and indices that have variable lengths.
vl_values = ([[1., 2.], [3., 4.]], [[3., 4.]], [[2., 1.]], [[0., 0.]],
[[3., 1.], [2., 1.]], [[2., 1.]])
vl_indices = ([1, 2], [2], [3], [4], [5, 6], [6])
indexed_slices = []
for i, d in enumerate(devices):
idx = i + start_i
indexed_slices.append(
_make_indexed_slices(
vl_values[idx] if variable_length else values[idx],
vl_indices[idx] if variable_length else indices[idx], dense_shape,
d))
if as_per_replica:
per_replica = value_lib.PerReplica(indexed_slices)
return per_replica
else:
return indexed_slices
def _test_reduce_indexed_slices(self,
task_type,
task_id,
num_gpus,
communication,
batch_reduce,
variable_length,
local_mode=False):
collective_all_reduce, devices, master_target = self._get_test_objects(
task_type,
task_id,
num_gpus,
communication=communication,
local_mode=local_mode)
if local_mode:
num_workers = 1
worker_device = None
else:
num_workers = len(self._cluster_spec.get("chief", [])) + len(
self._cluster_spec.get("worker", []))
worker_device = "/job:%s/task:%d" % (task_type, task_id)
config = config_pb2.ConfigProto()
# Disable coordination service for all tasks except task 0 to avoid
# parallel init of the service singleton.
if task_id != 0:
config.experimental.coordination_config.service_type = ""
with ops.Graph().as_default(), \
ops.device(worker_device), \
self.cached_session(target=master_target, config=config) as sess:
per_replica = self._get_indexed_slices(devices,
(task_id or 0) * max(num_gpus, 1),
variable_length)
if batch_reduce:
result = collective_all_reduce.batch_reduce(
reduce_util.ReduceOp.SUM, [(per_replica, per_replica)])[0]
else:
result = collective_all_reduce.reduce(reduce_util.ReduceOp.SUM,
per_replica, per_replica)
if num_gpus > 1:
self.assertIsInstance(result, value_lib.Mirrored)
run_options = config_pb2.RunOptions()
run_options.experimental.collective_graph_key = 7
if num_gpus > 1:
result = sess.run([ops.convert_to_tensor(v) for v in result.values],
options=run_options)[0]
else:
result = sess.run(ops.convert_to_tensor(result), options=run_options)
# Reduce the same indexed slices on CPU locally as our expected results.
devices_cpu = [(worker_device or "") + "/device:CPU:0"] * (
max(num_gpus, 1) * num_workers)
per_replica_on_cpu = self._get_indexed_slices(
devices_cpu, 0, variable_length, as_per_replica=False)
expected_result = cross_device_utils.aggregate_tensors_or_indexed_slices(
per_replica_on_cpu)
expected_result = sess.run(ops.convert_to_tensor(expected_result))
self.assertAllEqual(expected_result, result)
@combinations.generate(
combinations.combine(
mode=["graph"],
required_gpus=[0, 1, 2],
use_strategy_object=[True, False],
bytes_per_pack=[0, 1, 4]))
def testReductionDistributed(self, required_gpus, use_strategy_object,
bytes_per_pack):
hints = collective_util.Hints(bytes_per_pack=bytes_per_pack)
self._run_between_graph_clients(
self._test_reduction,
self._cluster_spec,
required_gpus,
communication=CollectiveCommunication.RING,
use_strategy_object=use_strategy_object,
hints=hints)
@combinations.generate(
combinations.combine(
mode=["graph"],
required_gpus=[0, 1, 2],
variable_length=[True, False]))
def testReduceIndexedSlicesDistributed(self, required_gpus, variable_length):
self._run_between_graph_clients(
self._test_reduce_indexed_slices,
self._cluster_spec,
required_gpus,
communication=CollectiveCommunication.RING,
batch_reduce=True,
variable_length=variable_length)
# Collective ops doesn't support strategy with one device.
@combinations.generate(
combinations.combine(
mode=["graph"],
required_gpus=2,
communication=[
CollectiveCommunication.NCCL, CollectiveCommunication.RING
],
use_strategy_object=[True, False]))
def testReductionLocal(self, required_gpus, communication,
use_strategy_object):
self._test_reduction(
None,
None,
required_gpus,
communication=communication,
use_strategy_object=use_strategy_object,
local_mode=True)
@combinations.generate(
combinations.combine(
mode=["graph"],
required_gpus=2,
batch_reduce=[True, False],
variable_length=[True, False],
communication=[
CollectiveCommunication.NCCL, CollectiveCommunication.RING
]))
def testReduceIndexedSlicesLocal(self, required_gpus, batch_reduce,
variable_length, communication):
self._test_reduce_indexed_slices(
None,
None,
required_gpus,
communication=communication,
batch_reduce=batch_reduce,
variable_length=variable_length,
local_mode=True)
@combinations.generate(
combinations.combine(
required_gpus=2,
mode="eager",
communication=[
CollectiveCommunication.NCCL, CollectiveCommunication.RING
]))
def testEagerMultiThread(self, communication):
collective, devices, _ = self._get_test_objects(
None,
None,
num_gpus=2,
communication=communication,
use_strategy_object=False,
local_mode=True)
# We would like to simulate the following sequence:
# thread-0 device0 device1
# thread-1 device0 device1
# If the kernel launch sequence is as-is the program will deadlock since
# NCCL requires the launch order to be same on each device.
v0 = _make_per_replica([1.0 for _ in devices], devices)
v1 = _make_per_replica([2.0 for _ in devices], devices)
# Add a delay to collective_ops.all_reduce according to the input tensors
# index in `sequence.`
sequence = [v0.values[0], v1.values[0], v1.values[1], v0.values[1]]
all_reduce = collective_ops.all_reduce
def delayed_all_reduce(input_tensor, *args, **kwargs):
for idx, v in enumerate(sequence):
if input_tensor is v:
time.sleep(idx)
break
return all_reduce(input_tensor, *args, **kwargs)
with test.mock.patch.object(collective_ops, "all_reduce",
delayed_all_reduce):
# We only use NCCL for batch reduce with two or more values, so we use two
# values here.
def thread_fn():
reduced = collective.batch_reduce(reduce_util.ReduceOp.SUM, [(v0, v0),
(v0, v0)])
self.assertAllEqual(reduced[0].values, [2.0, 2.0])
self.assertAllEqual(reduced[1].values, [2.0, 2.0])
t = threading.Thread(target=thread_fn)
t.start()
reduced = collective.batch_reduce(reduce_util.ReduceOp.SUM, [(v1, v1),
(v1, v1)])
self.assertAllEqual(reduced[0].values, [4.0, 4.0])
self.assertAllEqual(reduced[1].values, [4.0, 4.0])
t.join()
if __name__ == "__main__":
# Set default inter op thread pool size to one to ensure we don't exhaust the
# thread pool with the additional executors to run collectives in eager.
os.environ["TF_NUM_INTEROP_THREADS"] = "1"
test.main()
tensorflow/python/distribute/v1/input_lib.py
+420
View File
@@ -0,0 +1,420 @@
# Copyright 2021 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.
# ==============================================================================
"""Various classes representing distributed inputs."""
from tensorflow.python.data.experimental.ops import cardinality as cardinality_lib
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import multi_device_iterator_ops
from tensorflow.python.data.ops import optional_ops
from tensorflow.python.distribute import input_lib
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.types import data as data_types
from tensorflow.python.util.deprecation import deprecated
class DistributedDatasetV1(input_lib.DistributedDataset):
"""Distributed dataset that supports prefetching to multiple devices."""
def __init__(self,
dataset,
input_workers,
strategy,
num_replicas_in_sync=None,
input_context=None,
options=None):
self._input_workers = input_workers
super(DistributedDatasetV1, self).__init__(
input_workers,
strategy,
dataset,
num_replicas_in_sync=num_replicas_in_sync,
input_context=input_context,
options=options)
def make_one_shot_iterator(self):
"""Get a one time use iterator for DistributedDatasetV1.
Note: This API is deprecated. Please use `for ... in dataset:` to iterate
over the dataset or `iter` to create an iterator.
Returns:
A DistributedIteratorV1 instance.
"""
return self._make_one_shot_iterator()
def _make_one_shot_iterator(self):
"""Get an iterator for DistributedDatasetV1."""
# Graph mode with one shot iterator is disabled because we have to call
# `initialize` on the iterator which is only required if we are using a
# tf.distribute strategy.
if not context.executing_eagerly():
raise ValueError("Cannot create a one shot iterator. Please use "
"`make_initializable_iterator()` instead.")
return self._get_iterator()
def make_initializable_iterator(self):
"""Get an initializable iterator for DistributedDatasetV1.
Note: This API is deprecated. Please use
`tf.compat.v1.data.make_initializable_iterator(dataset)` to create an
initializable iterator.
Returns:
A DistributedIteratorV1 instance.
"""
return self._make_initializable_iterator()
def _make_initializable_iterator(self, shared_name=None): # pylint: disable=unused-argument
"""Get an initializable iterator for DistributedDatasetV1."""
# Eager mode generates already initialized iterators. Hence we cannot create
# an initializable iterator.
if context.executing_eagerly():
raise ValueError("Cannot create initializable iterator in Eager mode. "
"Please use `iter()` instead.")
return self._get_iterator()
def _get_iterator(self):
worker_iterators = _create_iterators_per_worker(self._cloned_datasets,
self._input_workers,
self._options)
cardinality = input_lib._cardinality(self._cloned_datasets[0]) # pylint: disable=protected-access
iterator = DistributedIteratorV1(self._input_workers, worker_iterators,
self._strategy, cardinality,
self._enable_get_next_as_optional)
iterator._element_spec = self.element_spec # pylint: disable=protected-access
# When async eager is enabled, sometimes the iterator may not finish
# initialization before passing to a multi device function, add a sync point
# here to make sure all underlying iterators are initialized.
if context.executing_eagerly():
context.async_wait()
return iterator
# pylint: disable=non-iterator-returned
def __iter__(self):
if (ops.executing_eagerly_outside_functions() or
ops.get_default_graph().building_function):
return self._get_iterator()
raise RuntimeError("__iter__() is only supported inside of tf.function "
"or when eager execution is enabled.")
# pylint: enable=non-iterator-returned
class DistributedDatasetsFromFunctionV1(
input_lib.DistributedDatasetsFromFunction):
"""Inputs created from dataset function."""
def _make_initializable_iterator(self, shared_name=None):
"""Get an initializable iterator for DistributedDatasetsFromFunctionV1."""
del shared_name # Unused
# Eager mode generates already initialized iterators. Hence we cannot create
# an initializable iterator.
if context.executing_eagerly():
raise ValueError("Cannot create initializable iterator in Eager mode. "
"Please use `iter()` instead.")
return self._get_iterator()
def _make_one_shot_iterator(self):
"""Get an iterator for iterating over DistributedDatasetsFromFunctionV1."""
# Graph mode with one shot iterator is disabled because we have to call
# `initialize` on the iterator which is only required if we are using a
# tf.distribute strategy.
if not context.executing_eagerly():
raise ValueError("Cannot create a one shot iterator. Please use "
"`make_initializable_iterator()` instead.")
return self._get_iterator()
def _get_iterator(self):
iterators = _create_iterators_per_worker(self._datasets,
self._input_workers, self._options)
cardinality = input_lib._cardinality(self._datasets[0]) # pylint: disable=protected-access
iterator = DistributedIteratorV1(self._input_workers, iterators,
self._strategy, cardinality,
self._enable_get_next_as_optional)
iterator._element_spec = self._element_spec # pylint: disable=protected-access
# When async eager is enabled, sometimes the iterator may not finish
# initialization before passing to a multi device function, add a sync point
# here to make sure all underlying iterators are initialized.
if context.executing_eagerly():
context.async_wait()
return iterator
# pylint: disable=non-iterator-returned
def __iter__(self):
if (ops.executing_eagerly_outside_functions() or
ops.get_default_graph().building_function):
return self._get_iterator()
raise RuntimeError("__iter__() is only supported inside of tf.function "
"or when eager execution is enabled.")
# pylint: enable=non-iterator-returned
class DistributedIteratorV1(input_lib.DistributedIteratorBase):
"""Input Iterator for a distributed dataset."""
# We need a private initializer method for re-initializing multidevice
# iterators when used with Keras training loops. If we don't reinitialize the
# iterator we run into memory leak issues (b/123315763).
@property
def _initializer(self):
init_ops = []
for it in self._iterators:
init_ops.extend(it.initialize())
return control_flow_ops.group(init_ops)
@deprecated(None, "Use the iterator's `initializer` property instead.")
def initialize(self):
"""Initialize underlying iterators.
Returns:
A list of any initializer ops that should be run.
"""
return self._initializer
@property
def initializer(self):
"""Returns a list of ops that initialize the iterator."""
return self.initialize()
# TODO(priyag): Remove when we switch to using `MultiDeviceIterator` for TPUs.
@property
def output_classes(self):
return self._iterators[0].output_classes
# TODO(priyag): Remove when we switch to using `MultiDeviceIterator` for TPUs.
@property
def output_shapes(self):
return self._iterators[0].output_shapes
# TODO(priyag): Remove when we switch to using `MultiDeviceIterator` for TPUs.
@property
def output_types(self):
return self._iterators[0].output_types
# TODO(priyag): Remove when we switch to using `MultiDeviceIterator` for TPUs.
def get_iterator(self, worker):
for i, w in enumerate(self._input_workers.worker_devices):
if worker == w:
return self._iterators[i]
return None
@property
def element_spec(self):
"""The type specification of an element of this iterator."""
return self._element_spec
class DatasetIterator(DistributedIteratorV1):
"""Iterator created from input dataset."""
def __init__(self,
dataset,
input_workers,
strategy,
num_replicas_in_sync=None,
input_context=None):
"""Make an iterator for the dataset on given devices.
If `num_replicas_in_sync` is not None, we split each batch of the dataset
into `num_replicas_in_sync` smaller batches, to be distributed among that
worker's replicas, so that the batch size for a global step (across all
workers and replicas) is as expected.
Args:
dataset: `tf.data.Dataset` that will be used as the input source.
input_workers: an `InputWorkers` object.
strategy: a `tf.distribute.Strategy` object, used to run all-reduce to
handle last partial batch.
num_replicas_in_sync: Optional integer. If this is not None, the value is
used to decide how to rebatch datasets into smaller batches so that the
total batch size for each step (across all workers and replicas) adds up
to `dataset`'s batch size.
input_context: `InputContext` for sharding. Only pass this in for between
graph multi-worker cases where there is only one `input_worker`. In
these cases, we will shard based on the `input_pipeline_id` and
`num_input_pipelines` in the `InputContext`.
"""
dist_dataset = DistributedDatasetV1(
dataset,
input_workers,
strategy,
num_replicas_in_sync=num_replicas_in_sync,
input_context=input_context)
# pylint: disable=protected-access
worker_iterators = _create_iterators_per_worker(
dist_dataset._cloned_datasets, input_workers)
super(DatasetIterator,
self).__init__(input_workers, worker_iterators, strategy,
dist_dataset.cardinality,
dist_dataset._enable_get_next_as_optional)
self._element_spec = dist_dataset.element_spec
# pylint: enable=protected-access
class InputFunctionIterator(DistributedIteratorV1):
"""Iterator created from input function."""
def __init__(self, input_fn, input_workers, input_contexts, strategy):
"""Make an iterator for input provided via an input function.
Currently implements PER_WORKER mode, in which the `input_fn` is called
once on each worker.
TODO(priyag): Add other replication modes.
Args:
input_fn: Input function that returns a `tf.data.Dataset` object.
input_workers: an `InputWorkers` object.
input_contexts: A list of `InputContext` instances to be passed to call(s)
to `input_fn`. Length and order should match worker order in
`worker_device_pairs`.
strategy: a `tf.distribute.Strategy` object, used to run all-reduce to
handle last partial batch.
"""
assert isinstance(input_workers, input_lib.InputWorkers)
if input_workers.num_workers != len(input_contexts):
raise ValueError("Number of input workers (%d) is not same as number of "
"input_contexts (%d)" %
(input_workers.num_workers, len(input_contexts)))
iterators = []
for i, ctx in enumerate(input_contexts):
worker = input_workers.worker_devices[i]
with ops.device(worker):
result = input_fn(ctx)
devices = input_workers.compute_devices_for_worker(i)
if isinstance(result, data_types.DatasetV2):
iterator = _SingleWorkerDatasetIterator(result, worker, devices)
elif callable(result):
iterator = _SingleWorkerCallableIterator(result, worker, devices)
else:
raise ValueError(
"input_fn must return a tf.data.Dataset or a callable.")
iterators.append(iterator)
super(InputFunctionIterator, self).__init__(
input_workers,
iterators,
strategy,
cardinality=cardinality_lib.UNKNOWN,
enable_get_next_as_optional=False)
self._enable_get_next_as_optional = False
class _SingleWorkerDatasetIterator(input_lib._SingleWorkerDatasetIteratorBase): # pylint: disable=protected-access
"""Iterator for a single DistributedDatasetV1 instance."""
def _make_iterator(self):
"""Make appropriate iterator on the dataset."""
with ops.device(self._worker):
if self._options is not None:
self._iterator = multi_device_iterator_ops.MultiDeviceIterator(
self._dataset,
self._devices,
max_buffer_size=self._options.experimental_per_replica_buffer_size,
prefetch_buffer_size=self._options
.experimental_per_replica_buffer_size)
else:
self._iterator = multi_device_iterator_ops.MultiDeviceIterator(
self._dataset,
self._devices,
)
def initialize(self):
"""Initialize underlying iterator.
In eager execution, this simply recreates the underlying iterator.
In graph execution, it returns the initializer ops for the underlying
iterator.
Returns:
A list of any initializer ops that should be run.
"""
if ops.executing_eagerly_outside_functions():
self._iterator._eager_reset() # pylint: disable=protected-access
return []
else:
return [self._iterator.initializer]
@property
def output_classes(self):
return dataset_ops.get_legacy_output_classes(self._iterator)
@property
def output_shapes(self):
return dataset_ops.get_legacy_output_shapes(self._iterator)
@property
def output_types(self):
return dataset_ops.get_legacy_output_types(self._iterator)
class _SingleWorkerCallableIterator(object):
"""Iterator for a single tensor-returning callable."""
def __init__(self, fn, worker, devices):
self._fn = fn
self._worker = worker
self._devices = devices
def get_next(self, device, name=None):
"""Get next element for the given device from the callable."""
del device, name
with ops.device(self._worker):
return self._fn()
def get_next_as_list(self, name=None):
"""Get next element from the callable."""
del name
with ops.device(self._worker):
data_list = [self._fn() for _ in self._devices]
return data_list
def get_next_as_optional_list(self):
with ops.device(self._worker):
data_list = [
optional_ops.Optional.from_value(self._fn()) for _ in self._devices
]
return data_list
def initialize(self):
# TODO(petebu) Should this throw an exception instead?
return []
def _create_iterators_per_worker(worker_datasets, input_workers, options=None):
"""Create a multidevice iterator on each of the workers."""
assert isinstance(input_workers, input_lib.InputWorkers)
assert len(worker_datasets) == len(input_workers.worker_devices)
iterators = []
for i, worker in enumerate(input_workers.worker_devices):
with ops.device(worker):
worker_devices = input_workers.compute_devices_for_worker(i)
iterator = _SingleWorkerDatasetIterator(
worker_datasets[i], # pylint: disable=protected-access
worker,
worker_devices,
options)
iterators.append(iterator)
return iterators