Add gpu gloo tests (pytorch#1668)

Summary:

Add GPU tests for gloo.

Reviewed By: gnahzg

Differential Revision: D53207076

torchrec/distributed/tests/test_model_parallel_gloo_gpu.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torchrec.distributed.test_utils.test_model_parallel import ModelParallelBase
+
+# GPU tests for Gloo.
+
+
+class ModelParallelTestGloo(ModelParallelBase):
+    def setUp(self, backend: str = "gloo") -> None:
+        super().setUp(backend=backend)

torchrec/distributed/tests/test_model_parallel_gloo_gpu_single_rank.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torchrec.distributed.test_utils.test_model_parallel_base import (
+    ModelParallelSparseOnlyBase,
+    ModelParallelStateDictBase,
+)
+
+# Single rank GPU tests for Gloo.
+
+
+class ModelParallelStateDictTestGloo(ModelParallelStateDictBase):
+    def setUp(self, backend: str = "gloo") -> None:
+        super().setUp(backend=backend)
+
+
+class ModelParallelSparseOnlyTestGloo(ModelParallelSparseOnlyBase):
+    def setUp(self, backend: str = "gloo") -> None:
+        super().setUp(backend=backend)