refactor(module): rename Module to TunableModule

docs/api/tuning/module.md → 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

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])

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)
 ```
 

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)
 ```
 

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)
 ```
 

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)
 ```
 

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)
 ```

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

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)
 

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,15 +66,15 @@ 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(
                 False
             ),  # just one strategy, there is no prefill/decode distinction
         )
     elif cache == "static":
-        model = Module(
+        model = TunableModule(
             model,
             model.__class__.__name__,
             strategies=[

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 []

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),

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,

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(

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
 

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.

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()))"
    ]
   },
   {

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()))"
    ]
   },
   {

notebooks/wrapper_module_demo.ipynb → 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\")"
    ]
   },
   {

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

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)