def test_bluefog_allreduce_cpu(self):
"""Test on CPU that the allreduce correctly sums 1D, 2D, 3D tensors."""
size = bf.size()
dtypes = [tf.int32, tf.int64, tf.float32, tf.float64]
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
with tf.device("/cpu:0"):
tensor = random_uniform([17] * dim, -100, 100, dtype=dtype)
summed = bf.allreduce(tensor, average=False)
multiplied = tensor * size
max_difference = tf.reduce_max(tf.abs(summed - multiplied))
# Threshold for floating point equality depends on number of
# ranks, since we're comparing against precise multiplication.
if size <= 3 or dtype in [tf.int32, tf.int64]:
threshold = 0
elif size < 10:
threshold = 1e-4
elif size < 15:
threshold = 5e-4
else:
break
diff = self.evaluate(max_difference)
self.assertTrue(diff <= threshold,
"bf.allreduce produces incorrect results")
def test_bluefog_size(self):
"""Test that the size returned by bf.size() is correct."""
_, true_size = mpi_env_rank_and_size()
bf.init()
size = bf.size()
# print("Size: ", true_size, size)
assert true_size == size
def test_bluefog_allgather(self):
"""Test that the allgather correctly gathers 1D, 2D, 3D tensors."""
rank = bf.rank()
size = bf.size()
dtypes = [tf.int32, tf.int64, tf.float32, tf.float64, tf.bool]
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = tf.ones([17] * dim) * rank
if dtype == tf.bool:
tensor = tensor % 2
tensor = tf.cast(tensor, dtype=dtype)
gathered = bf.allgather(tensor)
gathered_tensor = self.evaluate(gathered)
self.assertEqual(list(gathered_tensor.shape),
[17 * size] + [17] * (dim - 1))
for i in range(size):
rank_tensor = tf.slice(gathered_tensor,
[i * 17] + [0] * (dim - 1),
[17] + [-1] * (dim - 1))
self.assertEqual(list(rank_tensor.shape), [17] * dim)
# tf.equal() does not support tf.uint16 as of TensorFlow 1.2,
# so need to cast rank_tensor to tf.int32.
if dtype == tf.bool:
value = i % 2
else:
value = i
self.assertTrue(
self.evaluate(
tf.reduce_all(
tf.equal(tf.cast(rank_tensor, tf.int32), value))),
"bf.allgather produces incorrect gathered tensor")
def test_bluefog_broadcast(self):
"""Test that the broadcast correctly broadcasts 1D, 2D, 3D tensors."""
rank = bf.rank()
size = bf.size()
# This test does not apply if there is only one worker.
if size == 1:
return
dtypes = [tf.int32, tf.int64, tf.float32, tf.float64, tf.bool]
dims = [1, 2, 3]
root_ranks = list(range(size))
for dtype, dim, root_rank in itertools.product(dtypes, dims,
root_ranks):
tensor = tf.ones([17] * dim) * rank
root_tensor = tf.ones([17] * dim) * root_rank
if dtype == tf.bool:
tensor = tensor % 2
root_tensor = root_tensor % 2
tensor = tf.cast(tensor, dtype=dtype)
root_tensor = tf.cast(root_tensor, dtype=dtype)
broadcasted_tensor = bf.broadcast(tensor, root_rank)
self.assertTrue(
self.evaluate(
tf.reduce_all(
tf.equal(tf.cast(root_tensor, tf.int32),
tf.cast(broadcasted_tensor, tf.int32)))),
"bf.broadcast produces incorrect broadcasted tensor")
def test_bluefog_allreduce_grad_cpu(self):
"""Test the correctness of the allreduce gradient on CPU."""
size = bf.size()
# As of TensorFlow v1.9, gradients are not supported on
# integer tensors
dtypes = [tf.float32, tf.float64]
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
with tf.device("/cpu:0"):
if _executing_eagerly():
tensor = self.tfe.Variable(
random_uniform([5] * dim, -100, 100, dtype=dtype))
with tf.GradientTape() as tape:
summed = bf.allreduce(tensor, average=False)
else:
tensor = random_uniform([5] * dim, -100, 100, dtype=dtype)
summed = bf.allreduce(tensor, average=False)
grad_ys = tf.ones([5] * dim)
if _executing_eagerly():
grad_out = tape.gradient(summed, tensor, grad_ys)
else:
grad = tf.gradients(summed, tensor, grad_ys)[0]
grad_out = self.evaluate(grad)
expected = np.ones([5] * dim) * size
err = np.linalg.norm(expected - grad_out)
self.assertLess(
err, 0.00000001, "gradient %s differs from expected %s, "
"error: %s" % (grad_out, expected, str(err)))
def test_bluefog_allreduce_gpu(self):
"""Test that the allreduce works on GPUs."""
# Only do this test if there are GPUs available.
if not tf.test.is_gpu_available(cuda_only=True):
return
local_rank = bf.local_rank()
size = bf.size()
dtypes = [tf.int32, tf.int64, tf.float32, tf.float64]
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
with tf.device("/gpu:%d" % local_rank):
tensor = random_uniform([17] * dim, -100, 100, dtype=dtype)
summed = bf.allreduce(tensor, average=False)
multiplied = tensor * size
max_difference = tf.reduce_max(tf.abs(summed - multiplied))
# Threshold for floating point equality depends on number of
# ranks, since we're comparing against precise multiplication.
if size <= 3 or dtype in [tf.int32, tf.int64]:
threshold = 0
elif size < 10:
threshold = 1e-4
elif size < 15:
threshold = 5e-4
else:
return
diff = self.evaluate(max_difference)
self.assertTrue(diff <= threshold,
"bf.allreduce on GPU produces incorrect results")
def test_bluefog_broadcast_grad_cpu(self):
"""Test the correctness of the broadcast gradient on CPU."""
rank = bf.rank()
size = bf.size()
# This test does not apply if there is only one worker.
if size == 1:
return
# As of TensorFlow v1.9, gradients are not supported on
# integer tensors
dtypes = [tf.float32, tf.float64]
dims = [1, 2, 3]
root_ranks = list(range(size))
for dtype, dim, root_rank in itertools.product(dtypes, dims,
root_ranks):
if _executing_eagerly():
tensor = self.tfe.Variable(tf.ones([5] * dim) * rank)
else:
tensor = tf.ones([5] * dim) * rank
if dtype == tf.bool:
tensor = tensor % 2
if _executing_eagerly():
with tf.GradientTape() as tape:
tensor = tf.cast(tensor, dtype=dtype)
broadcasted_tensor = bf.broadcast(tensor, root_rank)
with tf.device("/cpu:0"):
grad_out = tape.gradient(broadcasted_tensor, tensor)
else:
tensor = tf.cast(tensor, dtype=dtype)
broadcasted_tensor = bf.broadcast(tensor, root_rank)
grad_ys = tf.ones([5] * dim)
with tf.device("/cpu:0"):
grad = tf.gradients(broadcasted_tensor, tensor, grad_ys)[0]
grad_out = self.evaluate(grad)
c = size if rank == root_rank else 0
expected = np.ones([5] * dim) * c
err = np.linalg.norm(expected - grad_out)
self.assertLess(
err, 0.00000001, "gradient %s differs from expected %s, "
"error: %s" % (grad_out, expected, str(err)))