diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index ce0c12804b..f2a16fb29d 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -41,6 +41,8 @@ "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/distributed_reshape.cc" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/distributed_expand.cc" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/distributed_reduce.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/distributed_unsqueeze.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/distributed_squeeze.cc" ) endif() # add using ONNXRUNTIME_ROOT so they show up under the 'contrib_ops' folder in Visual Studio diff --git a/cmake/onnxruntime_rocm_hipify.cmake b/cmake/onnxruntime_rocm_hipify.cmake index 4140eeee0d..ce27bf756f 100644 --- a/cmake/onnxruntime_rocm_hipify.cmake +++ b/cmake/onnxruntime_rocm_hipify.cmake @@ -115,6 +115,8 @@ if (NOT onnxruntime_USE_NCCL) list(APPEND contrib_ops_excluded_files "collective/distributed_reshape.cc") list(APPEND contrib_ops_excluded_files "collective/distributed_expand.cc") list(APPEND contrib_ops_excluded_files "collective/distributed_reduce.cc") + list(APPEND contrib_ops_excluded_files "collective/distributed_unsqueeze.cc") + list(APPEND contrib_ops_excluded_files "collective/distributed_squeeze.cc") endif() set(provider_excluded_files diff --git a/onnxruntime/contrib_ops/cuda/collective/distributed_reshape.cc b/onnxruntime/contrib_ops/cuda/collective/distributed_reshape.cc index a0ac40defb..e413ccf580 100644 --- a/onnxruntime/contrib_ops/cuda/collective/distributed_reshape.cc +++ b/onnxruntime/contrib_ops/cuda/collective/distributed_reshape.cc @@ -472,8 +472,8 @@ std::tuple ComputeNativeSpecForTwoAxisFusion( } else if (src_spec.HasShard() && (src_spec.GetPartitionAxis() < fused_axis_in_src || src_spec.GetPartitionAxis() > fused_axis_in_src + 1)) { // It's two-axis fusion but the fused axes is not sharded. // Example: S[0]RR, shape=[2, 3, 5], device_mesh=[0, 1] -> S[0]R, shape = [2, 15], device_mesh=[0, 1] - auto dst_spec = TensorPartitionSpec::CreateByDropOneAxis( - src_spec, fused_axis_in_src + 1); + auto dst_spec = TensorPartitionSpec::CreateByDropAxes( + src_spec, {fused_axis_in_src + 1}); return std::make_tuple(true, dst_spec); } else { return std::make_tuple(false, TensorPartitionSpec()); diff --git a/onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.cc b/onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.cc new file mode 100644 index 0000000000..c3cae2d0bf --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.cc @@ -0,0 +1,96 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Distributed computation. +#include "distributed_squeeze.h" +#include "mpi_include.h" + +// ORT system. +#include "core/providers/cuda/cuda_check_memory.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#if defined(ORT_USE_NCCL) +template +DistributedSqueeze::DistributedSqueeze(const OpKernelInfo& info) : DistributedKernel(info) { +} + +template +Status DistributedSqueeze::ComputeInternal(OpKernelContext* context) const { + auto input_tensor = context->Input(0); + auto axes_tensor = context->Input(1); + auto axes_span = axes_tensor->DataAsSpan(); + + const TensorPartitionSpec& input_spec = input_shard_specs_[0]; + const TensorPartitionSpec& axes_spec = input_shard_specs_[1]; + const TensorPartitionSpec& output_spec = output_shard_specs_[0]; + + ORT_ENFORCE(axes_spec.HasNoShard(), "Axes tensor cannot be sharded."); + + // Non-negative collection of axes to drop. + std::vector axes; + for (const auto axis : axes_span) { + axes.push_back(axis >= 0 ? axis : axis + input_tensor->Shape().NumDimensions()); + } + // Shape after dropping axes. + auto dims = input_tensor->Shape().AsShapeVector(); + // Sort in descending order so that we can drop axes from the end. + std::sort(axes.begin(), axes.end(), [](Tind a, Tind b) { return a > b; }); + for (const auto axis : axes) { + ORT_ENFORCE(input_tensor->Shape()[axis] == 1, "Cannot squeeze non-singleton dimension."); + dims.erase(dims.begin() + axis); + } + auto native_output_spec = TensorPartitionSpec::CreateByDropAxes( + input_spec, + axes); + ORT_ENFORCE( + output_spec == native_output_spec, + "Re-sharding is required but not supported yet for this case."); + auto output_tensor = context->Output(0, dims); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( + output_tensor->MutableDataRaw(), + input_tensor->DataRaw(), + input_tensor->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream(context))); + return Status::OK(); +} + +ONNX_OPERATOR_TYPED_KERNEL_EX( + DistributedSqueeze, + kMSDomain, + 1, + float, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + DistributedSqueeze); + +ONNX_OPERATOR_TYPED_KERNEL_EX( + DistributedSqueeze, + kMSDomain, + 1, + MLFloat16, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + DistributedSqueeze); + +ONNX_OPERATOR_TYPED_KERNEL_EX( + DistributedSqueeze, + kMSDomain, + 1, + int64_t, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + DistributedSqueeze); + +#endif + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.h b/onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.h new file mode 100644 index 0000000000..5b81d9c479 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.h @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "sharding.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#if defined(ORT_USE_NCCL) + +template +class DistributedSqueeze final : public DistributedKernel { + public: + explicit DistributedSqueeze(const OpKernelInfo& info); + + Status ComputeInternal(OpKernelContext* context) const override; +}; + +#endif + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.cc b/onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.cc new file mode 100644 index 0000000000..a78f19101b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.cc @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Distributed computation. +#include "distributed_unsqueeze.h" +#include "mpi_include.h" + +// ORT system. +#include "core/providers/cuda/cuda_check_memory.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#if defined(ORT_USE_NCCL) +template +DistributedUnsqueeze::DistributedUnsqueeze(const OpKernelInfo& info) : DistributedKernel(info) { +} + +template +Status DistributedUnsqueeze::ComputeInternal(OpKernelContext* context) const { + auto input_tensor = context->Input(0); + auto axes_tensor = context->Input(1); + auto axes_span = axes_tensor->DataAsSpan(); + + const TensorPartitionSpec& input_spec = input_shard_specs_[0]; + const TensorPartitionSpec& axes_spec = input_shard_specs_[1]; + const TensorPartitionSpec& output_spec = output_shard_specs_[0]; + + ORT_ENFORCE(axes_spec.HasNoShard(), "Axes tensor cannot be sharded."); + + std::vector axes(axes_span.begin(), axes_span.end()); + std::sort(axes.begin(), axes.end()); + auto dims = input_tensor->Shape().AsShapeVector(); + auto native_output_spec = input_spec; + for (auto axis : axes) { + if (axis < 0) { + axis += input_tensor->Shape().NumDimensions() + 1; + } + dims.insert(dims.begin() + axis, 1); + native_output_spec = TensorPartitionSpec::CreateByInsertOneAxis( + native_output_spec, + axis); + } + ORT_ENFORCE( + output_spec == native_output_spec, + "Re-sharding is required but not supported yet for this case. ", + "Specified: ", output_spec.ToString(), + " Actual: ", native_output_spec.ToString()); + auto output_tensor = context->Output(0, dims); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( + output_tensor->MutableDataRaw(), + input_tensor->DataRaw(), + input_tensor->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream(context))); + return Status::OK(); +} + +ONNX_OPERATOR_TYPED_KERNEL_EX( + DistributedUnsqueeze, + kMSDomain, + 1, + float, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + DistributedUnsqueeze); + +ONNX_OPERATOR_TYPED_KERNEL_EX( + DistributedUnsqueeze, + kMSDomain, + 1, + MLFloat16, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + DistributedUnsqueeze); + +ONNX_OPERATOR_TYPED_KERNEL_EX( + DistributedUnsqueeze, + kMSDomain, + 1, + int64_t, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + DistributedUnsqueeze); + +#endif + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.h b/onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.h new file mode 100644 index 0000000000..005093ef78 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.h @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "sharding.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +#if defined(ORT_USE_NCCL) + +template +class DistributedUnsqueeze final : public DistributedKernel { + public: + explicit DistributedUnsqueeze(const OpKernelInfo& info); + + Status ComputeInternal(OpKernelContext* context) const override; +}; + +#endif + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/collective/sharding_spec.h b/onnxruntime/contrib_ops/cuda/collective/sharding_spec.h index 5185c41e68..5abc50a61c 100644 --- a/onnxruntime/contrib_ops/cuda/collective/sharding_spec.h +++ b/onnxruntime/contrib_ops/cuda/collective/sharding_spec.h @@ -257,15 +257,25 @@ class TensorPartitionSpec { return TensorPartitionSpec::Create(axis_specs, device_mesh); } - static TensorPartitionSpec CreateByDropOneAxis( - const TensorPartitionSpec& TensorPartitionSpec, const size_t axis_to_drop) { + static TensorPartitionSpec CreateByDropAxes( + const TensorPartitionSpec& spec, const std::vector& axes_to_drop) { std::vector axis_specs; - for (size_t i = 0; i < TensorPartitionSpec.axis_specs.size(); ++i) { - if (i != axis_to_drop) { - axis_specs.push_back(TensorPartitionSpec.axis_specs[i]); + for (size_t i = 0; i < spec.axis_specs.size(); ++i) { + if (std::find(axes_to_drop.begin(), axes_to_drop.end(), i) != axes_to_drop.end()) { + // This axis, i, is in axes_to_drop. Let's not copy its spec. + continue; } + axis_specs.push_back(spec.axis_specs[i]); } - return TensorPartitionSpec::Create(axis_specs, TensorPartitionSpec.device_mesh); + return TensorPartitionSpec::Create(axis_specs, spec.device_mesh); + } + + static TensorPartitionSpec CreateByInsertOneAxis( + const TensorPartitionSpec& spec, + const size_t axis_to_insert) { + std::vector axis_specs(spec.axis_specs); + axis_specs.insert(axis_specs.begin() + axis_to_insert, AxisPartitionSpec::CreateReplica()); + return TensorPartitionSpec::Create(axis_specs, spec.device_mesh); } // Helper to debug and generate error message; e.g., diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 8e157da6cb..6162895d00 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -184,6 +184,14 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, DistributedReduceMean); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, DistributedReduceMean); + +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int64_t, DistributedUnsqueeze); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, DistributedUnsqueeze); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, DistributedUnsqueeze); + +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int64_t, DistributedSqueeze); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, DistributedSqueeze); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, DistributedSqueeze); #endif template <> @@ -372,6 +380,14 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, #endif }; diff --git a/onnxruntime/core/graph/contrib_ops/collective_defs.cc b/onnxruntime/core/graph/contrib_ops/collective_defs.cc index 8b5b561c1a..59adfc523c 100644 --- a/onnxruntime/core/graph/contrib_ops/collective_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/collective_defs.cc @@ -396,6 +396,80 @@ void RegisterCollectiveOps() { OpSchema::NonDifferentiable) .Output(0, "output", "Output tensor", "T", OpSchema::Single, true, 1, OpSchema::Differentiable) .TypeConstraint("T", OpSchema::all_tensor_types_ir4(), "Constrain input and output types to all tensors."); + + ONNX_CONTRIB_OPERATOR_SCHEMA(DistributedUnsqueeze) + .SetDomain(kMSDomain) + .SinceVersion(1) + .Attr("input_device_mesh_elements", + "device_mesh_elements[i] defines the device mesh's value for the i-th input. " + "E.g., device_mesh_elements=[\"[0, 1]\", \"[0, 1]\"] means the 1st and the 2nd " + " inputs are stored on the 0-th and the 1st devices, respectively.", + AttributeProto::STRINGS) + .Attr("input_device_mesh_shapes", + "device_mesh_shape[i] defines the device mesh's shape for the i-th input.", + AttributeProto::STRINGS) + .Attr("input_shard_specs", + "The sharding spec of inputs. " + "E.g., if input_shard_specs[i] is \"RRR\", the i-th input is a unsharded 3-D tensor.", + AttributeProto::STRINGS) + .Attr("output_device_mesh_elements", + "Similar to input_device_mesh_elments but for outputs.", + AttributeProto::STRINGS) + .Attr("output_device_mesh_shapes", + "Similar to input_device_mesh_shapes but for outputs.", + AttributeProto::STRINGS) + .Attr("output_shard_specs", + "Similar to input_shard_specs but for outputs.", + AttributeProto::STRINGS) + .Input(0, "input", "Input tensor", "T", OpSchema::Single, true, 1, OpSchema::Differentiable) + .Input( + 1, + "axes", + "A 1-D tensor indicates the axes to add.", + "tensor(int64)", + OpSchema::Single, + true, + 1, + OpSchema::NonDifferentiable) + .Output(0, "output", "Output tensor", "T", OpSchema::Single, true, 1, OpSchema::Differentiable) + .TypeConstraint("T", OpSchema::all_tensor_types_ir4(), "Constrain input and output types to all tensors."); + + ONNX_CONTRIB_OPERATOR_SCHEMA(DistributedSqueeze) + .SetDomain(kMSDomain) + .SinceVersion(1) + .Attr("input_device_mesh_elements", + "device_mesh_elements[i] defines the device mesh's value for the i-th input. " + "E.g., device_mesh_elements=[\"[0, 1]\", \"[0, 1]\"] means the 1st and the 2nd " + " inputs are stored on the 0-th and the 1st devices, respectively.", + AttributeProto::STRINGS) + .Attr("input_device_mesh_shapes", + "device_mesh_shape[i] defines the device mesh's shape for the i-th input.", + AttributeProto::STRINGS) + .Attr("input_shard_specs", + "The sharding spec of inputs. " + "E.g., if input_shard_specs[i] is \"RRR\", the i-th input is a unsharded 3-D tensor.", + AttributeProto::STRINGS) + .Attr("output_device_mesh_elements", + "Similar to input_device_mesh_elments but for outputs.", + AttributeProto::STRINGS) + .Attr("output_device_mesh_shapes", + "Similar to input_device_mesh_shapes but for outputs.", + AttributeProto::STRINGS) + .Attr("output_shard_specs", + "Similar to input_shard_specs but for outputs.", + AttributeProto::STRINGS) + .Input(0, "input", "Input tensor", "T", OpSchema::Single, true, 1, OpSchema::Differentiable) + .Input( + 1, + "axes", + "A 1-D tensor indicates the axes to add.", + "tensor(int64)", + OpSchema::Single, + true, + 1, + OpSchema::NonDifferentiable) + .Output(0, "output", "Output tensor", "T", OpSchema::Single, true, 1, OpSchema::Differentiable) + .TypeConstraint("T", OpSchema::all_tensor_types_ir4(), "Constrain input and output types to all tensors."); } } // namespace contrib diff --git a/onnxruntime/test/python/onnxruntime_test_distributed.py b/onnxruntime/test/python/onnxruntime_test_distributed.py index 6f69197218..de70478761 100644 --- a/onnxruntime/test/python/onnxruntime_test_distributed.py +++ b/onnxruntime/test/python/onnxruntime_test_distributed.py @@ -852,6 +852,235 @@ class TestDistributedExpand(unittest.TestCase): ) +class TestDistributedUnsqueeze(unittest.TestCase): + def _check_distributed_unsqueeze( + self, + shape: Tuple[int, ...], + axes: Tuple[int, ...], + input_device_meshes: np.ndarray, + input_shard_specs: Tuple[str, ...], + output_device_meshes: np.ndarray, + output_shard_specs: Tuple[str, ...], + ): + assert len(input_device_meshes) == len(input_shard_specs) + assert len(output_device_meshes) == len(output_shard_specs) + + input_device_mesh_shapes, input_device_mesh_elements = translate_all_device_meshes(input_device_meshes) + output_device_mesh_shapes, output_device_mesh_elements = translate_all_device_meshes(output_device_meshes) + + @onnxscript.script() + def distributed_unsqueeze_instance(data_tensor: FLOAT, axes_tensor: INT64): + return MICROSOFT_OPSET.DistributedUnsqueeze( + data_tensor, + axes_tensor, + input_device_mesh_shapes=input_device_mesh_shapes, + input_device_mesh_elements=input_device_mesh_elements, + input_shard_specs=input_shard_specs, + output_device_mesh_shapes=output_device_mesh_shapes, + output_device_mesh_elements=output_device_mesh_elements, + output_shard_specs=output_shard_specs, + ) + + rank = comm.Get_rank() + data_tensor = np.arange(np.prod(shape), dtype=np.float32).reshape(*shape) + axes_tensor = np.array(axes, dtype=np.int64) + + local_data_tensor = shard_tensor_per_spec(data_tensor, rank, input_shard_specs[0], input_device_meshes[0]) + assert "S" not in input_shard_specs[1], "Shape should not be sharded." + + expected = data_tensor.copy() + for axis in sorted(axes): + expected = np.expand_dims(expected, axis=axis) + + local_expected = shard_tensor_per_spec(expected, rank, output_shard_specs[0], output_device_meshes[0]) + + onnx_model = distributed_unsqueeze_instance.to_model_proto( + input_types=[FLOAT[tuple(local_data_tensor.shape)], INT64[tuple(axes_tensor.shape)]], + output_types=[FLOAT[tuple(local_expected.shape)]], + ) + + # Each MPI process owns a sharded model. + sess = ort.InferenceSession( + onnx_model.SerializeToString(), + providers=["CUDAExecutionProvider"], + provider_options=[{"device_id": str(rank)}], + ) + + # Each MPI process executes its sharded model. + # The result is `local` tensor stored on a specific MPI rank + # instead of `logical` tensor. + result = sess.run( + None, + { + "data_tensor": local_data_tensor, + "axes_tensor": axes_tensor, + }, + ) + + # Compare local tensor and the corresponding logical sub-tensor + # obtained by sharding logical tensor following output's sharding spec. + np.testing.assert_allclose(result[0], local_expected, rtol=1e-5, atol=1e-8) + + def test_unsqueeze_sharded(self): + # data: shape=[8,1], spec=(RR, [0,1]) + # shape: shape=[2], spec=(R, [0,1]), value=[1,4] + # output: shape=[8,4], spec=(RS, [0,1]) + self._check_distributed_unsqueeze( + shape=( + 8, + 2, + ), + axes=(1,), + input_device_meshes=[np.array([0, 1])] * 2, + input_shard_specs=("RS[0]", "R"), + output_device_meshes=[np.array([0, 1])], + output_shard_specs=("RRS[0]",), + ) + + def test_unsqueeze_descending_axes(self): + # data: shape=[8,1], spec=(RR, [0,1]) + # shape: shape=[2], spec=(R, [0,1]), value=[1,4] + # output: shape=[8,4], spec=(RS, [0,1]) + self._check_distributed_unsqueeze( + shape=( + 8, + 2, + ), + axes=( + 4, + 1, + 0, + ), + input_device_meshes=[np.array([0, 1])] * 2, + input_shard_specs=("RS[0]", "R"), + output_device_meshes=[np.array([0, 1])], + output_shard_specs=("RRRS[0]R",), + ) + + def test_unsqueeze_not_sharded(self): + # data: shape=[8,1], spec=(RR, [0,1]) + # shape: shape=[2], spec=(R, [0,1]), value=[1,4] + # output: shape=[8,4], spec=(RS, [0,1]) + self._check_distributed_unsqueeze( + shape=( + 8, + 2, + ), + axes=(2,), + input_device_meshes=[np.array([0, 1])] * 2, + input_shard_specs=("RR", "R"), + output_device_meshes=[np.array([0, 1])], + output_shard_specs=("RRR",), + ) + + +class TestDistributedSqueeze(unittest.TestCase): + def _check_distributed_squeeze( + self, + shape: Tuple[int, ...], + axes: Tuple[int, ...], + input_device_meshes: np.ndarray, + input_shard_specs: Tuple[str, ...], + output_device_meshes: np.ndarray, + output_shard_specs: Tuple[str, ...], + ): + assert len(input_device_meshes) == len(input_shard_specs) + assert len(output_device_meshes) == len(output_shard_specs) + + input_device_mesh_shapes, input_device_mesh_elements = translate_all_device_meshes(input_device_meshes) + output_device_mesh_shapes, output_device_mesh_elements = translate_all_device_meshes(output_device_meshes) + + @onnxscript.script() + def distributed_squeeze_instance(data_tensor: FLOAT, axes_tensor: INT64): + return MICROSOFT_OPSET.DistributedSqueeze( + data_tensor, + axes_tensor, + input_device_mesh_shapes=input_device_mesh_shapes, + input_device_mesh_elements=input_device_mesh_elements, + input_shard_specs=input_shard_specs, + output_device_mesh_shapes=output_device_mesh_shapes, + output_device_mesh_elements=output_device_mesh_elements, + output_shard_specs=output_shard_specs, + ) + + rank = comm.Get_rank() + data_tensor = np.arange(np.prod(shape), dtype=np.float32).reshape(*shape) + axes_tensor = np.array(axes, dtype=np.int64) + + local_data_tensor = shard_tensor_per_spec(data_tensor, rank, input_shard_specs[0], input_device_meshes[0]) + assert "S" not in input_shard_specs[1], "Shape should not be sharded." + + expected = data_tensor.copy() + for axis in sorted(axes, reverse=True): + expected = np.squeeze(expected, axis=axis) + + local_expected = shard_tensor_per_spec(expected, rank, output_shard_specs[0], output_device_meshes[0]) + + onnx_model = distributed_squeeze_instance.to_model_proto( + input_types=[FLOAT[tuple(local_data_tensor.shape)], INT64[tuple(axes_tensor.shape)]], + output_types=[FLOAT[tuple(local_expected.shape)]], + ) + + # Each MPI process owns a sharded model. + sess = ort.InferenceSession( + onnx_model.SerializeToString(), + providers=["CUDAExecutionProvider"], + provider_options=[{"device_id": str(rank)}], + ) + + # Each MPI process executes its sharded model. + # The result is `local` tensor stored on a specific MPI rank + # instead of `logical` tensor. + result = sess.run( + None, + { + "data_tensor": local_data_tensor, + "axes_tensor": axes_tensor, + }, + ) + + # Compare local tensor and the corresponding logical sub-tensor + # obtained by sharding logical tensor following output's sharding spec. + np.testing.assert_allclose(result[0], local_expected, rtol=1e-5, atol=1e-8) + + def test_squeeze_sharded(self): + # data: shape=[8,1], spec=(RR, [0,1]) + # shape: shape=[2], spec=(R, [0,1]), value=[1,4] + # output: shape=[8,4], spec=(RS, [0,1]) + self._check_distributed_squeeze( + shape=( + 1, + 2, + ), + axes=(0,), + input_device_meshes=[np.array([0, 1])] * 2, + input_shard_specs=("RS[0]", "R"), + output_device_meshes=[np.array([0, 1])], + output_shard_specs=("S[0]",), + ) + + def test_squeeze_not_sharded(self): + # data: shape=[8,1], spec=(RR, [0,1]) + # shape: shape=[2], spec=(R, [0,1]), value=[1,4] + # output: shape=[8,4], spec=(RS, [0,1]) + self._check_distributed_squeeze( + shape=( + 8, + 1, + 1, + 1, + ), + axes=( + 1, + 3, + ), + input_device_meshes=[np.array([0, 1])] * 2, + input_shard_specs=("RRRR", "R"), + output_device_meshes=[np.array([0, 1])], + output_shard_specs=("RR",), + ) + + class TestDistributedReduce(unittest.TestCase): def _check_distributed_reduce( self,