diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index bd0152b348..5099eca81c 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -514,8 +514,8 @@ Do not modify directly.*
|||[1, 10]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)|
|GatherElements|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|13+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)|
|||[11, 12]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)|
-|GatherND|*in* data:**T**
*in* indices:**tensor(int64)**
*out* output:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int64)
**Tind** = tensor(int64)|
-|||12|**T** = tensor(double), tensor(float), tensor(float16), tensor(int64)
**Tind** = tensor(int64)|
+|GatherND|*in* data:**T**
*in* indices:**tensor(int64)**
*out* output:**T**|13+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int64)
**Tind** = tensor(int64)|
+|||12|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int64)
**Tind** = tensor(int64)|
|Gemm|*in* A:**T**
*in* B:**T**
*in* C:**T**
*out* Y:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|||[11, 12]|**T** = tensor(double), tensor(float), tensor(float16)|
|||[9, 10]|**T** = tensor(double), tensor(float), tensor(float16)|
diff --git a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc
index 64e1f958b2..214a2f527b 100644
--- a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc
+++ b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc
@@ -107,6 +107,7 @@ Status GatherNDBase::PrepareCompute(
DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType(), \
+ DataTypeImpl::GetTensorType(), \
}) \
.TypeConstraint("Tind", DataTypeImpl::GetTensorType()), \
GatherND);
@@ -117,15 +118,17 @@ Status GatherNDBase::PrepareCompute(
DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType(), \
+ DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType() }
-#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, BFloat16
+#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, BFloat16, bool
#else
#define GATHER_ND_T_TENSOR_TYPES \
{ DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType(), \
+ DataTypeImpl::GetTensorType(), \
DataTypeImpl::GetTensorType() }
-#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t
+#define GATHER_ND_T_DATA_TYPES float, MLFloat16, double, int64_t, bool
#endif
#define REGISTER_KERNEL_TYPED_GATHER_ND(TIndex, ver) \
diff --git a/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu b/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu
index 3f0275547c..272d3a3886 100644
--- a/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu
+++ b/onnxruntime/core/providers/cuda/tensor/gather_nd_impl.cu
@@ -107,6 +107,7 @@ void GatherNDImpl(
SPECIALIZED_COMPUTE_SLICE_OFFSETS_IMPL(int32_t)
SPECIALIZED_COMPUTE_SLICE_OFFSETS_IMPL(int64_t)
+SPECIALIZED_IMPL(bool)
SPECIALIZED_IMPL(float)
SPECIALIZED_IMPL(int64_t)
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 600
diff --git a/onnxruntime/core/providers/rocm/tensor/gather_nd_impl.cu b/onnxruntime/core/providers/rocm/tensor/gather_nd_impl.cu
index 5d70cbe7d4..fed7d1b544 100644
--- a/onnxruntime/core/providers/rocm/tensor/gather_nd_impl.cu
+++ b/onnxruntime/core/providers/rocm/tensor/gather_nd_impl.cu
@@ -107,6 +107,7 @@ void GatherNDImpl(
SPECIALIZED_COMPUTE_SLICE_OFFSETS_IMPL(int32_t)
SPECIALIZED_COMPUTE_SLICE_OFFSETS_IMPL(int64_t)
+SPECIALIZED_IMPL(bool)
SPECIALIZED_IMPL(float)
SPECIALIZED_IMPL(int64_t)
SPECIALIZED_IMPL(half)
diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py
index e7c7886108..736ef00b9e 100755
--- a/onnxruntime/python/tools/symbolic_shape_infer.py
+++ b/onnxruntime/python/tools/symbolic_shape_infer.py
@@ -191,6 +191,7 @@ class SymbolicShapeInference:
'aten::embedding': self._infer_Gather,
'aten::diagonal': self._infer_aten_diagonal,
'aten::max_pool2d_with_indices': self._infer_aten_pool2d,
+ 'aten::multinomial': self._infer_aten_multinomial,
'aten::unfold': self._infer_aten_unfold,
'aten::argmax': self._infer_aten_argmax,
'aten::avg_pool2d': self._infer_aten_pool2d,
@@ -1104,6 +1105,19 @@ class SymbolicShapeInference:
helper.make_tensor_value_info(node.output[0], self.known_vi_[node.input[0]].type.tensor_type.elem_type,
get_shape_from_sympy_shape(new_shape)))
+ def _infer_aten_multinomial(self, node):
+ sympy_shape = self._get_sympy_shape(node, 0)
+ rank = len(sympy_shape)
+ assert rank in [1,2]
+ num_samples = self._try_get_value(node, 1)
+ di = rank - 1
+ last_dim = num_samples if num_samples else str(self._new_symbolic_dim_from_output(node, 0, di))
+ output_shape = sympy_shape[:-1] + [last_dim]
+ vi = self.known_vi_[node.output[0]]
+ vi.CopyFrom(
+ helper.make_tensor_value_info(node.output[0], onnx.TensorProto.INT64,
+ get_shape_from_sympy_shape(output_shape)))
+
def _infer_aten_pool2d(self, node):
sympy_shape = self._get_sympy_shape(node, 0)
assert len(sympy_shape) == 4
diff --git a/orttraining/orttraining/core/framework/gradient_graph_builder.cc b/orttraining/orttraining/core/framework/gradient_graph_builder.cc
index df9fbb4a81..2f09afb113 100644
--- a/orttraining/orttraining/core/framework/gradient_graph_builder.cc
+++ b/orttraining/orttraining/core/framework/gradient_graph_builder.cc
@@ -123,8 +123,8 @@ NodeSet GradientGraphBuilder::BFSWithStopGradient(const std::unordered_set nodes = graph_->GetConsumerNodes(name);
for (const Node* node : nodes) {
int input_index = graph_utils::GetNodeInputIndexFromInputName(*node, name);
- auto it = STOP_GRADIENT_EDGES.find(node->OpType());
- if (it != STOP_GRADIENT_EDGES.end() && it->second.count(input_index)) {
+ const std::unordered_set* edges = GetStopGradientEdges(*node);
+ if (edges != nullptr && edges->count(input_index)) {
continue;
}
queue.push_back(node);
@@ -139,8 +139,8 @@ NodeSet GradientGraphBuilder::BFSWithStopGradient(const std::unordered_setOutputEdgesBegin(); edge_it != n->OutputEdgesEnd(); ++edge_it) {
const Node& node = edge_it->GetNode();
- auto it = STOP_GRADIENT_EDGES.find(node.OpType());
- if (it != STOP_GRADIENT_EDGES.end() && it->second.count(edge_it->GetDstArgIndex())) {
+ const std::unordered_set* edges = GetStopGradientEdges(node);
+ if (edges != nullptr && edges->count(edge_it->GetDstArgIndex())) {
continue;
}
@@ -163,8 +163,8 @@ NodeSet GradientGraphBuilder::ReverseBFSWithStopGradient(const NodeSet& nodes) c
queue.pop_front();
for (auto edge_it = n->InputEdgesBegin(); edge_it != n->InputEdgesEnd(); ++edge_it) {
- auto it = STOP_GRADIENT_EDGES.find(n->OpType());
- if (it != STOP_GRADIENT_EDGES.end() && it->second.count(edge_it->GetDstArgIndex())) {
+ const std::unordered_set* edges = GetStopGradientEdges(*n);
+ if (edges != nullptr && edges->count(edge_it->GetDstArgIndex())) {
LOGS(logger_, INFO) << "Skip building gradient for input_" << edge_it->GetDstArgIndex()
<< " of node: " << n->Name();
continue;
@@ -200,6 +200,22 @@ Status GradientGraphBuilder::CheckNodeArgsReachable() const {
return Status::OK();
}
+const std::unordered_set* GradientGraphBuilder::GetStopGradientEdges(const Node& node) const {
+ std::string op_type = node.OpType();
+
+ if (op_type == "ATenOp") {
+ std::string key = GetGradientDefinitionKeyByNode(node);
+ return GradientDefinitionRegistry::Instance().GetStopGradientEdgesForNode(key);
+ } else {
+ auto it = STOP_GRADIENT_EDGES.find(op_type);
+ if (it == STOP_GRADIENT_EDGES.end()) {
+ return nullptr;
+ }
+
+ return &it->second;
+ }
+}
+
Status GradientGraphBuilder::Build(const std::unordered_set* p_initializer_names_to_preserve) {
auto opt_ret = graph_transformation_mgr_.ApplyTransformers(*graph_, TransformerLevel::Level2, logger_);
ORT_RETURN_IF_ERROR(opt_ret);
@@ -233,8 +249,8 @@ Status GradientGraphBuilder::Build(const std::unordered_set* p_init
if (!IsReachable(&next_node)) continue;
- auto it = STOP_GRADIENT_EDGES.find(next_node.OpType());
- if (it != STOP_GRADIENT_EDGES.end() && it->second.count(edge_it->GetDstArgIndex())) {
+ const std::unordered_set* edges = GetStopGradientEdges(next_node);
+ if (edges != nullptr && edges->count(edge_it->GetDstArgIndex())) {
LOGS(logger_, WARNING) << "Skip building gradient for input_" << edge_it->GetDstArgIndex()
<< " of node: " << next_node.Name();
continue;
diff --git a/orttraining/orttraining/core/framework/gradient_graph_builder.h b/orttraining/orttraining/core/framework/gradient_graph_builder.h
index 9b952a4bc1..9fcc71f46b 100644
--- a/orttraining/orttraining/core/framework/gradient_graph_builder.h
+++ b/orttraining/orttraining/core/framework/gradient_graph_builder.h
@@ -138,6 +138,8 @@ class GradientGraphBuilder {
// Tracks tensors that are stashed in the forward pass for later use in backward pass.
std::unordered_set stashed_tensors_;
+ const std::unordered_set* GetStopGradientEdges(const Node& node) const;
+
/**
Performs a BFS on the graph with STOP_GRADIENT_EDGES constrain
It will skip traversing over the edges defined in STOP_GRADIENT_EDGES map.
diff --git a/orttraining/orttraining/core/graph/gradient_definition_registry.h b/orttraining/orttraining/core/graph/gradient_definition_registry.h
index 9de17218eb..f4caea8c3b 100644
--- a/orttraining/orttraining/core/graph/gradient_definition_registry.h
+++ b/orttraining/orttraining/core/graph/gradient_definition_registry.h
@@ -51,8 +51,22 @@ class GradientDefinitionRegistry {
definitions_.emplace(key, definition);
}
+ void SetStopGradientEdgesForNode(const std::string& key, const std::unordered_set edges) {
+ custom_stop_gradient_edges_.emplace(key, edges);
+ }
+
+ const std::unordered_set* GetStopGradientEdgesForNode(const std::string& key) {
+ auto it = custom_stop_gradient_edges_.find(key);
+ if (it == custom_stop_gradient_edges_.end()) {
+ return nullptr;
+ }
+
+ return &it->second;
+ }
+
private:
std::unordered_map> definitions_;
+ std::unordered_map> custom_stop_gradient_edges_;
};
} // namespace training
diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc
index 9cfeb5078e..91a70f07b8 100644
--- a/orttraining/orttraining/python/orttraining_pybind_state.cc
+++ b/orttraining/orttraining/python/orttraining_pybind_state.cc
@@ -749,6 +749,11 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
[](const std::string& key, const std::vector& gradient_def) -> void {
GradientDefinitionRegistry::Instance().Register(key, gradient_def);
});
+
+ m.def("register_custom_stop_gradient_edges",
+ [](const std::string& key, const std::unordered_set edges) -> void {
+ GradientDefinitionRegistry::Instance().SetStopGradientEdgesForNode(key, edges);
+ });
}
} // namespace python
diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py
index e3225678cd..251bbef8bd 100644
--- a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py
+++ b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py
@@ -56,17 +56,24 @@ def _to_gradient_definition(gradient):
class CustomGradientRegistry:
_GRADIENTS = {}
+ _STOP_GRADIENT_EDGES = {}
@classmethod
def register(cls, domain, name, attributes, fn):
key = '::'.join([domain, name] + list(attributes))
cls._GRADIENTS[key] = _to_gradient_definition(fn())
+ @classmethod
+ def register_custom_stop_gradient_edges(cls, edges, domain, name, *attributes):
+ key = '::'.join([domain, name] + list(attributes))
+ cls._STOP_GRADIENT_EDGES[key] = set(edges)
+
@classmethod
def register_all(cls):
for key, value in cls._GRADIENTS.items():
C.register_gradient_definition(key, value)
-
+ for key, value in cls._STOP_GRADIENT_EDGES.items():
+ C.register_custom_stop_gradient_edges(key, value)
def register_gradient(domain, name, *attributes):
def gradient_wrapper(fn):
@@ -125,3 +132,5 @@ def adaptive_avg_pool2d_gradient():
(('ATenOp', 'com.microsoft'), ['GO(0)', 'I(0)'], [
'GI(0)'], {'name': {'value': 'aten::_adaptive_avg_pool2d_backward', 'dtype': 'string'}}),
]
+
+CustomGradientRegistry.register_custom_stop_gradient_edges([0], 'com.microsoft', 'ATenOp', 'aten::multinomial', '')
\ No newline at end of file
diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py
index 954ddc1718..a29e03300d 100644
--- a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py
+++ b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py
@@ -78,6 +78,12 @@ def diagonal(g, self, offset, dim1, dim2):
return g.op("com.microsoft::ATenOp", self, offset, dim1, dim2,
name_s='aten::diagonal')
+@register_symbolic('multinomial')
+def multinomial(g, self, num_samples, replacement=False, generator=None):
+ if generator is not None and not sym_help._is_none(generator):
+ raise RuntimeError("Unsupported: ONNX does not support generator for multinomial")
+ return g.op("com.microsoft::ATenOp", self, num_samples, replacement, generator,
+ name_s='aten::multinomial')
@register_symbolic('max_pool2d')
def max_pool2d(g, self, kernel_size, stride, padding, dilation, ceil_mode):
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
index b9cb2c6af4..2ba544ae9f 100644
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
@@ -975,6 +975,44 @@ def test_gradient_correctness_argmax_diagonal(offset, dim1, dim2):
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
_test_helpers.assert_values_are_close(ort_input.grad, pt_input.grad)
+# Since multinomial is a generator function, we do not have to test for gradient
+# Two consecutive calls on the torch.multinomail on a probability distribution with more
+# than one index with non-zero probability(eg, [0, 10, 3, 0]) will not result in
+# the same output. Thus we reset the seed before each call to the op torch.multinomial.
+@pytest.mark.parametrize("input_shape", ([5], [2,5]))
+@pytest.mark.parametrize("num_samples, replacement", ((1, False), (2, True)))
+def test_aten_multinomial(input_shape, num_samples, replacement):
+ class NeuralNetDiagonal(torch.nn.Module):
+ def __init__(self, num_samples, replacement):
+ super(NeuralNetDiagonal, self).__init__()
+ self.num_samples = num_samples
+ self.replacement = replacement
+
+ def forward(self, input):
+ return torch.multinomial(input, self.num_samples, self.replacement)
+
+ torch.backends.cudnn.deterministic = True
+ device = 'cuda'
+ pt_model = NeuralNetDiagonal(num_samples, replacement).to(device)
+ ort_model = ORTModule(copy.deepcopy(pt_model))
+
+ def run_step(model, input):
+ # reset manual seed to reset the generator
+ torch.manual_seed(5032)
+ prediction = model(input)
+ return prediction
+
+ pt_input = torch.rand(input_shape, dtype=torch.float, device=device)
+ ort_input = copy.deepcopy(pt_input)
+ pt_prediction = run_step(pt_model, pt_input)
+ ort_prediction = run_step(ort_model, ort_input)
+ # run the ort prediction again since the first call involves export
+ # and run step, which means the torch.multinomial is called twice in a row without
+ # resetting the generator in between, which will result in a different output
+ ort_prediction = run_step(ort_model, ort_input)
+
+ _test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
+
def test_module_with_non_differential_output():
device = 'cuda'
N, D_in, H, D_out = 32, 128, 64, 10