Support for primitive types in ortmodule (#7588)

orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py

@@ -229,7 +229,7 @@ class GraphExecutionManager(ABC):
         # Therefore, deepcopy only the data component of the input tensors for export.
         sample_inputs_copy, sample_kwargs_copy = _io.deepcopy_model_input(*inputs, **kwargs)
         # NOTE: Flattening the input will change the 'input schema', resulting in a re-export
-        sample_inputs_as_tuple = tuple(self._input_info.flatten(sample_inputs_copy, sample_kwargs_copy))
+        sample_inputs_as_tuple = tuple(self._input_info.flatten(sample_inputs_copy, sample_kwargs_copy, self._device))
         # Ops behaving differently under train/eval mode need to exported with the
         # correct training flag to reflect the expected behavior.
         # For example, the Dropout node in a model is dropped under eval mode.

orttraining/orttraining/python/training/ortmodule/_inference_manager.py

@@ -94,7 +94,8 @@ class InferenceManager(GraphExecutionManager):
                                                                              self._input_info,
                                                                              self._flattened_module.named_buffers(),
                                                                              inputs,
-                                                                             kwargs))
+                                                                             kwargs,
+                                                                             self._device))
 
         return _io.unflatten_user_output(self._module_output_schema,
                                          self._graph_info.user_output_names,

orttraining/orttraining/python/training/ortmodule/_io.py

@@ -9,6 +9,22 @@ import inspect
 import torch
 import warnings
 
+class _PrimitiveType(object):
+    _primitive_types = {int, bool, float}
+    @staticmethod
+    def is_primitive_type(value):
+        return type(value) in _PrimitiveType._primitive_types
+
+    @staticmethod
+    def get_tensor(value, device):
+        return torch.tensor(value, device=device)
+
+    @staticmethod
+    def get_primitive_dtype(value):
+        # If `value` is a boolean, save the value of the boolean in dtype.
+        # This way, if the value changes from one forward call to the next, the schema will mismatch,
+        # and the model will be re-exported.
+        return f"{str(type(value))}_{value}" if isinstance(value, bool) else str(type(value))
 
 class _InputInfo(object):
     def __init__(self,
@@ -40,11 +56,13 @@ class _InputInfo(object):
             \t#Positionals (non-None): {self.num_positionals_non_none}
             \tKeyword names:           {self.keyword_names}'''
 
-    def flatten(self, args, kwargs):
+    def flatten(self, args, kwargs, device):
         '''Flatten args and kwargs in a single tuple of tensors with strict ordering'''
 
-        ret = list(args)
-        ret += [kwargs[name] for name in self.names if name in kwargs]
+        ret = [_PrimitiveType.get_tensor(arg, device) if _PrimitiveType.is_primitive_type(arg) else arg for arg in args]
+        ret += [_PrimitiveType.get_tensor(kwargs[name], device) if _PrimitiveType.is_primitive_type(kwargs[name])
+            else kwargs[name] for name in self.names if name in kwargs]
+
         return ret
 
     def unflatten(self, flat_args):
@@ -55,7 +73,7 @@ class _InputInfo(object):
             if name in self.keyword_names}
         return args, kwargs
 
-def _combine_input_buffers_initializers(param_names, onnx_input_names, input_info, buffer_names, inputs, kwargs):
+def _combine_input_buffers_initializers(param_names, onnx_input_names, input_info, buffer_names, inputs, kwargs, device):
     '''Creates forward `*inputs` list from user input and PyTorch initializers
 
     ONNX Runtime forward requires an ordered list of:
@@ -88,6 +106,8 @@ def _combine_input_buffers_initializers(param_names, onnx_input_names, input_inf
                 raise KeyError(f'Registered buffer name {name} not found.')
 
         if inp is not None:
+            if _PrimitiveType.is_primitive_type(inp):
+                inp = _PrimitiveType.get_tensor(inp, device)
             result.append(inp)
         else:
             raise RuntimeError(f'Input is present in ONNX graph but not provided: {name}.')
@@ -201,6 +221,8 @@ def _extract_schema(data):
 
     if data is None:
         return None
+    elif _PrimitiveType.is_primitive_type(data):
+        return _TensorStub(dtype=_PrimitiveType.get_primitive_dtype(data), shape_dims=0)
     # Depth first traversal to iterate over the data to replace every tensor with a stub
     elif isinstance(data, torch.Tensor):
         return _TensorStub(dtype=str(data.dtype), shape_dims=len(data.size()))
@@ -305,10 +327,17 @@ def parse_inputs_for_onnx_export(all_input_parameters, onnx_graph, inputs, kwarg
         return dynamic_axes
 
     def _add_input(name, input, onnx_graph, onnx_graph_input_names):
-        if input is not None and (onnx_graph is None or name in onnx_graph_input_names):
+        if input is None:
+            # Drop all None inputs.
+            return
+
+        # InputInfo should contain all the names irrespective of whether they are
+        # a part of the onnx graph or not.
+        input_names.append(name)
+
+        if (onnx_graph is None or name in onnx_graph_input_names) and isinstance(input, torch.Tensor):
             if input.requires_grad:
                 input_names_require_grad.append(name)
-            input_names.append(name)
             dynamic_axes.update(_add_dynamic_shape(name, input))
             input_shape.append(list(input.size()))
 

orttraining/orttraining/python/training/ortmodule/_training_manager.py

@@ -73,7 +73,9 @@ class TrainingManager(GraphExecutionManager):
                                                       kwargs)
 
         # Reinitialize graph builder if the inputs or initializers requiring gradient have changed.
-        build_gradient_graph = build_gradient_graph or self._reinitialize_graph_builder(input_info)
+        # Order of or operation is important here because we always need to call
+        # _reinitialize_graph_builder irrespective of the value of build_gradient_graph.
+        build_gradient_graph = self._reinitialize_graph_builder(input_info) or build_gradient_graph
 
         # Build the gradient graph
         if build_gradient_graph:
@@ -191,7 +193,8 @@ class TrainingManager(GraphExecutionManager):
                                                 self._input_info,
                                                 self._flattened_module.named_buffers(),
                                                 inputs,
-                                                kwargs)))
+                                                kwargs,
+                                                self._device)))
 
     def _build_graph(self):
         """Build an optimized gradient graph using the module_graph_builder"""

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -2252,3 +2252,79 @@ def test_forward_call_lots_None():
                  a, b, c, d, e, f, y, z)
         run_step(a.item() + b.item() + c.item() + d.item() + e.item() + f.item() + y.item() + z.item(),
                  **{'a': a, 'b': b, 'c': c, 'd': d, 'e': e, 'f': f, 'y': y, 'z': z})
+
+@pytest.mark.parametrize("bool_argument", [True, False])
+@pytest.mark.parametrize("int_argument", [100, 100000, 100000000, -100, -100000, -100000000])
+@pytest.mark.parametrize("float_argument", [1.23, 11209123.12452, 12093702935.1249863, -1.23, -11209123.12452, -12093702935.1249863])
+def test_primitive_inputs(bool_argument, int_argument, float_argument):
+    class PrimitiveTypesInputNet(torch.nn.Module):
+        def __init__(self, input_size, hidden_size, num_classes):
+            super(PrimitiveTypesInputNet, self).__init__()
+
+            self.fc1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu = torch.nn.ReLU()
+            self.fc2 = torch.nn.Linear(hidden_size, num_classes)
+
+        def forward(self, input1, bool_argument, int_argument, float_argument):
+            input1 = input1 + int_argument + float_argument
+            if bool_argument:
+                out = self.fc1(input1)
+                out = self.relu(out)
+                out = self.fc2(out)
+            else:
+                out = self.fc1(input1)
+                out = self.fc2(out)
+                out = self.relu(out)
+            return out
+
+    assert type(bool_argument) is bool
+    assert type(int_argument) is int
+    assert type(float_argument) is float
+
+    device = 'cuda'
+    N, D_in, H, D_out = 32, 784, 500, 10
+    pt_model = PrimitiveTypesInputNet(D_in, H, D_out).to(device)
+    ort_model = ORTModule(copy.deepcopy(pt_model))
+
+    input1 = torch.randn(N, D_in, device=device)
+    pt_out = pt_model(input1, bool_argument, int_argument, float_argument)
+    ort_out = ort_model(input1, bool_argument, int_argument, float_argument)
+    assert torch.equal(pt_out, ort_out)
+
+@pytest.mark.parametrize("bool_arguments", [(True, False), (False, True)])
+def test_changing_bool_input_re_exports_model(bool_arguments):
+    class PrimitiveTypesInputNet(torch.nn.Module):
+        def __init__(self, input_size, hidden_size, num_classes):
+            super(PrimitiveTypesInputNet, self).__init__()
+
+            self.fc1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu = torch.nn.ReLU()
+            self.fc2 = torch.nn.Linear(hidden_size, num_classes)
+
+        def forward(self, input1, bool_argument):
+            if bool_argument:
+                out = self.fc1(input1)
+                out = self.relu(out)
+                out = self.fc2(out)
+            else:
+                out = self.fc1(input1)
+                out = self.fc2(out)
+                out = self.relu(out)
+            return out
+
+    assert type(bool_arguments[0]) is bool
+    assert type(bool_arguments[1]) is bool
+
+    device = 'cuda'
+    N, D_in, H, D_out = 32, 784, 500, 10
+    pt_model = PrimitiveTypesInputNet(D_in, H, D_out).to(device)
+    ort_model = ORTModule(pt_model)
+
+    input1 = torch.randn(N, D_in, device=device)
+    ort_model(input1, bool_arguments[0])
+    exported_model1 = ort_model._execution_manager(ort_model._is_training())._onnx_model
+
+    ort_model(input1, bool_arguments[1])
+    exported_model2 = ort_model._execution_manager(ort_model._is_training())._onnx_model
+
+    assert exported_model1 != exported_model2