Move import to when needed to avoid circular dependency error (micros…

…oft#19579) (#35) ### Description Move import to when needed to avoid circular dependency error ### Motivation and Context Fixes dependency error described here: microsoft/DeepSpeed#5140 --------- Co-authored-by: jingyanwangms <[email protected]> Co-authored-by: Thiago Crepaldi <[email protected]>
ROCm · May 9, 2024 · 1caf382 · 1caf382
1 parent 9c748d7
commit 1caf382
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
@@ -20,7 +20,6 @@
 from onnxruntime.capi import _pybind_state as C
 from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
 from onnxruntime.training.utils import ORTModelInputOutputSchemaType, PTable, onnx_dtype_to_pytorch_dtype
-from onnxruntime.training.utils.hooks import configure_ort_compatible_zero_stage3
 
 from . import _are_deterministic_algorithms_enabled, _io, _logger, _onnx_models, _utils
 from ._fallback import (
@@ -144,6 +143,9 @@ def __init__(
 
         self._zero_stage3_param_map = {}
         if self._runtime_options.enable_zero_stage3_support:
+            # Move import to here to avoid circular dependency error
+            from onnxruntime.training.utils.hooks import configure_ort_compatible_zero_stage3  # type: ignore[import]
+
             # Cannot toggle feature enabling/disabling after the first time enabled.
 
             configure_ort_compatible_zero_stage3(debug=False, stats_output_dir="ort_output", stats_overwrite=True)