Move exported model to cpu before save.

jkosek · jkosek · commit 4130fb2e8c4e · 2025-03-31T04:18:17.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -28,6 +28,7 @@ limitations under the License.
 - change: Install the TensorRT package for architectures other than x86_64
 - change: Disable conversion fallback for TensorRT paths and expose control option in custom config
 - fix: Correctness command relative tolerance formula
+- fix: Memory management during export and conversion process for Torch
 
 ## 0.13.0
 
diff --git a/model_navigator/commands/data_dump/samples.py b/model_navigator/commands/data_dump/samples.py
@@ -25,6 +25,7 @@
 from model_navigator.core.tensor import TensorMetadata
 from model_navigator.core.workspace import Workspace
 from model_navigator.frameworks import Framework
+from model_navigator.frameworks.memory import offload_model_to_cpu
 from model_navigator.runners.base import NavigatorRunner
 
 
@@ -165,6 +166,7 @@ class FetchOutputModelData(Command, is_required=True):
 
     def _run(
         self,
+        framework: Framework,
         workspace: Workspace,
         model: Any,
         runner_cls: Type[NavigatorRunner],
@@ -194,6 +196,8 @@ def _run(
         output_data_path = workspace.path / "model_output"
         output_data_path.mkdir(parents=True, exist_ok=True)
 
+        offload_model_to_cpu(model, framework)
+
         runner_kwargs = runner_config.to_dict() if runner_config is not None else {}
         runner = runner_cls(
             model=model, input_metadata=input_metadata, output_metadata=output_metadata, **runner_kwargs
@@ -211,6 +215,8 @@ def _run(
                 sample_path = output_data_path / sample_name
                 samples_to_npz(outputs, sample_path, batch_dim, raise_on_error=raise_on_error, num_samples=len(samples))
 
+        runner.deactivate()
+
         return CommandOutput(
             status=CommandStatus.OK,
         )
diff --git a/model_navigator/commands/export/exporters/torch2dynamo_onnx.py b/model_navigator/commands/export/exporters/torch2dynamo_onnx.py
@@ -70,6 +70,7 @@ def export(
         device_max_batch_size: Maximum batch size that fits on the device. Defaults to None.
     """
     model = get_model()
+    model.to(target_device)
 
     if not navigator_workspace:
         navigator_workspace = pathlib.Path.cwd()
diff --git a/model_navigator/commands/export/exporters/torch2exportedprogram.py b/model_navigator/commands/export/exporters/torch2exportedprogram.py
@@ -22,6 +22,7 @@
 from model_navigator.core.dataloader import expand_sample, load_samples
 from model_navigator.core.tensor import TensorMetadata
 from model_navigator.exceptions import ModelNavigatorRuntimeError
+from model_navigator.frameworks.torch.utils import offload_torch_model_to_cpu
 
 
 def get_model() -> torch.nn.Module:
@@ -133,4 +134,7 @@ def export(
     exported_model_path = pathlib.Path(exported_model_path)
     if not exported_model_path.is_absolute():
         exported_model_path = navigator_workspace / exported_model_path
+
     torch.export.save(exported_model, exported_model_path.as_posix())
+
+    offload_torch_model_to_cpu(exported_model.module())
diff --git a/model_navigator/commands/export/exporters/torch2onnx.py b/model_navigator/commands/export/exporters/torch2onnx.py
@@ -63,7 +63,7 @@ def export(
             For available arguments check PyTorch documentation: https://pytorch.org/docs/stable/onnx.html#torch.onnx.export
     """
     model = get_model()
-    model = model.to(export_device)
+    model.to(export_device)
 
     if not navigator_workspace:
         navigator_workspace = pathlib.Path.cwd()
diff --git a/model_navigator/commands/export/exporters/torch2quantized_onnx.py b/model_navigator/commands/export/exporters/torch2quantized_onnx.py
@@ -26,6 +26,7 @@
 from model_navigator.core.dataloader import load_samples
 from model_navigator.core.logger import LOGGER
 from model_navigator.core.tensor import TensorMetadata
+from model_navigator.frameworks.torch.utils import offload_torch_model_to_cpu
 from model_navigator.utils.common import numpy_to_torch_dtype
 
 
@@ -105,7 +106,7 @@ def export(
     model_copy = deepcopy(original_model)
 
     # Offload original model to CPU
-    original_model.to("cpu")
+    offload_torch_model_to_cpu(original_model)
 
     try:
         # Move model copy to target device
@@ -226,10 +227,8 @@ def forward_loop(model):
         LOGGER.info("Quantized ONNX export completed successfully")
 
         # Clean up
-        del model_copy
-        del quantized_model
-        torch.cuda.empty_cache()
-
+        offload_torch_model_to_cpu(model_copy)
+        offload_torch_model_to_cpu(quantized_model)
     except Exception as e:
         LOGGER.error(f"Error during quantized ONNX export: {str(e)}")
         raise
diff --git a/model_navigator/commands/export/exporters/torch2torchscript.py b/model_navigator/commands/export/exporters/torch2torchscript.py
@@ -61,6 +61,7 @@ def export(
             For available arguments check PyTorch documentation: https://pytorch.org/docs/stable/jit.html#torch.jit.trace
     """
     model = get_model()
+
     target_jit_type = JitType(target_jit_type)
 
     if not navigator_workspace:
diff --git a/model_navigator/commands/export/torch.py b/model_navigator/commands/export/torch.py
@@ -27,6 +27,7 @@
 from model_navigator.core.tensor import TensorMetadata
 from model_navigator.core.workspace import Workspace
 from model_navigator.exceptions import ModelNavigatorConfigurationError
+from model_navigator.frameworks.torch.utils import offload_torch_model_to_cpu
 from model_navigator.utils.common import parse_kwargs_to_cmd
 
 
@@ -93,7 +94,7 @@ def _run(
 
         # Keep model on CPU after operation
         def on_exit():
-            model.to("cpu")
+            offload_torch_model_to_cpu(model)
 
         with ExecutionContext(
             workspace=workspace,
@@ -197,7 +198,7 @@ def _run(
 
         # Keep model on CPU after operation
         def on_exit():
-            model.to("cpu")
+            offload_torch_model_to_cpu(model)
 
         with ExecutionContext(
             workspace=workspace,
@@ -283,7 +284,7 @@ def _run(
 
         # Keep model on CPU after operation
         def on_exit():
-            model.to("cpu")
+            offload_torch_model_to_cpu(model)
 
         with ExecutionContext(
             workspace=workspace,
@@ -375,7 +376,7 @@ def _run(
 
         # Keep model on CPU after operation
         def on_exit():
-            model.to("cpu")
+            offload_torch_model_to_cpu(model)
 
         if dynamo_dynamic_shapes is None:
             dynamic_shapes = batch_dim is not None or dynamic_axes
@@ -478,7 +479,7 @@ def _run(
 
         # Keep model on CPU after operation
         def on_exit():
-            model.to("cpu")
+            offload_torch_model_to_cpu(model)
 
         with ExecutionContext(
             workspace=workspace,
diff --git a/model_navigator/commands/infer_metadata.py b/model_navigator/commands/infer_metadata.py
@@ -33,6 +33,7 @@
 from model_navigator.core.workspace import Workspace
 from model_navigator.exceptions import ModelNavigatorUserInputError
 from model_navigator.frameworks import Framework, is_torch_available
+from model_navigator.frameworks.memory import offload_model_to_cpu
 from model_navigator.frameworks.onnx.utils import get_onnx_io_names
 from model_navigator.frameworks.tensorrt.utils import get_tensorrt_io_names
 from model_navigator.runners.base import NavigatorRunner
@@ -290,6 +291,8 @@ def _run(
         else:
             temp_output_metadata = None
 
+        offload_model_to_cpu(model, framework)
+
         runner_kwargs = runner_config.to_dict() if runner_config is not None else {}
         runner = runner_cls(
             model=model,
@@ -320,6 +323,8 @@ def _run(
 
         output_metadata = _get_metadata_from_axes_shapes(pytree_metadata, axes_shapes, batch_dim, output_dtypes)
 
+        runner.deactivate()
+
         return CommandOutput(
             status=CommandStatus.OK,
             output={"output_metadata": output_metadata},
diff --git a/model_navigator/frameworks/memory.py b/model_navigator/frameworks/memory.py
@@ -0,0 +1,31 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Memory management utilities for frameworks."""
+
+from typing import Any
+
+from model_navigator.frameworks import Framework, is_torch_available
+
+
+def offload_model_to_cpu(model: Any, framework: Framework):
+    """Offload model to CPU.
+
+    Args:
+        model: Model to offload.
+        framework: Framework of model to offload.
+    """
+    if is_torch_available() and framework == Framework.TORCH:
+        from model_navigator.frameworks.torch.utils import offload_torch_model_to_cpu
+
+        offload_torch_model_to_cpu(model)
diff --git a/model_navigator/frameworks/torch/utils.py b/model_navigator/frameworks/torch/utils.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Torch utils."""
 
+import gc
 from typing import Optional
 
 from model_navigator.utils.module import lazy_import
@@ -33,3 +34,18 @@ def get_module_device(module: "torch.nn.Module") -> Optional["torch.device"]:
         return next(module.parameters()).device
     except StopIteration:
         return None
+
+
+def offload_torch_model_to_cpu(model: "torch.nn.Module"):
+    """Offload PyTorch model to CPU.
+
+    Args:
+        model: PyTorch model to offload.
+    """
+    model.to("cpu")
+
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+
+    gc.collect()
diff --git a/model_navigator/runners/torch.py b/model_navigator/runners/torch.py
@@ -275,6 +275,7 @@ def deactivate_impl(self):
         if not self._inplace:
             self.model.to(self._input_module_device)
         torch.cuda.empty_cache()
+        gc.collect()
 
 
 class TorchCPURunner(_BaseTorchRunner):
diff --git a/model_navigator/torch/__init__.py b/model_navigator/torch/__init__.py
@@ -46,7 +46,7 @@
     verify_builder,
 )
 from model_navigator.pipelines.builders.find_device_max_batch_size import find_device_max_batch_size_builder
-from model_navigator.pipelines.builders.torch import torch_dynamo_onnx_builder, torch_exportedprogram_builder
+from model_navigator.pipelines.builders.torch import torch_exportedprogram_builder
 from model_navigator.pipelines.wrappers.optimize import optimize_pipeline
 from model_navigator.runners.base import NavigatorRunner
 from model_navigator.runners.utils import default_runners, filter_runners
@@ -145,7 +145,6 @@ def optimize(
         torch_export_builder,
         find_device_max_batch_size_builder,
         torch_exportedprogram_builder,
-        torch_dynamo_onnx_builder,
         torch_conversion_builder,
         torch_tensorrt_conversion_builder,
         tensorrt_conversion_builder,
diff --git a/tests/unit/base/test_python_optimize.py b/tests/unit/base/test_python_optimize.py
@@ -76,11 +76,11 @@ def test_python_package_return_valid_runner():
         }
 
         with runner:
-            input = {"input__0": dataloader[0]}
-            output = runner.infer(input)
+            inputs = {"input__0": dataloader[0]}
+            output = runner.infer(inputs)
             assert output is not None
             assert output["output__0"] is not None
-            assert np.array_equal(output["output__0"], input["input__0"])
+            assert np.array_equal(output["output__0"], inputs["input__0"])
 
 
 def test_export_formats_return_empty_list_for_framework_none():
