Fixes for GatherND, Multinomial (#9143)

* register gathernd kernel, aten multinomial

* fix CI, add test

* review comments

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)<br/> **Tind** = tensor(int32), tensor(int64)|
 |GatherElements|*in* data:**T**<br> *in* indices:**Tind**<br> *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)<br/> **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)<br/> **Tind** = tensor(int32), tensor(int64)|
-|GatherND|*in* data:**T**<br> *in* indices:**tensor(int64)**<br> *out* output:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int64)<br/> **Tind** = tensor(int64)|
-|||12|**T** = tensor(double), tensor(float), tensor(float16), tensor(int64)<br/> **Tind** = tensor(int64)|
+|GatherND|*in* data:**T**<br> *in* indices:**tensor(int64)**<br> *out* output:**T**|13+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int64)<br/> **Tind** = tensor(int64)|
+|||12|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int64)<br/> **Tind** = tensor(int64)|
 |Gemm|*in* A:**T**<br> *in* B:**T**<br> *in* C:**T**<br> *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)|

onnxruntime/core/providers/cuda/tensor/gather_nd.cc

@@ -107,6 +107,7 @@ Status GatherNDBase::PrepareCompute(
                               DataTypeImpl::GetTensorType<double>(),        \
                               DataTypeImpl::GetTensorType<MLFloat16>(),     \
                               DataTypeImpl::GetTensorType<int64_t>(),       \
+                              DataTypeImpl::GetTensorType<bool>(),          \
                           })                                                \
           .TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()),   \
       GatherND<TIndex>);
@@ -117,15 +118,17 @@ Status GatherNDBase::PrepareCompute(
     DataTypeImpl::GetTensorType<double>(),    \
     DataTypeImpl::GetTensorType<MLFloat16>(), \
     DataTypeImpl::GetTensorType<BFloat16>(),  \
+    DataTypeImpl::GetTensorType<bool>(),      \
     DataTypeImpl::GetTensorType<int64_t>() }
-#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<float>(),     \
     DataTypeImpl::GetTensorType<double>(),    \
     DataTypeImpl::GetTensorType<MLFloat16>(), \
+    DataTypeImpl::GetTensorType<bool>(),      \
     DataTypeImpl::GetTensorType<int64_t>() }
-#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)                      \

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

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)

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

orttraining/orttraining/core/framework/gradient_graph_builder.cc

@@ -123,8 +123,8 @@ NodeSet GradientGraphBuilder::BFSWithStopGradient(const std::unordered_set<std::
     std::vector<const Node*> 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<size_t>* 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_set<std::
     for (auto edge_it = n->OutputEdgesBegin(); 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<size_t>* 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<size_t>* 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<size_t>* 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<std::string>* 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<std::string>* 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<size_t>* 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;

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<std::string> stashed_tensors_;
 
+  const std::unordered_set<size_t>* 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.

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<size_t> edges) {
+    custom_stop_gradient_edges_.emplace(key, edges);
+  }
+
+  const std::unordered_set<size_t>* 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<std::string, std::vector<GradientNodeDefinition>> definitions_;
+  std::unordered_map<std::string, std::unordered_set<size_t>> custom_stop_gradient_edges_;
 };
 
 }  // namespace training

orttraining/orttraining/python/orttraining_pybind_state.cc

@@ -749,6 +749,11 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
         [](const std::string& key, const std::vector<GradientNodeDefinition>& gradient_def) -> void {
           GradientDefinitionRegistry::Instance().Register(key, gradient_def);
         });
+  
+  m.def("register_custom_stop_gradient_edges",
+        [](const std::string& key, const std::unordered_set<size_t> edges) -> void {
+          GradientDefinitionRegistry::Instance().SetStopGradientEdgesForNode(key, edges);
+        });
 }
 
 }  // namespace python

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', '')

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

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