test/distributed/test_multi_threaded_pg.py - platform/external/pytorch - Git at Google

 # Owner(s): ["oncall: distributed"]

 import operator
 import os
 import sys
 import threading
 from functools import reduce
 from unittest import skip, SkipTest

 import torch
 import torch.autograd
 import torch.distributed as dist
 from torch._C._distributed_c10d import ReduceOp


 if not dist.is_available():
     print("Distributed not available, skipping tests", file=sys.stderr)
     sys.exit(0)

 from torch.testing._internal.common_distributed import (
     MultiThreadedTestCase,
     skip_if_lt_x_gpu,
     spawn_threads_and_init_comms,
 )
 from torch.testing._internal.common_utils import IS_SANDCASTLE, run_tests, TestCase


 DEFAULT_WORLD_SIZE = 4


 class TestCollectivesWithWrapper(TestCase):
     @spawn_threads_and_init_comms(world_size=4)
     def test_broadcast_object_list(self):
         val = 99 if dist.get_rank() == 0 else None
         object_list = [val] * dist.get_world_size()

         dist.broadcast_object_list(object_list=object_list)
         self.assertEqual(99, object_list[0])

     def test_collective_error_on_rank_zero(self):
         @spawn_threads_and_init_comms(world_size=4)
         def _test_method(self):
             input_tensor = torch.ones(3, 3) * dist.get_rank()  # perform 1st all gather
             output_tensors = [
                 torch.empty_like(input_tensor) for _ in range(dist.get_world_size())
             ]
             dist.all_gather(output_tensors, input_tensor)

             if dist.get_rank() == 0:
                 raise AssertionError("Mimic real test failure.")  # fail on rank 0

             dist.all_gather(output_tensors, input_tensor)  # perform 2nd all gather

         with self.assertRaises(RuntimeError):
             _test_method(self)

     def test_collective_error_on_rank_non_zero(self):
         @spawn_threads_and_init_comms(world_size=4)
         def _test_method(self):
             input_tensor = torch.ones(3, 3) * dist.get_rank()  # perform 1st all gather
             output_tensors = [
                 torch.empty_like(input_tensor) for _ in range(dist.get_world_size())
             ]
             dist.all_gather(output_tensors, input_tensor)

             if dist.get_rank() == 1:
                 raise AssertionError("Mimic real test failure.")  # fail on rank 1

             dist.all_gather(output_tensors, input_tensor)  # perform 2nd all gather

         with self.assertRaises(RuntimeError):
             _test_method(self)

     def test_collective_error_on_rank_non_zero_all(self):
         @spawn_threads_and_init_comms(world_size=4)
         def _test_method(self):
             input_tensor = torch.ones(3, 3) * dist.get_rank()  # perform 1st all gather
             output_tensors = [
                 torch.empty_like(input_tensor) for _ in range(dist.get_world_size())
             ]
             dist.all_gather(output_tensors, input_tensor)

             if dist.get_rank() > 0:
                 raise AssertionError(
                     "Mimic real test failure."
                 )  # fail on all non-zero rank

             dist.all_gather(output_tensors, input_tensor)  # perform 2nd all gather

         with self.assertRaises(RuntimeError):
             _test_method(self)

     def test_skip(self):
         @spawn_threads_and_init_comms(world_size=4)
         @skip("check if skip exception can be captured correctly.")
         def _test_method(self):
             pass

         if not IS_SANDCASTLE:
             with self.assertRaises(SkipTest):
                 _test_method(self)

     @spawn_threads_and_init_comms(world_size=4)
     def test_all_to_all_single_tensor(self):
         rank = dist.get_rank()
         world_size = dist.get_world_size()
         send = torch.full((world_size, 2), rank)
         sizes = torch.ones(world_size, dtype=torch.int64)

         out = torch.zeros(world_size, 2, dtype=send.dtype)
         dist.all_to_all_single(out, send, sizes, sizes)
         self.assertEqual(out.tolist(), list(zip(range(world_size), range(world_size))))

     @spawn_threads_and_init_comms(world_size=4)
     def test_all_to_all_single_list(self):
         rank = dist.get_rank()
         world_size = dist.get_world_size()
         send = torch.full((world_size, 2), rank)
         sizes = [1] * world_size

         out = torch.zeros(world_size, 2, dtype=send.dtype)
         dist.all_to_all_single(out, send, sizes, sizes)
         self.assertEqual(out.tolist(), list(zip(range(world_size), range(world_size))))

     @spawn_threads_and_init_comms(world_size=4)
     def test_all_to_all_single_none(self):
         rank = dist.get_rank()
         world_size = dist.get_world_size()
         send = torch.full((world_size, 2), rank)

         out = torch.zeros(world_size, 2, dtype=send.dtype)
         dist.all_to_all_single(out, send)
         self.assertEqual(out.tolist(), list(zip(range(world_size), range(world_size))))


 class TestCollectivesWithBaseClass(MultiThreadedTestCase):
     @property
     def world_size(self):
         return 4

     def setUp(self):
         os.environ["TORCH_DIST_INIT_BARRIER"] = "1"
         super().setUp()
         self._spawn_threads()

     def tearDown(self):
         super().tearDown()
         os.environ["TORCH_DIST_INIT_BARRIER"] = "0"

     def test_allgather(self):
         input_tensor = torch.ones(3, 3) * dist.get_rank()
         output_tensors = [
             torch.empty_like(input_tensor) for _ in range(self.world_size)
         ]
         dist.all_gather(output_tensors, input_tensor)
         for rank, out_tensor in enumerate(output_tensors):
             self.assertEqual(out_tensor, torch.ones(3, 3) * rank)

     def test_broadcast(self):
         input_tensor = torch.ones(3, 3) * dist.get_rank()
         for rank in range(self.world_size):
             cloned_input = input_tensor.clone()
             dist.broadcast(cloned_input, src=rank)
             self.assertEqual(cloned_input, torch.ones(3, 3) * rank)

     def test_scatter(self):
         if dist.get_rank() == 0:
             scatter_list = [torch.ones(3, 3) * rank for rank in range(self.world_size)]
         else:
             scatter_list = None
         output_tensor = torch.empty(3, 3)

         dist.scatter(output_tensor, scatter_list)
         self.assertEqual(output_tensor, torch.ones(3, 3) * dist.get_rank())

     def test_reduce_scatter(self):
         to_reduce_scatter = [torch.ones(3, 3) * rank for rank in range(self.world_size)]
         output_tensor = torch.empty(3, 3)

         dist.reduce_scatter(output_tensor, to_reduce_scatter)
         expected_tensor = torch.ones(3, 3) * dist.get_rank() * self.world_size
         self.assertEqual(output_tensor, expected_tensor)

         output_tensor = torch.empty(3, 3)
         dist.reduce_scatter(output_tensor, to_reduce_scatter, op=dist.ReduceOp.AVG)
         expected_tensor = torch.ones(3, 3) * dist.get_rank()
         self.assertEqual(output_tensor, expected_tensor)

     def test_broadcast_object_list(self):
         val = 99 if dist.get_rank() == 0 else None
         object_list = [val] * dist.get_world_size()
         print(f"{dist.get_rank()} -> {dist.get_world_size()}")

         dist.broadcast_object_list(object_list=object_list)
         self.assertEqual(99, object_list[0])

     def test_all_reduce(self):
         output = torch.ones(3, 3) * dist.get_rank()
         dist.all_reduce(output)
         res_num = ((0 + self.world_size - 1) * self.world_size) / 2
         self.assertEqual(output, torch.ones(3, 3) * res_num)

     def test_all_to_all(self):
         rank = self.rank
         world_size = self.world_size
         input_tensor_list = [
             torch.ones(3, 3) * x
             for x in range(rank * world_size, (rank + 1) * world_size)
         ]
         output_tensor_list = [torch.empty_like(tensor) for tensor in input_tensor_list]
         dist.all_to_all(output_tensor_list, input_tensor_list)
         expected_tensor_list = [
             torch.ones(3, 3) * x
             for x in range(rank, world_size * world_size, world_size)
         ]
         self.assertEqual(expected_tensor_list, output_tensor_list)

     def test_all_reduce_ops(self):
         tensor = torch.tensor([dist.get_rank() + 1])
         dist.all_reduce(tensor, op=ReduceOp.PRODUCT)
         expected = reduce(operator.mul, range(1, self.world_size + 1))
         self.assertEqual(expected, tensor.item())

         tensor = torch.tensor([dist.get_rank() + 1])
         dist.all_reduce(tensor, op=ReduceOp.MIN)
         self.assertEqual(1, tensor.item())

         tensor = torch.tensor([dist.get_rank() + 1])
         dist.all_reduce(tensor, op=ReduceOp.MAX)
         self.assertEqual(self.world_size, tensor.item())

         tensor = torch.tensor([dist.get_rank() + 1])
         dist.all_reduce(tensor, op=ReduceOp.BAND)
         expected = reduce(operator.and_, range(1, self.world_size + 1))
         self.assertEqual(expected, tensor.item())

         tensor = torch.tensor([dist.get_rank() + 1])
         dist.all_reduce(tensor, op=ReduceOp.BOR)
         expected = reduce(operator.or_, range(1, self.world_size + 1))
         self.assertEqual(expected, tensor.item())

         tensor = torch.tensor([dist.get_rank() + 1])
         dist.all_reduce(tensor, op=ReduceOp.BXOR)
         expected = reduce(operator.xor, range(1, self.world_size + 1))
         self.assertEqual(expected, tensor.item())

     def test_assert_equal_on_rank(self):
         # RNG is shared across threads. So instead of asserting on all threads
         # we only assert on rank 0
         self_tensor = torch.rand(3, 3)
         rank_0_tensor = self_tensor.clone()
         dist.broadcast(rank_0_tensor, src=0)
         self.assertEqualOnRank(rank_0_tensor, self_tensor, rank=0)
         self.assertNotEqualOnRank(rank_0_tensor, self_tensor, rank=1)

     def test_subpg(self):
         subpg0 = dist.new_group([0, 1])
         subpg1 = dist.new_group([2, 3])
         current_rank = dist.get_rank()
         output = torch.ones(3, 3) * current_rank

         # call all_reduce on subpg0 and subpg1 concurrently
         if current_rank in [0, 1]:
             dist.all_reduce(output, group=subpg0)
         else:
             dist.all_reduce(output, group=subpg1)

         if current_rank in [0, 1]:
             self.assertEqual(output, torch.ones(3, 3) * 1)
         else:
             self.assertEqual(output, torch.ones(3, 3) * 5)

     def test_using_pg_from_another_thread(self):
         def stuff_in_other_thread(pg):
             x = torch.rand(4, requires_grad=True)
             dist.all_reduce(x, group=pg)

         t = threading.Thread(target=stuff_in_other_thread, args=(dist.group.WORLD,))
         t.start()
         t.join()

     def test_gather(self):
         if dist.get_rank() == 0:
             gather_list = [torch.empty(3, 3) for _ in range(self.world_size)]
         else:
             gather_list = None
         input_tensor = torch.ones(3, 3) * dist.get_rank()

         dist.gather(input_tensor, gather_list)
         if dist.get_rank() == 0:
             for i in range(self.world_size):
                 self.assertEqual(gather_list[i], torch.ones(3, 3) * i)

     def test_all_reduce_coalesced(self):
         t0 = torch.ones(3, 3) * dist.get_rank()
         t1 = torch.ones(3, 3) * dist.get_rank() * 2
         dist.all_reduce_coalesced([t0, t1])
         res_num = ((0 + self.world_size - 1) * self.world_size) / 2
         self.assertEqual(t0, torch.ones(3, 3) * res_num)
         self.assertEqual(t1, torch.ones(3, 3) * (res_num * 2))

     @skip_if_lt_x_gpu(1)
     def test_bwd_sees_fwd_pg(self):
         fwd_tid = threading.current_thread().ident

         class MyFunc(torch.autograd.Function):
             @staticmethod
             def forward(ctx, rank):
                 result = rank * 2

                 ctx.save_for_backward(result, rank)
                 assert int(rank.item()) == dist.get_rank()
                 return result

             @staticmethod
             def backward(ctx, grad_output):
                 result, rank = ctx.saved_tensors
                 bwd_tid = threading.current_thread().ident

                 self.assertEqual(
                     fwd_tid,
                     bwd_tid,
                     f"bwd not running in the same thread a fwd for rank {rank.item()}",
                 )
                 self.assertTrue(dist.is_initialized())
                 self.assertEqual(int(rank.item()), dist.get_rank())
                 dist.all_reduce(result)
                 self.assertEqual(int(result.item()), 12)  # (0 + 1 + 2 + 3) * 2

                 return grad_output * result

         x = torch.tensor(
             [dist.get_rank()], dtype=torch.float, device="cuda", requires_grad=True
         )
         x = MyFunc.apply(x)
         x.sum().backward()


 if __name__ == "__main__":
     run_tests()
	# Owner(s): ["oncall: distributed"]

	import operator
	import os
	import sys
	import threading
	from functools import reduce
	from unittest import skip, SkipTest

	import torch
	import torch.autograd
	import torch.distributed as dist
	from torch._C._distributed_c10d import ReduceOp


	if not dist.is_available():
	print("Distributed not available, skipping tests", file=sys.stderr)
	sys.exit(0)

	from torch.testing._internal.common_distributed import (
	MultiThreadedTestCase,
	skip_if_lt_x_gpu,
	spawn_threads_and_init_comms,
	)
	from torch.testing._internal.common_utils import IS_SANDCASTLE, run_tests, TestCase


	DEFAULT_WORLD_SIZE = 4


	class TestCollectivesWithWrapper(TestCase):
	@spawn_threads_and_init_comms(world_size=4)
	def test_broadcast_object_list(self):
	val = 99 if dist.get_rank() == 0 else None
	object_list = [val] * dist.get_world_size()

	dist.broadcast_object_list(object_list=object_list)
	self.assertEqual(99, object_list[0])

	def test_collective_error_on_rank_zero(self):
	@spawn_threads_and_init_comms(world_size=4)
	def _test_method(self):
	input_tensor = torch.ones(3, 3) * dist.get_rank() # perform 1st all gather
	output_tensors = [
	torch.empty_like(input_tensor) for _ in range(dist.get_world_size())
	]
	dist.all_gather(output_tensors, input_tensor)

	if dist.get_rank() == 0:
	raise AssertionError("Mimic real test failure.") # fail on rank 0

	dist.all_gather(output_tensors, input_tensor) # perform 2nd all gather

	with self.assertRaises(RuntimeError):
	_test_method(self)

	def test_collective_error_on_rank_non_zero(self):
	@spawn_threads_and_init_comms(world_size=4)
	def _test_method(self):
	input_tensor = torch.ones(3, 3) * dist.get_rank() # perform 1st all gather
	output_tensors = [
	torch.empty_like(input_tensor) for _ in range(dist.get_world_size())
	]
	dist.all_gather(output_tensors, input_tensor)

	if dist.get_rank() == 1:
	raise AssertionError("Mimic real test failure.") # fail on rank 1

	dist.all_gather(output_tensors, input_tensor) # perform 2nd all gather

	with self.assertRaises(RuntimeError):
	_test_method(self)

	def test_collective_error_on_rank_non_zero_all(self):
	@spawn_threads_and_init_comms(world_size=4)
	def _test_method(self):
	input_tensor = torch.ones(3, 3) * dist.get_rank() # perform 1st all gather
	output_tensors = [
	torch.empty_like(input_tensor) for _ in range(dist.get_world_size())
	]
	dist.all_gather(output_tensors, input_tensor)

	if dist.get_rank() > 0:
	raise AssertionError(
	"Mimic real test failure."
	) # fail on all non-zero rank

	dist.all_gather(output_tensors, input_tensor) # perform 2nd all gather

	with self.assertRaises(RuntimeError):
	_test_method(self)

	def test_skip(self):
	@spawn_threads_and_init_comms(world_size=4)
	@skip("check if skip exception can be captured correctly.")
	def _test_method(self):
	pass

	if not IS_SANDCASTLE:
	with self.assertRaises(SkipTest):
	_test_method(self)

	@spawn_threads_and_init_comms(world_size=4)
	def test_all_to_all_single_tensor(self):
	rank = dist.get_rank()
	world_size = dist.get_world_size()
	send = torch.full((world_size, 2), rank)
	sizes = torch.ones(world_size, dtype=torch.int64)

	out = torch.zeros(world_size, 2, dtype=send.dtype)
	dist.all_to_all_single(out, send, sizes, sizes)
	self.assertEqual(out.tolist(), list(zip(range(world_size), range(world_size))))

	@spawn_threads_and_init_comms(world_size=4)
	def test_all_to_all_single_list(self):
	rank = dist.get_rank()
	world_size = dist.get_world_size()
	send = torch.full((world_size, 2), rank)
	sizes = [1] * world_size

	out = torch.zeros(world_size, 2, dtype=send.dtype)
	dist.all_to_all_single(out, send, sizes, sizes)
	self.assertEqual(out.tolist(), list(zip(range(world_size), range(world_size))))

	@spawn_threads_and_init_comms(world_size=4)
	def test_all_to_all_single_none(self):
	rank = dist.get_rank()
	world_size = dist.get_world_size()
	send = torch.full((world_size, 2), rank)

	out = torch.zeros(world_size, 2, dtype=send.dtype)
	dist.all_to_all_single(out, send)
	self.assertEqual(out.tolist(), list(zip(range(world_size), range(world_size))))


	class TestCollectivesWithBaseClass(MultiThreadedTestCase):
	@property
	def world_size(self):
	return 4

	def setUp(self):
	os.environ["TORCH_DIST_INIT_BARRIER"] = "1"
	super().setUp()
	self._spawn_threads()

	def tearDown(self):
	super().tearDown()
	os.environ["TORCH_DIST_INIT_BARRIER"] = "0"

	def test_allgather(self):
	input_tensor = torch.ones(3, 3) * dist.get_rank()
	output_tensors = [
	torch.empty_like(input_tensor) for _ in range(self.world_size)
	]
	dist.all_gather(output_tensors, input_tensor)
	for rank, out_tensor in enumerate(output_tensors):
	self.assertEqual(out_tensor, torch.ones(3, 3) * rank)

	def test_broadcast(self):
	input_tensor = torch.ones(3, 3) * dist.get_rank()
	for rank in range(self.world_size):
	cloned_input = input_tensor.clone()
	dist.broadcast(cloned_input, src=rank)
	self.assertEqual(cloned_input, torch.ones(3, 3) * rank)

	def test_scatter(self):
	if dist.get_rank() == 0:
	scatter_list = [torch.ones(3, 3) * rank for rank in range(self.world_size)]
	else:
	scatter_list = None
	output_tensor = torch.empty(3, 3)

	dist.scatter(output_tensor, scatter_list)
	self.assertEqual(output_tensor, torch.ones(3, 3) * dist.get_rank())

	def test_reduce_scatter(self):
	to_reduce_scatter = [torch.ones(3, 3) * rank for rank in range(self.world_size)]
	output_tensor = torch.empty(3, 3)

	dist.reduce_scatter(output_tensor, to_reduce_scatter)
	expected_tensor = torch.ones(3, 3) * dist.get_rank() * self.world_size
	self.assertEqual(output_tensor, expected_tensor)

	output_tensor = torch.empty(3, 3)
	dist.reduce_scatter(output_tensor, to_reduce_scatter, op=dist.ReduceOp.AVG)
	expected_tensor = torch.ones(3, 3) * dist.get_rank()
	self.assertEqual(output_tensor, expected_tensor)

	def test_broadcast_object_list(self):
	val = 99 if dist.get_rank() == 0 else None
	object_list = [val] * dist.get_world_size()
	print(f"{dist.get_rank()} -> {dist.get_world_size()}")

	dist.broadcast_object_list(object_list=object_list)
	self.assertEqual(99, object_list[0])

	def test_all_reduce(self):
	output = torch.ones(3, 3) * dist.get_rank()
	dist.all_reduce(output)
	res_num = ((0 + self.world_size - 1) * self.world_size) / 2
	self.assertEqual(output, torch.ones(3, 3) * res_num)

	def test_all_to_all(self):
	rank = self.rank
	world_size = self.world_size
	input_tensor_list = [
	torch.ones(3, 3) * x
	for x in range(rank * world_size, (rank + 1) * world_size)
	]
	output_tensor_list = [torch.empty_like(tensor) for tensor in input_tensor_list]
	dist.all_to_all(output_tensor_list, input_tensor_list)
	expected_tensor_list = [
	torch.ones(3, 3) * x
	for x in range(rank, world_size * world_size, world_size)
	]
	self.assertEqual(expected_tensor_list, output_tensor_list)

	def test_all_reduce_ops(self):
	tensor = torch.tensor([dist.get_rank() + 1])
	dist.all_reduce(tensor, op=ReduceOp.PRODUCT)
	expected = reduce(operator.mul, range(1, self.world_size + 1))
	self.assertEqual(expected, tensor.item())

	tensor = torch.tensor([dist.get_rank() + 1])
	dist.all_reduce(tensor, op=ReduceOp.MIN)
	self.assertEqual(1, tensor.item())

	tensor = torch.tensor([dist.get_rank() + 1])
	dist.all_reduce(tensor, op=ReduceOp.MAX)
	self.assertEqual(self.world_size, tensor.item())

	tensor = torch.tensor([dist.get_rank() + 1])
	dist.all_reduce(tensor, op=ReduceOp.BAND)
	expected = reduce(operator.and_, range(1, self.world_size + 1))
	self.assertEqual(expected, tensor.item())

	tensor = torch.tensor([dist.get_rank() + 1])
	dist.all_reduce(tensor, op=ReduceOp.BOR)
	expected = reduce(operator.or_, range(1, self.world_size + 1))
	self.assertEqual(expected, tensor.item())

	tensor = torch.tensor([dist.get_rank() + 1])
	dist.all_reduce(tensor, op=ReduceOp.BXOR)
	expected = reduce(operator.xor, range(1, self.world_size + 1))
	self.assertEqual(expected, tensor.item())

	def test_assert_equal_on_rank(self):
	# RNG is shared across threads. So instead of asserting on all threads
	# we only assert on rank 0
	self_tensor = torch.rand(3, 3)
	rank_0_tensor = self_tensor.clone()
	dist.broadcast(rank_0_tensor, src=0)
	self.assertEqualOnRank(rank_0_tensor, self_tensor, rank=0)
	self.assertNotEqualOnRank(rank_0_tensor, self_tensor, rank=1)

	def test_subpg(self):
	subpg0 = dist.new_group([0, 1])
	subpg1 = dist.new_group([2, 3])
	current_rank = dist.get_rank()
	output = torch.ones(3, 3) * current_rank

	# call all_reduce on subpg0 and subpg1 concurrently
	if current_rank in [0, 1]:
	dist.all_reduce(output, group=subpg0)
	else:
	dist.all_reduce(output, group=subpg1)

	if current_rank in [0, 1]:
	self.assertEqual(output, torch.ones(3, 3) * 1)
	else:
	self.assertEqual(output, torch.ones(3, 3) * 5)

	def test_using_pg_from_another_thread(self):
	def stuff_in_other_thread(pg):
	x = torch.rand(4, requires_grad=True)
	dist.all_reduce(x, group=pg)

	t = threading.Thread(target=stuff_in_other_thread, args=(dist.group.WORLD,))
	t.start()
	t.join()

	def test_gather(self):
	if dist.get_rank() == 0:
	gather_list = [torch.empty(3, 3) for _ in range(self.world_size)]
	else:
	gather_list = None
	input_tensor = torch.ones(3, 3) * dist.get_rank()

	dist.gather(input_tensor, gather_list)
	if dist.get_rank() == 0:
	for i in range(self.world_size):
	self.assertEqual(gather_list[i], torch.ones(3, 3) * i)

	def test_all_reduce_coalesced(self):
	t0 = torch.ones(3, 3) * dist.get_rank()
	t1 = torch.ones(3, 3) * dist.get_rank() * 2
	dist.all_reduce_coalesced([t0, t1])
	res_num = ((0 + self.world_size - 1) * self.world_size) / 2
	self.assertEqual(t0, torch.ones(3, 3) * res_num)
	self.assertEqual(t1, torch.ones(3, 3) * (res_num * 2))

	@skip_if_lt_x_gpu(1)
	def test_bwd_sees_fwd_pg(self):
	fwd_tid = threading.current_thread().ident

	class MyFunc(torch.autograd.Function):
	@staticmethod
	def forward(ctx, rank):
	result = rank * 2

	ctx.save_for_backward(result, rank)
	assert int(rank.item()) == dist.get_rank()
	return result

	@staticmethod
	def backward(ctx, grad_output):
	result, rank = ctx.saved_tensors
	bwd_tid = threading.current_thread().ident

	self.assertEqual(
	fwd_tid,
	bwd_tid,
	f"bwd not running in the same thread a fwd for rank {rank.item()}",
	)
	self.assertTrue(dist.is_initialized())
	self.assertEqual(int(rank.item()), dist.get_rank())
	dist.all_reduce(result)
	self.assertEqual(int(result.item()), 12) # (0 + 1 + 2 + 3) * 2

	return grad_output * result

	x = torch.tensor(
	[dist.get_rank()], dtype=torch.float, device="cuda", requires_grad=True
	)
	x = MyFunc.apply(x)
	x.sum().backward()


	if __name__ == "__main__":
	run_tests()