[Debuggability] Add feature to ORTTrainer Frontend (#5124)

* add option, feature to orttrainer and test

* address comments

* minor fixes

* further address comments

* minor changes

Co-authored-by: Rayan Krishnan <t-rakr@OrtDevTest2v100.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>

orttraining/orttraining/python/training/orttrainer.py

@@ -285,6 +285,40 @@ class ORTTrainer(object):
         with open(path, "wb") as f:
             f.write(self._onnx_model.SerializeToString())
 
+    def _debug_model_export(self, input):
+        from onnx import helper, TensorProto, numpy_helper
+        import numpy as np
+        from numpy.testing import assert_allclose
+        import _test_helpers
+        onnx_model_copy = copy.deepcopy(self._onnx_model)
+        
+        # Mute the dropout nodes
+        dropout_nodes = [n for n in onnx_model_copy.graph.node if n.op_type == 'Dropout']
+        for node in dropout_nodes:
+            ratio_node = [n for n in onnx_model_copy.graph.node if node.input[1] in n.output][0]
+            training_mode_node = [n for n in onnx_model_copy.graph.node if node.input[2] in n.output][0]
+            
+            training_mode_node.attribute.pop()
+            ratio_node.attribute.pop()
+            new_training_mode_arr = np.array(False, dtype=bool)
+            new_ratio_arr = np.array(0.0, dtype=np.float32)
+            new_training_mode = numpy_helper.from_array(new_training_mode_arr)
+            new_ratio = numpy_helper.from_array(new_ratio_arr)
+            training_mode_node.attribute.add().t.CopyFrom(new_training_mode)
+            ratio_node.attribute.add().t.CopyFrom(new_ratio)
+            training_mode_node.attribute[0].type = 4
+            ratio_node.attribute[0].type = 4
+            training_mode_node.attribute[0].name = "value"
+            ratio_node.attribute[0].name = "value"
+            
+        _inference_sess = ort.InferenceSession(onnx_model_copy.SerializeToString())
+        inf_inputs = {}
+        for i, input_elem in enumerate(input):
+            inf_inputs[_inference_sess.get_inputs()[i].name] = input_elem.cpu().numpy()
+        _inference_outs = _inference_sess.run(None, inf_inputs)
+        for torch_item, ort_item in zip(self.torch_sample_outputs, _inference_outs):
+            assert_allclose(torch_item, ort_item, rtol=1e-2, atol=1e-6)
+
     def train_step(self, *args, **kwargs):
         r"""Train step method
 
@@ -303,6 +337,11 @@ class ORTTrainer(object):
         if self._onnx_model is None:
             sample_input = self._prepare_model_input(self.model_desc.inputs, None, None, *args, **kwargs)
             self._init_onnx_model(sample_input)
+            
+            # Debug Model Export if indicated
+            if self.options.debug.check_model_export:
+                self._debug_model_export(sample_input)
+
 
         # Prepare inputs+lr and output descriptions
         inputs_desc = self._model_desc_inputs_with_lr
@@ -461,6 +500,7 @@ class ORTTrainer(object):
                 warnings.warn("This model cannot be deep copied (or pickled), which is a required step for stateful models to be properly exported to ONNX."
                               " Compute will continue, but unexpected results may occur!")
             sample_outputs = model_copy(*sample_inputs_copy)
+            self.torch_sample_outputs = sample_outputs
         model.train()
 
         if isinstance(sample_outputs, torch.Tensor):

orttraining/orttraining/python/training/orttrainer_options.py

@@ -144,6 +144,10 @@ class ORTTrainerOptions(object):
                             'type' : 'boolean',
                             'default' : False
                         },
+                        'check_model_export' : {
+                            'type' : 'boolean',
+                            'default' : False
+                        },
                     }
                 },
                 '_internal_use' : {
@@ -225,6 +229,9 @@ class ORTTrainerOptions(object):
             debug options
         debug.deterministic_compute (bool, default is False)
             forces compute to be deterministic accross runs
+        debug.check_model_export (bool, default is False)
+            compares PyTorch outputs with InferenceSession outputs before the first
+            train step to ensure successful model export
         _internal_use (dict):
             internal options, possibly undocumented, that might be removed without notice
         _internal_use.enable_internal_postprocess (bool, default is True):
@@ -450,6 +457,10 @@ _ORTTRAINER_OPTIONS_SCHEMA = {
                 'type': 'boolean',
                 'default': False
             },
+            'check_model_export': {
+                'type': 'boolean',
+                'default': False
+            },
         }
     },
     '_internal_use': {

orttraining/orttraining/test/python/orttraining_test_debuggability.py

@@ -0,0 +1,52 @@
+import inspect
+import onnx
+import os
+import pytest
+import torch
+import torchvision
+
+from numpy.testing import assert_allclose
+
+from onnxruntime import set_seed
+from onnxruntime.capi.ort_trainer import IODescription as Legacy_IODescription,\
+                                         ModelDescription as Legacy_ModelDescription,\
+                                         LossScaler as Legacy_LossScaler,\
+                                         ORTTrainer as Legacy_ORTTrainer
+from onnxruntime.training import _utils, amp, optim, orttrainer, TrainStepInfo,\
+                                      model_desc_validation as md_val,\
+                                      orttrainer_options as orttrainer_options
+
+from orttraining_test_orttrainer_frontend import _load_pytorch_transformer_model
+
+import _test_helpers
+
+
+###############################################################################
+# Testing starts here #########################################################
+###############################################################################
+
+
+@pytest.mark.parametrize("seed, device", [
+    (24, 'cuda'),
+])
+def testORTTransformerModelExport(seed, device):
+    # Common setup
+    optim_config = optim.LambConfig()
+    opts = orttrainer.ORTTrainerOptions({
+        'debug' : {
+            'check_model_export': True,
+        },
+        'device' : {
+            'id' : device,
+        }
+    })
+
+    # Setup for the first ORTTRainer run
+    torch.manual_seed(seed)
+    set_seed(seed)
+    model, model_desc, my_loss, batcher_fn, train_data, val_data, _ = _load_pytorch_transformer_model(device)
+    first_trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=my_loss, options=opts)
+    data, targets = batcher_fn(train_data, 0)
+    _ = first_trainer.train_step(data, targets)
+    assert first_trainer._onnx_model is not None
+

orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py

@@ -100,7 +100,8 @@ def testORTTrainerOptionsDefaultValues(test_input):
             'invertible_layer_norm_gradient': False,
         },
         'debug': {
-            'deterministic_compute': False
+            'deterministic_compute': False,
+            'check_model_export': False
         },
         '_internal_use': {
             'enable_internal_postprocess': True,