diff --git a/docs/api/tuning/module.md b/docs/api/tuning/tunable_module.md similarity index 80% rename from docs/api/tuning/module.md rename to docs/api/tuning/tunable_module.md index cadc0673..d672c020 100644 --- a/docs/api/tuning/module.md +++ b/docs/api/tuning/tunable_module.md @@ -5,11 +5,11 @@ SPDX-License-Identifier: Apache-2.0 NOTE: This file has been modified by Intel Corporation. --> -# Module API +# TunableModule API -## Module +## TunableModule -::: torch_tweak.torch.Module +::: torch_tweak.torch.TunableModule options: show_root_heading: true heading_level: 3 diff --git a/docs/guides/advanced/tuning_workflow.md b/docs/guides/advanced/tuning_workflow.md index 48c9f1dc..e8aa53e9 100644 --- a/docs/guides/advanced/tuning_workflow.md +++ b/docs/guides/advanced/tuning_workflow.md @@ -30,7 +30,7 @@ When you call `tt.tune()`, Torch Tweak first enters the sample gathering phase: import torch_tweak.torch as tt # Wrap your module -module = tt.Module(model, "my_model") +module = tt.TunableModule(model, "my_model") # Tune with different batch sizes tt.tune(module, dataset, batch_sizes=[1, 2, 4, 8]) diff --git a/docs/guides/backends/openvino_backend.md b/docs/guides/backends/openvino_backend.md index 30c261f2..5bed0d2e 100644 --- a/docs/guides/backends/openvino_backend.md +++ b/docs/guides/backends/openvino_backend.md @@ -23,7 +23,7 @@ config = OpenVINOBackendConfig(use_dynamo=False, performance_hint="THROUGHPUT") backend = OpenVINOBackend(config) strategy = tt.OneBackendStrategy(backend=backend) -model = tt.Module(model, "my-model", strategy=strategy) +model = tt.TunableModule(model, "my-model", strategy=strategy) tt.tune(model, input_data) ``` diff --git a/docs/guides/backends/torch_inductor_backend.md b/docs/guides/backends/torch_inductor_backend.md index 33a75438..b8c024ed 100644 --- a/docs/guides/backends/torch_inductor_backend.md +++ b/docs/guides/backends/torch_inductor_backend.md @@ -34,7 +34,7 @@ from torch_tweak.torch.tune_strategy import OneBackendStrategy strategy = tt.OneBackendStrategy(backend=backend) -model = tt.Module(model, "my-model", strategy=strategy) +model = tt.TunableModule(model, "my-model", strategy=strategy) tt.tune(model, input_data) ``` diff --git a/docs/guides/backends/torchao_backend.md b/docs/guides/backends/torchao_backend.md index c2a2bb04..02d8af0c 100644 --- a/docs/guides/backends/torchao_backend.md +++ b/docs/guides/backends/torchao_backend.md @@ -30,7 +30,7 @@ backend = TorchAOBackend(config) # Use in tuning strategy = tt.OneBackendStrategy(backend=backend) -model = tt.Module(model, "my-model", strategy=strategy) +model = tt.TunableModule(model, "my-model", strategy=strategy) tt.tune(model, input_data) ``` diff --git a/docs/guides/tune_strategies/tune_strategies.md b/docs/guides/tune_strategies/tune_strategies.md index 2fe37577..617c4616 100644 --- a/docs/guides/tune_strategies/tune_strategies.md +++ b/docs/guides/tune_strategies/tune_strategies.md @@ -61,7 +61,7 @@ backend = TorchInductorBackend(config) strategy = tt.OneBackendStrategy(backend=backend) # Use in tuning -model = tt.Module(model, "my-model", strategy=strategy) +model = tt.TunableModule(model, "my-model", strategy=strategy) tt.tune(model, input_data) ``` @@ -106,7 +106,7 @@ backends = [ strategy = tt.FirstWinsStrategy(backends=backends) # Use in tuning -model = tt.Module(model, "my-model", strategy=strategy) +model = tt.TunableModule(model, "my-model", strategy=strategy) tt.tune(model, input_data) ``` @@ -167,7 +167,7 @@ strategy = tt.HighestThroughputStrategy( ) # Use in tuning -model = tt.Module(model, "my-model", strategy=strategy) +model = tt.TunableModule(model, "my-model", strategy=strategy) tt.tune(model, input_data) ``` diff --git a/docs/guides/tuning.md b/docs/guides/tuning.md index 6be1f9de..5e51362d 100644 --- a/docs/guides/tuning.md +++ b/docs/guides/tuning.md @@ -92,8 +92,8 @@ If you would like to have more control over picking the modules, you can manuall pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers") pipe.to("xpu") -pipe.unet = tt.Module(pipe.unet, strategy=strategy_for_unet) -pipe.transformer = tt.Module(pipe.transformer, strategy=strategy_for_transformer) +pipe.unet = tt.TunableModule(pipe.unet, strategy=strategy_for_unet) +pipe.transformer = tt.TunableModule(pipe.transformer, strategy=strategy_for_transformer) ``` @@ -229,7 +229,7 @@ from torch import nn import torch_tweak.torch as tt model = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 8)).eval() -module = tt.Module(model, name="mlp") +module = tt.TunableModule(model, name="mlp") sample = torch.randn(32) tt.tune(module, sample, batch_sizes=[1, 4], device="xpu", dry_run=True) ``` diff --git a/docs/learn/core_functionalities.md b/docs/learn/core_functionalities.md index 7f740a1a..e8f4541e 100644 --- a/docs/learn/core_functionalities.md +++ b/docs/learn/core_functionalities.md @@ -43,7 +43,7 @@ import torch model = SimpleModel() # Wrap the model -model = tt.Module(model) +model = tt.TunableModule(model) # Define inference function diff --git a/examples/FLUX/flux/tune.py b/examples/FLUX/flux/tune.py index 5c4061ff..ffdb0ace 100644 --- a/examples/FLUX/flux/tune.py +++ b/examples/FLUX/flux/tune.py @@ -70,7 +70,7 @@ def tune_model( ) strategy.enable_find_max_batch_size(enable=False) - # Wrap all modules with Torch Tweak Module + # Wrap all modules with TunableModule modules = modules_info.get_modules() pipe = tt.wrap(pipe, modules, strategy=strategy) diff --git a/examples/LLM/llm/tune.py b/examples/LLM/llm/tune.py index ee5decf2..e71cf039 100644 --- a/examples/LLM/llm/tune.py +++ b/examples/LLM/llm/tune.py @@ -12,7 +12,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline from llm.cmd_args import get_tune_parser -from torch_tweak.torch import Module, OneBackendStrategy +from torch_tweak.torch import OneBackendStrategy, TunableModule from torch_tweak.torch import tune as torch_tweak from torch_tweak.torch.backend import TorchEagerBackend, TorchInductorBackend, TorchInductorBackendConfig @@ -66,7 +66,7 @@ def tune_model(model, tokenizer, cache="static"): """ if cache == "dynamic": - model = Module( + model = TunableModule( model, model.__class__.__name__, strategy=OneBackendStrategy(TorchEagerBackend()).enable_find_max_batch_size( @@ -74,7 +74,7 @@ def tune_model(model, tokenizer, cache="static"): ), # just one strategy, there is no prefill/decode distinction ) elif cache == "static": - model = Module( + model = TunableModule( model, model.__class__.__name__, strategies=[ diff --git a/examples/Minimal/demo.py b/examples/Minimal/demo.py index 1c2a8a71..e3b055cb 100644 --- a/examples/Minimal/demo.py +++ b/examples/Minimal/demo.py @@ -32,7 +32,7 @@ def main() -> None: ], ).enable_find_max_batch_size(False) - module = tt.Module(model, name="mlp", strategy=strategy) + module = tt.TunableModule(model, name="mlp", strategy=strategy) tt.tune(module, sample, batch_sizes=[4], device=device) results = strategy.results[-1].highest_throughput_results if strategy.results else [] diff --git a/examples/ParakeetCTC/parakeet_ctc/model.py b/examples/ParakeetCTC/parakeet_ctc/model.py index 96ff160c..e8b148ee 100644 --- a/examples/ParakeetCTC/parakeet_ctc/model.py +++ b/examples/ParakeetCTC/parakeet_ctc/model.py @@ -12,7 +12,7 @@ import torch from nemo.core.classes.common import typecheck -from torch_tweak.torch import Module, TuneStrategy +from torch_tweak.torch import TunableModule, TuneStrategy def get_model(model_name: str = "nvidia/parakeet-ctc-0.6b") -> nemo.collections.asr.models.EncDecCTCModel: @@ -64,13 +64,13 @@ def wrap_pipeline(name: str, pipeline, strategy: TuneStrategy | None = None): Returns: ASR model """ - pipeline.encoder = Module( + pipeline.encoder = TunableModule( pipeline.encoder, name=f"{name}-encoder", strategy=deepcopy(strategy), ) - pipeline.decoder = Module( + pipeline.decoder = TunableModule( pipeline.decoder, name=f"{name}-decoder", strategy=deepcopy(strategy), diff --git a/examples/ParakeetRNNT/parakeet_rnnt/model.py b/examples/ParakeetRNNT/parakeet_rnnt/model.py index d9403e88..869cad3a 100644 --- a/examples/ParakeetRNNT/parakeet_rnnt/model.py +++ b/examples/ParakeetRNNT/parakeet_rnnt/model.py @@ -69,13 +69,13 @@ def wrap_pipeline(name: str, pipeline: nemo.collections.asr.models.EncDecRNNTBPE Returns: ASR model """ - pipeline.encoder = tt.Module( + pipeline.encoder = tt.TunableModule( pipeline.encoder, name=f"{name}-encoder", strategy=strategy, ) - pipeline.decoder.prediction["dec_rnn"] = tt.Module( + pipeline.decoder.prediction["dec_rnn"] = tt.TunableModule( pipeline.decoder.prediction["dec_rnn"], name=f"{name}-decoder-rnn", strategy=strategy, diff --git a/examples/ResNet/resnet/tune.py b/examples/ResNet/resnet/tune.py index 87997620..d992e1c6 100644 --- a/examples/ResNet/resnet/tune.py +++ b/examples/ResNet/resnet/tune.py @@ -13,7 +13,7 @@ from resnet.cmd_args import get_parser from resnet.model import get_model, get_transform -from torch_tweak.torch import HighestThroughputStrategy, LocalTorchStorage, Module, save, tune +from torch_tweak.torch import HighestThroughputStrategy, LocalTorchStorage, TunableModule, save, tune from torch_tweak.torch.backend import ( TorchAOBackend, TorchAOBackendConfig, @@ -51,7 +51,7 @@ def tune_model( module_name = f"example-{model_name}" - module = Module( + module = TunableModule( model, module_name, strategy=HighestThroughputStrategy( diff --git a/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py b/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py index f23b1f89..20043fc0 100644 --- a/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py +++ b/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py @@ -20,7 +20,7 @@ from torch_tweak.torch import inspect, tune, wrap from torch_tweak.torch.config import torch_tweak_cache_dir from torch_tweak.torch.module.tensor_spec import InfoLevel -from torch_tweak.torch.module.wrapper_module import ModuleState +from torch_tweak.torch.module.tunable_module import ModuleState from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.utils.xpu import synchronize as device_synchronize diff --git a/notebooks/README.md b/notebooks/README.md index 743201c3..51123751 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -26,8 +26,8 @@ This notebook demonstrates how to inspect modules in a pipeline, pick modules fo ## 2. SampleMetadata Walkthrough This notebook provides a comprehensive tutorial on the `SampleMetadata` class, which is a core component of the Torch Tweak library for tracking and managing tensor shapes in PyTorch models. -## 3. Wrapper Module Demo -This notebook demonstrates how to wrap a module and tune it. +## 3. Tunable Module Demo +This notebook demonstrates how to wrap a module in `TunableModule` and tune it. ## 4. Storage Demo This notebook demonstrates how to use the `Storage` class to save and load checkpoints. diff --git a/notebooks/saving_loading_checkpoints.ipynb b/notebooks/saving_loading_checkpoints.ipynb index 04cf05d0..f0c3d8f5 100644 --- a/notebooks/saving_loading_checkpoints.ipynb +++ b/notebooks/saving_loading_checkpoints.ipynb @@ -47,7 +47,7 @@ "\n", "from torch_tweak.torch.backend import TorchInductorBackend\n", "from torch_tweak.torch.module_registry import MODULE_REGISTRY\n", - "from torch_tweak.torch import load, save, tune, LocalTorchStorage, OneBackendStrategy, Module" + "from torch_tweak.torch import load, save, tune, LocalTorchStorage, OneBackendStrategy, TunableModule" ] }, { @@ -115,7 +115,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### Wrap the whole model with Module\n", + "### Wrap the whole model with TunableModule\n", "Below we are wrapping the whole model.\n" ] }, @@ -127,7 +127,7 @@ "source": [ "model = SimpleNet(5, 10, 2).to(device)\n", "original_model = copy.deepcopy(model) # let's keep original model for comparison\n", - "module = Module(model, \"demo-simple\", strategy=OneBackendStrategy(TorchInductorBackend()))" + "module = TunableModule(model, \"demo-simple\", strategy=OneBackendStrategy(TorchInductorBackend()))" ] }, { @@ -319,8 +319,8 @@ "source": [ "model = SimpleNet(5, 10, 2).to(device)\n", "original_model = copy.deepcopy(model) # let's keep original model for comparison\n", - "model.fc1 = Module(model.fc1, \"demo-simple1\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", - "model.fc3 = Module(model.fc3, \"demo-simple2\", strategy=OneBackendStrategy(TorchInductorBackend()))" + "model.fc1 = TunableModule(model.fc1, \"demo-simple1\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", + "model.fc3 = TunableModule(model.fc3, \"demo-simple2\", strategy=OneBackendStrategy(TorchInductorBackend()))" ] }, { @@ -476,8 +476,8 @@ "source": [ "pipeline = Pipeline(5, 10, 2).to(device)\n", "original_pipeline = copy.deepcopy(pipeline) # let's keep original pipeline for comparison\n", - "pipeline.fc1 = Module(pipeline.fc1, \"demo-simple\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", - "pipeline.fc3 = Module(pipeline.fc3, \"demo-simple2\", strategy=OneBackendStrategy(TorchInductorBackend()))" + "pipeline.fc1 = TunableModule(pipeline.fc1, \"demo-simple\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", + "pipeline.fc3 = TunableModule(pipeline.fc3, \"demo-simple2\", strategy=OneBackendStrategy(TorchInductorBackend()))" ] }, { @@ -634,9 +634,9 @@ "metadata": {}, "outputs": [], "source": [ - "pipeline.net.fc1 = Module(pipeline.net.fc1, \"demo-pipecomplex-fc1\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", - "pipeline.net.fc2 = Module(pipeline.net.fc2, \"demo-pipecomplex-fc2\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", - "pipeline.net.fc3 = Module(pipeline.net.fc3, \"demo-pipecomplex-fc3\", strategy=OneBackendStrategy(TorchInductorBackend()))" + "pipeline.net.fc1 = TunableModule(pipeline.net.fc1, \"demo-pipecomplex-fc1\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", + "pipeline.net.fc2 = TunableModule(pipeline.net.fc2, \"demo-pipecomplex-fc2\", strategy=OneBackendStrategy(TorchInductorBackend()))\n", + "pipeline.net.fc3 = TunableModule(pipeline.net.fc3, \"demo-pipecomplex-fc3\", strategy=OneBackendStrategy(TorchInductorBackend()))" ] }, { diff --git a/notebooks/storage.ipynb b/notebooks/storage.ipynb index fbb86b35..7b02eb48 100644 --- a/notebooks/storage.ipynb +++ b/notebooks/storage.ipynb @@ -48,7 +48,7 @@ "import torch.nn as nn\n", "from torch_tweak.torch import LocalTorchStorage\n", "from torch_tweak.torch.backend import TorchInductorBackend, OpenVINOBackend\n", - "from torch_tweak.torch import save, tune, LocalTorchStorage, OneBackendStrategy, Module" + "from torch_tweak.torch import save, tune, LocalTorchStorage, OneBackendStrategy, TunableModule" ] }, { @@ -261,9 +261,9 @@ "metadata": {}, "outputs": [], "source": [ - "net.fc1 = Module(net.fc1, \"fc1\", strategy=OneBackendStrategy(backend=TorchInductorBackend()))\n", - "net.fc2 = Module(net.fc2, \"fc2\", strategy=OneBackendStrategy(backend=OpenVINOBackend()))\n", - "net.fc3 = Module(net.fc3, \"fc3\", strategy=OneBackendStrategy(backend=OpenVINOBackend()))" + "net.fc1 = TunableModule(net.fc1, \"fc1\", strategy=OneBackendStrategy(backend=TorchInductorBackend()))\n", + "net.fc2 = TunableModule(net.fc2, \"fc2\", strategy=OneBackendStrategy(backend=OpenVINOBackend()))\n", + "net.fc3 = TunableModule(net.fc3, \"fc3\", strategy=OneBackendStrategy(backend=OpenVINOBackend()))" ] }, { diff --git a/notebooks/wrapper_module_demo.ipynb b/notebooks/tunable_module_demo.ipynb similarity index 96% rename from notebooks/wrapper_module_demo.ipynb rename to notebooks/tunable_module_demo.ipynb index 63eea9e9..76d938a8 100644 --- a/notebooks/wrapper_module_demo.ipynb +++ b/notebooks/tunable_module_demo.ipynb @@ -12,7 +12,7 @@ "NOTE: This file has been modified by Intel Corporation.\n", "-->\n", "\n", - "# Module wrapper demo\n", + "# TunableModule wrapper demo\n", "The purpose of this notebook is to walk through the Torch Tweak module wrapper usage." ] }, @@ -40,7 +40,7 @@ "\n", "from torch_tweak.global_context import BATCH_SIZE_KEY, global_context\n", "from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend\n", - "from torch_tweak.torch.module.wrapper_module import Module\n", + "from torch_tweak.torch.module.tunable_module import TunableModule\n", "from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy\n", "\n", "log_level = os.environ.get(\"TORCH_TWEAK_LOG_LEVEL\", \"INFO\")\n", @@ -82,7 +82,7 @@ "outputs": [], "source": [ "model = Identity()\n", - "module = Module(model, \"demo-identity\")\n", + "module = TunableModule(model, \"demo-identity\")\n", "device = torch.device(\"xpu\" if torch.xpu.is_available() else \"cpu\")\n", "print(f\"Using device: {device}\")" ] @@ -247,7 +247,7 @@ "\n", "As we have seen module detects dynamic axes just by looking at different input data. However for checking performance and sanity Torch Tweak needs to detect batch axis - which is special type of dynamic axis i.e. grows with a batch size. One example could be LLM pipeline where we have batch and sequence length dynamic axes `B` and `L` but only `B` grows with a batch size.\n", "\n", - "In order to detect batch dimension Module wrapper has to have information about current global batch size. This is because the `global batch size` can be different than `local batch size` e.g. SDXL model stacks input tensor and with `global bs=1` a module can receive `local bs=2`.\n", + "In order to detect batch dimension TunableModule has to have information about current global batch size. This is because the `global batch size` can be different than `local batch size` e.g. SDXL model stacks input tensor and with `global bs=1` a module can receive `local bs=2`.\n", "\n", "When you perform tuning, the Torch Tweak takes dataloader and populates information about `batch size` so that module wrapper can obtain it. Since we are demonstrating module wrapper in isolation, we have to provide that information ourselves." ] @@ -399,7 +399,7 @@ "metadata": {}, "outputs": [], "source": [ - "module = Module(model, \"demo-resnet50\")" + "module = TunableModule(model, \"demo-resnet50\")" ] }, { diff --git a/properdocs.yml b/properdocs.yml index fe1998df..a6110c05 100644 --- a/properdocs.yml +++ b/properdocs.yml @@ -106,7 +106,7 @@ nav: - StableDiffusion: examples/StableDiffusion/README.md - API: - Tuning API: - - Module: api/tuning/module.md + - TunableModule: api/tuning/tunable_module.md - Tuning: api/tuning/tuning.md - Config: api/tuning/config.md - Backend: api/tuning/backend.md diff --git a/tests/functional/pytorch/test_torch_tweak_resnet_openvino_backend.py b/tests/functional/pytorch/test_torch_tweak_resnet_openvino_backend.py index d8360d21..b7a0c79d 100644 --- a/tests/functional/pytorch/test_torch_tweak_resnet_openvino_backend.py +++ b/tests/functional/pytorch/test_torch_tweak_resnet_openvino_backend.py @@ -10,7 +10,7 @@ import torch_tweak.torch as tt from torch_tweak.torch.backend import OpenVINOBackend, OpenVINOBackendConfig -from torch_tweak.torch.module import Module +from torch_tweak.torch.module import TunableModule from torch_tweak.torch.tune_strategy import OneBackendStrategy @@ -22,7 +22,7 @@ def test_openvino(use_dynamo: bool): backend = OpenVINOBackend(config=OpenVINOBackendConfig(use_dynamo=use_dynamo)) strategy = OneBackendStrategy(backend) - module = Module(model, "test-resnet18", strategy=strategy) + module = TunableModule(model, "test-resnet18", strategy=strategy) batch_sizes = [1, 10] tt.tune(module, samples, batch_sizes=batch_sizes) diff --git a/tests/functional/pytorch/test_torch_tweak_tune_llm.py b/tests/functional/pytorch/test_torch_tweak_tune_llm.py index a5a4b9f0..ad4d8261 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_llm.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_llm.py @@ -11,7 +11,7 @@ import torch from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline -from torch_tweak.torch import Module, OneBackendStrategy +from torch_tweak.torch import OneBackendStrategy, TunableModule from torch_tweak.torch import tune as torch_tweak from torch_tweak.torch.backend import TorchEagerBackend, TorchInductorBackend @@ -65,7 +65,7 @@ def tune_model(model, tokenizer, cache="static"): """ if cache == "dynamic": - model = Module( + model = TunableModule( model, model.__class__.__name__, strategy=OneBackendStrategy(TorchEagerBackend()).enable_find_max_batch_size( @@ -73,7 +73,7 @@ def tune_model(model, tokenizer, cache="static"): ), # just one strategy, there is no prefill/decode distinction ) elif cache == "static": - model = Module( + model = TunableModule( model, model.__class__.__name__, strategies=[ diff --git a/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_compile_backend.py b/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_compile_backend.py index fc63bd6d..f61cdfb1 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_compile_backend.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_compile_backend.py @@ -11,7 +11,7 @@ from tests.utilities.helpers import requires_xpu from torch_tweak.torch import tune from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend, TorchInductorBackendConfig -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy # max-autotune-no-cudagraphs exhausts XPU memory on current CI hardware. @@ -35,7 +35,7 @@ def test_tune_resnet_torch_inductor(mode: str, torch_device: torch.device): expected_probs = torch.nn.functional.softmax(out[0], dim=0) backend = TorchInductorBackend(config=TorchInductorBackendConfig(mode=mode)) - module = Module(model, "functional-resnet18", strategy=OneBackendStrategy(backend)) + module = TunableModule(model, "functional-resnet18", strategy=OneBackendStrategy(backend)) tune(module, data, batch_sizes=[2, 1], device=torch_device) out = module(data.unsqueeze(0)) diff --git a/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_eager_backend.py b/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_eager_backend.py index 8b89c32b..666f1afd 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_eager_backend.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_resnet_torch_eager_backend.py @@ -12,7 +12,7 @@ from torch_tweak.torch import tune from torch_tweak.torch.backend.torch_eager import TorchEagerBackend, TorchEagerBackendConfig -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy @@ -32,7 +32,7 @@ def do_test(backend: TorchEagerBackend): out = model(data.unsqueeze(0)) expected_probs = torch.nn.functional.softmax(out[0], dim=0) - module = Module(model, "functional-resnet18", strategy=OneBackendStrategy(backend)) + module = TunableModule(model, "functional-resnet18", strategy=OneBackendStrategy(backend)) # when tune(module, data, batch_sizes=[2, 1], dry_run=False, disable_external_logging=False) # then - verify tuning diff --git a/tests/functional/pytorch/test_torch_tweak_tune_resnet_torchao_backend.py b/tests/functional/pytorch/test_torch_tweak_tune_resnet_torchao_backend.py index e56fd749..782c0c67 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_resnet_torchao_backend.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_resnet_torchao_backend.py @@ -12,7 +12,7 @@ import torch from torch_tweak.torch.backend.torchao_backend import TorchAOBackend, TorchAOBackendConfig -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy from torch_tweak.torch.tuning import tune @@ -36,7 +36,7 @@ def do_test(backend: TorchAOBackend, dtype: torch.dtype): expected_probs = torch.nn.functional.softmax(out[0], dim=0) - module = Module( + module = TunableModule( model, "functional-resnet18", strategy=OneBackendStrategy(backend).enable_find_max_batch_size(False) ) # when diff --git a/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_forward_hooks.py b/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_forward_hooks.py index 90763f80..4802efcf 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_forward_hooks.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_forward_hooks.py @@ -10,7 +10,7 @@ import torch from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy @@ -39,7 +39,7 @@ def post_hook(module, input, output): # noqa: A002 expected_arg_max = model() # notice: not argument - it will be added by pre hook # when - module = Module(model, "functional-resnet50") + module = TunableModule(model, "functional-resnet50") # then - verify recording module() # notice: not argument - it will be added by pre hook diff --git a/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_module_device.py b/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_module_device.py index b548fce2..db7d8397 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_module_device.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_resnet_with_module_device.py @@ -10,7 +10,7 @@ import torch from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy @@ -39,7 +39,7 @@ def post_hook(module, input, output): # noqa: A002 expected_arg_max = model() # notice: not argument - it will be added by pre hook # when - module = Module(model, "functional-resnet50") + module = TunableModule(model, "functional-resnet50") # then - verify recording module() # notice: not argument - it will be added by pre hook diff --git a/tests/functional/pytorch/test_torch_tweak_tune_save_load.py b/tests/functional/pytorch/test_torch_tweak_tune_save_load.py index c6315811..ce7130f7 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_save_load.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_save_load.py @@ -13,7 +13,7 @@ import torch from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy from torch_tweak.torch.tuning import load, save, tune @@ -62,12 +62,12 @@ def _tune_save_load_helper(model_factory, samples, output_dir, wrap_model_fn): def tune_save_load_whole_model(model_factory, samples, output_dir, backend_factory): - """Test tuning, saving and loading a whole model wrapped in a single Module.""" + """Test tuning, saving and loading a whole model wrapped in a single TunableModule.""" def wrap_whole_model(model): strategy = OneBackendStrategy(backend_factory()) strategy.enable_find_max_batch_size(enable=False) - return Module(model, "demo-simple", strategy=strategy) + return TunableModule(model, "demo-simple", strategy=strategy) _tune_save_load_helper(model_factory, samples, output_dir, wrap_whole_model) @@ -78,8 +78,8 @@ def tune_save_load_part_of_model(model_factory, samples, output_dir, backend_fac def wrap_partial_model(model): strategy = OneBackendStrategy(backend_factory()) strategy.enable_find_max_batch_size(enable=False) - model.layer1 = Module(model.layer1, "demo-simple1", strategy=strategy) - model.layer2 = Module(model.layer2, "demo-simple2", strategy=strategy) + model.layer1 = TunableModule(model.layer1, "demo-simple1", strategy=strategy) + model.layer2 = TunableModule(model.layer2, "demo-simple2", strategy=strategy) return model _tune_save_load_helper(model_factory, samples, output_dir, wrap_partial_model) diff --git a/tests/functional/pytorch/test_torch_tweak_tune_save_load_first_module_only.py b/tests/functional/pytorch/test_torch_tweak_tune_save_load_first_module_only.py index 4d505eb4..3e605753 100644 --- a/tests/functional/pytorch/test_torch_tweak_tune_save_load_first_module_only.py +++ b/tests/functional/pytorch/test_torch_tweak_tune_save_load_first_module_only.py @@ -14,7 +14,7 @@ from diffusers import StableDiffusionPipeline from torch_tweak.torch.backend import TorchEagerBackend, TorchInductorBackend -from torch_tweak.torch.module import Module +from torch_tweak.torch.module import TunableModule from torch_tweak.torch.tune_strategy import FirstWinsStrategy from torch_tweak.torch.tuning import load, save, tune @@ -41,7 +41,7 @@ def _get_pipeline(model_name: str = "stable-diffusion-v1-5/stable-diffusion-v1-5 def _tune_and_save(save_path: pathlib.Path, device: str = "xpu"): pipeline = _get_pipeline(device=device) - pipeline.text_encoder = Module( + pipeline.text_encoder = TunableModule( pipeline.text_encoder, "text_encoder", strategy=FirstWinsStrategy(backends=[TorchEagerBackend(), TorchInductorBackend()]), diff --git a/tests/functional/pytorch/test_torch_tweak_wrap_module_resnet.py b/tests/functional/pytorch/test_torch_tweak_wrap_module_resnet.py index 19f1b8c8..904d0b53 100644 --- a/tests/functional/pytorch/test_torch_tweak_wrap_module_resnet.py +++ b/tests/functional/pytorch/test_torch_tweak_wrap_module_resnet.py @@ -8,7 +8,7 @@ import timm import torch -from torch_tweak.torch import Module, OneBackendStrategy +from torch_tweak.torch import OneBackendStrategy, TunableModule from torch_tweak.torch.backend import TorchInductorBackend @@ -27,7 +27,7 @@ def test_resnet50(): expected_probs = torch.nn.functional.softmax(out[0], dim=0) # when - module = Module(model, "functional-resnet50") + module = TunableModule(model, "functional-resnet50") # then - verify recording module(data) diff --git a/tests/unit/torch/backend/test_xpu_graph_backend.py b/tests/unit/torch/backend/test_xpu_graph_backend.py index eda892d5..82547f12 100644 --- a/tests/unit/torch/backend/test_xpu_graph_backend.py +++ b/tests/unit/torch/backend/test_xpu_graph_backend.py @@ -16,7 +16,7 @@ from torch_tweak.torch.module.graph_spec import GraphSpec from torch_tweak.torch.module.recording_module import Sample from torch_tweak.torch.module.sample_metadata import SampleMetadata -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.tune_strategy import FirstWinsStrategy, HighestThroughputStrategy, OneBackendStrategy from torch_tweak.torch.tuning import tune @@ -323,7 +323,7 @@ def test_tune_module(torch_device, tmp_path, eager_calls): # tune routes later calls of module through the selected backend. reference = copy.deepcopy(module) sample = torch.randn(16, device=torch_device) - wrapped = Module( + wrapped = TunableModule( module, "xpu-graph-mlp", strategy=OneBackendStrategy(XPUGraphBackend()).enable_find_max_batch_size(False), diff --git a/tests/unit/torch/checkpoint/test_torch_checkpoint.py b/tests/unit/torch/checkpoint/test_torch_checkpoint.py index 202ddabc..023f9cc0 100644 --- a/tests/unit/torch/checkpoint/test_torch_checkpoint.py +++ b/tests/unit/torch/checkpoint/test_torch_checkpoint.py @@ -17,7 +17,7 @@ from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend from torch_tweak.torch.checkpoint.local_torch_storage import LocalTorchStorage from torch_tweak.torch.checkpoint.torch_checkpoint import TorchCheckpoint -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy from torch_tweak.torch.tuning import tune @@ -95,12 +95,12 @@ def _tune_save_load_helper(model_factory, samples, wrap_model_fn, checkpoint, de def _wrap_whole_model(model, backend=None): - return Module(model, "demo-simple", strategy=OneBackendStrategy(backend or TorchInductorBackend())) + return TunableModule(model, "demo-simple", strategy=OneBackendStrategy(backend or TorchInductorBackend())) def _wrap_linear_layers(model): - model.linear1 = Module(model.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend())) - model.linear2 = Module(model.linear2, "demo-simple2", strategy=OneBackendStrategy(TorchInductorBackend())) + model.linear1 = TunableModule(model.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend())) + model.linear2 = TunableModule(model.linear2, "demo-simple2", strategy=OneBackendStrategy(TorchInductorBackend())) return model @@ -138,7 +138,7 @@ def __call__(self, x): ], ) # cover both compiled and eager backends def test_tune_save_load_whole_model(model_factory, sample, backend, checkpoint): - """Test tuning, saving and loading a whole model wrapped in a single Module.""" + """Test tuning, saving and loading a whole model wrapped in a single TunableModule.""" _tune_save_load_helper(model_factory, sample, lambda model: _wrap_whole_model(model, backend), checkpoint) @@ -281,7 +281,7 @@ def test_tune_save_load_pipeline_with_nested_tuned_modules_rejects_leaf_name_in_ @requires_xpu def test_state_dict_from_pipeline_stores_only_nested_tuned_modules(complex_pipeline_factory, sample): pipeline = complex_pipeline_factory() - pipeline.net.linear1 = Module( + pipeline.net.linear1 = TunableModule( pipeline.net.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend()) ) tune(pipeline, sample, batch_sizes=[1, 2]) @@ -308,7 +308,7 @@ def test_tune_save_load_pipeline_restores_aliased_tuned_module_as_one_module( restored = checkpoint.load(_alias_net_linear1(complex_pipeline_factory()), "checkpoint.tt") - assert isinstance(restored.linear1, Module) + assert isinstance(restored.linear1, TunableModule) assert restored.linear1 is restored.net.linear1 @@ -325,7 +325,7 @@ def test_tune_save_load_pipeline_restores_tuned_module_aliased_by_its_parent( restored = checkpoint.load(restored, "checkpoint.tt") - assert isinstance(restored.net.linear1_alias, Module) + assert isinstance(restored.net.linear1_alias, TunableModule) assert restored.net.linear1_alias is restored.net.linear1 @@ -337,7 +337,7 @@ def test_load_state_dict_for_pipeline_restores_duplicate_alias_keys_as_one_modul TorchCheckpoint.load_state_dict_for_pipeline(restored, state_dict, device_map={}) - assert isinstance(restored.linear1, Module) + assert isinstance(restored.linear1, TunableModule) assert restored.linear1 is restored.net.linear1 @@ -364,8 +364,12 @@ def test_load_state_dict_for_pipeline_rejects_missing_path(complex_pipeline_fact @requires_xpu def test_state_dict_from_pipeline_warns_about_tuned_module_outside_module_attributes(sample, caplog): pipeline = _PipelineWithLayerList() - pipeline.linear1 = Module(pipeline.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend())) - pipeline.layers[0] = Module(pipeline.layers[0], "demo-simple2", strategy=OneBackendStrategy(TorchInductorBackend())) + pipeline.linear1 = TunableModule( + pipeline.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend()) + ) + pipeline.layers[0] = TunableModule( + pipeline.layers[0], "demo-simple2", strategy=OneBackendStrategy(TorchInductorBackend()) + ) tune(pipeline, sample, batch_sizes=[1, 2]) # `tune` calls `setup_logging`, which removes the `caplog` handler from the root logger. logging.getLogger().addHandler(caplog.handler) @@ -415,7 +419,7 @@ def test_get_pipeline_modules_does_not_descend_into_module_twice(): @requires_xpu def test_state_dict_from_pipeline(pipeline_factory): pipeline = pipeline_factory() - pipeline.linear1 = Mock(spec=Module) + pipeline.linear1 = Mock(spec=TunableModule) # second module is still torch.nn.Module i.e. not tuned state_dict = TorchCheckpoint.state_dict_from_pipeline(pipeline) assert isinstance(state_dict, dict) @@ -433,7 +437,7 @@ def test_load_state_dict_for_pipeline_unknown_module(pipeline_factory): def test_load_state_dict_for_pipeline_non_module_attribute(pipeline_factory): pipeline = pipeline_factory() pipeline.config = "not a module" - with pytest.raises(ValueError, match="Module config is not a torch.nn.Module or WrapperModule"): + with pytest.raises(ValueError, match="Module config is not a torch.nn.Module or TunableModule"): TorchCheckpoint.load_state_dict_for_pipeline(pipeline, {"config": {}}, device_map={}) @@ -441,6 +445,6 @@ def test_load_state_dict_for_pipeline_non_module_attribute(pipeline_factory): def test_state_dict_from_pipeline_no_tuned_modules(pipeline_factory): """Test that ValueError is raised when trying to get state dict from a pipeline with no tuned modules.""" pipeline = pipeline_factory() - # Both modules are regular torch.nn.Module, not wrapped in Module + # Both modules are regular torch.nn.Module, not wrapped in TunableModule with pytest.raises(ValueError, match="No tuned modules found in the pipeline"): TorchCheckpoint.state_dict_from_pipeline(pipeline) diff --git a/tests/unit/torch/inspecting/test_wrapping.py b/tests/unit/torch/inspecting/test_wrapping.py index d29d6968..67d65d31 100644 --- a/tests/unit/torch/inspecting/test_wrapping.py +++ b/tests/unit/torch/inspecting/test_wrapping.py @@ -15,7 +15,7 @@ from torch_tweak.torch.inspecting.module_inspector import DictOfModulesInfo, ListOfModulesInfo from torch_tweak.torch.inspecting.wrapping import wrap from torch_tweak.torch.module.sample_metadata import SampleMetadata -from torch_tweak.torch.module.wrapper_module import Module, ModuleState +from torch_tweak.torch.module.tunable_module import ModuleState, TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.tune_strategy import FirstWinsStrategy, OneBackendStrategy from torch_tweak.torch.tuning import tune @@ -53,8 +53,8 @@ def test_wrap_root_module(simple_model): ], ) - # Then the model should be wrapped in a Module - assert isinstance(wrapped_model, Module) + # Then the model should be wrapped in a TunableModule + assert isinstance(wrapped_model, TunableModule) assert wrapped_model._self_name == "Sequential" @@ -70,9 +70,9 @@ def test_wrap_submodule(nested_model): ) # Then the submodule should be wrapped - assert not isinstance(wrapped_model, Module) + assert not isinstance(wrapped_model, TunableModule) - assert isinstance(wrapped_model.linear, Module) + assert isinstance(wrapped_model.linear, TunableModule) assert wrapped_model.linear._self_name == "linear" # And the original model should be returned @@ -94,8 +94,8 @@ def test_wrap_multiple_modules(nested_model): ) # Then both submodules should be wrapped - assert isinstance(nested_model.linear, Module) - assert isinstance(nested_model.relu, Module) + assert isinstance(nested_model.linear, TunableModule) + assert isinstance(nested_model.relu, TunableModule) assert nested_model.linear._self_name == "linear" assert nested_model.relu._self_name == "relu" @@ -117,7 +117,7 @@ def test_wrap_with_strategy(simple_model): ) # Then the module should be wrapped with the strategy - assert isinstance(wrapped_model, Module) + assert isinstance(wrapped_model, TunableModule) assert wrapped_model._self_strategy.__class__ == strategy.__class__ assert wrapped_model._self_strategy_list is None assert wrapped_model._self_strategy_map is None @@ -140,7 +140,7 @@ def test_wrap_with_strategies_list(simple_model): ) # Then the module should be wrapped with the strategy - assert isinstance(wrapped_model, Module) + assert isinstance(wrapped_model, TunableModule) assert wrapped_model._self_strategy is None assert wrapped_model._self_strategy_map is None @@ -167,7 +167,7 @@ def test_wrap_with_strategies_map(simple_model): ) # Then the module should be wrapped with the strategy - assert isinstance(wrapped_model, Module) + assert isinstance(wrapped_model, TunableModule) assert wrapped_model._self_strategy is None assert wrapped_model._self_strategy_list is None assert set(wrapped_model._self_strategy_map.keys()) == set(strategies.keys()) @@ -179,7 +179,7 @@ def test_wrap_empty_modules_list(simple_model): wrapped_model = wrap(simple_model, []) # Then the original model should be returned unchanged - assert not isinstance(wrapped_model, Module) + assert not isinstance(wrapped_model, TunableModule) assert wrapped_model is simple_model @@ -229,7 +229,7 @@ def __call__(self, x): # Then verify the results assert isinstance(wrapped_obj, CustomObject) - assert isinstance(wrapped_obj.model, Module) + assert isinstance(wrapped_obj.model, TunableModule) def test_wrap_custom_object_nested(): @@ -262,7 +262,7 @@ def __call__(self, x): # Then verify the results assert isinstance(wrapped_obj, CustomObject) - assert isinstance(wrapped_obj.layer, Module) + assert isinstance(wrapped_obj.layer, TunableModule) def test_wrap_custom_object_multiple_modules(): @@ -292,8 +292,8 @@ def __call__(self, x): # Then verify the results assert isinstance(wrapped_obj, CustomObject) - assert isinstance(wrapped_obj.model1, Module) - assert isinstance(wrapped_obj.model2, Module) + assert isinstance(wrapped_obj.model1, TunableModule) + assert isinstance(wrapped_obj.model2, TunableModule) def test_wrap_custom_object_without_modules(): @@ -315,7 +315,7 @@ def __call__(self, x): # Then verify the results assert isinstance(wrapped_obj, CustomObject) - assert not isinstance(wrapped_obj.data, Module) + assert not isinstance(wrapped_obj.data, TunableModule) def test_wrap_custom_object_with_attributes(): @@ -343,11 +343,11 @@ def __call__(self, x): # Then verify the results assert isinstance(wrapped_obj, CustomObject) - assert isinstance(wrapped_obj.model, Module) + assert isinstance(wrapped_obj.model, TunableModule) - assert not isinstance(wrapped_obj.data, Module) - assert not isinstance(wrapped_obj.list_data, Module) - assert not isinstance(wrapped_obj.dict_data, Module) + assert not isinstance(wrapped_obj.data, TunableModule) + assert not isinstance(wrapped_obj.list_data, TunableModule) + assert not isinstance(wrapped_obj.dict_data, TunableModule) def test_wrap_module_in_list(): @@ -356,10 +356,10 @@ def test_wrap_module_in_list(): parent = ListOfModulesInfo(module=model) wrapped_model = wrap(model, [ModuleInfo(name="0", module=model[0], parent=parent, object_path=".0")]) - assert not isinstance(wrapped_model, Module) - assert isinstance(wrapped_model[0], Module) - assert not isinstance(wrapped_model[1], Module) - assert not isinstance(wrapped_model[2], Module) + assert not isinstance(wrapped_model, TunableModule) + assert isinstance(wrapped_model[0], TunableModule) + assert not isinstance(wrapped_model[1], TunableModule) + assert not isinstance(wrapped_model[2], TunableModule) def test_wrap_module_in_dict(): @@ -370,10 +370,10 @@ def test_wrap_module_in_dict(): wrapped_model = wrap( model, [ModuleInfo(name="linear", module=model["linear"], parent=parent, object_path=".linear")] ) - assert not isinstance(wrapped_model, Module) - assert isinstance(wrapped_model["linear"], Module) - assert not isinstance(wrapped_model["relu"], Module) - assert not isinstance(wrapped_model["linear2"], Module) + assert not isinstance(wrapped_model, TunableModule) + assert isinstance(wrapped_model["linear"], TunableModule) + assert not isinstance(wrapped_model["relu"], TunableModule) + assert not isinstance(wrapped_model["linear2"], TunableModule) def test_wrap_nested_diffrent_types(): @@ -399,10 +399,10 @@ def forward(self, x): ] wrapped_model = wrap(model, modules) - assert not isinstance(wrapped_model, Module) - assert isinstance(wrapped_model.linear, Module) - assert isinstance(wrapped_model.dict["linear"], Module) - assert isinstance(wrapped_model.list[0], Module) + assert not isinstance(wrapped_model, TunableModule) + assert isinstance(wrapped_model.linear, TunableModule) + assert isinstance(wrapped_model.dict["linear"], TunableModule) + assert isinstance(wrapped_model.list[0], TunableModule) class _Block(nn.Module): diff --git a/tests/unit/torch/module/test_passthrough_module.py b/tests/unit/torch/module/test_passthrough_module.py index 4e3b6b46..775a18d3 100644 --- a/tests/unit/torch/module/test_passthrough_module.py +++ b/tests/unit/torch/module/test_passthrough_module.py @@ -11,7 +11,7 @@ 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.module.tunable_module import TunableModule from torch_tweak.torch.utils.xpu import is_available as is_xpu_available devices = ["cpu", "xpu"] if is_xpu_available() else ["cpu"] @@ -27,7 +27,7 @@ def test_call(device): def test_module_without_parameters_keeps_inputs_on_their_device(): - module = Module(nn.Flatten(), "flatten") + module = TunableModule(nn.Flatten(), "flatten") module.enable_passthrough() output = module(torch.empty(2, 3, 4, device="meta")) diff --git a/tests/unit/torch/module/test_wrapper_module.py b/tests/unit/torch/module/test_tunable_module.py similarity index 87% rename from tests/unit/torch/module/test_wrapper_module.py rename to tests/unit/torch/module/test_tunable_module.py index 8f1160d5..c7f6a698 100644 --- a/tests/unit/torch/module/test_wrapper_module.py +++ b/tests/unit/torch/module/test_tunable_module.py @@ -2,7 +2,7 @@ # Copyright (c) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 # NOTE: This file has been modified by Intel Corporation. -"""Test for wrapper module.""" +"""Tests for the tunable module.""" import inspect from collections import OrderedDict @@ -19,8 +19,8 @@ from torch_tweak.torch.config import torch_tweak_cache_dir from torch_tweak.torch.module.graph_spec import GraphSpec from torch_tweak.torch.module.sample_metadata import SampleMetadata +from torch_tweak.torch.module.tunable_module import ModuleState, TunableModule from torch_tweak.torch.module.tuned_module import TunedModule -from torch_tweak.torch.module.wrapper_module import Module, ModuleState from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.tune_strategy import ( OneBackendStrategy, @@ -50,12 +50,12 @@ def simulate_prefill_decode(self, x): def module(): model = Identity() - return Module(model, TEST_MODULE_NAME, strategy=DummyTuneStrategy()) + return TunableModule(model, TEST_MODULE_NAME, strategy=DummyTuneStrategy()) def test_init(): model = Identity() - module = Module(model, TEST_MODULE_NAME) + module = TunableModule(model, TEST_MODULE_NAME) assert module._self_name == TEST_MODULE_NAME assert module.state == ModuleState.INIT assert module.device == torch.device("cpu") @@ -68,7 +68,7 @@ def test_init(): def test_module_registration(): assert MODULE_REGISTRY.modules == {} - module = Module(Identity(), TEST_MODULE_NAME) + module = TunableModule(Identity(), TEST_MODULE_NAME) assert MODULE_REGISTRY.modules == {TEST_MODULE_NAME: module} @@ -125,7 +125,9 @@ def test_tune_dry_run(module, torch_device): strategy = Mock() global_config.strict_mode = True - with pytest.raises(ValueError, match="Module: 'demo-identity' has not recorded any samples. Cannot tune it."): + with pytest.raises( + ValueError, match="TunableModule: 'demo-identity' has not recorded any samples. Cannot tune it." + ): module.tune(strategy=strategy, dry_run=True, device=torch_device) # record a sample @@ -169,7 +171,9 @@ def test_tune(module, torch_device): strategy = Mock() global_config.strict_mode = True - with pytest.raises(ValueError, match="Module: 'demo-identity' has not recorded any samples. Cannot tune it."): + with pytest.raises( + ValueError, match="TunableModule: 'demo-identity' has not recorded any samples. Cannot tune it." + ): module.tune(strategy=strategy, dry_run=True, device=torch_device) # record a sample @@ -210,7 +214,7 @@ def test_tune(module, torch_device): def test_tune_with_provided_strategy(torch_device): model = Identity() strategy = Mock() - module = Module(model, TEST_MODULE_NAME, strategy=strategy) + module = TunableModule(model, TEST_MODULE_NAME, strategy=strategy) module(1, a=1) module(2, b=2) @@ -224,7 +228,7 @@ def test_tune_with_list_of_strategies(torch_device): model = Identity() strategy1 = Mock() strategy2 = Mock() - module = Module(model, TEST_MODULE_NAME, strategies=[strategy1, strategy2]) + module = TunableModule(model, TEST_MODULE_NAME, strategies=[strategy1, strategy2]) module(1, a=1) module(2, b=2) @@ -234,7 +238,7 @@ def test_tune_with_list_of_strategies(torch_device): module._self_strategy_list[1].tune.assert_called_once() assert module.state == ModuleState.TUNED - with pytest.raises(RuntimeError, match="Module is already tuned. Use force=True to reset tuned module."): + with pytest.raises(RuntimeError, match="TunableModule is already tuned. Use force=True to reset tuned module."): module.enable_recording() module.enable_recording(force=True) @@ -253,7 +257,7 @@ def test_tune_with_dict_of_strategies(torch_device): strategy2 = Mock() graph1 = SampleMetadata.from_inputs(args=(1,), kwargs={}, strict=True) graph2 = SampleMetadata.from_inputs(args=(2,), kwargs={}, strict=True) - module = Module(model, TEST_MODULE_NAME, strategies={graph1: strategy1, graph2: strategy2}) + module = TunableModule(model, TEST_MODULE_NAME, strategies={graph1: strategy1, graph2: strategy2}) module(1) module(2) @@ -263,7 +267,7 @@ def test_tune_with_dict_of_strategies(torch_device): module._self_strategy_map[graph2].tune.assert_called_once() assert module.state == ModuleState.TUNED - with pytest.raises(RuntimeError, match="Module is already tuned. Use force=True to reset tuned module."): + with pytest.raises(RuntimeError, match="TunableModule is already tuned. Use force=True to reset tuned module."): module.enable_recording() module.enable_recording(force=True) @@ -323,7 +327,7 @@ def test_tune_with_torch_compile_backend_simulate_prefill_decode(module, torch_d def test_serialization(torch_device): """Test converting module to dictionary.""" model = Identity() - module = Module(model, TEST_MODULE_NAME) + module = TunableModule(model, TEST_MODULE_NAME) # Record samples and tune the module module(torch.tensor(1)) @@ -335,13 +339,13 @@ def test_serialization(torch_device): state_dict = module.to_dict(prefix="test") assert isinstance(state_dict, OrderedDict) module_data = state_dict["test"] - assert module_data[Module.TYPE_KEY] == Module.__name__ - assert module_data[Module.NAME_KEY] == TEST_MODULE_NAME - assert Module.TUNED_MODULE_KEY in module_data + assert module_data[TunableModule.TYPE_KEY] == TunableModule.__name__ + assert module_data[TunableModule.NAME_KEY] == TEST_MODULE_NAME + assert TunableModule.TUNED_MODULE_KEY in module_data # Test from_dict model = Identity() - module = Module.from_dict(model, state_dict["test"], device=torch_device) + module = TunableModule.from_dict(model, state_dict["test"], device=torch_device) assert module.state == ModuleState.TUNED assert module._self_name == TEST_MODULE_NAME assert isinstance(module._self_wrapper, TunedModule) @@ -349,15 +353,15 @@ def test_serialization(torch_device): def test_to_dict_raises_error_when_not_tuned(): """Test that to_dict raises error when module is not tuned.""" - module = Module(Identity(), TEST_MODULE_NAME) - with pytest.raises(RuntimeError, match="Module is not tuned. Cannot save state_dict."): + module = TunableModule(Identity(), TEST_MODULE_NAME) + with pytest.raises(RuntimeError, match="TunableModule is not tuned. Cannot save state_dict."): module.to_dict() def test_from_dict_raises_error_for_invalid_state_dict(torch_device): """Test that from_dict raises error for invalid state dict.""" - with pytest.raises(ValueError, match=f"Invalid dictionary format for {Module.__class__.__name__}"): - Module.from_dict(Identity(), {"type": "invalid"}, device=torch_device) + with pytest.raises(ValueError, match=f"Invalid dictionary format for {TunableModule.__class__.__name__}"): + TunableModule.from_dict(Identity(), {"type": "invalid"}, device=torch_device) def test_deepcopy(): @@ -368,7 +372,7 @@ def test_deepcopy(): a backend. """ model = Identity() - module = Module(model, TEST_MODULE_NAME) + module = TunableModule(model, TEST_MODULE_NAME) module._restore_original_forward() deepcopy(model) # this must not fail @@ -387,7 +391,7 @@ def forward(self, x, y, z, pos=4): original_model = TestNet() expected_keys = {"x", "y", "z", "pos"} - model = Module(original_model, TEST_MODULE_NAME) + model = TunableModule(original_model, TEST_MODULE_NAME) assert model.state == ModuleState.INIT assert set(inspect.signature(original_model.forward).parameters.keys()) == expected_keys model(1, 2, 3) @@ -403,7 +407,7 @@ def forward(self, x, y, z, pos=4): def test_sleep_backend_keeps_original_module_after_tune(): """Backends that run the original module must not have it offloaded to meta.""" model = Identity() - module = Module(model, TEST_MODULE_NAME) + module = TunableModule(model, TEST_MODULE_NAME) sample = torch.randn(2, 5) module(sample) diff --git a/tests/unit/torch/task/test_profiling.py b/tests/unit/torch/task/test_profiling.py index 36866936..1fb2a2c2 100644 --- a/tests/unit/torch/task/test_profiling.py +++ b/tests/unit/torch/task/test_profiling.py @@ -8,7 +8,7 @@ import pytest from tests.toy_models.torch_models import ToyTorchModel -from torch_tweak.torch import Module +from torch_tweak.torch import TunableModule from torch_tweak.torch.backend import TorchEagerBackend from torch_tweak.torch.task.profiling.config import ProfilingConfig from torch_tweak.torch.task.profiling.events import ProfilingResultEvent @@ -194,7 +194,7 @@ def test_measurement_stop_strategy_stable_window_stability(window_size: int, sta def test_profile_toy_model(): profile_config = ProfilingConfig(batch_sizes=[2**n for n in range(1, 5)]) - model = Module(ToyTorchModel()) + model = TunableModule(ToyTorchModel()) dataset = model.samples() result = profile(model, dataset, profile_config) @@ -205,7 +205,7 @@ def test_profile_toy_model(): def test_profile_toy_model_no_batching(): profile_config = ProfilingConfig(batch_sizes=[2**n for n in range(1, 5)], batching=False) - model = Module(ToyTorchModel()) + model = TunableModule(ToyTorchModel()) dataset = model.samples() result = profile(model, dataset, profile_config) diff --git a/tests/unit/torch/test_torch_llm_support.py b/tests/unit/torch/test_torch_llm_support.py index 935b6e91..19786fcc 100644 --- a/tests/unit/torch/test_torch_llm_support.py +++ b/tests/unit/torch/test_torch_llm_support.py @@ -12,10 +12,10 @@ import pytest import torch -from torch_tweak.torch import Module, OneBackendStrategy +from torch_tweak.torch import OneBackendStrategy, TunableModule from torch_tweak.torch.backend import TorchInductorBackend from torch_tweak.torch.module.locator import Locator -from torch_tweak.torch.module.wrapper_module import ModuleState +from torch_tweak.torch.module.tunable_module import ModuleState try: from tests.utilities.helpers import requires_xpu @@ -255,7 +255,7 @@ def test_llm_dynamic_cache(cache_type, registered_type, torch_device): model.eval() kv_cache = cache_type() - model = Module(model, "test_llm") # wrap the model with Torch Tweak + model = TunableModule(model, "test_llm") # wrap the model with Torch Tweak with torch.no_grad(): result, kv_cache = model.generate(prompt_token_ids, kv_cache, max_seq_len=MAX_SEQ_LEN) diff --git a/tests/unit/torch/tune_strategy/test_first_wins_strategy.py b/tests/unit/torch/tune_strategy/test_first_wins_strategy.py index 0577b409..dc3e8d3d 100644 --- a/tests/unit/torch/tune_strategy/test_first_wins_strategy.py +++ b/tests/unit/torch/tune_strategy/test_first_wins_strategy.py @@ -12,10 +12,10 @@ from tests.toy_backends import BuildFailsBackend, SleepBackend from tests.toy_models.torch_models import ToyTorchModel -from torch_tweak.torch import Module, tune +from torch_tweak.torch import TunableModule, tune from torch_tweak.torch.backend import Backend from torch_tweak.torch.module.graph_spec import GraphSpec -from torch_tweak.torch.module.wrapper_module import ModuleState +from torch_tweak.torch.module.tunable_module import ModuleState from torch_tweak.torch.tune_strategy.first_wins_strategy import FirstWinsStrategy @@ -176,7 +176,7 @@ def test_first_wins_strategy_find_max_batch_size_fails(torch_device): sample = model.sample().to(torch_device) assert model(sample) is not None - model = Module(model, strategy=strategy) + model = TunableModule(model, strategy=strategy) tune(model, sample, device=torch_device, disable_external_logging=False) assert model.state == ModuleState.PASSTHROUGH @@ -197,7 +197,7 @@ class TestOutOfMemoryException(Exception): model = ToyTorchModel().eval().to(torch_device) sample = model.sample().to(torch_device) - model = Module(model, strategy=strategy) + model = TunableModule(model, strategy=strategy) tune(model, sample, device=torch_device, disable_external_logging=False) diff --git a/tests/unit/torch/tune_strategy/test_highest_throughput_strategy.py b/tests/unit/torch/tune_strategy/test_highest_throughput_strategy.py index 97f703ac..9c0f4b6f 100644 --- a/tests/unit/torch/tune_strategy/test_highest_throughput_strategy.py +++ b/tests/unit/torch/tune_strategy/test_highest_throughput_strategy.py @@ -13,11 +13,11 @@ from tests.toy_backends import BuildFailsBackend, SleepBackend from tests.toy_models.torch_models import ToyTorchModel from tests.utilities.helpers import requires_xpu -from torch_tweak.torch import Module +from torch_tweak.torch import TunableModule from torch_tweak.torch.backend import Backend from torch_tweak.torch.backend.torch_eager import TorchEagerBackend from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend -from torch_tweak.torch.module.wrapper_module import ModuleState, get_object_name +from torch_tweak.torch.module.tunable_module import ModuleState, get_object_name from torch_tweak.torch.task.correctness import CorrectnessValueError from torch_tweak.torch.task.profiling import NumStepsMeasuringStopStrategy, StableWindowMeasuringStopStrategy from torch_tweak.torch.task.profiling.profiling_stop_strategy import ( @@ -118,7 +118,7 @@ def test_highest_throughput_strategy_num_steps_all_samples(torch_device): strategy.set_find_max_batch_size_profiling_config(find_profiling_config) strategy.enable_correctness_check(True) - model = Module(ToyTorchModel().eval().to(torch_device), strategy=strategy) + model = TunableModule(ToyTorchModel().eval().to(torch_device), strategy=strategy) sample = model.sample().to(torch_device) batch_sizes = [1, 2, 4, 8, 16] n_backends = len(strategy._backends) @@ -158,7 +158,7 @@ def test_highest_throughput_strategy_stable_window(torch_device): measurement_stop_strategy=StableWindowMeasuringStopStrategy(window_size=10, stability_percentage=90), profiling_stop_strategy=ThroughputSaturatedProfilingStopStrategy(throughput_cutoff_threshold=0.99), ).enable_find_max_batch_size(False) - model = Module(ToyTorchModel().eval().to(torch_device), strategy=strategy) + model = TunableModule(ToyTorchModel().eval().to(torch_device), strategy=strategy) sample = model.sample().to(torch_device) batch_sizes = list(range(1, 17)) @@ -208,7 +208,7 @@ def test_highest_throughput_strategy_fails_backend_if_all_of_backends_fails(torc batch_sizes = list(range(1, 17)) - model = Module(model, strategy=strategy) + model = TunableModule(model, strategy=strategy) tune(model, sample, batch_sizes=batch_sizes, device=torch_device, disable_external_logging=False) assert model.state == ModuleState.PASSTHROUGH @@ -231,7 +231,7 @@ def test_highest_throughput_strategy_select_backend_if_one_of_backends_succeeds( batch_sizes = list(range(1, 17)) - model = Module(model, strategy=strategy) + model = TunableModule(model, strategy=strategy) tune(model, sample, batch_sizes=batch_sizes, device=torch_device, disable_external_logging=False) assert model(sample.unsqueeze(0)) is not None # Note: does not work without unsqueeze @@ -255,7 +255,7 @@ def test_highest_throughput_strategy_find_max_batch_size_fails(torch_device): batch_sizes = list(range(1, 17)) - model = Module(model, strategy=strategy) + model = TunableModule(model, strategy=strategy) tune(model, sample, batch_sizes=batch_sizes, device=torch_device, disable_external_logging=False) assert model.state == ModuleState.PASSTHROUGH diff --git a/torch_tweak/torch/__init__.py b/torch_tweak/torch/__init__.py index d3bf2cdd..6d627f33 100644 --- a/torch_tweak/torch/__init__.py +++ b/torch_tweak/torch/__init__.py @@ -10,7 +10,7 @@ from torch_tweak.torch.checkpoint.local_torch_storage import LocalTorchStorage from torch_tweak.torch.dataloader import DataLoaderFactory from torch_tweak.torch.inspecting import inspect, wrap -from torch_tweak.torch.module import Module +from torch_tweak.torch.module import TunableModule from torch_tweak.torch.tune_strategy import ( FirstWinsStrategy, HighestThroughputStrategy, @@ -27,7 +27,7 @@ "tune", "load", "save", - "Module", + "TunableModule", "TuneStrategy", "OneBackendStrategy", "FirstWinsStrategy", diff --git a/torch_tweak/torch/checkpoint/checkpoint.py b/torch_tweak/torch/checkpoint/checkpoint.py index f2ffd95c..411f550b 100644 --- a/torch_tweak/torch/checkpoint/checkpoint.py +++ b/torch_tweak/torch/checkpoint/checkpoint.py @@ -10,7 +10,7 @@ import torch from torch.nn import Module -from torch_tweak.torch.module.wrapper_module import Module as WrapperModule +from torch_tweak.torch.module.tunable_module import TunableModule class Checkpoint(ABC): @@ -20,7 +20,7 @@ class Checkpoint(ABC): """ @abstractmethod - def save(self, module: Module | WrapperModule, path: str | Path) -> None: + def save(self, module: Module | TunableModule, path: str | Path) -> None: """Save a module to a checkpoint. Args: @@ -38,7 +38,7 @@ def load( module: Module, path: str | Path, device_map: dict[str, str | torch.device] | None = None, - ) -> Module | WrapperModule: + ) -> Module | TunableModule: """Load a module from a checkpoint. Args: diff --git a/torch_tweak/torch/checkpoint/torch_checkpoint.py b/torch_tweak/torch/checkpoint/torch_checkpoint.py index a2ab8409..ad169600 100644 --- a/torch_tweak/torch/checkpoint/torch_checkpoint.py +++ b/torch_tweak/torch/checkpoint/torch_checkpoint.py @@ -14,8 +14,7 @@ from torch_tweak.torch.checkpoint.checkpoint import Checkpoint from torch_tweak.torch.checkpoint.storage import Storage -from torch_tweak.torch.module.wrapper_module import Module as WrapperModule -from torch_tweak.torch.module.wrapper_module import ModuleState +from torch_tweak.torch.module.tunable_module import ModuleState, TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.utils.device import get_device @@ -45,7 +44,7 @@ def save( This method handles both individual torch modules and pipelines containing modules. For individual modules, it directly uses the module's state_dict() method. - For a pipeline, it saves each tuned module (WrapperModule instance) under its dotted path, + For a pipeline, it saves each tuned module (TunableModule instance) under its dotted path, e.g. `"vae.decoder"`. Modules that are not tuned are not saved. Args: @@ -67,7 +66,7 @@ def load( module_or_pipeline: TorchModule | Any, path: str | Path, device_map: dict[str, str | torch.device] | None = None, - ) -> TorchModule | WrapperModule: + ) -> TorchModule | TunableModule: """Load a module or pipeline from a saved state dictionary. This method handles both individual torch modules and pipelines containing modules. @@ -76,7 +75,7 @@ def load( For pipelines, it loads state dictionaries into the corresponding modules in the pipeline. The pipeline is expected to be an object with attributes that are torch modules or - tuned modules (WrapperModule instances). + tuned modules (TunableModule instances). Args: module_or_pipeline: The module or pipeline to load the state into. This can be @@ -90,7 +89,7 @@ def load( Raises: ValueError: If a module in the state_dict is not found in the pipeline or - if a matched attribute is not a torch.nn.Module or WrapperModule. + if a matched attribute is not a torch.nn.Module or TunableModule. """ # Convert device_map to torch.device @@ -110,7 +109,7 @@ def load( return module_or_pipeline @staticmethod - def get_pipeline_modules(pipeline: Any) -> dict[str, TorchModule | WrapperModule]: + def get_pipeline_modules(pipeline: Any) -> dict[str, TorchModule | TunableModule]: """Get the modules from the pipeline. Args: @@ -120,14 +119,14 @@ def get_pipeline_modules(pipeline: Any) -> dict[str, TorchModule | WrapperModule dict: Maps the dotted path of every module the pipeline holds, e.g. `"vae.decoder"`, to the module. The search does not descend into tuned modules or into a module it has already visited. """ - modules: dict[str, TorchModule | WrapperModule] = {} + modules: dict[str, TorchModule | TunableModule] = {} visited_module_ids = set() - def add_module(path: str, module: TorchModule | WrapperModule) -> None: + def add_module(path: str, module: TorchModule | TunableModule) -> None: # Every path to a module is kept, so that `load_state_dict_for_pipeline` replaces each alias. # Children are listed under the first path only: replacing a child on a shared owner covers all paths. modules[path] = module - if isinstance(module, WrapperModule) or id(module) in visited_module_ids: + if isinstance(module, TunableModule) or id(module) in visited_module_ids: return visited_module_ids.add(id(module)) # `named_children()` yields a module held twice by one parent only once. @@ -138,7 +137,7 @@ def add_module(path: str, module: TorchModule | WrapperModule) -> None: # TODO: search lists, dicts and plain objects as `ModuleInspector` does, so that every module # `tune` tunes is also saved. for name, attr in vars(pipeline).items(): - if isinstance(attr, (TorchModule, WrapperModule)): + if isinstance(attr, (TorchModule, TunableModule)): add_module(name, attr) return modules @@ -159,7 +158,7 @@ def state_dict_from_pipeline(pipeline: Any) -> dict[str, Any]: saved_module_ids = set() for name, module in TorchCheckpoint.get_pipeline_modules(pipeline).items(): - if isinstance(module, WrapperModule) and id(module) not in saved_module_ids: + if isinstance(module, TunableModule) and id(module) not in saved_module_ids: state_dict[name] = module.state_dict() saved_module_ids.add(id(module)) @@ -180,11 +179,11 @@ def state_dict_from_pipeline(pipeline: Any) -> dict[str, Any]: @staticmethod def load_state_dict_for_module( - module: TorchModule | WrapperModule, + module: TorchModule | TunableModule, state_dict: dict, device_map: dict[str, torch.device], module_name: str = "", - ) -> TorchModule | WrapperModule: + ) -> TorchModule | TunableModule: """Load a state dictionary into a module. This method loads a state dictionary into a module, replacing regular torch modules @@ -194,7 +193,7 @@ def load_state_dict_for_module( The methods has two phases: 1. Traverse the module hierarchy and identify which parts of the state dictionary correspond to tuned modules. When a match is found, it replaces the original module - with a wrapper module initialized from the state dictionary. Corresponding keys are + with a tunable module initialized from the state dictionary. Corresponding keys are removed from the state_dict. The tuned module is compiled when necessary. 2. Handle original torch modules not tuned with `load_state_dict(state_dict, strict=False)`, strict=False because wrapped module keys were removed and those modules have been already loaded. @@ -204,13 +203,13 @@ def load_state_dict_for_module( This can be a regular torch module or a module that contains tuned submodules. state_dict (dict): The state dictionary containing the tuned module states. The dictionary should have keys that correspond to the module structure, - with tuned module states stored in a format recognized by WrapperModule. + with tuned module states stored in a format recognized by TunableModule. device_map (dict): The device map to load module to. module_name (str): The name of the module for device map matching. Returns: torch.nn.Module: The module with tuned components loaded from the state dictionary. - If the top-level module was tuned, returns the new wrapper module. + If the top-level module was tuned, returns the new tunable module. Otherwise, returns the original module with tuned submodules replaced. Note: @@ -240,18 +239,18 @@ def replace_tuned_module( module_name: The name of the module. """ stored_object = local_state_dict.get(prefix, {}) - if WrapperModule.is_state_dict_valid(stored_object): - # This is a tuned module, replace with wrapper module + if TunableModule.is_state_dict_valid(stored_object): + # This is a tuned module, replace with tunable module device = device_map.pop(module_name, None) - wrapper_module = WrapperModule.from_dict(module, stored_object, device) + tunable_module = TunableModule.from_dict(module, stored_object, device) if parent_module is None: # We replace top module nonlocal top_module - top_module = wrapper_module + top_module = tunable_module else: - # We replace a child module with wrapper module - setattr(parent_module, module_name, wrapper_module) + # We replace a child module with tunable module + setattr(parent_module, module_name, tunable_module) # Remove state_dict key - we already handled it del local_state_dict[prefix] else: @@ -301,7 +300,7 @@ def load_state_dict_for_pipeline(pipeline, state_dict, device_map: dict[str, tor Raises: ValueError: If a module in the state_dict is not found in the pipeline or - if a matched attribute is not a torch.nn.Module or WrapperModule. + if a matched attribute is not a torch.nn.Module or TunableModule. """ modules = TorchCheckpoint.get_pipeline_modules(pipeline) matched_modules = set() @@ -310,7 +309,7 @@ def load_state_dict_for_pipeline(pipeline, state_dict, device_map: dict[str, tor module = modules.get(name) if module is None: if hasattr(pipeline, name): - raise ValueError(f"Module {name} is not a torch.nn.Module or WrapperModule") + raise ValueError(f"Module {name} is not a torch.nn.Module or TunableModule") continue matched_modules.add(name) diff --git a/torch_tweak/torch/inspecting/wrapping.py b/torch_tweak/torch/inspecting/wrapping.py index 9fb2529a..198daee8 100644 --- a/torch_tweak/torch/inspecting/wrapping.py +++ b/torch_tweak/torch/inspecting/wrapping.py @@ -7,7 +7,7 @@ from logging import getLogger from torch_tweak.torch.inspecting.module_info import ModuleInfo -from torch_tweak.torch.module.wrapper_module import Module, StrategyList, StrategyMap +from torch_tweak.torch.module.tunable_module import StrategyList, StrategyMap, TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.tune_strategy.tune_strategy import TuneStrategy from torch_tweak.utils.logging import setup_logging @@ -39,11 +39,11 @@ def wrap( for module_info in modules: logger.info("Wrapping module: %s", module_info.object_path or module_info.name or obj.__class__.__name__) if module_info.parent is None: - return Module(obj, name=obj.__class__.__name__, strategy=strategy, strategies=strategies) + return TunableModule(obj, name=obj.__class__.__name__, strategy=strategy, strategies=strategies) # `module_info.name` repeats across siblings (`layers.0.mlp`, `layers.1.mlp`); the path does not. name = module_info.object_path.lstrip(".") - wrapped = Module(module_info.module, name=name, strategy=strategy, strategies=strategies) + wrapped = TunableModule(module_info.module, name=name, strategy=strategy, strategies=strategies) module_info.parent.set_wrapped(module_info.name, wrapped) return obj diff --git a/torch_tweak/torch/module/__init__.py b/torch_tweak/torch/module/__init__.py index 28d5a16e..5d1965d1 100644 --- a/torch_tweak/torch/module/__init__.py +++ b/torch_tweak/torch/module/__init__.py @@ -4,6 +4,6 @@ # NOTE: This file has been modified by Intel Corporation. """Torch wrappers module.""" -from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module.tunable_module import TunableModule -__all__ = ["Module"] +__all__ = ["TunableModule"] diff --git a/torch_tweak/torch/module/wrapper_module.py b/torch_tweak/torch/module/tunable_module.py similarity index 91% rename from torch_tweak/torch/module/wrapper_module.py rename to torch_tweak/torch/module/tunable_module.py index 61b6e36a..b18c2b58 100644 --- a/torch_tweak/torch/module/wrapper_module.py +++ b/torch_tweak/torch/module/tunable_module.py @@ -2,7 +2,7 @@ # Copyright (c) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 # NOTE: This file has been modified by Intel Corporation. -"""Torch Tweak wrapper module.""" +"""Torch Tweak tunable module.""" import gc from collections import OrderedDict @@ -38,7 +38,7 @@ class ModuleState(Enum): - """Possible states of the Module class.""" + """Possible states of the TunableModule class.""" INIT = "init" PASSTHROUGH = "passthrough" @@ -50,7 +50,7 @@ class ModuleState(Enum): StrategyMap = dict[SampleMetadata, TuneStrategy] -class Module(wrapt.CallableObjectProxy): +class TunableModule(wrapt.CallableObjectProxy): """Torch Tweak module wrapper. This class wraps a torch module and provides tuning functionality. @@ -67,7 +67,7 @@ class Module(wrapt.CallableObjectProxy): Example: >>> model = torch.nn.Linear(10, 10) - >>> model = Module(model, "name") + >>> model = TunableModule(model, "name") """ TYPE_KEY = "type" @@ -151,7 +151,7 @@ def graph_specs(self) -> list[GraphSpec]: elif self._self_state == ModuleState.TUNED and self._self_prev_recording: return self._self_prev_recording.graph_specs else: - raise ValueError("Module has not recorded any samples. Cannot collect graph specs.") + raise ValueError("TunableModule has not recorded any samples. Cannot collect graph specs.") @property def state(self) -> ModuleState: @@ -174,7 +174,7 @@ def enable_passthrough(self, force: bool = False): self._reset() self.enable_passthrough() else: - raise RuntimeError("Module is already tuned. Use force=True to reset tuned module.") + raise RuntimeError("TunableModule is already tuned. Use force=True to reset tuned module.") elif self._self_state == ModuleState.RECORDING: self._self_prev_recording = self._self_wrapper @@ -196,7 +196,7 @@ def enable_recording(self, force: bool = False): self._reset() self.enable_recording() else: - raise RuntimeError("Module is already tuned. Use force=True to reset tuned module.") + raise RuntimeError("TunableModule is already tuned. Use force=True to reset tuned module.") elif self._self_state in [ModuleState.PASSTHROUGH, ModuleState.INIT]: self._deactivate_wrapper() self._self_state = ModuleState.RECORDING @@ -217,7 +217,7 @@ def activate(self): elif self._self_state == ModuleState.PASSTHROUGH: return else: - raise RuntimeError("Module is not tuned. Cannot activate backends.") + raise RuntimeError("TunableModule is not tuned. Cannot activate backends.") def deactivate(self): """Deactivates the module backends.""" @@ -227,34 +227,34 @@ def deactivate(self): @staticmethod def is_state_dict_valid(state_dict: dict): - """Check if the state dictionary has a wrapper module.""" - return state_dict and Module.__name__ in state_dict[Module.TYPE_KEY] + """Check if the state dictionary has a tunable module.""" + return state_dict and TunableModule.__name__ in state_dict[TunableModule.TYPE_KEY] @staticmethod def from_dict(module: torch.nn.Module, state_dict: dict, device: torch.device | None = None): - """Create a wrapper module from a state dictionary.""" - if not Module.is_state_dict_valid(state_dict): - raise ValueError(f"Invalid dictionary format for {Module.__class__.__name__}") + """Create a tunable module from a state dictionary.""" + if not TunableModule.is_state_dict_valid(state_dict): + raise ValueError(f"Invalid dictionary format for {TunableModule.__class__.__name__}") - self = Module(module, state_dict[Module.NAME_KEY]) + self = TunableModule(module, state_dict[TunableModule.NAME_KEY]) self._self_state = ModuleState.TUNED - self._self_wrapper = TunedModule.from_dict(module, state_dict[Module.TUNED_MODULE_KEY]) + self._self_wrapper = TunedModule.from_dict(module, state_dict[TunableModule.TUNED_MODULE_KEY]) self._deploy_wrapper(device) return self def to_dict(self, *args, destination=None, prefix="", keep_vars=False): - """Convert the wrapper module to a state dictionary. + """Convert the tunable module to a state dictionary. The method follows torch.state_dict format so that it could be called recursively from torch.nn.Module.state_dict. The name of the function follows torch_tweak convention of from/to dict. Special alias is created to match torch convention i.e. state_dict. """ if self._self_state != ModuleState.TUNED: - raise RuntimeError("Module is not tuned. Cannot save state_dict.") + raise RuntimeError("TunableModule is not tuned. Cannot save state_dict.") destination = OrderedDict() if destination is None else destination destination[prefix] = { - self.TYPE_KEY: Module.__name__, # self.__class__ returns wrapped class, use direct class object instead + self.TYPE_KEY: TunableModule.__name__, # self.__class__ returns wrapped class, use direct class object instead self.NAME_KEY: self._self_name, self.TUNED_MODULE_KEY: self._self_wrapper.to_dict(), } @@ -279,9 +279,9 @@ def tune( if module has already defined strategy/strategies those will take precedence over the provided one. """ if self._self_state == ModuleState.TUNED: - raise ValueError(f"Module: '{self._self_name}' has already been tuned. Reset it to do it again.") + raise ValueError(f"TunableModule: '{self._self_name}' has already been tuned. Reset it to do it again.") 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.") + raise ValueError(f"TunableModule: '{self._self_name}' has not recorded any samples. Cannot tune it.") device = device or self._self_device or torch.device("cpu") @@ -430,7 +430,7 @@ def _get_strategies_for_graph_specs( ) -> list[TuneStrategy]: """Returns strategies for given graph specs. - Module level strategies (list/dict) take precedence over provided strategy. + TunableModule level strategies (list/dict) take precedence over provided strategy. The function checks if there is sufficient strategies (list/dict). """ if strategy is not None: diff --git a/torch_tweak/torch/module_registry.py b/torch_tweak/torch/module_registry.py index d33aa51b..80f7d855 100644 --- a/torch_tweak/torch/module_registry.py +++ b/torch_tweak/torch/module_registry.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from torch_tweak.torch.module.wrapper_module import Module + from torch_tweak.torch.module.tunable_module import TunableModule class ModuleRegistry: @@ -17,9 +17,9 @@ class ModuleRegistry: def __init__(self) -> None: """Initialize ModuleRegistry.""" - self._registry: OrderedDict[str, Module] = OrderedDict() + self._registry: OrderedDict[str, TunableModule] = OrderedDict() - def register(self, name: str, module: "Module") -> None: + def register(self, name: str, module: "TunableModule") -> None: """Register a module.""" self._registry[name] = module @@ -34,14 +34,14 @@ def clear(self) -> None: gc.collect() @property - def modules(self) -> dict[str, "Module"]: + def modules(self) -> dict[str, "TunableModule"]: """All registered modules.""" return self._registry - def get(self, name: str) -> "Module": + def get(self, name: str) -> "TunableModule": """Get a module.""" if name not in self._registry: - raise ValueError(f"Module {name} not found in registry.") + raise ValueError(f"TunableModule {name} not found in registry.") return self._registry[name] diff --git a/torch_tweak/torch/tuning.py b/torch_tweak/torch/tuning.py index a64ea0ef..09d2fc24 100644 --- a/torch_tweak/torch/tuning.py +++ b/torch_tweak/torch/tuning.py @@ -19,7 +19,7 @@ from torch_tweak.torch.config import DEFAULT_DEVICE, torch_tweak_cache_dir from torch_tweak.torch.dataloader import DataLoaderFactory, DatasetLike, samples_generator from torch_tweak.torch.module.tensor_spec import InfoLevel -from torch_tweak.torch.module.wrapper_module import Module, ModuleState +from torch_tweak.torch.module.tunable_module import ModuleState, TunableModule from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.utils.device import get_device from torch_tweak.torch.utils.xpu import synchronize @@ -131,7 +131,7 @@ def save( >>> doctest.ELLIPSIS_MARKER = "***" # using *** instead of ... to avoid doctest failure >>> import torch_tweak.torch as tt >>> from torch.nn import Linear - >>> model = tt.Module(Linear(10, 10), "model", strategy=tt.FirstWinsStrategy([tt.backend.TorchEagerBackend()])) + >>> model = tt.TunableModule(Linear(10, 10), "model", strategy=tt.FirstWinsStrategy([tt.backend.TorchEagerBackend()])) >>> dataset = torch.randn(10, 10) >>> tt.tune(model, dataset, batch_sizes=[1, 2], device="cpu") # doctest: +ELLIPSIS ***🎯 Tuning module: `model` (all graphs)*** @@ -152,7 +152,7 @@ def load( storage: Storage | None = None, device_map: dict[str, torch.device] | None = None, disable_external_logging: bool = True, -) -> torch.nn.Module | Module: +) -> torch.nn.Module | TunableModule: """Load the tuned module from a file. Args: @@ -203,7 +203,7 @@ def _activate_tuned_modules() -> None: logger.info("════════════════════════════════════════════════════════════════") -def _describe_module(module: Module): +def _describe_module(module: TunableModule): logger.info("📦 module: %s", module.name) if module.state == ModuleState.TUNED: for metadata, backend in module.module.backends.items():