saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

support allgather on different axis (#15610)

Browse source 
### Description
Extend the AllGather op to support perform allgather on different axis.
provide the implementation in nccl kernels.

### Motivation and Context
We hit some scenario in distributed inference that we need to support
gather on non-first axis.

---------

Co-authored-by: Cheng Tang <chenta@microsoft.com@orttrainingdev9.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
Co-authored-by: Wei-Sheng Chin <wschin@outlook.com>
This commit is contained in:
Tang, Cheng 2023-04-27 14:47:28 -07:00 committed by GitHub
parent 5bde1e8e37
commit 627f5c9767
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 94 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

75
onnxruntime/contrib_ops/cuda/collective/nccl_kernels.cc
View file

@ -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<int64_t>(1));
info.GetAttrOrDefault("axis", &axis_, static_cast<int64_t>(0));
cuda_ep_ = static_cast<const CUDAExecutionProvider*>(info.GetExecutionProvider());
}
Status AllGather::ComputeInternal(OpKernelContext* context) const {
@ -114,19 +115,67 @@ Status AllGather::ComputeInternal(OpKernelContext* context) const {
auto input_tensor = context->Input<Tensor>(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<size_t> 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<int64_t> 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<char*>(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();
}

2
onnxruntime/contrib_ops/cuda/collective/nccl_kernels.h
View file

@ -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 {

9
onnxruntime/core/graph/contrib_ops/collective_defs.cc
View file

@ -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<int64_t>(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<int64_t>(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<int>(axis)) * group_size;
*shape.mutable_dim(static_cast<int>(axis)) = dim;
*output_type->mutable_tensor_type()->mutable_shape() = shape;
}
});

30
onnxruntime/test/python/onnxruntime_test_collective.py
View file

@ -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))