More element types in AllGather and AllToAll (#16941)

Two things done in this PR.
- [2nd commit] More tensor element types are supported because in
distributed computation, we need to re-shard tensors in many different
types.
- [1st commit] We now specify opset version in test models. Without this
change, those models will have opset=20 with latest ONNX and results
test errors.
- [3rd commit] Tests are modified to test `AllGather` and `AllToAll` for
boolean tensors. Several graph patterns are tried for tests. We found
that `int64_tensor -> Cast -> bool_tensor -> AllToAll -> bool_tensor ->
Cast -> int64_tensor` always generate random results. My guess is that
`AllToAll` needs to synchronize all GPUs before calling `ncclSend` and
`ncclRecv` since `AllGather` doesn't hit this problem. For reproducing
the error, search for `TODO` in this PR. Note that this PR doesn't fix
it.

onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc

@@ -14,6 +14,9 @@ namespace cuda {
 static ncclDataType_t GetNcclDataType(onnxruntime::MLDataType type) {
   if (type == DataTypeImpl::GetType<uint8_t>()) {
     return ncclUint8;
+  } else if (type == DataTypeImpl::GetType<bool>()) {
+    // CUDA bool is 8-bit large.
+    return ncclUint8;
   } else if (type == DataTypeImpl::GetType<int8_t>()) {
     return ncclInt8;
   } else if (type == DataTypeImpl::GetType<int32_t>()) {
@@ -226,7 +229,7 @@ ONNX_OPERATOR_KERNEL_EX(
     kCudaExecutionProvider,
     (*KernelDefBuilder::Create())
         .AllocateInputsContiguously()
-        .TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes()),
+        .TypeConstraint("T", DataTypeImpl::AllTensorTypes()),
     AllGather);
 
 ONNX_OPERATOR_KERNEL_EX(
@@ -237,7 +240,7 @@ ONNX_OPERATOR_KERNEL_EX(
     (*KernelDefBuilder::Create())
         .VariadicAlias(0, 0)  // outputs and inputs are mapped one to one
         .AllocateInputsContiguously()
-        .TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes()),
+        .TypeConstraint("T", DataTypeImpl::AllTensorTypes()),
     AllToAll);
 
 }  // namespace cuda

onnxruntime/core/graph/contrib_ops/collective_defs.cc

@@ -42,8 +42,8 @@ void RegisterCollectiveOps() {
       .Output(0, "output", "gathered tensors", "T", OpSchema::Variadic)
       .TypeConstraint(
           "T",
-          {"tensor(float16)", "tensor(float)", "tensor(double)"},
-          "Constrain to float, float16 and double tensors.")
+          {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(int64)", "tensor(bool)"},
+          "Constrain to bool, float, float16 and double tensors.")
       .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
         auto group_size = getAttribute(ctx, "group_size", 1);
         auto axis = getAttribute(ctx, "axis", 0);
@@ -74,8 +74,8 @@ void RegisterCollectiveOps() {
       .Output(0, "output", "collected tensors", "T", OpSchema::Variadic)
       .TypeConstraint(
           "T",
-          {"tensor(float16)", "tensor(float)", "tensor(double)"},
-          "Constrain to float, float16 and double tensors.")
+          {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(int64)", "tensor(bool)"},
+          "Constrain to bool, float, float16 and double tensors.")
       .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
         propagateShapeAndTypeFromFirstInput(ctx);
       });

onnxruntime/test/python/onnxruntime_test_collective.py

@@ -3,17 +3,30 @@
 import unittest
 
 import numpy as np
-import onnx  # noqa: F401
 from mpi4py import MPI
-from onnx import AttributeProto, GraphProto, TensorProto, helper  # noqa: F401
+from onnx import TensorProto, helper
 
 import onnxruntime as ort
 
 
 class ORTBertPretrainTest(unittest.TestCase):
-    def _create_allreduce_ut_model(self, shape):
-        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)  # noqa: N806
-        Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, shape)  # noqa: N806
+    @staticmethod
+    def _create_model_with_opsets(
+        graph_def,
+    ):
+        opset_imports = [
+            helper.make_operatorsetid("", 18),
+            helper.make_operatorsetid(".com.microsoft", 1),
+        ]
+        return helper.make_model(
+            graph_def,
+            producer_name="ORTBertPretrainTest in onnxruntime_test_collective.py",
+            opset_imports=opset_imports,
+        )
+
+    def _create_allreduce_ut_model(self, shape, elem_type: TensorProto.DataType = TensorProto.FLOAT):
+        X = helper.make_tensor_value_info("X", elem_type, shape)  # noqa: N806
+        Y = helper.make_tensor_value_info("Y", elem_type, shape)  # noqa: N806
         node_def = helper.make_node("AllReduce", ["X"], ["Y"], domain="com.microsoft")
         graph_def = helper.make_graph(
             [node_def],
@@ -21,69 +34,176 @@ class ORTBertPretrainTest(unittest.TestCase):
             [X],
             [Y],
         )
-        return helper.make_model(graph_def, producer_name="ort-distributed-inference-unittest")
+        return ORTBertPretrainTest._create_model_with_opsets(graph_def)
 
     def _get_rank_size(self):
         comm = MPI.COMM_WORLD
         return comm.Get_rank(), comm.Get_size()
 
-    def _create_allgather_ut_model(self, shape, axis):
-        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)  # noqa: N806
-        rank, group_size = self._get_rank_size()
+    def _create_allgather_ut_model(
+        self,
+        shape,
+        axis,
+        # Element type for AllGather's input and output.
+        elem_type: TensorProto.DataType = TensorProto.FLOAT,
+        # Element type for Model's input and output.
+        communication_elem_type: TensorProto.DataType = TensorProto.FLOAT,
+    ):
+        X = helper.make_tensor_value_info("X", elem_type, shape)  # noqa: N806
+        _, group_size = self._get_rank_size()
         output_shape = [s * group_size if axis_index == axis else s for axis_index, s in enumerate(shape)]
-        Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape)  # noqa: N806
-        node_def = helper.make_node("AllGather", ["X"], ["Y"], domain="com.microsoft", group_size=group_size, axis=axis)
+        Y = helper.make_tensor_value_info("Y", elem_type, output_shape)  # noqa: N806
+        if elem_type != communication_elem_type:
+            # With elem_type and external_element_type, we use the pattern
+            #   model input type -> Cast -> elem_type -> AllGather -> elem_type -> Cast -> model output type
+            # so that we can test boolean tensors and other special types.
+            node_defs = [
+                helper.make_node("Cast", ["X"], ["X_casted"], to=communication_elem_type),
+                helper.make_node(
+                    "AllGather", ["X_casted"], ["Y_casted"], domain="com.microsoft", group_size=group_size, axis=axis
+                ),
+                helper.make_node("Cast", ["Y_casted"], ["Y"], to=elem_type),
+            ]
+        else:
+            # When elem_type == external_element_type, the pattern
+            #   model input type -> Cast -> elem_type -> AllGather -> elem_type -> Cast -> model output type
+            # is reduced to
+            #   model input type -> AllGather -> model output type
+            node_defs = [
+                helper.make_node("AllGather", ["X"], ["Y"], domain="com.microsoft", group_size=group_size, axis=axis),
+            ]
         graph_def = helper.make_graph(
-            [node_def],
+            node_defs,
             "",
             [X],
             [Y],
         )
-        return helper.make_model(graph_def, producer_name="ort-distributed-inference-unittest")
+        return ORTBertPretrainTest._create_model_with_opsets(graph_def)
 
-    def _create_alltoall_ut_model(self, shape):
-        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)  # noqa: N806
-        Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, shape)  # noqa: N806
+    def _create_alltoall_ut_model(
+        self,
+        shape,
+        elem_type: TensorProto.DataType = TensorProto.FLOAT,
+        communication_elem_type: TensorProto.DataType = TensorProto.FLOAT,
+    ):
+        X = helper.make_tensor_value_info("X", elem_type, shape)  # noqa: N806
+        Y = helper.make_tensor_value_info("Y", elem_type, shape)  # noqa: N806
         _, size = self._get_rank_size()
-        node_def = helper.make_node("AllToAll", ["X"], ["Y"], domain="com.microsoft", group_size=size)
+        # Explanation is in comments for model creation in _create_allgather_ut_model.
+        # Basically, ORT Python API doesn't support bool tensor yet, so we need to feed int64
+        # tensor and cast it to bool before running communication op.
+        if elem_type != communication_elem_type:
+            node_defs = [
+                helper.make_node("Cast", ["X"], ["X_casted"], to=communication_elem_type),
+                helper.make_node(
+                    "AllToAll",
+                    ["X_casted"],
+                    ["Y_casted"],
+                    domain="com.microsoft",
+                    group_size=size,
+                ),
+                helper.make_node("Cast", ["Y_casted"], ["Y"], to=elem_type),
+            ]
+        else:
+            node_defs = [
+                helper.make_node("AllToAll", ["X"], ["Y"], domain="com.microsoft", group_size=size),
+            ]
         graph_def = helper.make_graph(
-            [node_def],
+            node_defs,
             "",
             [X],
             [Y],
         )
-        return helper.make_model(graph_def, producer_name="ort-distributed-inference-unittest")
+        return ORTBertPretrainTest._create_model_with_opsets(graph_def)
+
+    def _create_alltoall_ut_model_for_boolean_tensor(
+        self,
+        shape,
+        # Tuple or list of bool values; e.g., [True, False] if
+        # shape is (2,) or [[True, False], [False, True]] if
+        # shape is (2, 2).
+        # It's input of AllToAll.
+        value,
+    ):
+        Y = helper.make_tensor_value_info("Y", TensorProto.BOOL, shape)  # noqa: N806
+        _, size = self._get_rank_size()
+        # Explanation is in comments for model creation in _create_allgather_ut_model.
+        # Basically, ORT Python API doesn't support bool tensor yet, so we need to feed int64
+        # tensor and cast it to bool before running communication op.
+        x_const_value = helper.make_tensor("condition", TensorProto.BOOL, shape, value)
+        node_defs = [
+            helper.make_node(
+                "Constant",
+                [],
+                ["X"],
+                value=x_const_value,
+            ),
+            helper.make_node(
+                "AllToAll",
+                ["X"],
+                ["Y"],
+                domain="com.microsoft",
+                group_size=size,
+            ),
+        ]
+        graph_def = helper.make_graph(
+            node_defs,
+            "",
+            [],
+            [Y],
+        )
+        return ORTBertPretrainTest._create_model_with_opsets(graph_def)
 
     def test_all_reduce(self):
-        model = self._create_allreduce_ut_model((128, 128))
-        rank, size = self._get_rank_size()
-        ort_sess = ort.InferenceSession(
-            model.SerializeToString(),
-            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
-            provider_options=[{"device_id": str(rank)}, {}],
-        )
+        for np_elem_type, elem_type in ((np.float32, TensorProto.FLOAT),):
+            model = self._create_allreduce_ut_model((128, 128), elem_type)
+            rank, size = self._get_rank_size()
+            ort_sess = ort.InferenceSession(
+                model.SerializeToString(),
+                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+                provider_options=[{"device_id": str(rank)}, {}],
+            )
 
-        input = np.ones((128, 128), dtype=np.float32)
-        outputs = ort_sess.run(None, {"X": input})
-        assert np.allclose(outputs[0], size * input)
+            input = np.ones((128, 128), dtype=np_elem_type)
+            outputs = ort_sess.run(None, {"X": input})
+            assert np.allclose(outputs[0], size * input)
 
     def test_all_gather(self):
-        model = self._create_allgather_ut_model((128, 128), 0)
+        for np_elem_type, elem_type, communication_elem_type in ((np.float32, TensorProto.FLOAT, TensorProto.FLOAT),):
+            model = self._create_allgather_ut_model((128, 128), 0, elem_type, communication_elem_type)
+            rank, size = self._get_rank_size()
+            ort_sess = ort.InferenceSession(
+                model.SerializeToString(),
+                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+                provider_options=[{"device_id": str(rank)}, {}],
+            )
+
+            input = np.ones((128, 128), dtype=np.float32) * rank
+            outputs = ort_sess.run(None, {"X": input})
+
+            expected_output = np.zeros((128, 128), dtype=np_elem_type)
+            for _ in range(size - 1):
+                expected_output = np.concatenate((expected_output, np.ones((128, 128), dtype=np_elem_type) * (_ + 1)))
+            np.testing.assert_allclose(outputs[0], expected_output, err_msg="all gather on axis0: result mismatch")
+
+    def test_all_gather_bool(self):
+        model = self._create_allgather_ut_model((4,), 0, TensorProto.INT64, TensorProto.INT64)
         rank, size = self._get_rank_size()
+        print(f"rank: {rank}, size: {size}")
         ort_sess = ort.InferenceSession(
             model.SerializeToString(),
             providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
             provider_options=[{"device_id": str(rank)}, {}],
         )
 
-        input = np.ones((128, 128), dtype=np.float32) * rank
-        outputs = ort_sess.run(None, {"X": input})
+        x = np.array([True, True, False, False]).astype(np.int64)
+        y = ort_sess.run(None, {"X": x})[0]
 
-        expected_output = np.zeros((128, 128), dtype=np.float32)
-        for _ in range(size - 1):
-            expected_output = np.concatenate((expected_output, np.ones((128, 128), dtype=np.float32) * (_ + 1)))
+        y_expected = np.array(
+            [True, True, False, False] * 4,
+        ).astype(np.int64)
 
-        np.testing.assert_allclose(outputs[0], expected_output, err_msg="all gather on axis0: result mismatch")
+        np.testing.assert_allclose(y, y_expected)
 
     def test_all_gather_axis1(self):
         model = self._create_allgather_ut_model((128, 128), 1)
@@ -104,24 +224,58 @@ class ORTBertPretrainTest(unittest.TestCase):
         np.testing.assert_allclose(outputs[0], expected_output, err_msg="all gather on axis1: result mismatch")
 
     def test_all_to_all(self):
-        model = self._create_alltoall_ut_model((128, 128))
-        rank, size = self._get_rank_size()
+        for np_elem_type, elem_type, communication_elem_type in (
+            (np.float32, TensorProto.FLOAT, TensorProto.FLOAT),
+            (np.int64, TensorProto.INT64, TensorProto.INT64),
+            # TODO: Fix the following case, which triggers random number-mismatch error.
+            # (np.float32, TensorProto.INT64, TensorProto.BOOL),
+        ):
+            model = self._create_alltoall_ut_model(
+                (128, 128), elem_type=elem_type, communication_elem_type=communication_elem_type
+            )
+            rank, size = self._get_rank_size()
+            ort_sess = ort.InferenceSession(
+                model.SerializeToString(),
+                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+                provider_options=[{"device_id": str(rank)}, {}],
+            )
+
+            input = np.ones((128, 128), dtype=np_elem_type) * rank
+            outputs = ort_sess.run(None, {"X": input})
+
+            expected_output = np.zeros((int(128 / size), 128), dtype=np_elem_type)
+            for _ in range(size - 1):
+                expected_output = np.concatenate(
+                    (expected_output, np.ones((int(128 / size), 128), dtype=np_elem_type) * (_ + 1))
+                )
+
+            print("outputs[0]: ", outputs[0] - expected_output)
+
+            assert np.allclose(outputs[0], expected_output)
+
+    def test_all_to_all_bool(self):
+        rank, _ = self._get_rank_size()
+
+        if rank == 0:
+            x = [True, True, True, True]
+        else:
+            x = [False, False, False, False]
+
+        model = self._create_alltoall_ut_model_for_boolean_tensor((4,), x)
+
         ort_sess = ort.InferenceSession(
             model.SerializeToString(),
             providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
             provider_options=[{"device_id": str(rank)}, {}],
         )
 
-        input = np.ones((128, 128), dtype=np.float32) * rank
-        outputs = ort_sess.run(None, {"X": input})
+        y = ort_sess.run(None, {})
 
-        expected_output = np.zeros((int(128 / size), 128), dtype=np.float32)
-        for _ in range(size - 1):
-            expected_output = np.concatenate(
-                (expected_output, np.ones((int(128 / size), 128), dtype=np.float32) * (_ + 1))
-            )
+        y_expected = np.array(
+            [True, False, False, False],
+        ).astype(np.int64)
 
-        assert np.allclose(outputs[0], expected_output)
+        np.testing.assert_allclose(y[0], y_expected)
 
 
 if __name__ == "__main__":