diff --git a/ChangeLog.md b/ChangeLog.md index 55ed3b3c..bbca75d6 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -7,6 +7,8 @@ SPDX-License-Identifier: Apache-2.0 ## Unreleased +- fix: a passthrough module without parameters keeps its inputs on their device instead of moving them to CPU. + - feat (breaking): `inspect` measures module depth in executed modules. Containers such as `nn.ModuleList` and modules that never run no longer add a level. `min_depth` is renamed to `depth`. diff --git a/tests/unit/torch/module/test_passthrough_module.py b/tests/unit/torch/module/test_passthrough_module.py index a36c9f14..4e3b6b46 100644 --- a/tests/unit/torch/module/test_passthrough_module.py +++ b/tests/unit/torch/module/test_passthrough_module.py @@ -6,9 +6,12 @@ """Test for passhrough module.""" import pytest +import torch +import torch.nn as nn from tests.toy_models import ToyTorchModel from torch_tweak.torch.module.passthrough_module import PassthroughModule +from torch_tweak.torch.module.wrapper_module import Module from torch_tweak.torch.utils.xpu import is_available as is_xpu_available devices = ["cpu", "xpu"] if is_xpu_available() else ["cpu"] @@ -21,3 +24,12 @@ def test_call(device): module = PassthroughModule(model, device=device) module(x) + + +def test_module_without_parameters_keeps_inputs_on_their_device(): + module = Module(nn.Flatten(), "flatten") + module.enable_passthrough() + + output = module(torch.empty(2, 3, 4, device="meta")) + + assert output.device.type == "meta" diff --git a/torch_tweak/torch/module/passthrough_module.py b/torch_tweak/torch/module/passthrough_module.py index fc18032d..8f2fba7d 100644 --- a/torch_tweak/torch/module/passthrough_module.py +++ b/torch_tweak/torch/module/passthrough_module.py @@ -6,45 +6,39 @@ from typing import Any -from torch_tweak.torch.module.sample_metadata import SampleMetadata +from torch_tweak.torch.module.locator import Locator class PassthroughModule: """Module that passes through the original module. - Before passing the data this module moves the data to the same device as the module. - + When a device is given, inputs are moved to it before the call. + When no device is given, inputs stay where they are. """ def __init__( self, module, - device: str, + device: str | None, ) -> None: """Initializes module. Args: module: module to be tuned. - device: Device on which tuned module has to be executed. + device: Device on which tuned module has to be executed. None leaves the module and its + inputs where they are, for modules without parameters. """ super().__init__() self._forward_call = module.__call__ self._device = device - module.to(self._device) + if self._device is not None: + module.to(self._device) def __call__(self, *args: Any, **kwargs: Any) -> Any: """Native inference on wrapped module.""" - args, kwargs = self._prepare_inputs(*args, **kwargs) + if self._device is not None: + for locator, tensor in Locator.find_leaves(args, only_tensors=True): + args = locator.set_value(args, tensor.to(self._device)) + for locator, tensor in Locator.find_leaves(kwargs, only_tensors=True): + kwargs = locator.set_value(kwargs, tensor.to(self._device)) return self._forward_call(*args, **kwargs) - - def _prepare_inputs(self, *args, **kwargs): - """Prepare inputs for inplace inference and place them on the same device if are not.""" - sample_metadata = SampleMetadata.from_inputs(args, kwargs) - - for locator, tensor_spec in sample_metadata.tensor_data: - if tensor_spec.name.startswith("args"): - args = locator.set_value(args, locator.get_value(args).to(self._device)) - else: - kwargs = locator.set_value(kwargs, locator.get_value(kwargs).to(self._device)) - - return args, kwargs diff --git a/torch_tweak/torch/module/wrapper_module.py b/torch_tweak/torch/module/wrapper_module.py index badb7a57..61b6e36a 100644 --- a/torch_tweak/torch/module/wrapper_module.py +++ b/torch_tweak/torch/module/wrapper_module.py @@ -98,7 +98,7 @@ def __init__( module.forward = self._self_proxy_forward self._setup_strategies(strategy, strategies) - self._self_device = get_module_device(module) or torch.device("cpu") + self._self_device = get_module_device(module) self._self_state = ModuleState.INIT self._self_wrapper = None self._self_prev_recording = None @@ -283,7 +283,7 @@ def tune( elif self._self_state == ModuleState.INIT or self._self_state == ModuleState.PASSTHROUGH: raise ValueError(f"Module: '{self._self_name}' has not recorded any samples. Cannot tune it.") - device = device or self._self_device + device = device or self._self_device or torch.device("cpu") recording = cast(RecordingModule, self._self_wrapper) backends: OrderedDict[SampleMetadata, Backend] = OrderedDict()