diff --git a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc index 30cb474d03..a89b861c8e 100644 --- a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "nccl_kernels.h" #include "mpi_include.h" +#include "core/providers/cuda/tensor/transpose.h" namespace onnxruntime { namespace contrib { @@ -99,14 +100,14 @@ Status AllReduce::ComputeInternal(OpKernelContext* context) const { void* output_data = context->Output(0, in_shape)->MutableDataRaw(); ncclDataType_t dtype = GetNcclDataType(input_tensor->DataType()); -#ifdef ORT_USE_NCCL NCCL_RETURN_IF_ERROR(ncclAllReduce(input_data, output_data, input_count, dtype, ncclSum, comm, Stream(context))); -#endif return Status::OK(); } AllGather::AllGather(const OpKernelInfo& info) : NcclKernel(info) { info.GetAttrOrDefault("group_size", &group_size_, static_cast(1)); + info.GetAttrOrDefault("axis", &axis_, static_cast(0)); + cuda_ep_ = static_cast(info.GetExecutionProvider()); } Status AllGather::ComputeInternal(OpKernelContext* context) const { @@ -114,19 +115,67 @@ Status AllGather::ComputeInternal(OpKernelContext* context) const { auto input_tensor = context->Input(0); const void* input_data = input_tensor->DataRaw(); - const auto in_shape = input_tensor->Shape(); + const auto& in_shape = input_tensor->Shape(); int64_t input_count = in_shape.Size(); - // construct output shape - TensorShape out_shape(in_shape); - out_shape[0] = group_size_ * out_shape[0]; - void* output_data = context->Output(0, out_shape)->MutableDataRaw(); + if (axis_ > 0) { + // Need transpose + // TODO: fuse transpose with allgather + std::vector permutation(in_shape.NumDimensions()); + AllocatorPtr alloc; + auto status = context->GetTempSpaceAllocator(&alloc); + if (!status.IsOK()) + return status; - ncclDataType_t dtype = GetNcclDataType(input_tensor->DataType()); -#ifdef ORT_USE_NCCL - NCCL_RETURN_IF_ERROR(ncclAllGather(input_data, output_data, input_count, dtype, comm, Stream(context))); -#endif - return Status::OK(); + std::iota(std::begin(permutation), std::end(permutation), 0); + + // swap rank 0 and rank axis + permutation[axis_] = 0; + permutation[0] = axis_; + std::vector transposed_input_dims; + transposed_input_dims.reserve(in_shape.NumDimensions()); + for (auto e : permutation) { + transposed_input_dims.push_back(in_shape[e]); + } + + // Allocate a temporary tensor to hold transposed input + auto temp_input = Tensor::Create(input_tensor->DataType(), TensorShape(transposed_input_dims), alloc); + + // Perform the transpose + ORT_RETURN_IF_ERROR(onnxruntime::cuda::Transpose::DoTranspose(cuda_ep_->GetDeviceProp(), + Stream(context), + GetCublasHandle(context), + permutation, *input_tensor, *temp_input)); + // Allocate a tempoarary buffer for all gather + TensorShape all_gather_out_shape(transposed_input_dims); + all_gather_out_shape[0] = group_size_ * all_gather_out_shape[0]; + auto all_gather_output = Tensor::Create(temp_input->DataType(), all_gather_out_shape, alloc); + ncclDataType_t dtype = GetNcclDataType(temp_input->DataType()); + NCCL_RETURN_IF_ERROR(ncclAllGather(temp_input->DataRaw(), + all_gather_output->MutableDataRaw(), + input_count, dtype, comm, Stream(context))); + // release temp_input + temp_input.release(); + // transpose to output + TensorShape out_shape(in_shape); + out_shape[axis_] = group_size_ * out_shape[axis_]; + auto* output_tensor = context->Output(0, out_shape); + + return onnxruntime::cuda::Transpose::DoTranspose(cuda_ep_->GetDeviceProp(), + Stream(context), + GetCublasHandle(context), + permutation, *all_gather_output, *output_tensor); + } else { + // construct output shape + TensorShape out_shape(in_shape); + out_shape[axis_] = group_size_ * out_shape[axis_]; + + void* output_data = context->Output(0, out_shape)->MutableDataRaw(); + + ncclDataType_t dtype = GetNcclDataType(input_tensor->DataType()); + NCCL_RETURN_IF_ERROR(ncclAllGather(input_data, output_data, input_count, dtype, comm, Stream(context))); + return Status::OK(); + } } AllToAll::AllToAll(const OpKernelInfo& info) : NcclKernel(info) { @@ -146,7 +195,6 @@ Status AllToAll::ComputeInternal(OpKernelContext* context) const { char* output_data = static_cast(context->Output(0, out_shape)->MutableDataRaw()); -#ifdef ORT_USE_NCCL NCCL_RETURN_IF_ERROR(ncclGroupStart()); for (int32_t r = 0; r < group_size_; r++) { NCCL_RETURN_IF_ERROR(ncclSend(input_data, rank_stride, dtype, r, comm, Stream(context))); @@ -155,7 +203,6 @@ Status AllToAll::ComputeInternal(OpKernelContext* context) const { output_data += (rank_stride * element_size); } NCCL_RETURN_IF_ERROR(ncclGroupEnd()); -#endif return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h index 4f7093dd49..24df69ea50 100644 --- a/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h +++ b/onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h @@ -67,6 +67,8 @@ class AllGather final : public NcclKernel { private: int64_t group_size_ = -1; + int64_t axis_ = -1; + const CUDAExecutionProvider* cuda_ep_; }; class AllToAll final : public NcclKernel { diff --git a/onnxruntime/core/graph/contrib_ops/collective_defs.cc b/onnxruntime/core/graph/contrib_ops/collective_defs.cc index c4815a75db..db2c209771 100644 --- a/onnxruntime/core/graph/contrib_ops/collective_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/collective_defs.cc @@ -30,6 +30,10 @@ void RegisterCollectiveOps() { ONNX_CONTRIB_OPERATOR_SCHEMA(AllGather) .SetDomain(kMSDomain) .SinceVersion(1) + .Attr("axis", + "the axis to gather on.", + AttributeProto::INT, + static_cast(1)) .Attr("group_size", "total size in the group that need to be gathered.", AttributeProto::INT, @@ -42,6 +46,7 @@ void RegisterCollectiveOps() { "Constrain to float, float16 and double tensors.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { auto group_size = getAttribute(ctx, "group_size", 1); + auto axis = getAttribute(ctx, "axis", 0); assert(group_size >= static_cast(1)); // propagate type for output propagateElemTypeFromInputToOutput(ctx, 0, 0); @@ -52,8 +57,8 @@ void RegisterCollectiveOps() { auto input_type = ctx.getInputType(0); if (hasShape(*input_type)) { auto shape = input_type->tensor_type().shape(); - auto dim = shape.dim(0) * group_size; - *shape.mutable_dim(0) = dim; + auto dim = shape.dim(static_cast(axis)) * group_size; + *shape.mutable_dim(static_cast(axis)) = dim; *output_type->mutable_tensor_type()->mutable_shape() = shape; } }); diff --git a/onnxruntime/test/python/onnxruntime_test_collective.py b/onnxruntime/test/python/onnxruntime_test_collective.py index a6e420fe47..16f5136235 100644 --- a/onnxruntime/test/python/onnxruntime_test_collective.py +++ b/onnxruntime/test/python/onnxruntime_test_collective.py @@ -27,12 +27,12 @@ class ORTBertPretrainTest(unittest.TestCase): comm = MPI.COMM_WORLD return comm.Get_rank(), comm.Get_size() - def _create_allgather_ut_model(self, shape): + def _create_allgather_ut_model(self, shape, axis): X = helper.make_tensor_value_info("X", TensorProto.FLOAT, shape) # noqa: N806 - rank, size = self._get_rank_size() - output_shape = [s * size if _ == 0 else s for _, s in enumerate(shape)] + rank, 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=size) + node_def = helper.make_node("AllGather", ["X"], ["Y"], domain="com.microsoft", group_size=group_size, axis=axis) graph_def = helper.make_graph( [node_def], "", @@ -68,7 +68,7 @@ class ORTBertPretrainTest(unittest.TestCase): assert np.allclose(outputs[0], size * input) def test_all_gather(self): - model = self._create_allgather_ut_model((128, 128)) + model = self._create_allgather_ut_model((128, 128), 0) rank, size = self._get_rank_size() ort_sess = ort.InferenceSession( model.SerializeToString(), @@ -83,7 +83,25 @@ class ORTBertPretrainTest(unittest.TestCase): for _ in range(size - 1): expected_output = np.concatenate((expected_output, np.ones((128, 128), dtype=np.float32) * (_ + 1))) - assert np.allclose(outputs[0], expected_output) + np.testing.assert_allclose(outputs[0], expected_output, err_msg="all gather on axis0: result mismatch") + + def test_all_gather_axis1(self): + model = self._create_allgather_ut_model((128, 128), 1) + 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.float32) + for _ in range(size - 1): + expected_output = np.concatenate((expected_output, np.ones((128, 128), dtype=np.float32) * (_ + 1)), axis=1) + + 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))