Support nested sequence and mapping types in ORTModule (#6791)

orttraining/orttraining/python/training/_ortmodule_output_transformation.py

@@ -0,0 +1,220 @@
+from collections import abc
+import copy
+import functools
+import torch
+import warnings
+
+def deepcopy_model_input(*inputs, **kwargs):
+    sample_inputs_copy = []
+    for model_input in inputs:
+        sample_inputs_copy.append(model_input.data if isinstance(model_input, torch.Tensor) else model_input)
+    sample_inputs_copy = copy.deepcopy(tuple(sample_inputs_copy))
+
+    sample_kwargs_copy = {}
+    for name, model_input in kwargs.items():
+        sample_kwargs_copy[name] = model_input.data if isinstance(model_input, torch.Tensor) else model_input
+    sample_kwargs_copy = copy.deepcopy(sample_kwargs_copy)
+
+    return sample_inputs_copy, sample_kwargs_copy
+
+class _TensorStub:
+    # Stub for a torch.Tensor value to be used to formulate the output schema
+    pass
+
+def populate_user_output_from_schema_and_outputs(output_schema, output_names, outputs):
+    """Follows the schema to generate an output that is expected by the user"""
+
+    def _replace_stub_with_tensor_value(user_output, outputs, output_idx):
+        # Recursively traverse across user_output and replace all _TensorStub
+        # with torch.Tensor values from outputs following output_idx
+
+        if isinstance(user_output, _TensorStub):
+            output_idx[0] += 1
+            return outputs[output_idx[0]-1]
+
+        if isinstance(user_output, abc.Sequence):
+            sequence_type = type(user_output)
+            user_output = list(user_output)
+            for idx in range(len(user_output)):
+                user_output[idx] = _replace_stub_with_tensor_value(user_output[idx], outputs, output_idx)
+            try:
+                # namedtuple can be created by passing the list sequence to method _make
+                user_output = sequence_type._make(user_output)
+            except AttributeError:
+                # If attribute error encountered, create the sequence directly
+                user_output = sequence_type(user_output)
+        elif isinstance(user_output, abc.Mapping):
+            for key in sorted(user_output):
+                user_output[key] = _replace_stub_with_tensor_value(user_output[key], outputs, output_idx)
+        else:
+            raise TypeError(f'ORTModule does not support the following model output type {type(user_output)}.')
+
+        return user_output
+
+    # Order the outputs according to the names so that the traversal order is consistent
+    outputs = [x for _, x in sorted(zip(output_names, outputs))]
+
+    # Replace every _TensorStub value in the schema with the torch.Tensor outputs calculated
+    output_schema_copy = copy.deepcopy(output_schema)
+    output_idx = [0]
+    user_output = _replace_stub_with_tensor_value(output_schema_copy, outputs, output_idx)
+
+    return user_output
+
+def _extract_output_schema(output):
+    """Extract the output schema by replacing every torch.Tensor value with _TensorStub"""
+
+    # Depth first traversal to iterate over the output to replace every tensor with a stub
+    if isinstance(output, torch.Tensor):
+        return _TensorStub()
+
+    if isinstance(output, abc.Sequence):
+        sequence_type = type(output)
+        output = list(output)
+        for idx in range(len(output)):
+            output[idx] = _extract_output_schema(output[idx])
+        try:
+            # namedtuple can be created by passing the list sequence to method _make
+            output = sequence_type._make(output)
+        except AttributeError:
+            # If attribute error encountered, create the sequence directly
+            output = sequence_type(output)
+    elif isinstance(output, abc.Mapping):
+        for key in sorted(output):
+            output[key] = _extract_output_schema(output[key])
+    else:
+        raise TypeError(f'ORTModule does not support the following model output type {type(output)}')
+
+    return output
+
+def _parse_outputs_and_extract_names_and_dynamic_axes(module_output):
+    """Parses through the module output and returns output names and dynamic axes"""
+
+    def _populate_output_names_and_dynamic_axes(output, output_names, output_dynamic_axes, output_idx):
+        # Depth first traversal to traverse through the entire output collecting output names and dynamic axes
+
+        if isinstance(output, torch.Tensor):
+            output_name = f'output{output_idx[0]}'
+            output_idx[0] += 1
+            output_names.append(output_name)
+            output_dynamic_axes[output_name] = {}
+            for dim_idx in range(len(output.shape)):
+                output_dynamic_axes[output_name].update({dim_idx: f'{output_name}_dim{dim_idx}'})
+            return
+
+        if isinstance(output, abc.Sequence):
+            for value in output:
+                _populate_output_names_and_dynamic_axes(value, output_names, output_dynamic_axes, output_idx)
+        elif isinstance(output, abc.Mapping):
+            for _, value in sorted(output.items()):
+                _populate_output_names_and_dynamic_axes(value, output_names, output_dynamic_axes, output_idx)
+        else:
+            raise TypeError(f'ORTModule does not support the following model output type {type(output)}')
+
+    output_names = []
+    output_dynamic_axes = {}
+    output_idx = [0]
+    _populate_output_names_and_dynamic_axes(module_output, output_names, output_dynamic_axes, output_idx)
+
+    return output_names, output_dynamic_axes
+
+def get_flattened_output_module(original_module):
+    """Returns a torch.nn.Module that flattens the output of the original module in its forward method"""
+
+    def _transform_output_to_flat_tuple(output):
+        """Converts the output to a flat tuple by iterating over the entire output structure"""
+
+        def _flatten_output(output, flat_output):
+            # Recursively traverse over the output and populate the flat_output with torch.Tensors
+
+            if isinstance(output, torch.Tensor):
+                flat_output.append(output)
+            elif isinstance(output, abc.Sequence):
+                for value in output:
+                    _flatten_output(value, flat_output)
+            elif isinstance(output, abc.Mapping):
+                for _, value in sorted(output.items()):
+                    _flatten_output(value, flat_output)
+            else:
+                raise TypeError(f'ORTModule does not support the following output type {type(output)}.')
+
+        flat_output = []
+        _flatten_output(output, flat_output)
+        return tuple(flat_output)
+
+    class FlattenedOutputModule(torch.nn.Module):
+        def __init__(self, module):
+            super(FlattenedOutputModule, self).__init__()
+            self._base_module = module
+
+            def _forward(self, *args, **kwargs):
+                return _transform_output_to_flat_tuple(self._base_module(*args, **kwargs))
+
+            # Exporter does not support use of **kwargs in the forward method.
+            # Work around it by making the signature of the forward method to resemble that of the
+            # original model
+            # Copy the forward signature from the original PyTorch module.
+            self.forward = _forward.__get__(self)
+            functools.update_wrapper(self.forward.__func__, module.forward.__func__)
+
+    return FlattenedOutputModule(original_module)
+
+def parse_inputs_for_onnx_export(all_input_names, onnx_graph, *inputs, **kwargs):
+    # Ignore optional inputs explicitly specified as None
+    # ONNX exporter may remove unused inputs
+    onnx_graph_input_names = []
+    if onnx_graph is not None:
+        onnx_graph_input_names = set([inp.name for inp in onnx_graph.graph.input])
+
+    input_names = []
+    dynamic_axes = {}
+    input_names_require_grad = []
+    input_shape = []
+
+    for input_idx, name in enumerate(all_input_names):
+        inp = None
+        if input_idx < len(inputs) and inputs[input_idx] is not None:
+            inp = inputs[input_idx]
+        elif name in kwargs and kwargs[name] is not None:
+            inp = kwargs[name]
+        if inp is not None and (onnx_graph is None or name in onnx_graph_input_names):
+            if inp.requires_grad:
+                # input_names_require_grad holds all input tensors that have requires_grad
+                input_names_require_grad.append(name)
+
+            input_names.append(name)
+            dynamic_axes[name] = {}
+            for dim_idx in range(len(inp.shape)):
+                dynamic_axes[name].update({dim_idx : f'input{input_idx}_dim{dim_idx}'})
+
+            input_shape.append(list(inp.size()))
+    return input_names, dynamic_axes, input_names_require_grad, input_shape
+
+def parse_outputs_for_onnx_export_and_extract_output_schema(module, inputs, kwargs):
+
+    #   Do an inference to grab outputs
+    is_train_mode = module.training
+    module.eval()
+    output_names = None
+    output_dynamic_axes = None
+    with torch.no_grad():
+        # Deepcopy inputs, since input values may change after model run.
+        sample_inputs_copy, sample_kwargs_copy = deepcopy_model_input(*inputs, **kwargs)
+        try:
+            # Deepcopy model, in case model is stateful and changes after model run.
+            model_copy = copy.deepcopy(module)
+        except Exception:
+            model_copy = module
+            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, **sample_kwargs_copy)
+
+        # Parse the output and extract the output_names and output_dynamic_axes to be used for onnx export
+        output_names, output_dynamic_axes = \
+            _parse_outputs_and_extract_names_and_dynamic_axes(sample_outputs)
+    if is_train_mode:
+        module.train()
+
+    # Return output names, output dynamic axes and output schema
+    return output_names, output_dynamic_axes, _extract_output_schema(sample_outputs)

orttraining/orttraining/python/training/ortmodule.py

@@ -1,24 +1,19 @@
-import copy
 import io
 import logging
 import onnx
 import onnxruntime
-import os
 import torch
-import warnings
-import numpy as np
 from inspect import signature
 
 from torch.utils.dlpack import from_dlpack
 from torch.utils.cpp_extension import load_inline
-from collections import abc
 
 # Needed to re-implement PyTorch's cpu,cuda,to methods
 from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict
 
 from onnxruntime.capi import _pybind_state as C
 from onnxruntime.training import register_custom_ops_pytorch_exporter
-from . import _utils
+from . import _utils, _ortmodule_output_transformation
 
 
 ONNX_OPSET_VERSION = 12
@@ -41,19 +36,6 @@ def _create_iobinding(io_binding, inputs, model, device):
         io_binding.bind_output(value_info.name, device.type,
                                device_id=_utils.get_device_index(device))
 
-def _deepcopy_model_input(*inputs, **kwargs):
-    sample_inputs_copy = []
-    for model_input in inputs:
-        sample_inputs_copy.append(model_input.data if isinstance(model_input, torch.Tensor) else model_input)
-    sample_inputs_copy = copy.deepcopy(tuple(sample_inputs_copy))
-
-    sample_kwargs_copy = {}
-    for name, model_input in kwargs.items():
-        sample_kwargs_copy[name] = model_input.data if isinstance(model_input, torch.Tensor) else model_input
-    sample_kwargs_copy = copy.deepcopy(sample_kwargs_copy)
-
-    return sample_inputs_copy, sample_kwargs_copy
-
 def _onnx_value_info_to_buffer_tensor(value_info, device):
     '''Create a torch zeroed tensor with the same shape and type of `value_info`'''
 
@@ -61,109 +43,6 @@ def _onnx_value_info_to_buffer_tensor(value_info, device):
     dtype = _utils.dtype_onnx_to_torch(value_info.type.tensor_type.elem_type)
     return torch.zeros(shape, device=device, dtype=dtype)
 
-def _parse_inputs_for_onnx_export(all_input_names, onnx_graph, *inputs, **kwargs):
-    # Ignore optional inputs explicitly specified as None
-    # ONNX exporter may remove unused inputs
-    onnx_graph_input_names = []
-    if onnx_graph is not None:
-        onnx_graph_input_names = set([inp.name for inp in onnx_graph.graph.input])
-
-    input_names = []
-    dynamic_axes = {}
-    input_names_require_grad = []
-    input_shape = []
-
-    for input_idx, name in enumerate(all_input_names):
-        inp = None
-        if input_idx < len(inputs) and inputs[input_idx] is not None:
-            inp = inputs[input_idx]
-        elif name in kwargs and kwargs[name] is not None:
-            inp = kwargs[name]
-        if inp is not None and (onnx_graph is None or name in onnx_graph_input_names):
-            if inp.requires_grad:
-                # input_names_require_grad holds all input tensors that have requires_grad
-                input_names_require_grad.append(name)
-
-            input_names.append(name)
-            dynamic_axes[name] = {}
-            for dim_idx in range(len(inp.shape)):
-                dynamic_axes[name].update({dim_idx : f'input{input_idx}_dim{dim_idx}'})
-
-            input_shape.append(list(inp.size()))
-    return input_names, dynamic_axes, input_names_require_grad, input_shape
-
-def _parse_outputs_for_onnx_export(module, inputs, kwargs):
-
-    def _create_output_dim_names_from_mapping(output):
-        output_names, dynamic_axes = [], {}
-        for name, value in output.items():
-            if not isinstance(value, torch.Tensor):
-                raise TypeError('ORTModule does not support the following model output type {} within a Mapping'.format(type(value)))
-            output_names.append(name)
-            dynamic_axes[name] = {}
-            for dim_idx in range(len(value.shape)):
-                dynamic_axes[name].update({dim_idx: '{}_dim{}'.format(name, dim_idx)})
-        return output_names, dynamic_axes
-
-    def _create_output_dim_names(output, output_idx, from_sequence):
-        if from_sequence and not isinstance(output, torch.Tensor):
-            raise TypeError('ORTModule does not support the following model output type {} within a Sequence'.format(type(output)))
-        output_names, dynamic_axes = [], {}
-        name = 'output{}'.format(output_idx)
-        output_names.append(name)
-        dynamic_axes[name] = {}
-        for dim_idx in range(len(output.shape)):
-            dynamic_axes[name].update({dim_idx : '{}_dim{}'.format(name, dim_idx)})
-        return output_names, dynamic_axes
-
-    #   Do an inference to grab outputs
-    is_train_mode = module.training
-    module.eval()
-    output_names = []
-    output_dynamic_axes = {}
-    sample_output_type = None
-    with torch.no_grad():
-        # Deepcopy inputs, since input values may change after model run.
-        sample_inputs_copy, sample_kwargs_copy = _deepcopy_model_input(*inputs, **kwargs)
-        try:
-            # Deepcopy model, in case model is stateful and changes after model run.
-            model_copy = copy.deepcopy(module)
-        except Exception:
-            model_copy = module
-            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, **sample_kwargs_copy)
-        sample_output_type = type(sample_outputs)
-        if isinstance(sample_outputs, torch.Tensor):
-            output_names, output_dynamic_axes = _create_output_dim_names(sample_outputs, 0, False)
-        elif isinstance(sample_outputs, abc.Mapping):
-            output_names, output_dynamic_axes = _create_output_dim_names_from_mapping(sample_outputs)
-        elif isinstance(sample_outputs, abc.Sequence):
-            for idx, out in enumerate(sample_outputs):
-                tmp_output_names, tmp_output_dynamic_axes = _create_output_dim_names(out, idx, True)
-                output_names += tmp_output_names
-                output_dynamic_axes.update(tmp_output_dynamic_axes)
-        else:
-            raise TypeError('ORTModule does not support the following model output type {}'.format(type(sample_outputs)))
-    if is_train_mode:
-        module.train()
-    return output_names, output_dynamic_axes, sample_output_type
-
-def _populate_user_output(user_output_type, user_output_names, user_outputs):
-    if issubclass(user_output_type, Mapping):
-        key_value_pairs = [(user_output_names[i], user_outputs[i]) for i in range(len(user_output_names))]
-        return user_output_type(key_value_pairs)
-    elif issubclass(user_output_type, tuple):
-        try:
-            # Try constructing the user named tuple from the output tuple
-            return user_output_type(*user_outputs)
-        except TypeError:
-            # The expected output type is not a namedtuple, but is a regular tuple type
-            pass
-
-    return user_outputs
-
 # TODO: PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8,
 # and convert the config to torch.uint8 tensor duing from_dlpack(). So a boolean tensor
 # from forward graph outputs will be converted to torch.uint8 tensor. When this tensor
@@ -206,6 +85,9 @@ class ORTModule(torch.nn.Module):
 
         # User module is wrapped to use its initializers and save computed gradients
         self._original_module = module
+        # Get the module that flattens the output from the original module into a tuple
+        self._flattened_output_module = \
+            _ortmodule_output_transformation.get_flattened_output_module(self._original_module)
         sig = signature(self._original_module.forward)
         self._original_module_input_names = sig.parameters.keys()
         self._onnx_inference = None
@@ -215,7 +97,7 @@ class ORTModule(torch.nn.Module):
         self._current_input_shape = None
         self._module_gradient_graph_builder = None
         self._input_names_require_grad = None
-        self._original_module_output_type = None
+        self._original_module_output_schema = None
 
         # Training model
         self._onnx_training = None
@@ -239,7 +121,7 @@ class ORTModule(torch.nn.Module):
 
     def _initialize_module_gradient_graph_builder(self):
         # TODO: PyTorch exporter bug: changes the initializer order in ONNX model
-        initializer_names = [p[0] for p in self._original_module.named_parameters()]
+        initializer_names = [p[0] for p in self._flattened_output_module.named_parameters()]
         onnx_initializer_names = [p.name for p in self._onnx_inference.graph.initializer]
         initializer_names = [p for p in initializer_names if p in onnx_initializer_names]
 
@@ -298,11 +180,11 @@ class ORTModule(torch.nn.Module):
 
     def eval(self: T) -> T:
         self._is_training = False
-        self._original_module.eval()
+        self._flattened_output_module.eval()
 
     def train(self: T, mode: bool = True) -> T:
         self._is_training = mode
-        self._original_module.train(mode)
+        self._flattened_output_module.train(mode)
 
     def forward(self, *inputs, **kwargs):
         '''Forward pass starts here and continues at `_ORTModuleFunction.forward`
@@ -324,7 +206,9 @@ class ORTModule(torch.nn.Module):
                     raise RuntimeError('A device must be specified in the model or data!')
             self._get_inference_graph_and_init_gradient_graph_builder(*inputs, **kwargs)
 
-        _, _, input_names_require_grad, new_input_shape = _parse_inputs_for_onnx_export(self._original_module_input_names, self._onnx_inference, *inputs, **kwargs)
+        _, _, input_names_require_grad, new_input_shape = \
+            _ortmodule_output_transformation.parse_inputs_for_onnx_export(
+                self._original_module_input_names, self._onnx_inference, *inputs, **kwargs)
         # If inputs requiring gradient change from one call to forward to the next, the module_gradient_graph_builder
         # needs to be reinitialized so it can compute the backward output for the new inputs that require_grad
         if input_names_require_grad != self._input_names_require_grad:
@@ -362,7 +246,7 @@ class ORTModule(torch.nn.Module):
                 # Run and return module outputs.
                 user_outputs = tuple(_ort_output_to_torch_tensor(forward_output) \
                     for forward_output in self._training_session.run_forward(self._training_io_binding, self._run_options))
-                return user_outputs[0] if len(user_outputs) == 1 else user_outputs
+                return user_outputs
 
             @staticmethod
             def backward(ctx, *grad_output):
@@ -398,7 +282,8 @@ class ORTModule(torch.nn.Module):
                             for backward_output in backward_outputs[num_user_input_grads:]]
                 return tuple(results)
 
-        return _populate_user_output(self._original_module_output_type, self._onnx_graphs_info.user_output_names,
+        return _ortmodule_output_transformation.populate_user_output_from_schema_and_outputs(self._original_module_output_schema,
+            self._onnx_graphs_info.user_output_names,
             _ORTModuleFunction.apply(*self._convert_training_graph_input_to_list(*inputs, **kwargs)))
 
     @_utils.timeit(enabled=__TEMP_ENABLE_METHOD_TIMING__)
@@ -423,7 +308,7 @@ class ORTModule(torch.nn.Module):
                 result.append(inp)
 
         # Initializers
-        for param in self._original_module.named_parameters():
+        for param in self._flattened_output_module.named_parameters():
             result.append(param[1])
 
         return result
@@ -436,8 +321,12 @@ class ORTModule(torch.nn.Module):
         '''
 
         # Setup dynamic axes for onnx model
-        input_names, dynamic_axes, self._input_names_require_grad, _ = _parse_inputs_for_onnx_export(self._original_module_input_names, None, *inputs, **kwargs)
-        output_names, output_dynamic_axes, self._original_module_output_type = _parse_outputs_for_onnx_export(self._original_module, inputs, kwargs)
+        input_names, dynamic_axes, self._input_names_require_grad, _ = \
+            _ortmodule_output_transformation.parse_inputs_for_onnx_export(
+                self._original_module_input_names, None, *inputs, **kwargs)
+        output_names, output_dynamic_axes, self._original_module_output_schema = \
+            _ortmodule_output_transformation.parse_outputs_for_onnx_export_and_extract_output_schema(
+                self._original_module, inputs, kwargs)
         dynamic_axes.update(output_dynamic_axes)
 
         # Export torch.nn.Module to ONNX
@@ -446,11 +335,12 @@ class ORTModule(torch.nn.Module):
         # Deepcopy inputs, since input values may change after model run.
         # NOTE: Inputs may contain tensors that have attributes preventing their deepcopy (example grad_fn).
         # Therefore, deepcopy only the data component of the input tensors for export.
-        sample_inputs_copy, sample_kwargs_copy = _deepcopy_model_input(*inputs, **kwargs)
+        sample_inputs_copy, sample_kwargs_copy = \
+            _ortmodule_output_transformation.deepcopy_model_input(*inputs, **kwargs)
 
         try:
             with torch.no_grad():
-                torch.onnx.export(self._original_module,
+                torch.onnx.export(self._flattened_output_module,
                                 sample_inputs_copy + (sample_kwargs_copy, ),
                                 f,
                                 input_names=input_names,

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -103,16 +103,18 @@ class NeuralNetSimplePositionalAndKeywordArguments(torch.nn.Module):
             return torch.mean(self.a) + 3 * y
         return torch.mean(self.a) + x
 
-def _get_bert_for_sequence_classification_model(device):
+def _get_bert_for_sequence_classification_model(device, output_attentions = False, \
+    output_hidden_states = False, return_dict = True):
     """Returns the BertForSequenceClassification pretrained model"""
 
     config = AutoConfig.from_pretrained(
             "bert-base-uncased",
             num_labels=2,
             num_hidden_layers=1,
-            output_attentions = False,
-            output_hidden_states = False,
+            output_attentions = output_attentions,
+            output_hidden_states = output_hidden_states,
     )
+    config.return_dict = return_dict
 
     model = BertForSequenceClassification.from_pretrained(
         "bert-base-uncased",
@@ -383,11 +385,9 @@ def test_gpu_reserved_memory_with_torch_no_grad():
 
     torch.cuda.empty_cache()
     model_with_no_grad = ORTModule(model_with_no_grad)
-    mem_reserved_before_export = torch.cuda.memory_reserved(device)
-    model_with_no_grad(x, y, None, None, None, None, z)
+    model_with_no_grad(x, attention_mask=y, labels=z)
     mem_reserved_after_export_with_torch_no_grad = torch.cuda.memory_reserved(device)
     del model_with_no_grad
-    mem_reserved_after_cache_empty = torch.cuda.memory_reserved(device)
 
     # Create another model and get the memory_reserved when torch.no_grad has not been enabled after export.
     model_without_no_grad = _get_bert_for_sequence_classification_model(device)
@@ -395,7 +395,7 @@ def test_gpu_reserved_memory_with_torch_no_grad():
     mem_reserved_after_export_without_torch_no_grad = 0
 
     with patch('torch.no_grad'):
-        model_without_no_grad(x, y, None, None, None, None, z)
+        model_without_no_grad(x, attention_mask=y, labels=z)
         mem_reserved_after_export_without_torch_no_grad = torch.cuda.memory_reserved(device)
 
     assert mem_reserved_after_export_with_torch_no_grad < mem_reserved_after_export_without_torch_no_grad
@@ -473,9 +473,46 @@ def test_dict_of_tuple_return_value_module(device):
     y = torch.randn(N, D_in, device=device)
     z = torch.randn(N, D_in, device=device)
 
-    with pytest.raises(TypeError) as type_error:
-        model(x, y, z)
-    assert 'ORTModule does not support the following model output type' in str(type_error.value)
+    output = model(x, y, z)
+    assert 'loss' in output
+    assert len(output['loss']) == 3
+
+@pytest.mark.parametrize("device", ['cuda', 'cpu'])
+def test_tuple_of_tuple_return_value_module(device):
+    class NeuralNetTupleOfTuplesOutput(torch.nn.Module):
+        def __init__(self, input_size, hidden_size, num_classes):
+            super(NeuralNetTupleOfTuplesOutput, self).__init__()
+
+            self.fc1_1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu1 = torch.nn.ReLU()
+            self.fc1_2 = torch.nn.Linear(hidden_size, num_classes)
+
+            self.fc2_1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu2 = torch.nn.ReLU()
+            self.fc2_2 = torch.nn.Linear(hidden_size, num_classes)
+
+            self.fc3_1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu3 = torch.nn.ReLU()
+            self.fc3_2 = torch.nn.Linear(hidden_size, num_classes)
+
+        def forward(self, input1, input2, input3):
+            out1 = self.fc1_2(self.relu1(self.fc1_1(input1)))
+            out2 = self.fc2_2(self.relu2(self.fc2_1(input2)))
+            out3 = self.fc3_2(self.relu3(self.fc3_1(input3)))
+            return ((out1, out2), out3)
+
+    N, D_in, H, D_out = 64, 784, 500, 10
+    model = NeuralNetTupleOfTuplesOutput(D_in, H, D_out).to(device)
+    model = ORTModule(model)
+    x = torch.randn(N, D_in, device=device)
+    y = torch.randn(N, D_in, device=device)
+    z = torch.randn(N, D_in, device=device)
+
+    output = model(x, y, z)
+    assert len(output) == 2
+    assert isinstance(output[0], tuple)
+    assert len(output[0]) == 2
+    assert isinstance(output[1], torch.Tensor)
 
 @pytest.mark.parametrize("device", ['cpu', 'cuda'])
 def test_named_tuple_return_value_module(device):
@@ -572,7 +609,7 @@ def test_dynamic_axes_config():
     model_with_no_grad = _get_bert_for_sequence_classification_model(device)
     model_with_no_grad = ORTModule(model_with_no_grad)
     x, y, z = _get_bert_for_sequence_classification_sample_data(device)
-    output = model_with_no_grad(x, y, None, None, None, None, z)
+    output = model_with_no_grad(x, attention_mask=y, labels=z)
     assert output is not None
     assert _test_helpers.is_dynamic_axes(model_with_no_grad)
 
@@ -730,3 +767,78 @@ def test_wrap_ortmodule_and_change_device():
 
     # Checking training finished normally
     assert y_pred is not None and loss is not None
+
+@pytest.mark.parametrize("return_dict", [True, False])
+def test_hf_model_output_with_tuples(return_dict):
+    device = 'cuda'
+
+    model = _get_bert_for_sequence_classification_model(device, output_attentions=True,
+        output_hidden_states=True, return_dict=return_dict)
+    x, y, z = _get_bert_for_sequence_classification_sample_data(device)
+
+    model = ORTModule(model)
+    output = model(x, attention_mask=y, labels=z)
+
+    if return_dict:
+        assert isinstance(output, SequenceClassifierOutput)
+        assert 'loss' in output and 'logits' in output and \
+            'attentions' in output and 'hidden_states' in output
+        assert isinstance(output['loss'], torch.Tensor)
+        assert isinstance(output['logits'], torch.Tensor)
+        assert isinstance(output['attentions'], tuple)
+        assert isinstance(output['hidden_states'], tuple)
+    else:
+        assert isinstance(output, tuple)
+        assert isinstance(output[0], torch.Tensor)
+        assert isinstance(output[1], torch.Tensor)
+        assert isinstance(output[2], tuple)
+        assert isinstance(output[3], tuple)
+
+@pytest.mark.parametrize("device", ['cuda', 'cpu'])
+def test_nested_return_value_module(device):
+    class NeuralNetNestedOutput(torch.nn.Module):
+        def __init__(self, input_size, hidden_size, num_classes):
+            super(NeuralNetNestedOutput, self).__init__()
+
+            self.fc1_1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu1 = torch.nn.ReLU()
+            self.relu = torch.nn.ReLU()
+            self.fc1_2 = torch.nn.Linear(hidden_size, num_classes)
+
+            self.fc2_1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu2 = torch.nn.ReLU()
+            self.fc2_2 = torch.nn.Linear(hidden_size, num_classes)
+
+            self.fc3_1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu3 = torch.nn.ReLU()
+            self.fc3_2 = torch.nn.Linear(hidden_size, num_classes)
+
+        def forward(self, input1, input2, input3):
+            out1 = self.fc1_2(self.relu1(self.fc1_1(input1)))
+            out2 = self.fc2_2(self.relu2(self.fc2_1(input2)))
+            out3 = self.fc3_2(self.relu(self.relu3(self.fc3_1(input3))))
+            return {
+                'a': {
+                    'b': {
+                        'c': out1
+                    },
+                    'd': (out2, out3)
+                }
+            }
+
+    N, D_in, H, D_out = 64, 784, 500, 10
+    model = NeuralNetNestedOutput(D_in, H, D_out).to(device)
+    model = ORTModule(model)
+    model._save_onnx = True
+    model._save_onnx_prefix = 'nested_model_output'
+    x = torch.randn(N, D_in, device=device)
+    y = torch.randn(N, D_in, device=device)
+    z = torch.randn(N, D_in, device=device)
+
+    output = model(x, y, z)
+    assert 'a' in output and 'b' in output['a'] and 'c' in output['a']['b']
+    assert isinstance(output['a']['b']['c'], torch.Tensor)
+
+    assert 'd' in output['a']
+    assert isinstance(output['a']['d'], tuple)
+    assert len(output['a']['d']) == 2