saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ORTModule] Add Env Variable to Control Disabling Custom AutoGrad Function Support (#13430)

Browse source 
Add env variable to control disabling custom autogard function support.
When using ORTModule, if the torch model has torch.nn.Function, if user
confirms that it can be exported to ONNX (for example, by inline
PythonOp) and the backward implementation is matched to the forward
impl, user can export "ORTMODULE_DISABLE_CUSTOM_AUTOGRAD_SUPPORT=1" to
disable the custom autograd support so that it won't use ORT's PythonOp
to fallback to PyTorch. Exporting to ONNX sometimes can leverage some
graph optimizations in ORT so that perf is better.
This commit is contained in:
Vincent Wang 2022-10-25 16:58:04 +08:00 committed by GitHub
parent ea1bdb162f
commit b6a3562ffb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
orttraining/orttraining/python/training/ortmodule/_custom_autograd_function.py
View file

@ -86,6 +86,9 @@ def enable_custom_autograd_support(to_enable=True):
from onnxruntime.capi._pybind_state import is_torch_interop_default_on
from onnxruntime.training import ortmodule
# Enable the custom autograd by default when PythonOp backend support is enabled during build.
enable_custom_autograd_support(is_torch_interop_default_on())
enable_custom_autograd_support(
not ortmodule._defined_from_envvar("ORTMODULE_DISABLE_CUSTOM_AUTOGRAD_SUPPORT", 0) and is_torch_interop_default_on()
)

8
orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py
View file

@ -4,7 +4,7 @@
# pylint: disable=missing-docstring
# pylint: disable=C0103
from distutils.version import LooseVersion
from packaging.version import Version
import pytest
import torch
@ -15,14 +15,14 @@ from torch.nn.parameter import Parameter
# Import external libraries.
import onnxruntime
from onnxruntime.training.ortmodule import DebugOptions, LogLevel, ORTModule
from onnxruntime.training.ortmodule import ORTModule
torch.manual_seed(1)
onnxruntime.set_seed(1)
def torch_version_lower_than(v):
return LooseVersion(torch.__version__) < LooseVersion(v)
return Version(torch.__version__) < Version(v)
@pytest.fixture(scope="session", autouse=True)
@ -1106,6 +1106,8 @@ def test_non_differentiable_autograd_function():
run()
# There is bug in exporter side since 1.13 that will throw "RuntimeError: _Map_base::at" for this test.
@pytest.mark.skipif(Version(torch.__version__) >= Version("1.13.0"), reason="PyTorch 1.13+ incompatible")
def test_checkpoint_function():
class A(torch.nn.Module):
# A supported module.