Only wrap sub-modules which can be wrapped as ORTModule (#9021)

cmake/onnxruntime_python.cmake

@@ -274,6 +274,9 @@ if (onnxruntime_ENABLE_TRAINING)
   file(GLOB onnxruntime_python_ortmodule_experimental_json_config_srcs CONFIGURE_DEPENDS
     "${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/experimental/json_config/*.py"
   )
+  file(GLOB onnxruntime_python_ortmodule_experimental_hierarchical_srcs CONFIGURE_DEPENDS
+    "${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/experimental/hierarchical_ortmodule/*.py"
+  )
   file(GLOB onnxruntime_python_ortmodule_torch_cpp_ext_srcs CONFIGURE_DEPENDS
     "${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/*.py"
   )
@@ -517,6 +520,7 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/experimental
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/experimental/json_config
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/experimental/hierarchical_ortmodule
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/aten_op_executor
     COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_interop_utils
@@ -542,6 +546,9 @@ if (onnxruntime_ENABLE_TRAINING)
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_ortmodule_experimental_json_config_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/experimental/json_config/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_ortmodule_experimental_hierarchical_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/experimental/hierarchical_ortmodule/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_ortmodule_torch_cpp_ext_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/

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

@@ -5,6 +5,7 @@
 
 import sys
 import torch
+import torch.utils.checkpoint
 import warnings
 from torch.onnx import symbolic_helper
 
@@ -12,6 +13,16 @@ from onnxruntime.capi._pybind_state import register_torch_autograd_function
 from ._fallback import _FallbackManager, ORTModuleONNXModelException, ORTModuleTorchModelException, wrap_exception
 from . import _logger
 
+# Some autograd.Function's shouldn't be exported as PythonOp.
+# If CheckpointFunction is exported as PythonOp, the checkpointed computation
+# may be computed by Pytorch, not ORT. This situation is especially important
+# for big models such as GPT-2. Exporting CheckpointFunction as PythonOp means
+# every transformer would be computed by Pytorch and ORT doesn't contribute
+# at all.
+BANNED_AUTOGRAD_FUNCTION_NAMES = set(
+    [torch.utils.checkpoint.CheckpointFunction.__name__])
+
+
 def _export(g, n, *args, **kwargs):
     '''
     This function exports PythonOp (input: "n") into a graph
@@ -20,6 +31,10 @@ def _export(g, n, *args, **kwargs):
     '''
     try:
         name = kwargs['name']
+        if name in BANNED_AUTOGRAD_FUNCTION_NAMES:
+            raise Exception(f'The autograd.Function {name} should not be exported to ONNX. '
+                            'Please replace ORTModule with HierarchalORTModule to only'
+                            'wrap exportable sub-nn.Module\'s as ORTModule.')
         inplace = kwargs['inplace']
         training_mode = symbolic_helper._training_mode
         cconv = n.cconv()

orttraining/orttraining/python/training/ortmodule/experimental/hierarchical_ortmodule/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# __init__.py
+
+from ._hierarchical_ortmodule import HierarchicalORTModule

orttraining/orttraining/python/training/ortmodule/experimental/hierarchical_ortmodule/_hierarchical_ortmodule.py

@@ -0,0 +1,198 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# debug_options.py
+import tempfile
+import torch
+from ... import ORTModule
+from ... import ONNX_OPSET_VERSION
+from ...debug_options import DebugOptions
+
+
+class HierarchicalORTModule(torch.nn.Module):
+    '''
+    Recursively wraps submodules of `module` as ORTModule whenever possible
+    Similarly to ORTModule, the actual wrapping happens in its first `forward` call during Pytorch-to-ONNX export.
+    Supported computation is delegated to ONNX Runtime and unsupported computation is still done by PyTorch.
+    Args:
+        module (torch.nn.Module): User's PyTorch module that HierarchicalORTModule specializes.
+
+    Example::
+
+        import torch
+        from torch.utils.checkpoint import checkpoint
+        from onnxruntime.training.ortmodule.experimental.hierarchical_ortmodule import HierarchicalORTModule
+
+
+        class Foo(torch.nn.Module):
+            def __init__(self):
+                super(Foo, self).__init__()
+                self.l1 = torch.nn.Linear(2, 2)
+                self.l2 = torch.nn.Linear(2, 2)
+                self.l3 = torch.nn.Linear(2, 2)
+
+            def forward(self, x):
+                def custom():
+                    def custom_forward(x_):
+                        return self.l2(x_)
+                    return custom_forward
+                z = self.l1(checkpoint(custom(), self.l3(x)))
+                return z
+
+        x = torch.rand(2)
+        m = HierarchicalORTModule(Foo())
+        y = m(x)
+
+    '''
+
+    def __init__(self, module, debug_options=None):
+        self._initialized = False
+        super(HierarchicalORTModule, self).__init__()
+        self._original_module = module
+        if not debug_options:
+            self._debug_options = DebugOptions()
+
+    def _initialize(self, *args, **kwargs):
+        handle_pool = []
+        module_arg_pool = {}
+
+        # A forward pre-hook to record inputs for each nn.Module.
+        def record_args(module, args):
+            if module in module_arg_pool:
+                module_arg_pool[module].append(args)
+            else:
+                module_arg_pool[module] = [args]
+
+        # Recursively hook "record_args" to module and all its sub-modules.
+        # The function "record_args" records the inputs for each nn.Module,
+        # and later we will try exporting those nn.Module's with their recorded
+        # inputs.
+        def recursive_hook(module):
+            handle_pool.append(module.register_forward_pre_hook(record_args))
+            for name, sub in module._modules.items():
+                if isinstance(sub, torch.nn.ModuleList):
+                    for name1, sub1 in sub._modules.items():
+                        recursive_hook(sub1)
+                else:
+                    recursive_hook(sub)
+
+        exportable_list = {}
+
+        # Fill "exportable_list". exportable_list[module] = True means
+        # "module" can be wrapped as ORTModule. Otherwise, "module" is
+        # not exportable to ONNX.
+        def check_exportable(module):
+            sub_dict = module._modules
+            if not sub_dict:
+                # No sub-module exists, so this module is a leaf
+                # module in overall model hierarchy.
+                exportable = True
+                # Check if this leaf module is exportable.
+                for args in module_arg_pool[module]:
+                    try:
+                        with tempfile.NamedTemporaryFile(prefix='sub-module') as temp:
+                            torch.onnx.export(
+                                module, args, temp, opset_version=ONNX_OPSET_VERSION,
+                                do_constant_folding=False, export_params=False,
+                                keep_initializers_as_inputs=True,
+                                training=torch.onnx.TrainingMode.TRAINING)
+                    except Exception as e:
+                        exportable = False
+
+                exportable_list[module] = exportable
+                return exportable
+            else:
+                sub_exportable = True
+                for name, sub_module in sub_dict.items():
+                    if isinstance(sub_module, torch.nn.ModuleList):
+                        for name1, sub_module1 in sub_module._modules.items():
+                            sub_exportable1 = check_exportable(sub_module1)
+                            sub_exportable = sub_exportable and sub_exportable1
+                    else:
+                        sub_exportable1 = check_exportable(sub_module)
+                        sub_exportable = sub_exportable and sub_exportable1
+
+                if sub_exportable is False:
+                    # At least one existing sub-module is not exportable,
+                    # so is the entire module.
+                    exportable_list[module] = sub_exportable
+                    return sub_exportable
+                else:
+                    # Now, we know all sub-modules are exportable, so
+                    # we are going to check if the composition of them
+                    # is still exportable at this module level.
+                    module_exportable = True
+                    for args in module_arg_pool[module]:
+                        try:
+                            with tempfile.NamedTemporaryFile(prefix='sub-module') as temp:
+                                torch.onnx.export(
+                                    module, args, temp, opset_version=ONNX_OPSET_VERSION,
+                                    do_constant_folding=False, export_params=False,
+                                    keep_initializers_as_inputs=True,
+                                    training=torch.onnx.TrainingMode.TRAINING)
+                        except Exception as e:
+                            # If this module is not exportable for one arg
+                            # group, we say this module is not exportable.
+                            module_exportable = False
+                            # Already found a broken case.
+                            # No need to check next case.
+                            break
+
+                    exportable_list[module] = module_exportable
+                    return exportable_list[module]
+
+        # Add a hook to record forward's input for all modules.
+        recursive_hook(self._original_module)
+
+        # Run forward with actual input to record all possible
+        # inputs for all invoked modules.
+        _ = self._original_module(*args, **kwargs)
+
+        # We already have "supported_modules" so
+        # we no longer need those hooks in forward pass.
+        for h in handle_pool:
+            h.remove()
+
+        # Try exporter on all module-input pairs. If a module can be exported with
+        # all its recorded inputs, then it's exporable.
+        check_exportable(self._original_module)
+
+        # A naive way of determining if ORT can run nn.Module
+        def is_supported(module):
+            return module in exportable_list and exportable_list[module]
+
+        # Top-down wrapper to replace nn.Module's with ORTModule.
+        # Note that using bottom-up wrapper may lead to much
+        # ORTModule instances and each ORTModule owns a much smaller graph.
+        def recursive_wrap(module):
+            sub_dict = module._modules
+            for name, sub in sub_dict.items():
+                if isinstance(sub, torch.nn.ModuleList):
+                    # We encounter a list of sub-modules.
+                    # Let's wrap them one-by-one.
+                    for name1, sub1 in sub._modules.items():
+                        if is_supported(sub1):
+                            sub._modules[name1] = ORTModule(
+                                sub1,
+                                debug_options=self._debug_options)
+                        else:
+                            recursive_wrap(sub1)
+                else:
+                    if is_supported(sub):
+                        # Just wrap it as ORTModule when possible.
+                        sub_dict[name] = ORTModule(
+                            sub,
+                            debug_options=self._debug_options)
+                    else:
+                        # This sub-module is not exportable to ONNX
+                        # Let's check its sub-modules.
+                        recursive_wrap(sub)
+
+        recursive_wrap(self._original_module)
+        self._initialized = True
+
+    def forward(self, *inputs, **kwargs):
+        if not self._initialized:
+            self._initialize(*inputs, **kwargs)
+        # forward can be run only after initialization is done.
+        assert self._initialized
+        return self._original_module(*inputs, **kwargs)

orttraining/orttraining/test/python/orttraining_ortmodule_tests.py

@@ -95,6 +95,14 @@ def run_ortmodule_custom_autograd_tests(cwd, log):
     run_subprocess(command, cwd=cwd, log=log).check_returncode()
 
 
+def run_ortmodule_hierarchical_ortmodule_tests(cwd, log):
+    log.debug('Running: ORTModule-Hierarchical model tests')
+
+    command = [sys.executable, '-m', 'pytest', '-sv', 'orttraining_test_hierarchical_ortmodule.py']
+
+    run_subprocess(command, cwd=cwd, log=log).check_returncode()
+
+
 def run_ortmodule_experimental_json_config_tests(cwd, log):
     log.debug('Running: ORTModule Experimental Load Config tests')
 
@@ -129,6 +137,8 @@ def main():
 
     run_ortmodule_fallback_tests(cwd, log, args.transformers_cache)
 
+    run_ortmodule_hierarchical_ortmodule_tests(cwd, log,)
+
     return 0
 
 

orttraining/orttraining/test/python/orttraining_test_hierarchical_ortmodule.py

@@ -0,0 +1,204 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections.abc import Iterable
+from torch.utils.checkpoint import checkpoint
+from onnxruntime.training.ortmodule import ORTModule
+from onnxruntime.training.ortmodule.experimental.hierarchical_ortmodule import HierarchicalORTModule
+
+
+class A(nn.Module):
+    # A supported module.
+    def __init__(self):
+        super(A, self).__init__()
+        self.l1 = nn.Linear(2, 2)
+
+    def forward(self, x):
+        return self.l1(x)
+
+
+class B(nn.Module):
+    # This module is not exportable to ONNX because it
+    # uses gradient-checkpointing. However, its two sub-module's
+    # are exportable, so ORTModule should be used to compute them.
+    def __init__(self):
+        super(B, self).__init__()
+        self.l1 = nn.Linear(2, 2)
+        self.a = A()
+
+    def forward(self, x):
+
+        def custom():
+            def custom_forward(x_):
+                return self.a(x_)
+            return custom_forward
+        z = self.l1(checkpoint(custom(), x))
+        return z
+
+
+class C(nn.Module):
+    # A supported module.
+    def __init__(self):
+        super(C, self).__init__()
+        self.l1 = nn.Linear(2, 2)
+
+    def forward(self, x):
+        y = F.relu(self.l1(x))
+        return y
+
+
+class D(nn.Module):
+    # This module is not exportable to ONNX because it
+    # inner module self.b uses gradient-checkpointing.
+    def __init__(self):
+        super(D, self).__init__()
+        self.b = B()
+
+    def forward(self, x):
+        y = F.sigmoid(self.b(x))
+        return y
+
+
+class Main(nn.Module):
+    # Main module.
+    def __init__(self):
+        super(Main, self).__init__()
+        self.alpha = nn.Parameter(torch.tensor(0.941736), requires_grad=True)
+        self.a = A()
+        self.b = B()
+        self.c = C()
+        self.d = D()
+
+    def forward(self, x):
+        z = self.alpha * self.d(self.c(self.b(self.a(x))))
+        return z
+
+
+class MainWithNonTensorInput(nn.Module):
+    # Module for testing non-tensor input.
+    def __init__(self):
+        super(MainWithNonTensorInput, self).__init__()
+        self.alpha = nn.Parameter(torch.tensor(0.941736), requires_grad=True)
+        self.a = A()
+        self.b = B()
+        self.c = C()
+        self.d = D()
+
+    def forward(self, x, case):
+        if case == 'reverse':
+            z = self.alpha * self.a(self.b(self.c(self.d(x))))
+        else:
+            z = self.alpha * self.d(self.c(self.b(self.a(x))))
+        return z
+
+
+class E(nn.Module):
+    # Sub-modules are stored in nn.ModuleList.
+    def __init__(self):
+        super(E, self).__init__()
+        self.my_layers = nn.ModuleList([A(), B(), C(), D()])
+
+    def forward(self, x):
+        y = x
+        for layer in self.my_layers:
+            y = layer(y)
+        return y
+
+
+class MainWithModuleList(nn.Module):
+    # Sub-modules are stored in nn.ModuleList.
+    def __init__(self):
+        super(MainWithModuleList, self).__init__()
+        self.my_layers = nn.ModuleList([E(), E()])
+
+    def forward(self, x):
+        y = x
+        for layer in self.my_layers:
+            y = layer(y)
+        return y
+
+
+class MainWithMultiModuleOutputs(nn.Module):
+    # Module with repeated sub-modules and producing
+    # multiple outputs.
+    def __init__(self):
+        super(MainWithMultiModuleOutputs, self).__init__()
+        self.layer_list1 = nn.ModuleList([D(), A(), B()])
+        self.layer_list2 = nn.ModuleList([C(), B(), D()])
+
+    def forward(self, x):
+        y1 = x
+        for layer in self.layer_list1:
+            y1 = layer(y1)
+        y2 = x
+        for layer in self.layer_list2:
+            y2 = layer(y2)
+        return y1, y2
+
+
+def test_hierarchical_ortmodule():
+    def count_ortmodule(module):
+        n = 1 if isinstance(module, ORTModule) else 0
+        for sub in module._modules.values():
+            n = n + count_ortmodule(sub)
+        return n
+
+    def call_backward(y):
+        if isinstance(y, tuple):
+            for ele in y:
+                ele.sum().backward()
+        else:
+            y.sum().backward()
+
+    def call_allclose(y, y_ref):
+        assert type(y) == type(y_ref)
+        if isinstance(y, Iterable):
+            for ele, ele_ref in zip(y, y_ref):
+                torch.allclose(ele, ele_ref)
+        else:
+            torch.allclose(y, y_ref)
+
+    def trial(module_to_wrap, args, expected_num_ortmodule):
+        # Run baseline model.
+        m = module_to_wrap
+
+        y_ref = m(*args)
+        call_backward(y_ref)
+        g_ref = []
+        for param in m.parameters():
+            g_ref.append(param.grad.detach())
+
+        m.zero_grad()
+
+        # Run hierarchical ORTModule model.
+        m = HierarchicalORTModule(m)
+
+        y = m(*args)
+        call_backward(y)
+        g = []
+        for param in m.parameters():
+            g.append(param.grad.detach())
+
+        # Some sub-modules become ORTModule.
+        assert expected_num_ortmodule == count_ortmodule(m)
+
+        call_allclose(y, y_ref)
+        call_allclose(g, g_ref)
+
+    num_trials = 4
+    for _ in range(num_trials):
+        trial(Main(), [torch.rand(2).requires_grad_()], 6)
+        trial(MainWithModuleList(), [torch.rand(2).requires_grad_()], 12)
+        trial(MainWithMultiModuleOutputs(), [
+              torch.rand(2).requires_grad_()], 10)
+        trial(MainWithNonTensorInput(), [
+              torch.rand(2).requires_grad_(), 'reverse'], 6)
+        trial(MainWithNonTensorInput(), [
+              torch.rand(2).requires_grad_(), 'normal'], 6)
+
+
+if __name__ == '__main__':
+    test_hierarchical_ortmodule()

setup.py

@@ -327,6 +327,7 @@ if enable_training:
                      'onnxruntime.training.ortmodule',
                      'onnxruntime.training.ortmodule.experimental',
                      'onnxruntime.training.ortmodule.experimental.json_config',
+                     'onnxruntime.training.ortmodule.experimental.hierarchical_ortmodule',
                      'onnxruntime.training.ortmodule.torch_cpp_extensions',
                      'onnxruntime.training.ortmodule.torch_cpp_extensions.aten_op_executor',
                      'onnxruntime.training.ortmodule.torch_cpp_extensions.torch_interop_utils',