Fix pythonop training_mode in evaluation mode (#13514)

Customer reported this issue: they see many warnings when doing hte
evaluation using ORTModule.


![image](https://user-images.githubusercontent.com/10530022/199371757-5fed7d05-a951-4f1b-8f88-049c5ab89886.png)

After investigation, we found the `training_mode` is exported to a wrong
value in evaluation mode, it's value should be 0, but we found it is 1.

Fix: 
fix pythonop training mode

if training_mode's type is torch._C._onnx.TrainingMode, then not matter
it is EVAL or TRAINING, "if training_mode" will always be true

orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py

@@ -88,9 +88,19 @@ def _export_pt_1_10(g, n, *args, **kwargs):
         training_mode = None
         runtime_pytorch_version = version.parse(torch.__version__.split("+")[0])
         if runtime_pytorch_version >= version.parse("1.12"):
+            # FIXME: using privated modules
             from torch.onnx import _globals
 
-            training_mode = _globals.GLOBALS.training_mode
+            # before https://github.com/pytorch/pytorch/commit/c8b9b6266b505328e503b12f6a42fd88c56374f9, training_mode is still a bool type
+            if isinstance(_globals.GLOBALS.training_mode, bool):
+                training_mode = _globals.GLOBALS.training_mode
+            else:
+                if _globals.GLOBALS.training_mode not in [
+                    torch.onnx.TrainingMode.EVAL,
+                    torch.onnx.TrainingMode.TRAINING,
+                ]:
+                    raise Exception(f"Unexpected training mode {_globals.GLOBALS.training_mode}")
+                training_mode = _globals.GLOBALS.training_mode == torch.onnx.TrainingMode.TRAINING
         else:
             training_mode = symbolic_helper._training_mode
         cconv = n.cconv()

orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py

@@ -3,14 +3,14 @@
 
 # pylint: disable=missing-docstring
 # pylint: disable=C0103
-
-from packaging.version import Version
+# pylint: disable=W0212
 
 import pytest
 import torch
 
 # Import ORT modules.
 from _test_helpers import *
+from packaging.version import Version
 from torch.nn.parameter import Parameter
 
 # Import external libraries.
@@ -1210,3 +1210,65 @@ def test_skipped_autograd_function():
     assert not can_run
 
     del os.environ["ORTMODULE_SKIPPED_AUTOGRAD_FUNCTIONS"]
+
+
+def test_pythonop_training_mode():
+    def check_pythonop_training_mode(model, is_eval_mode):
+        ## make sure the ort's PythonOp's training_mode is correct
+        if is_eval_mode:
+            onnx_nodes = (
+                model._torch_module._execution_manager._inference_manager._onnx_models.exported_model.graph.node
+            )
+        else:
+            onnx_nodes = model._torch_module._execution_manager._training_manager._onnx_models.exported_model.graph.node
+
+        found_pythonop = False
+        for node in onnx_nodes:
+            if node.op_type == "PythonOp":
+                found_pythonop = True
+                for attr in node.attribute:
+                    if attr.name == "training_mode":
+                        if is_eval_mode:
+                            assert attr.i == 0, f"in eval mode, it shoule be 0, while it is {attr.i} now"
+                        else:
+                            assert attr.i == 1, f"in training mode, it should be 1, while it is {attr.i} now"
+
+        assert found_pythonop, "PythonOp should be found in the exported model"
+
+    class TestFunction(torch.autograd.Function):
+        @staticmethod
+        # bias is an optional argument
+        def forward(ctx, x):
+            ctx.save_for_backward(x)
+            return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            x = ctx.saved_tensors
+            tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+            ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
+            return ff * grad_output
+
+    class TestModel(torch.nn.Module):
+        def __init__(self, output_size):
+            super(TestModel, self).__init__()
+            self.custom_fn = TestFunction.apply
+            self.bias = Parameter(torch.empty(output_size, dtype=torch.float))
+
+            with torch.no_grad():
+                self.bias.uniform_()
+
+        def forward(self, model_input):
+            # model_input did not require_grad
+            out = self.custom_fn(model_input)
+            return out + self.bias
+
+    output_size = 1024
+    # 1 check traning mode
+    ortmodule = ORTModule(TestModel(output_size)).train()
+    _ = ortmodule(torch.randn(output_size, dtype=torch.float))
+    check_pythonop_training_mode(ortmodule, is_eval_mode=False)
+    # 2 check eval mode
+    ortmodule = ORTModule(TestModel(output_size)).eval()
+    _ = ortmodule(torch.randn(output_size, dtype=torch.float))
+    check_pythonop_training_mode(ortmodule, is_eval_mode=True)