# 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.
# ==============================================================================
"""Tests for Collective Operations."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import time
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import kernels
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import collective_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
class CollectiveOpTest(test.TestCase):
def setUp(self):
context._reset_context() # pylint: disable=protected-access
super(CollectiveOpTest, self).setUp()
def _testCollectiveReduce(self,
inputs,
expected,
set_graph_key,
communication_hint='auto',
fp16=False,
instance_key=1,
merge_op='Add',
final_op='Div',
timeout=0,
reported_group_size=None):
group_key = 1
group_size = len(inputs)
if reported_group_size is None:
reported_group_size = group_size
device_type = 'CPU'
config = config_pb2.ConfigProto(device_count={device_type: group_size})
devices = ['/{}:{}'.format(device_type, i) for i in range(group_size)]
with self.session(config=config) as sess:
colred = []
for i in range(group_size):
with ops.device(devices[i]):
tensor = constant_op.constant(inputs[i], dtype=(
dtypes.float16 if fp16 else dtypes.float32))
colred.append(
collective_ops.all_reduce(
tensor,
reported_group_size,
group_key,
instance_key,
merge_op,
final_op,
communication_hint=communication_hint,
timeout=timeout))
run_options = config_pb2.RunOptions()
if set_graph_key:
run_options.experimental.collective_graph_key = 1
results = sess.run(colred, options=run_options)
tolerance = 1e-3 if fp16 else 1e-5
for i in range(group_size):
logging.info('i {} result {} expected {}'.format(i, results[i], expected))
self.assertAllClose(results[i], expected, rtol=tolerance, atol=tolerance)
def _testMultipleConcurrentCollectiveReduce(self, t0, t1, expected):
group_key = 1
group_size = 2
num_instances = 2
all_reduces = []
config = config_pb2.ConfigProto(device_count={'CPU': group_size})
config.experimental.collective_deterministic_sequential_execution = True
with self.session(config=config) as sess:
for cpu in range(group_size):
with ops.device('/CPU:%d' % cpu):
in_tensor = constant_op.constant(t0 if cpu == 0 else t1)
for instance in range(num_instances):
all_reduces.append(collective_ops.all_reduce(
in_tensor, group_size, group_key, instance, 'Add', 'Div'))
results = sess.run(all_reduces)
for i in range(group_size * num_instances):
self.assertAllClose(results[i], expected, rtol=1e-5, atol=1e-5)
def testCollectiveReduce(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
set_graph_key=True)
def testCollectiveAutoGraphKey(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
set_graph_key=False)
def testFp16Reduce(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
set_graph_key=True,
fp16=True)
def testCollectiveMultipleConcurrentReduce(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testMultipleConcurrentCollectiveReduce(
[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2])
def testCollectiveTimeoutV1(self):
timeout = 4.5
kwargs = dict(
inputs=[[i + j + 0.1 for i in range(8)] for j in range(3)],
expected=[1 + i + 0.1 for i in range(8)],
set_graph_key=True,
timeout=timeout)
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(**kwargs)
start_time = time.time()
with ops.Graph().as_default():
with self.assertRaisesRegex(
errors.DeadlineExceededError,
'Collective has timed out waiting for other workers'):
self._testCollectiveReduce(
reported_group_size=len(kwargs['inputs']) + 1, **kwargs)
elapsed = time.time() - start_time
self.assertAllGreaterEqual(elapsed, timeout)
def testNcclHintFallbackToRingReduce(self):
"""Tests that setting `communication_hint=nccl` works on non-GPU builds."""
if kernels.get_registered_kernels_for_op('NcclAllReduce'):
self.skipTest('Run only on non-GPU environments')
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(
inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
set_graph_key=False,
communication_hint='nccl')
def _testWhile(self, num_vars, num_iterations, key_base):
group_size = 2
group_key = 1
instances = [(key_base + i) for i in range(num_vars)]
devices = ['CPU:{}'.format(i) for i in range(group_size)]
config = config_pb2.ConfigProto(device_count={'CPU': group_size})
rewrite_options = config.graph_options.rewrite_options
rewrite_options.scoped_allocator_optimization = (
rewriter_config_pb2.RewriterConfig.ON)
del rewrite_options.scoped_allocator_opts.enable_op[:]
rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
with self.session(config=config) as sess:
loop_vars = []
for device in devices:
with ops.device(device):
loop_vars.append(
[variables.VariableV1((1 << i) * 1.) for i in range(num_vars)])
# This variable controls number of iterations.
loop_vars.append(variables.VariableV1(0.))
def loop_body(dev0_tensors, dev1_tensors, loop_tensor):
return_ops = []
for i in range(len(devices)):
device = devices[i]
device_tensors = dev0_tensors if i == 0 else dev1_tensors
with ops.device(device):
device_collectives = []
for j in range(num_vars):
# NOTE(ayushd): we need the `cast` here to ensure that the input
# to `all_reduce` has an explicit device string. We don't use
# `identity` because `cast` is more resilient to getting optimized
# away by various optimization passes.
input_tensor = math_ops.cast(device_tensors[j], dtypes.float16)
collective_op = collective_ops.all_reduce(
input_tensor, group_size, group_key, instances[j],
'Add', 'Id')
output_tensor = math_ops.cast(collective_op, dtypes.float32)
device_collectives.append(output_tensor)
return_ops.append(device_collectives)
return_ops.append(math_ops.add(loop_tensor, 1.))
return return_ops
# Run until last variable exceeds number of iterations.
loop_cond = lambda d0, d1, i: math_ops.less(i, num_iterations)
sess.run(variables.global_variables_initializer())
results = sess.run(control_flow_ops.while_loop(loop_cond, loop_body,
loop_vars))
self.assertEqual(results[:-1], [
[((1 << (num_iterations + v)) * 1.) for v in range(num_vars)]
for _ in range(group_size)])
def testSimpleWhile(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testWhile(num_vars=1, num_iterations=4, key_base=20)
def testWhileMultipleAllReduce(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testWhile(num_vars=2, num_iterations=4, key_base=20)
def testWhileWithScopedAllocator(self):
group_size = 2
group_key = 1
instance_key0 = 1
instance_key1 = 2
config = config_pb2.ConfigProto(device_count={'CPU': group_size})
rewrite_options = config.graph_options.rewrite_options
rewrite_options.scoped_allocator_optimization = (
rewriter_config_pb2.RewriterConfig.ON)
del rewrite_options.scoped_allocator_opts.enable_op[:]
rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
with self.session(config=config) as sess:
run_ops = []
for i in range(group_size):
with ops.device('CPU:%d' % i):
constant = constant_op.constant(0.)
cond = lambda i: math_ops.less(i, 10.)
body = lambda i: math_ops.add(i, 1.)
input0 = control_flow_ops.while_loop(cond, body, [constant])
input1 = math_ops.add(constant, 5)
colred0 = collective_ops.all_reduce(input0, group_size, group_key,
instance_key0, 'Add', 'Id')
colred1 = collective_ops.all_reduce(input1, group_size, group_key,
instance_key1, 'Add', 'Id')
run_ops.append(math_ops.add_n([colred0, colred1]))
results = sess.run(run_ops)
self.assertEqual(results, [30., 30.])
def testCollectiveReduceScalar(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(inputs=[0.1, 0.3], expected=0.2,
set_graph_key=True)
def testCollectiveReduceMaximum(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(
inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
expected=[10., 20., 30., 40., 50.],
set_graph_key=True,
instance_key=30,
merge_op='Max',
final_op='Id')
def testCollectiveReduceMinimum(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveReduce(
inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
expected=[1., 2., 3., 4., 5.],
set_graph_key=True,
instance_key=40,
merge_op='Min',
final_op='Id')
def _testCollectiveBroadcast(self, in_val):
group_key = 1
instance_key = 1
with self.session(
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
with ops.device('/CPU:0'):
in0 = constant_op.constant(in_val)
out0 = collective_ops.broadcast_send(in0, in0.shape, in0.dtype,
2, group_key, instance_key)
with ops.device('/CPU:1'):
c1 = constant_op.constant(in_val)
out1 = collective_ops.broadcast_recv(c1.shape, c1.dtype,
2, group_key, instance_key)
run_options = config_pb2.RunOptions()
run_options.experimental.collective_graph_key = 1
results = sess.run([out0, out1], options=run_options)
self.assertAllClose(results[0], in_val, rtol=1e-5, atol=1e-5)
self.assertAllClose(results[1], in_val, rtol=1e-5, atol=1e-5)
def testCollectiveBroadcast(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveBroadcast([0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1])
def testCollectiveBroadcastBool(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveBroadcast([True, False])
def _testCollectiveGather(self, t0, t1, expected, set_graph_key):
group_key = 1
instance_key = 1
with self.session(
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
with ops.device('/CPU:0'):
in0 = constant_op.constant(t0)
c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
with ops.device('/CPU:1'):
in1 = constant_op.constant(t1)
c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
run_options = config_pb2.RunOptions()
if set_graph_key:
run_options.experimental.collective_graph_key = 1
results = sess.run([c0, c1], options=run_options)
self.assertAllClose(results[0], expected, rtol=1e-5, atol=1e-5)
self.assertAllClose(results[1], expected, rtol=1e-5, atol=1e-5)
def testCollectiveGather(self):
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
self._testCollectiveGather([0, 1, 2, 3, 4, 5, 6, 7],
[10, 11, 12, 13, 14, 15, 16, 17],
[0, 1, 2, 3, 4, 5, 6, 7,
10, 11, 12, 13, 14, 15, 16, 17],
True)
self._testCollectiveGather([[0, 1, 2, 3], [4, 5, 6, 7]],
[[10, 11, 12, 13], [14, 15, 16, 17]],
[[0, 1, 2, 3], [4, 5, 6, 7],
[10, 11, 12, 13], [14, 15, 16, 17]],
True)
self._testCollectiveGather([[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
[[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
[[[0, 1], [2, 3]], [[4, 5], [6, 7]],
[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
True)
def testCollectiveGatherShapeMismatch(self):
group_key = 1
instance_key = 1
t0 = [1, 2, 3, 4]
t1 = [5, 6, 7, 8]
t2 = [9, 10]
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
with self.session(
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
with ops.device('/CPU:0'):
in0 = constant_op.constant(t0)
c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
with ops.device('/CPU:1'):
in1 = constant_op.constant(t1)
in2 = constant_op.constant(t2)
c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
c2 = collective_ops.all_gather(in2, 2, group_key, instance_key)
run_options = config_pb2.RunOptions()
run_options.experimental.collective_graph_key = 1
sess.run([c0, c1], options=run_options)
with self.assertRaisesRegex(errors.InvalidArgumentError,
'Shape mismatch'):
sess.run([c0, c2], options=run_options)
def testCollectiveGatherShapeMismatchAcrossDevices(self):
group_key = 1
instance_key = 1
t0 = [1, 2, 3, 4]
t1 = [5, 6]
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
with self.session(
config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
with ops.device('/CPU:0'):
in0 = constant_op.constant(t0)
c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
with ops.device('/CPU:1'):
in1 = constant_op.constant(t1)
c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
run_options = config_pb2.RunOptions()
run_options.experimental.collective_graph_key = 1
with self.assertRaisesRegex(errors.InvalidArgumentError,
'Shape mismatch'):
sess.run([c0, c1], options=run_options)
def testCollectiveGatherPolymorphicShape(self):
t0 = [0, 1, 2, 3, 4, 5, 6, 7]
t1 = [10, 11, 12, 13, 14, 15, 16, 17]
group_size = 2
group_key = 1
instance_key = 123
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
with self.session(
config=config_pb2.ConfigProto(
device_count={'CPU': group_size})) as sess:
with ops.device('/CPU:0'):
in0 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
c0 = collective_ops.all_gather(in0, group_size, group_key,
instance_key)
with ops.device('/CPU:1'):
in1 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
c1 = collective_ops.all_gather(in1, group_size, group_key,
instance_key)
results = sess.run([c0, c1], feed_dict={in0: t0, in1: t1})
results_ = sess.run([c0, c1], feed_dict={in0: t0[1:], in1: t1[1:]})
expected_output = [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 15, 16, 17]
self.assertAllClose(results[0], expected_output, rtol=1e-5, atol=1e-5)
self.assertAllClose(results[1], expected_output, rtol=1e-5, atol=1e-5)
expected_output_ = [1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17]
self.assertAllClose(results_[0], expected_output_, rtol=1e-5, atol=1e-5)
self.assertAllClose(results_[1], expected_output_, rtol=1e-5, atol=1e-5)
@test_util.run_v2_only
@test_util.disable_tfrt(
'b/177270918: TFRT has dead lock when executing collective ops.')
def testCollectiveGroupSizeMismatch(self):
cpus = config.list_physical_devices('CPU')
self.assertEqual(len(cpus), 1)
config.set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
context.ensure_initialized()
@def_function.function
def run_all_reduce():
group_key = 10
instance_key = 20
t0 = [1, 2, 3, 4]
t1 = [5, 6, 7, 8]
with ops.device('/CPU:0'):
in0 = constant_op.constant(t0)
c0 = collective_ops.all_reduce(
in0, group_size=2, group_key=group_key, instance_key=instance_key,
merge_op='Add', final_op='Id')
with ops.device('/CPU:1'):
in1 = constant_op.constant(t1)
c1 = collective_ops.all_reduce(
in1, group_size=3, group_key=group_key, instance_key=instance_key,
merge_op='Add', final_op='Id')
return c0, c1
with self.assertRaisesRegex(errors.InternalError,
'but that group has size'):
run_all_reduce()
@test_util.run_v2_only
def testCollectiveTensorsHaveNoDeviceSpecified(self):
cpus = config.list_physical_devices('CPU')
self.assertEqual(len(cpus), 1)
config.set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
context.ensure_initialized()
group_size = 2
group_key = 1
instance_key = 1
@def_function.function
def fn(all_args):
results = []
# The inputs have no devices set. This is expected to be a trace-time
# check only.
self.assertEqual(all_args[0].device, '')
self.assertEqual(all_args[1].device, '')
with ops.device('/CPU:0'):
results.append(
collective_ops.all_reduce(all_args[0], group_size, group_key,
instance_key, 'Add', 'Div'))
with ops.device('/CPU:1'):
results.append(
collective_ops.all_reduce(all_args[1], group_size, group_key,
instance_key, 'Add', 'Div'))
return results
with ops.device('/CPU:0'):
in0 = constant_op.constant(1)
with ops.device('/CPU:1'):
in1 = constant_op.constant(3)
result = fn([in0, in1])
self.assertAllClose(result, [2, 2])
def testConstantWithScopedAllocator(self):
group_size = 2
group_key = 1
instance_key1 = 1
instance_key2 = 2
graph_options = config_pb2.GraphOptions(
optimizer_options=config_pb2.OptimizerOptions(do_constant_folding=True))
cfg = config_pb2.ConfigProto(device_count={'CPU': group_size},
graph_options=graph_options)
rewrite_options = cfg.graph_options.rewrite_options
rewrite_options.scoped_allocator_optimization = (
rewriter_config_pb2.RewriterConfig.ON)
del rewrite_options.scoped_allocator_opts.enable_op[:]
rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
# Tests that execute collectives need to be enclosed in graph or tf.function
with ops.Graph().as_default():
with self.session(config=cfg) as sess:
run_ops = []
for i in range(group_size):
with ops.device('CPU:%d' % i):
constant = constant_op.constant(i + 1.)
input_tensor1 = array_ops.identity(constant)
input_tensor2 = array_ops.identity(constant)
reduced_tensor1 = collective_ops.all_reduce(
input_tensor1, group_size, group_key, instance_key1, 'Add',
'Id')
reduced_tensor2 = collective_ops.all_reduce(
input_tensor2, group_size, group_key, instance_key2, 'Add',
'Id')
run_ops.append(array_ops.identity(reduced_tensor1))
run_ops.append(array_ops.identity(reduced_tensor2))
results = sess.run(run_ops)
self.assertEqual(results, [3., 3., 3., 3.])
if __name__ == '__main__':
test.main()