From 06b0891b59ce80284602e365f4d6f3de4057a45f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Thu, 22 Oct 2026 12:13:30 +0200 Subject: [PATCH] fix(module): keep inputs on their device in a passthrough module without parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A module without parameters had CPU as its device. When such a module failed to tune, its passthrough moved every input to CPU. The next module then got a CPU tensor, for example after `nn.Flatten` in ConvNeXt. A passthrough module without parameters now leaves its inputs where they are. When a device is set, tensor inputs are moved with Locator.find_leaves. Signed-off-by: Łukasz Ślusarczyk --- ChangeLog.md | 2 ++ .../torch/module/test_passthrough_module.py | 12 +++++++ .../torch/module/passthrough_module.py | 32 ++++++++----------- torch_tweak/torch/module/wrapper_module.py | 4 +-- 4 files changed, 29 insertions(+), 21 deletions(-) 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()