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

Distributed Squeeze and Distributed Unsqueeze (#18269)

Browse source 
Implementat DistributedSqueeze & DistributedUnsqueeze for llama 2.
This commit is contained in:
Wei-Sheng Chin 2023-11-06 20:11:35 -08:00 committed by GitHub
parent ad34c67a44
commit fb6737e893
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 596 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cmake/onnxruntime_providers_cuda.cmake
View file

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

2
cmake/onnxruntime_rocm_hipify.cmake
View file

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

4
onnxruntime/contrib_ops/cuda/collective/distributed_reshape.cc
View file

@ -472,8 +472,8 @@ std::tuple<bool, TensorPartitionSpec> 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());

96
onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.cc Normal file
View file

@ -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 <typename T, typename Tind>
DistributedSqueeze<T, Tind>::DistributedSqueeze(const OpKernelInfo& info) : DistributedKernel(info) {
}
template <typename T, typename Tind>
Status DistributedSqueeze<T, Tind>::ComputeInternal(OpKernelContext* context) const {
auto input_tensor = context->Input<Tensor>(0);
auto axes_tensor = context->Input<Tensor>(1);
auto axes_span = axes_tensor->DataAsSpan<Tind>();
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<Tind> 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<float>()),
DistributedSqueeze<float, int64_t>);
ONNX_OPERATOR_TYPED_KERNEL_EX(
DistributedSqueeze,
kMSDomain,
1,
MLFloat16,
kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.InputMemoryType(OrtMemTypeCPUInput, 1)
.TypeConstraint("T", DataTypeImpl::GetTensorType<MLFloat16>()),
DistributedSqueeze<MLFloat16, int64_t>);
ONNX_OPERATOR_TYPED_KERNEL_EX(
DistributedSqueeze,
kMSDomain,
1,
int64_t,
kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.InputMemoryType(OrtMemTypeCPUInput, 1)
.TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>()),
DistributedSqueeze<int64_t, int64_t>);
#endif
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

32
onnxruntime/contrib_ops/cuda/collective/distributed_squeeze.h Normal file
View file

@ -0,0 +1,32 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <algorithm>
#include <tuple>
#include <optional>
#include <string>
#include <nccl.h>
#include <sstream>
#include "sharding.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
#if defined(ORT_USE_NCCL)
template <typename T, typename Tind>
class DistributedSqueeze final : public DistributedKernel {
public:
explicit DistributedSqueeze(const OpKernelInfo& info);
Status ComputeInternal(OpKernelContext* context) const override;
};
#endif
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

95
onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.cc Normal file
View file

@ -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 <typename T, typename Tind>
DistributedUnsqueeze<T, Tind>::DistributedUnsqueeze(const OpKernelInfo& info) : DistributedKernel(info) {
}
template <typename T, typename Tind>
Status DistributedUnsqueeze<T, Tind>::ComputeInternal(OpKernelContext* context) const {
auto input_tensor = context->Input<Tensor>(0);
auto axes_tensor = context->Input<Tensor>(1);
auto axes_span = axes_tensor->DataAsSpan<Tind>();
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<int64_t> 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<float>()),
DistributedUnsqueeze<float, int64_t>);
ONNX_OPERATOR_TYPED_KERNEL_EX(
DistributedUnsqueeze,
kMSDomain,
1,
MLFloat16,
kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.InputMemoryType(OrtMemTypeCPUInput, 1)
.TypeConstraint("T", DataTypeImpl::GetTensorType<MLFloat16>()),
DistributedUnsqueeze<MLFloat16, int64_t>);
ONNX_OPERATOR_TYPED_KERNEL_EX(
DistributedUnsqueeze,
kMSDomain,
1,
int64_t,
kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.InputMemoryType(OrtMemTypeCPUInput, 1)
.TypeConstraint("T", DataTypeImpl::GetTensorType<int64_t>()),
DistributedUnsqueeze<int64_t, int64_t>);
#endif
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

32
onnxruntime/contrib_ops/cuda/collective/distributed_unsqueeze.h Normal file
View file

@ -0,0 +1,32 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <algorithm>
#include <tuple>
#include <optional>
#include <string>
#include <nccl.h>
#include <sstream>
#include "sharding.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
#if defined(ORT_USE_NCCL)
template <typename T, typename Tind>
class DistributedUnsqueeze final : public DistributedKernel {
public:
explicit DistributedUnsqueeze(const OpKernelInfo& info);
Status ComputeInternal(OpKernelContext* context) const override;
};
#endif
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

22
onnxruntime/contrib_ops/cuda/collective/sharding_spec.h
View file

@ -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<int64_t>& axes_to_drop) {
std::vector<AxisPartitionSpec> 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<AxisPartitionSpec> 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.,

16
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
View file

@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, DistributedReduceMean)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, DistributedReduceMean)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int64_t, DistributedUnsqueeze)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, DistributedUnsqueeze)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, DistributedUnsqueeze)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, int64_t, DistributedSqueeze)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, DistributedSqueeze)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, DistributedSqueeze)>,
#endif
};

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

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

229
onnxruntime/test/python/onnxruntime_test_distributed.py
View file

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