From 271a34302406fd14a138fa39b03d21335948d0ad Mon Sep 17 00:00:00 2001 From: baijumeswani Date: Wed, 2 Jun 2021 10:07:49 -0700 Subject: [PATCH] Cache initializers and avoid device check ot end of forward (#7905) --- .../python/training/ortmodule/_graph_execution_manager.py | 6 ++++++ .../python/training/ortmodule/_inference_manager.py | 3 +-- .../python/training/ortmodule/_training_manager.py | 6 +----- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 8c54e93ddb..53df873ac4 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -54,6 +54,7 @@ class GraphExecutionManager(ABC): self._graph_info = None self._graph_initializer_names = None self._graph_initializer_names_to_train = None + self._graph_initializers = None # TrainingAgent or InferenceAgent self._execution_agent = None @@ -320,3 +321,8 @@ class GraphExecutionManager(ABC): # a set (unordered_set in the backend) that does not require a copy on each reference. self._graph_initializer_names = set(initializer_names) self._graph_initializer_names_to_train = set(initializer_names_to_train) + + # Initializers can be cached and used since they are expected not to be re-instantiated + # between forward calls. + self._graph_initializers = [param for name, param in self._flattened_module.named_parameters() + if name in self._graph_initializer_names] diff --git a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py index 260dacf55b..ba688e0b7d 100644 --- a/orttraining/orttraining/python/training/ortmodule/_inference_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_inference_manager.py @@ -89,8 +89,7 @@ class InferenceManager(GraphExecutionManager): self._optimized_onnx_model, self._device, *_io._combine_input_buffers_initializers( - [param for name, param in self._flattened_module.named_parameters() - if name in self._graph_initializer_names], + self._graph_initializers, self._graph_info.user_input_names, self._input_info, self._flattened_module.named_buffers(), diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index 96c634105a..462000d580 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -47,9 +47,6 @@ class TrainingManager(GraphExecutionManager): execution_session.run_forward(forward_inputs, forward_outputs, state) user_outputs = tuple(_utils._ortvalue_to_torch_tensor(forward_output) for forward_output in forward_outputs) - # Assert that the outputs and model device match - _utils._check_same_device(device, "Output argument from forward", *user_outputs) - output_info = [(output.shape, output.device, output.dtype) for output in user_outputs] run_info = RunStateInfo(state, output_info) # Return user outputs and forward run information @@ -192,8 +189,7 @@ class TrainingManager(GraphExecutionManager): return _io.unflatten_user_output(self._module_output_schema, _ORTModuleFunction.apply( *_io._combine_input_buffers_initializers( - [param for name, param in self._flattened_module.named_parameters() - if name in self._graph_initializer_names], + self._graph_initializers, self._graph_info.user_input_names, self._input_info, self._flattened_module.named_buffers(),