From d7d7665023e29d686d462c9a304e5ca376904c71 Mon Sep 17 00:00:00 2001
From: mindest <30493312+mindest@users.noreply.github.com>
Date: Thu, 17 Mar 2022 11:56:25 +0800
Subject: [PATCH] restore random states after export_model (#10705)

* restore random states after export_model

* move get/set_random_states inside _export_model

* add comments for random state save/restore

* add unit test for random state check

* resolve comments

* fix error
---
 .../ortmodule/_graph_execution_manager.py     |  7 ++++
 .../python/training/ortmodule/_utils.py       | 18 +++++++++
 .../python/orttraining_test_ortmodule_api.py  | 40 +++++++++++++++++++
 3 files changed, 65 insertions(+)

diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
index 7a1f0a07bc..359f09114d 100644
--- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
+++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
@@ -311,6 +311,10 @@ class GraphExecutionManager(GraphExecutionInterface):
         #       Model is not re-exported when the model parameters change. This can happen when the model is a stateful model,
         #       or the user explicitly changed model parameters after the onnx export.
 
+        # Record random states here and restore later in case any of them gets changed during the export,
+        # e.g., some sympy functions in symbolic_shape_infer will change Python's random state.
+        random_states = _utils.get_random_states()
+
         schema = _io._extract_schema(
             {'args': copy.copy(inputs), 'kwargs': copy.copy(kwargs)})
         if self._onnx_models.exported_model and schema == self._input_info.schema and not self._original_model_has_changed:
@@ -329,6 +333,9 @@ class GraphExecutionManager(GraphExecutionInterface):
             self._onnx_models.exported_model = SymbolicShapeInference.infer_shapes(self._onnx_models.exported_model,
                                                                                    auto_merge=True, guess_output_rank=True)
 
+        # Restore the recorded random states
+        _utils.set_random_states(random_states)
+
         return True
 
     def _get_exported_model(self, input_schema, *inputs, **kwargs):
diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py
index 23dfea316a..534efd2804 100644
--- a/orttraining/orttraining/python/training/ortmodule/_utils.py
+++ b/orttraining/orttraining/python/training/ortmodule/_utils.py
@@ -25,6 +25,24 @@ from typing import List
 import types
 import warnings
 from distutils.version import LooseVersion
+import random
+import numpy as np
+
+def get_random_states():
+    r_state = random.getstate()
+    np_state = np.random.get_state()
+    torch_state = torch.get_rng_state()
+    torch_cuda_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None
+    return r_state, np_state, torch_state, torch_cuda_state
+
+def set_random_states(states):
+    r_state, np_state, torch_state, torch_cuda_state = states
+    random.setstate(r_state)
+    np.random.set_state(np_state)
+    torch.set_rng_state(torch_state)
+    if torch_cuda_state is not None:
+        torch.cuda.set_rng_state(torch_cuda_state)
+
 
 def _ortvalue_from_torch_tensor(torch_tensor):
     # TODO: Current DLPack doesn't support bool and PyTorch disables converting bool tensor to DLPack in recent commit.
diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
index 0ff12f9040..84ce8405c6 100644
--- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
+++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
@@ -4818,3 +4818,43 @@ def test_check_opset_is_default_opset_after_training():
     _test_helpers.assert_values_are_close(pt_loss, ort_loss)
     _test_helpers.assert_values_are_close(pt_x.grad, ort_x.grad)
     assert ortmodule_module.ONNX_OPSET_VERSION == DEFAULT_OPSET
+
+
+def test_random_states_unchanged_for_ortmodule():
+    import numpy
+
+    os.environ['ORTMODULE_FALLBACK_RETRY'] = 'False'
+
+    class NeuralNetSlice(torch.nn.Module):
+        def __init__(self):
+            super(NeuralNetSlice, self).__init__()
+            self.dim = 32
+
+        def forward(self, x):
+            # This slice operation will call sympy.Min() when exporting, which will change Python's random state
+            return x[:self.dim, :]
+
+    def random_state_equal(a, b):
+        assert type(a) == type(b)
+        if isinstance(a, tuple):
+            assert len(a) == len(b)
+            return all([random_state_equal(a_i, b_i) for a_i, b_i in zip(a, b)])
+        if isinstance(a, numpy.ndarray):
+            return numpy.array_equal(a, b)
+        if isinstance(a, torch.Tensor):
+            return torch.equal(a, b)
+        return a == b
+
+    model = NeuralNetSlice()
+    x = torch.randn(16, 16)
+
+    ori_random_states = _utils.get_random_states()
+
+    ort_model = ORTModule(model)
+    ort_model(x)
+
+    new_random_states = _utils.get_random_states()
+
+    assert random_state_equal(ori_random_states, new_random_states)
+
+    del os.environ['ORTMODULE_FALLBACK_RETRY']