From 6fa7a2854e5102ccc65fbd0e33b6e30babeada5f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Thu, 29 Oct 2026 15:45:47 +0100 Subject: [PATCH 1/5] feat(tune-strategy): tune with every backend for the device by default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A module without a strategy is now tuned with tt.default_strategy(device). It is HighestThroughputStrategy over default_backends(device). On XPU these are XPU graph, Torch Inductor, OpenVINO, and Torch Eager. On CPU these are Torch Inductor, OpenVINO, and Torch Eager. TorchAO is left out because quantization changes the results. The old default was one FirstWinsStrategy instance shared by every module. FirstWins keeps the first backend that builds, not the fastest one. The shared instance also carried state from one module to the next. Each graph now gets its own strategy, built for the device it is tuned on. FirstWinsStrategy and HighestThroughputStrategy without backends use the same default_backends(device). FirstWinsStrategy keeps its order, XPU graph first. HighestThroughputStrategy() now also measures OpenVINO. On CPU neither strategy tries XPU graph. Without backends, a device other than XPU or CPU raises ValueError. On an Arc Pro B70 the default strategy was the fastest or close to it on five of six models. OpenVINO wins on whole ViT and DistilBERT models and loses on small modules, so it is measured, not excluded. Signed-off-by: Łukasz Ślusarczyk --- ChangeLog.md | 6 +++ docs/guides/backends/xpu_graph_backend.md | 5 ++- .../guides/tune_strategies/tune_strategies.md | 17 +++++++- docs/learn/core_functionalities.md | 4 +- .../torch/backend/test_xpu_graph_backend.py | 11 +---- .../unit/torch/module/test_wrapper_module.py | 17 ++++++++ .../tune_strategy/test_default_strategy.py | 42 +++++++++++++++++++ .../tune_strategy/test_first_wins_strategy.py | 18 ++++++++ .../test_highest_throughput_strategy.py | 19 +++++++++ torch_tweak/torch/__init__.py | 2 + torch_tweak/torch/module/wrapper_module.py | 21 ++++------ torch_tweak/torch/tune_strategy/__init__.py | 11 ++++- .../torch/tune_strategy/default_backends.py | 27 ++++++++++++ .../torch/tune_strategy/default_strategy.py | 13 ++++++ .../tune_strategy/first_wins_strategy.py | 27 ++++-------- .../highest_throughput_strategy.py | 31 ++++++-------- 16 files changed, 203 insertions(+), 68 deletions(-) create mode 100644 tests/unit/torch/tune_strategy/test_default_strategy.py create mode 100644 torch_tweak/torch/tune_strategy/default_backends.py create mode 100644 torch_tweak/torch/tune_strategy/default_strategy.py diff --git a/ChangeLog.md b/ChangeLog.md index 3c0b29bc..47da1ec2 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -11,6 +11,12 @@ SPDX-License-Identifier: Apache-2.0 - fix: a tuned module that a pipeline holds under several paths is saved once and restored as one module. - feat: `save` warns about tuned modules it does not save, e.g. a module the pipeline holds in a plain list. - fix: a passthrough module without parameters keeps its inputs on their device instead of moving them to CPU. +- feat (breaking): a module without a strategy is tuned with `tt.default_strategy(device)` instead of `FirstWinsStrategy`. + Tuning measures every backend and keeps the fastest, so it takes longer. +- feat: `tt.default_strategy(device)` returns `HighestThroughputStrategy` over every backend that runs on the device. +- feat (breaking): `FirstWinsStrategy` and `HighestThroughputStrategy` without `backends` use `default_backends(device)` for the tuning device. + `HighestThroughputStrategy()` now also measures OpenVINO, and neither strategy tries XPU graph on CPU. + Tuning on a device other than XPU or CPU without `backends` raises `ValueError`. - feat (breaking): `inspect` measures module depth in executed modules. Containers such as `nn.ModuleList` and modules that never run no longer add a level. `min_depth` is renamed to `depth`. diff --git a/docs/guides/backends/xpu_graph_backend.md b/docs/guides/backends/xpu_graph_backend.md index 4195619d..2f60581c 100644 --- a/docs/guides/backends/xpu_graph_backend.md +++ b/docs/guides/backends/xpu_graph_backend.md @@ -25,8 +25,9 @@ from torch_tweak.torch.tune_strategy import HighestThroughputStrategy strategy = HighestThroughputStrategy(backends=[XPUGraphBackend()]) ``` -The backend is first in the default `FirstWinsStrategy` list, before Torch Inductor. -It is also in the default `HighestThroughputStrategy` list. +On XPU the backend is first in `default_backends(device)`, before Torch Inductor. +`FirstWinsStrategy`, `HighestThroughputStrategy`, and `tt.default_strategy(device)` use that list when no backends are given. +On CPU the list leaves this backend out. PyTorch records a native SYCL graph when it was built with a SYCL compiler of 2026.1 or newer, except on Intel Data Center GPU Max Series. An older build records the legacy SYCL command graph. diff --git a/docs/guides/tune_strategies/tune_strategies.md b/docs/guides/tune_strategies/tune_strategies.md index 2fe37577..f77952b3 100644 --- a/docs/guides/tune_strategies/tune_strategies.md +++ b/docs/guides/tune_strategies/tune_strategies.md @@ -18,6 +18,21 @@ Torch Tweak provides three built-in strategies: - **FirstWinsStrategy**: Tries backends in order, uses the first that succeeds - **HighestThroughputStrategy**: Profiles all backends, selects the fastest +## Default Strategy + +A module without a strategy is tuned with `tt.default_strategy(device)`. +It is `HighestThroughputStrategy` over `default_backends(device)`, every backend that runs on the device and keeps the results of the model: + +| Device | Backends | +|---|---| +| XPU | XPU graph, Torch Inductor, OpenVINO, Torch Eager | +| CPU | Torch Inductor, OpenVINO, Torch Eager | + +TorchAO is left out because quantization changes the results. +Every module gets its own strategy instance, built for the device it is tuned on. +`FirstWinsStrategy` and `HighestThroughputStrategy` without a backend list use the same `default_backends(device)`. +Without a strategy, tuning on any other device raises `ValueError`. + ## Why Backends Can Fail Not every backend can successfully tune every model. Each backend relies on a different compilation or export technology, and each has its own limitations: @@ -82,7 +97,7 @@ tt.tune(model, input_data) ## FirstWinsStrategy -Tries backends in priority order and returns the first one that successfully builds and validates. If a backend fails, the strategy moves on to the next candidate instead of aborting. If all backends fail, the error is caught and the original model is used as-is. The default order is XPU graph, Torch Inductor, OpenVINO, then eager. Put a backend earlier only when you want that one to win as soon as it builds. +Tries backends in priority order and returns the first one that successfully builds and validates. If a backend fails, the strategy moves on to the next candidate instead of aborting. If all backends fail, the error is caught and the original model is used as-is. Without a backend list it tries `default_backends(device)`: XPU graph, Torch Inductor, OpenVINO, then eager, with XPU graph skipped on CPU. Put a backend earlier only when you want that one to win as soon as it builds. ### Usage diff --git a/docs/learn/core_functionalities.md b/docs/learn/core_functionalities.md index 7f740a1a..3253fc7e 100644 --- a/docs/learn/core_functionalities.md +++ b/docs/learn/core_functionalities.md @@ -118,7 +118,7 @@ strategy = FirstWinsStrategy( ) ``` -That order is also the default when you construct `FirstWinsStrategy` with no backend list. The first backend that builds and validates is kept. +`FirstWinsStrategy` with no backend list tries `default_backends(device)` in that order. On CPU it skips XPU graph. The first backend that builds and validates is kept. ## OneBackendStrategy @@ -132,7 +132,7 @@ strategy = OneBackendStrategy(backend=OpenVINOBackend()) ## HighestThroughputStrategy -Profiles all compatible backends and selects the fastest. Use this when maximum throughput matters and you can afford longer tuning time. `HighestThroughputStrategy()` with no backend list compares Torch Inductor, XPU graph, and eager. +Profiles all compatible backends and selects the fastest. Use this when maximum throughput matters and you can afford longer tuning time. `HighestThroughputStrategy()` with no backend list compares `default_backends(device)`: XPU graph, Torch Inductor, OpenVINO, and eager on XPU, and the same without XPU graph on CPU. ```python from torch_tweak.torch.tune_strategy import HighestThroughputStrategy diff --git a/tests/unit/torch/backend/test_xpu_graph_backend.py b/tests/unit/torch/backend/test_xpu_graph_backend.py index eda892d5..61ebc6fc 100644 --- a/tests/unit/torch/backend/test_xpu_graph_backend.py +++ b/tests/unit/torch/backend/test_xpu_graph_backend.py @@ -18,7 +18,7 @@ from torch_tweak.torch.module.sample_metadata import SampleMetadata from torch_tweak.torch.module.wrapper_module import Module from torch_tweak.torch.module_registry import MODULE_REGISTRY -from torch_tweak.torch.tune_strategy import FirstWinsStrategy, HighestThroughputStrategy, OneBackendStrategy +from torch_tweak.torch.tune_strategy import OneBackendStrategy from torch_tweak.torch.tuning import tune @@ -294,15 +294,6 @@ def test_default_backend_has_no_recording_switch(): assert XPUGraphBackend().describe() == "XPUGraphBackend()" -def test_first_wins_tries_the_graph_backend_before_inductor(): - text = FirstWinsStrategy().describe() - assert text.index("XPUGraphBackend()") < text.index("TorchInductorBackend()") - - -def test_highest_throughput_lists_the_graph_backend(): - assert "XPUGraphBackend()" in HighestThroughputStrategy().describe() - - @requires_xpu def test_serialization(model, sample_data, tmp_path): backend = backend_build(XPUGraphBackend(), torch.float16, model, sample_data, tmp_path) diff --git a/tests/unit/torch/module/test_wrapper_module.py b/tests/unit/torch/module/test_wrapper_module.py index 8f1160d5..1c0be592 100644 --- a/tests/unit/torch/module/test_wrapper_module.py +++ b/tests/unit/torch/module/test_wrapper_module.py @@ -220,6 +220,23 @@ def test_tune_with_provided_strategy(torch_device): assert module.state == ModuleState.TUNED +def test_tune_without_strategy_builds_a_default_strategy_per_graph_for_the_device(torch_device, monkeypatch): + requested_devices = [] + + def default_strategy_stub(device): + requested_devices.append(device) + return Mock() + + monkeypatch.setattr("torch_tweak.torch.module.wrapper_module.default_strategy", default_strategy_stub) + module = Module(Identity(), TEST_MODULE_NAME) + module(1, a=1) + module(2, b=2) + + module.tune(device=torch_device) + + assert requested_devices == [torch_device, torch_device] + + def test_tune_with_list_of_strategies(torch_device): model = Identity() strategy1 = Mock() diff --git a/tests/unit/torch/tune_strategy/test_default_strategy.py b/tests/unit/torch/tune_strategy/test_default_strategy.py new file mode 100644 index 00000000..227c16df --- /dev/null +++ b/tests/unit/torch/tune_strategy/test_default_strategy.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the default backends and strategy.""" + +import pytest + +from torch_tweak.torch.backend import ( + OpenVINOBackend, + TorchEagerBackend, + TorchInductorBackend, + XPUGraphBackend, +) +from torch_tweak.torch.tune_strategy import default_backends, default_strategy + + +@pytest.mark.parametrize( + ("device", "expected"), + [ + ("xpu:0", [XPUGraphBackend, TorchInductorBackend, OpenVINOBackend, TorchEagerBackend]), + ("cpu", [TorchInductorBackend, OpenVINOBackend, TorchEagerBackend]), + ], +) +def test_default_backends_are_every_backend_that_runs_on_the_device(device, expected): + assert [type(backend) for backend in default_backends(device)] == expected + + +def test_default_backends_reject_an_unknown_device(): + with pytest.raises(ValueError, match="No default backends"): + default_backends("meta") + + +def test_default_strategy_on_cpu_measures_inductor_openvino_and_eager(): + strategy = default_strategy("cpu") + + assert strategy.describe().splitlines() == [ + "name: Highest Throughput Strategy", + "description: evaluate all backends, return backend with highest throughput", + "backends:", + " TorchInductorBackend()", + " OpenVINOBackend()", + " TorchEagerBackend()", + ] 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..0b4cd02e 100644 --- a/tests/unit/torch/tune_strategy/test_first_wins_strategy.py +++ b/tests/unit/torch/tune_strategy/test_first_wins_strategy.py @@ -5,6 +5,7 @@ # NOTE: This file has been modified by Intel Corporation. """Unit tests for FirstWinsStrategy.""" +from typing import cast from unittest.mock import MagicMock import pytest @@ -155,6 +156,23 @@ def test_tune_all_backends_fail(mock_module, mock_graph_spec, mock_sample, torch assert mock_module.to.call_count == 2 # Called after each failed backend +def test_without_backends_tunes_with_the_default_backends_for_the_device(torch_device, tmp_path, monkeypatch): + monkeypatch.setattr( + "torch_tweak.torch.tune_strategy.first_wins_strategy.default_backends", + lambda device: {torch_device: [SleepBackend(sleep_time=2e-5)]}[device], + ) + strategy = FirstWinsStrategy() + strategy.enable_find_max_batch_size(False) + strategy.enable_correctness_check(False) + model = ToyTorchModel().eval().to(torch_device) + sample = model.sample().unsqueeze(0).to(torch_device) + graph_spec = model.graph_spec(batch_sizes=[1, 2], device=torch_device) + + backend = strategy.tune(model, "test", graph_spec, [((sample,), {})], torch_device, tmp_path) + + assert cast(SleepBackend, backend).sleep_time == 2e-5 + + class InferenceFailsBackend(SleepBackend): """Backend that fails on inference.""" 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..771ad1b7 100644 --- a/tests/unit/torch/tune_strategy/test_highest_throughput_strategy.py +++ b/tests/unit/torch/tune_strategy/test_highest_throughput_strategy.py @@ -81,6 +81,25 @@ def forward_slow(x): assert highest_throughput_backend.sleep_time == faster.sleep_time +def test_without_backends_tunes_with_the_default_backends_for_the_device(torch_device, tmp_path, monkeypatch): + monkeypatch.setattr( + "torch_tweak.torch.tune_strategy.highest_throughput_strategy.default_backends", + lambda device: {torch_device: [SleepBackend(sleep_time=2e-5)]}[device], + ) + strategy = HighestThroughputStrategy() + strategy.set_find_max_batch_size_profiling_config( + TuneStrategyFindMaxBatchSizeExtension.default_profiling_config(max_batch_size=8) + ) + strategy.enable_correctness_check(False) + model = ToyTorchModel().eval().to(torch_device) + sample = model.sample().unsqueeze(0).to(torch_device) + graph_spec = model.graph_spec(batch_sizes=[1, 2], device=torch_device) + + backend = strategy.tune(model, "test", graph_spec, [((sample,), {})], torch_device, tmp_path) + + assert cast(SleepBackend, backend).sleep_time == 2e-5 + + def test_highest_throughput_strategy_max_batch_size_in_graph_spec(torch_device, tmp_path): profiling_config = TuneStrategyFindMaxBatchSizeExtension.default_profiling_config(max_batch_size=8) strategy = HighestThroughputStrategy(backends=[SleepBackend(sleep_time=1e-5)]) diff --git a/torch_tweak/torch/__init__.py b/torch_tweak/torch/__init__.py index d3bf2cdd..2a14edbf 100644 --- a/torch_tweak/torch/__init__.py +++ b/torch_tweak/torch/__init__.py @@ -16,6 +16,7 @@ HighestThroughputStrategy, OneBackendStrategy, TuneStrategy, + default_strategy, ) from torch_tweak.torch.tuning import load, save, tune @@ -32,6 +33,7 @@ "OneBackendStrategy", "FirstWinsStrategy", "HighestThroughputStrategy", + "default_strategy", "LocalTorchStorage", "DataLoaderFactory", ] diff --git a/torch_tweak/torch/module/wrapper_module.py b/torch_tweak/torch/module/wrapper_module.py index 61b6e36a..04f0bc17 100644 --- a/torch_tweak/torch/module/wrapper_module.py +++ b/torch_tweak/torch/module/wrapper_module.py @@ -23,9 +23,8 @@ from torch_tweak.torch.module.sample_metadata import SampleMetadata from torch_tweak.torch.module.tuned_module import TunedModule from torch_tweak.torch.module_registry import MODULE_REGISTRY -from torch_tweak.torch.tune_strategy import FirstWinsStrategy +from torch_tweak.torch.tune_strategy import default_strategy from torch_tweak.torch.tune_strategy.tune_strategy import ( - DummyTuneStrategy, TuneStrategy, ) from torch_tweak.torch.utils.module import get_module_device, offload @@ -34,8 +33,6 @@ logger = getLogger(__name__) -DEFAULT_STRATEGY = FirstWinsStrategy() - class ModuleState(Enum): """Possible states of the Module class.""" @@ -86,7 +83,8 @@ def __init__( Args: module: torch module to wrap. name: module name which differences from other tuned modules. - strategy: optional strategy, which will be used for all encountered graphs + strategy: optional strategy, which will be used for all encountered graphs. Without `strategy` and + `strategies`, every graph gets its own `default_strategy` for the tuning device. strategies: list or dict of strategies, each for a graph """ super().__init__(module) @@ -287,7 +285,7 @@ def tune( recording = cast(RecordingModule, self._self_wrapper) backends: OrderedDict[SampleMetadata, Backend] = OrderedDict() - strategies = self._get_strategies_for_graph_specs(strategy, recording.graph_specs, dry_run) + strategies = self._get_strategies_for_graph_specs(strategy, recording.graph_specs, device) for strategy, graph_spec in zip(strategies, recording.graph_specs, strict=True): cache_dir = self._create_graph_cache_dir(graph_spec) @@ -404,9 +402,6 @@ def _setup_strategies( strategies: StrategyList | StrategyMap | None, ): """Sets up strategy or strategy_map or strategy_list depending on input args.""" - if strategy is None and strategies is None: - strategy = DEFAULT_STRATEGY - if strategy is not None and strategies is not None: raise ValueError("Only one of strategy or strategies should be provided") self._self_strategy = strategy @@ -426,12 +421,13 @@ def _get_strategies_for_graph_specs( self, strategy: TuneStrategy | None, graph_specs: list[GraphSpec], - dry_run: bool, + device: torch.device | str, ) -> list[TuneStrategy]: """Returns strategies for given graph specs. Module level strategies (list/dict) take precedence over provided strategy. The function checks if there is sufficient strategies (list/dict). + Without any strategy, every graph gets its own `default_strategy(device)`. """ if strategy is not None: return [strategy] * len(graph_specs) @@ -456,10 +452,7 @@ def _get_strategies_for_graph_specs( if self._self_strategy is not None: return [self._self_strategy] * len(graph_specs) - if dry_run: - return [DummyTuneStrategy()] * len(graph_specs) - - raise RuntimeError("No strategy provided. Cannot tune module.") + return [default_strategy(device) for _ in graph_specs] def _offload(self, backends: OrderedDict[SampleMetadata, Backend]): """Offload the module to meta only when backends do not retain it.""" diff --git a/torch_tweak/torch/tune_strategy/__init__.py b/torch_tweak/torch/tune_strategy/__init__.py index 601f41e6..94adc5d0 100644 --- a/torch_tweak/torch/tune_strategy/__init__.py +++ b/torch_tweak/torch/tune_strategy/__init__.py @@ -4,9 +4,18 @@ # NOTE: This file has been modified by Intel Corporation. """Torch tweak strategy module.""" +from torch_tweak.torch.tune_strategy.default_backends import default_backends +from torch_tweak.torch.tune_strategy.default_strategy import default_strategy from torch_tweak.torch.tune_strategy.first_wins_strategy import FirstWinsStrategy from torch_tweak.torch.tune_strategy.highest_throughput_strategy import HighestThroughputStrategy from torch_tweak.torch.tune_strategy.one_backend_strategy import OneBackendStrategy from torch_tweak.torch.tune_strategy.tune_strategy import TuneStrategy -__all__ = ["FirstWinsStrategy", "OneBackendStrategy", "HighestThroughputStrategy", "TuneStrategy"] +__all__ = [ + "FirstWinsStrategy", + "OneBackendStrategy", + "HighestThroughputStrategy", + "TuneStrategy", + "default_backends", + "default_strategy", +] diff --git a/torch_tweak/torch/tune_strategy/default_backends.py b/torch_tweak/torch/tune_strategy/default_backends.py new file mode 100644 index 00000000..55dbe7f0 --- /dev/null +++ b/torch_tweak/torch/tune_strategy/default_backends.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Default backends for a device.""" + +import torch + +from torch_tweak.torch.backend import ( + Backend, + OpenVINOBackend, + TorchEagerBackend, + TorchInductorBackend, + XPUGraphBackend, +) + + +def default_backends(device: str | torch.device) -> list[Backend]: + """Every backend that runs on `device` and keeps the results of the model. + + `FirstWinsStrategy` tries them in this order. + TorchAO is left out because quantization changes the results. + """ + device_type = torch.device(device).type + if device_type == "xpu": + return [XPUGraphBackend(), TorchInductorBackend(), OpenVINOBackend(), TorchEagerBackend()] + if device_type == "cpu": + return [TorchInductorBackend(), OpenVINOBackend(), TorchEagerBackend()] + raise ValueError(f"No default backends for device {device_type}.") diff --git a/torch_tweak/torch/tune_strategy/default_strategy.py b/torch_tweak/torch/tune_strategy/default_strategy.py new file mode 100644 index 00000000..2f2d223f --- /dev/null +++ b/torch_tweak/torch/tune_strategy/default_strategy.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Default strategy for a device.""" + +import torch + +from torch_tweak.torch.tune_strategy.default_backends import default_backends +from torch_tweak.torch.tune_strategy.highest_throughput_strategy import HighestThroughputStrategy + + +def default_strategy(device: str | torch.device) -> HighestThroughputStrategy: + """`HighestThroughputStrategy` over `default_backends(device)`: every backend is measured and the fastest wins.""" + return HighestThroughputStrategy(backends=default_backends(device)) diff --git a/torch_tweak/torch/tune_strategy/first_wins_strategy.py b/torch_tweak/torch/tune_strategy/first_wins_strategy.py index ac3be7e0..55958caa 100644 --- a/torch_tweak/torch/tune_strategy/first_wins_strategy.py +++ b/torch_tweak/torch/tune_strategy/first_wins_strategy.py @@ -12,15 +12,10 @@ import torch import torch.nn as nn -from torch_tweak.torch.backend import ( - OpenVINOBackend, - TorchEagerBackend, - TorchInductorBackend, - XPUGraphBackend, -) from torch_tweak.torch.backend.backend import Backend from torch_tweak.torch.module.graph_spec import GraphSpec from torch_tweak.torch.module.recording_module import Sample +from torch_tweak.torch.tune_strategy.default_backends import default_backends from torch_tweak.torch.tune_strategy.extension import TuneStrategyFindMaxBatchSizeExtension from torch_tweak.utils.logging import control_output, log from torch_tweak.utils.timer import Timer @@ -29,13 +24,13 @@ class FirstWinsStrategy(TuneStrategyFindMaxBatchSizeExtension): """Runs backends in order and keeps the first one that builds and validates. - The default order is XPU graph, Torch Inductor, OpenVINO, then eager. + Without `backends` it tries `default_backends(device)` for the tuning device. """ def __init__(self, backends: Iterable[Backend] | None = None, **kwargs): """Initializes strategy.""" super().__init__(**kwargs) - self._backends = backends or self._default_backends() + self._backends = backends def _tune( self, @@ -57,7 +52,7 @@ def _tune( ) # Run backend by backend until first working backend is found - for backend in self._backends: + for backend in self._backends or default_backends(device): backend_cache_dir = cache_dir / backend.key() log_file = self._log_file(backend_cache_dir, "build.log") @@ -86,20 +81,14 @@ def _tune( raise RuntimeError(f"There is no valid backend for a module: {name}, graph_spec: {graph_spec}") - def _default_backends(self) -> list[Backend]: - """Returns default backends.""" - return [ - XPUGraphBackend(), - TorchInductorBackend(), - OpenVINOBackend(), - TorchEagerBackend(), - ] - def _describe_parts(self) -> list[str]: """Returns the parts of the description.""" + backends = ( + [backend.describe() for backend in self._backends] if self._backends else ["default_backends(device)"] + ) return [ "name: First Wins Strategy", "description: evaluate backends in order, return first working backend", "backends:", - *[f" {backend.describe()}" for backend in self._backends], + *[f" {backend}" for backend in backends], ] diff --git a/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py b/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py index 353de525..8a965964 100644 --- a/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py +++ b/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py @@ -20,12 +20,7 @@ import torch import torch.nn as nn -from torch_tweak.torch.backend import ( - Backend, - TorchEagerBackend, - TorchInductorBackend, - XPUGraphBackend, -) +from torch_tweak.torch.backend import Backend from torch_tweak.torch.config import ( DEFAULT_STABILITY_PERCENTAGE, DEFAULT_THROUGHPUT_BACKOFF_LIMIT, @@ -44,6 +39,7 @@ StableWindowMeasuringStopStrategy, ThroughputSaturatedProfilingStopStrategy, ) +from torch_tweak.torch.tune_strategy.default_backends import default_backends from torch_tweak.torch.tune_strategy.extension import TuneStrategyFindMaxBatchSizeExtension from torch_tweak.utils.logging import control_output, log from torch_tweak.utils.timer import Timer @@ -80,13 +76,13 @@ def __init__( """Initializes strategy. Args: - backends: List of backends to tune. + backends: List of backends to tune. Without it, `default_backends(device)` for the tuning device. measurement_stop_strategy: Measurement stop strategy. profiling_stop_strategy: Profiling stop strategy. kwargs: Additional arguments for the parent class """ super().__init__(**kwargs) - self._backends = backends or self._default_backends() + self._backends = backends self._measurement_stop_strategy = measurement_stop_strategy or StableWindowMeasuringStopStrategy( window_size=DEFAULT_WINDOW_SIZE, stability_percentage=DEFAULT_STABILITY_PERCENTAGE, @@ -128,8 +124,10 @@ def _tune( best_throughput = float("-inf") best_batch_size = 1 + backends = list(self._backends or default_backends(device)) + # Run all backends with given max batch size. - for backend in self._backends: + for backend in backends: backend_cache_dir = cache_dir / backend.key() log_file = self._log_file(backend_cache_dir, "build.log") @@ -183,7 +181,7 @@ def _tune( log( "ℹ️ No correct backend found with throughput %f > 0. Backends considered: %s", best_throughput, - ", ".join([backend.describe() for backend in self._backends]), + ", ".join([backend.describe() for backend in backends]), sink=self._sink, ) raise RuntimeError("No correct backend found with throughput > 0") @@ -218,19 +216,14 @@ def _get_profiling_config(self, batching: bool, max_batch_size: int) -> Profilin profiling_stop_strategy=self._profiling_stop_strategy, ) - def _default_backends(self) -> list[Backend]: - """Returns default backends.""" - return [ - TorchInductorBackend(), - XPUGraphBackend(), - TorchEagerBackend(), - ] - def _describe_parts(self) -> list[str]: """Returns the parts of the description.""" + backends = ( + [backend.describe() for backend in self._backends] if self._backends else ["default_backends(device)"] + ) return [ "name: Highest Throughput Strategy", "description: evaluate all backends, return backend with highest throughput", "backends:", - *[f" {backend.describe()}" for backend in self._backends], + *[f" {backend}" for backend in backends], ] From 9a9dc2242433682f01164ea83bdf4f9fa87ac475 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Fri, 23 Oct 2026 17:04:49 +0200 Subject: [PATCH 2/5] feat(xpu_graph): record graphs for a grid of batch sizes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A graph spec with a batch axis now records graphs for batch sizes 1, 2, 4, and every multiple of 8 up to the largest batch size seen during tuning. max_recorded_batch_size caps that size at 256 by default. A call pads its batch with zeros to the smallest recorded size that holds it. A larger batch replays in chunks. All graphs share one memory pool and one input buffer. Tuning runs eager once on the largest batch and compares every graph with it, full and partly padded. Modules whose rows depend on other rows are rejected. Signed-off-by: Łukasz Ślusarczyk --- .claude/skills/torch-tweak-assistant/SKILL.md | 1 + docs/guides/backends/xpu_graph_backend.md | 36 +- docs/learn/backends.md | 3 +- examples/Minimal/demo.py | 1 - .../torch/backend/test_xpu_graph_backend.py | 198 +++++++- .../torch/backend/xpu_graph_backend.py | 477 ++++++++++++++---- 6 files changed, 593 insertions(+), 123 deletions(-) diff --git a/.claude/skills/torch-tweak-assistant/SKILL.md b/.claude/skills/torch-tweak-assistant/SKILL.md index a013bb78..398fb3f4 100644 --- a/.claude/skills/torch-tweak-assistant/SKILL.md +++ b/.claude/skills/torch-tweak-assistant/SKILL.md @@ -71,6 +71,7 @@ Do NOT activate this skill when: ### XPU Graph Backend - **Best for**: replaying a module that was recorded once during tuning +- **Batch sizes**: recordings for 1, 2, 4, and every multiple of 8 up to `max_recorded_batch_size` (256); other sizes are padded - **Rejection**: failed, empty, or incorrect recordings are dropped for that module and input ### OpenVINO Backend diff --git a/docs/guides/backends/xpu_graph_backend.md b/docs/guides/backends/xpu_graph_backend.md index 2f60581c..8ba1df50 100644 --- a/docs/guides/backends/xpu_graph_backend.md +++ b/docs/guides/backends/xpu_graph_backend.md @@ -10,11 +10,41 @@ Recording is not repeated for each inference call. ## Overview -- **One recording**: warmup and capture run in `build`, while the module is being tuned +- **Recording**: warmup and capture run in `build`, while the module is being tuned - **Replay**: inference copies new tensor values into the recorded inputs and calls `replay` - **Rejection**: a failed recording, an empty recording, or outputs that disagree with eager drop this backend for that module and input -- **Static shapes**: a batch axis or another dynamic axis rejects the backend. Replay cannot change tensor shapes -- **Fallback**: an inference call whose inputs differ from the recorded call in shape, dtype, device, or a non-tensor value runs the original module +- **Dynamic axes**: an axis that changes without following the batch size, such as sequence length, rejects the backend +- **Fallback**: an inference call whose inputs differ from the recordings in shape, dtype, device, or a non-tensor value runs the original module + +## Batch Sizes + +A graph spec with static shapes records one graph. +A graph spec with a batch axis records graphs for batch sizes 1, 2, 4, and then every multiple of 8, up to the largest batch size seen during tuning. +`max_recorded_batch_size` caps that size. The default is 256, which records 35 graphs. + +A call pads its batch with zeros up to the smallest recorded size that holds it, replays that graph, and returns only the real rows. +A batch larger than the largest recorded size replays in chunks of that size. +Padding costs compute. A batch of 3 runs the graph recorded for 4, and a batch of 17 runs the graph recorded for 24. + +Padding is only correct when every row of the batch is computed on its own. +Tuning checks that: it runs eager once on the largest recorded batch size. +Every graph then replays a prefix of that batch, full and partly padded with zeros and with copies of a real row, and is compared with the matching rows of eager. +The backend is rejected when: + +- an output tensor has no batch axis, for example `x.sum()` +- one row of the output depends on another row, for example `x - x.mean(dim=0)` +- the module rearranges the rows of its batch, for example `torch.cat([y, y])` + +All graphs of one module share one memory pool and one input buffer sized for the largest recorded graph. +Peak memory follows the largest recorded graph. +Tuning time grows with the number of graphs: each one is warmed up, recorded, and replayed by the check. +An out-of-memory error while recording a size skips that size and keeps the smaller graphs. + +```python +from torch_tweak.torch.backend import XPUGraphBackend, XPUGraphBackendConfig + +backend = XPUGraphBackend(XPUGraphBackendConfig(max_recorded_batch_size=64)) +``` ## Quick Start diff --git a/docs/learn/backends.md b/docs/learn/backends.md index 2327d5ea..d6f3dd85 100644 --- a/docs/learn/backends.md +++ b/docs/learn/backends.md @@ -35,7 +35,8 @@ Guide: [Torch Inductor Backend](../guides/backends/torch_inductor_backend.md). ## XPU Graph Backend -The XPU graph backend records the module once while tuning and replays that recording for later inference. +The XPU graph backend records the module while tuning and replays that recording for later inference. +A module with a batch axis gets recordings for batch sizes 1, 2, 4, and every multiple of 8, and other batch sizes are padded. PyTorch chooses native or legacy recording. This backend uses whichever one that build records. A recording that fails, or that disagrees with eager, is rejected for that module and input. diff --git a/examples/Minimal/demo.py b/examples/Minimal/demo.py index 1c2a8a71..2990ba57 100644 --- a/examples/Minimal/demo.py +++ b/examples/Minimal/demo.py @@ -23,7 +23,6 @@ def main() -> None: model = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 8)).eval().to(device) sample = torch.randn(32, device=device) - # One batch size keeps every tensor shape static, so the XPU graph can be recorded. strategy = tt.HighestThroughputStrategy( backends=[ TorchInductorBackend(), diff --git a/tests/unit/torch/backend/test_xpu_graph_backend.py b/tests/unit/torch/backend/test_xpu_graph_backend.py index 61ebc6fc..fba22c01 100644 --- a/tests/unit/torch/backend/test_xpu_graph_backend.py +++ b/tests/unit/torch/backend/test_xpu_graph_backend.py @@ -18,7 +18,7 @@ from torch_tweak.torch.module.sample_metadata import SampleMetadata from torch_tweak.torch.module.wrapper_module import Module from torch_tweak.torch.module_registry import MODULE_REGISTRY -from torch_tweak.torch.tune_strategy import OneBackendStrategy +from torch_tweak.torch.tune_strategy import HighestThroughputStrategy, OneBackendStrategy from torch_tweak.torch.tuning import tune @@ -244,22 +244,192 @@ def test_empty_recording_is_rejected(torch_device, tmp_path): backend.build(module, static_graph_spec(module, [(args, {})]), [(args, {})], torch_device, tmp_path) +def seen_graph_spec(module, samples: list[tuple[int, Sample]]) -> GraphSpec: + """Builds the graph spec tuning records after it sees every (batch size, sample) pair.""" + spec = None + for batch_size, (args, kwargs) in samples: + with torch.no_grad(): + outputs = module(*args, **kwargs) + inputs = SampleMetadata.from_inputs(args, kwargs, batch_size=batch_size) + outputs = SampleMetadata.from_outputs(outputs, batch_size=batch_size) + if spec is None: + spec = GraphSpec(name="seen", input_spec=inputs, output_spec=outputs) + else: + spec.update_shapes_seen(inputs, outputs) + return spec + + +def batched_build(module, torch_device, tmp_path, largest=64, config=None, features=16): + """Builds a backend on a graph spec that saw batch sizes 1 and largest.""" + module = module.to(torch_device).eval() + small = ((torch.randn(1, features, device=torch_device),), {}) + large = ((torch.randn(largest, features, device=torch_device),), {}) + spec = seen_graph_spec(module, [(1, small), (largest, large)]) + return XPUGraphBackend(config).build(module, spec, [small, large], torch_device, tmp_path) + + +class RowMixing(nn.Module): + def forward(self, x): + return x - x.mean(dim=0, keepdim=True) + + +class Stacked(nn.Module): + """Takes a batch stacked twice along axis 0, like classifier-free guidance.""" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 8) + + def forward(self, x): + y = self.linear(x) + return torch.cat([y, y], dim=0) + + +@requires_xpu +def test_batch_axis_records_the_size_grid(torch_device, tmp_path): + module = LinearRelu() + backend = batched_build(module, torch_device, tmp_path) + assert backend.recorded_batch_sizes == [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64] + with torch.no_grad(): + for batch_size in (1, 3, 5, 8, 17, 64, 100): + x = torch.randn(batch_size, 16, device=torch_device) + torch.testing.assert_close(backend.infer(x), module(x)) + + +@requires_xpu +def test_grid_stops_at_the_first_size_that_holds_the_largest_batch(torch_device, tmp_path): + assert batched_build(LinearRelu(), torch_device, tmp_path, largest=20).recorded_batch_sizes == [1, 2, 4, 8, 16, 24] + config = XPUGraphBackendConfig(max_recorded_batch_size=12) + backend = batched_build(LinearRelu(), torch_device, tmp_path, config=config) + assert backend.recorded_batch_sizes == [1, 2, 4, 8, 12] + + +@requires_xpu +def test_graphs_share_one_input_buffer(torch_device, tmp_path): + backend = batched_build(LinearRelu(), torch_device, tmp_path) + pointers = {recording.static_leaves[0].data_ptr() for recording in backend._recordings.values()} + assert len(pointers) == 1 + + +@requires_xpu +def test_padding_does_not_leak_rows_of_a_larger_graph(torch_device, tmp_path): + module = LinearRelu() + backend = batched_build(module, torch_device, tmp_path) + with torch.no_grad(): + backend.infer(torch.full((64, 16), 1e4, device=torch_device)) + x = torch.randn(5, 16, device=torch_device) + torch.testing.assert_close(backend.infer(x), module(x)) + + +@requires_xpu +def test_max_recorded_batch_size_caps_recording_and_chunks(torch_device, tmp_path): + module = LinearRelu() + config = XPUGraphBackendConfig(max_recorded_batch_size=4) + backend = batched_build(module, torch_device, tmp_path, config=config) + assert backend.recorded_batch_sizes == [1, 2, 4] + assert backend.describe() == "XPUGraphBackend(max_recorded_batch_size=4)" + with torch.no_grad(): + x = torch.randn(11, 16, device=torch_device) + torch.testing.assert_close(backend.infer(x), module(x)) + + +@requires_xpu +def test_out_of_memory_keeps_smaller_graphs(torch_device, tmp_path, monkeypatch): + record_size = XPUGraphBackend._record_size + + def failing(self, size, *rest): + if size > 8: + raise torch.OutOfMemoryError("fake") + return record_size(self, size, *rest) + + monkeypatch.setattr(XPUGraphBackend, "_record_size", failing) + module = LinearRelu() + backend = batched_build(module, torch_device, tmp_path) + assert backend.recorded_batch_sizes == [1, 2, 4, 8] + with torch.no_grad(): + x = torch.randn(30, 16, device=torch_device) + torch.testing.assert_close(backend.infer(x), module(x)) + + +@requires_xpu +def test_output_without_batch_axis_is_rejected(torch_device, tmp_path): + with pytest.raises(RuntimeError, match="no batch axis"): + batched_build(Scaled(), torch_device, tmp_path, largest=4, features=5) + + @requires_xpu -def test_dynamic_shapes_are_rejected(torch_device, tmp_path): +def test_rows_that_depend_on_other_rows_are_rejected(torch_device, tmp_path): + with pytest.raises(RuntimeError, match="differs from eager"): + batched_build(RowMixing(), torch_device, tmp_path, largest=4) + + +@requires_xpu +def test_module_that_stacks_its_batch_is_rejected(torch_device, tmp_path): + with pytest.raises(RuntimeError, match="differs from eager"): + batched_build(Stacked(), torch_device, tmp_path, largest=8) + + +@requires_xpu +def test_input_with_two_rows_per_sample(torch_device, tmp_path): module = LinearRelu().to(torch_device).eval() - small = (torch.randn(1, 16, device=torch_device),) - large = (torch.randn(2, 16, device=torch_device),) + small = ((torch.randn(2, 16, device=torch_device),), {}) + large = ((torch.randn(16, 16, device=torch_device),), {}) + spec = seen_graph_spec(module, [(1, small), (8, large)]) + backend = XPUGraphBackend().build(module, spec, [small, large], torch_device, tmp_path) + assert backend.recorded_batch_sizes == [1, 2, 4, 8] with torch.no_grad(): - small_out = module(*small) - large_out = module(*large) - input_spec = SampleMetadata.from_inputs(small, {}, batch_size=1) - output_spec = SampleMetadata.from_outputs(small_out, batch_size=1) - input_spec.update_shapes_seen(SampleMetadata.from_inputs(large, {}, batch_size=2)) - output_spec.update_shapes_seen(SampleMetadata.from_outputs(large_out, batch_size=2)) - spec = GraphSpec(name="dynamic", input_spec=input_spec, output_spec=output_spec) - backend = XPUGraphBackend() - with pytest.raises(RuntimeError, match="static tensor shapes"): - backend.build(module, spec, [(large, {})], torch_device, tmp_path) + x = torch.randn(10, 16, device=torch_device) + torch.testing.assert_close(backend.infer(x), module(x)) + + +@requires_xpu +def test_dynamic_axis_is_rejected(torch_device, tmp_path): + module = LinearRelu().to(torch_device).eval() + short = ((torch.randn(1, 3, 16, device=torch_device),), {}) + long = ((torch.randn(2, 7, 16, device=torch_device),), {}) + spec = seen_graph_spec(module, [(1, short), (2, long)]) + with pytest.raises(RuntimeError, match="dynamic axis"): + XPUGraphBackend().build(module, spec, [short, long], torch_device, tmp_path) + + +@requires_xpu +def test_batch_axis_serialization(torch_device, tmp_path): + module = LinearRelu() + backend = batched_build(module, torch_device, tmp_path, largest=8) + TorchSaveTask().save(tmp_path, backend.to_dict()) + loaded = XPUGraphBackend.from_dict(module, TorchLoadTask().load(tmp_path)) + loaded.deploy(device=torch_device) + assert loaded.recorded_batch_sizes == [1, 2, 4, 8] + with torch.no_grad(): + x = torch.randn(6, 16, device=torch_device) + torch.testing.assert_close(loaded.infer(x), module(x)) + + +@requires_xpu +def test_tune_with_two_batch_sizes_records_the_batch_axis(torch_device, tmp_path): + module = LinearRelu().to(torch_device).eval() + reference = copy.deepcopy(module) + wrapped = Module( + module, + "xpu-graph-batched", + strategy=HighestThroughputStrategy(backends=[XPUGraphBackend()]).enable_find_max_batch_size(False), + ) + try: + tune( + wrapped, + torch.randn(16, device=torch_device), + batch_sizes=[1, 8], + device=torch_device, + ignore_failing_modules=False, + ) + (backend,) = wrapped._self_wrapper.backends.values() + assert isinstance(backend, XPUGraphBackend) + assert backend.recorded_batch_sizes == [1, 2, 4, 8] + x = torch.randn(5, 16, device=torch_device) + with torch.no_grad(): + torch.testing.assert_close(wrapped(x), reference(x)) + finally: + MODULE_REGISTRY.clear() @requires_xpu diff --git a/torch_tweak/torch/backend/xpu_graph_backend.py b/torch_tweak/torch/backend/xpu_graph_backend.py index ebc1f366..6eb67640 100644 --- a/torch_tweak/torch/backend/xpu_graph_backend.py +++ b/torch_tweak/torch/backend/xpu_graph_backend.py @@ -10,15 +10,20 @@ import torch import torch.nn as nn -from torch.utils._pytree import TreeSpec, tree_flatten, tree_map_only +from torch.utils._pytree import TreeSpec, tree_flatten, tree_map_only, tree_unflatten from torch_tweak.torch.backend.backend import Backend, BackendConfig, BackendState 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.tensor_spec import recovered_batch_dim_bounds from torch_tweak.utils.trace import annotate logger = getLogger(__name__) +# Maps a tensor axis to its batch multiplier. The axis holds multiplier * batch_size elements. +BatchAxes = dict[int, int] + @dataclass class XPUGraphBackendConfig(BackendConfig): @@ -28,10 +33,14 @@ class XPUGraphBackendConfig(BackendConfig): autocast_enabled: If True, enable autocast while recording. Replay runs the recorded casts. The autocast cache is disabled so later eager calls cannot reuse tensors captured in the graph. autocast_dtype: The dtype to use for autocast. + max_recorded_batch_size: The largest batch size recorded for a graph spec with a batch axis. + Graphs are recorded for batch sizes 1, 2, 4, and every multiple of 8 up to this size. + A larger batch replays in chunks of this size. """ autocast_enabled: bool = False autocast_dtype: torch.dtype | None = None + max_recorded_batch_size: int = 256 @classmethod def from_dict(cls, state_dict: dict[str, Any]) -> Self: @@ -39,8 +48,41 @@ def from_dict(cls, state_dict: dict[str, Any]) -> Self: return cls(**state_dict) +@dataclass +class _Recording: + """One recorded graph and the tensors it reads and writes.""" + + graph: torch.xpu.XPUGraph + static_leaves: list[Any] + captured_leaves: list[Any] + + +@dataclass +class _Layout: + """Batch axes of every flattened input and output leaf. Non-tensor leaves have no batch axes.""" + + input_spec: TreeSpec + input_axes: list[BatchAxes] + output_spec: TreeSpec + output_axes: list[BatchAxes] + + @property + def has_batch_axis(self) -> bool: + return any(self.input_axes) + + @property + def can_chunk(self) -> bool: + """Chunks are concatenated along one axis, so every batched leaf has exactly one batch axis.""" + return all(len(axes) <= 1 for axes in self.input_axes + self.output_axes) + + class XPUGraphBackend(Backend): - """Backend that records one XPU graph while tuning and replays it for inference.""" + """Backend that records XPU graphs while tuning and replays them for inference. + + A graph spec with static shapes records one graph. + A graph spec with a batch axis records graphs for batch sizes 1, 2, 4, and every multiple of 8. + A call pads its batch with zeros up to the nearest recorded size and gets the real rows back. + """ _devices: ClassVar[list[str]] = ["xpu"] requires_original_module: ClassVar[bool] = True @@ -50,6 +92,7 @@ class XPUGraphBackend(Backend): STATE_ORIG_MODULE = "orig_module" STATE_DATA = "data" STATE_DEVICE = "device" + STATE_INPUT_SPEC = "input_spec" def __init__(self, config: XPUGraphBackendConfig | None = None): """Initializes backend. @@ -61,10 +104,11 @@ def __init__(self, config: XPUGraphBackendConfig | None = None): self._config: XPUGraphBackendConfig = config or XPUGraphBackendConfig() self._orig_module: nn.Module | None = None self._data: list[Sample] | None = None - self._graph: torch.xpu.XPUGraph | None = None - self._static_args: tuple[list[Any], TreeSpec] | None = None - self._static_kwargs: tuple[list[Any], TreeSpec] | None = None - self._captured: Any = None + self._input_spec: SampleMetadata | None = None + self._layout: _Layout | None = None + # Keyed by batch size. A graph spec without a batch axis has one recording under key 1. + self._recordings: dict[int, _Recording] = {} + self._pool: Any = None self._warned_eager = False def key(self) -> str: @@ -75,35 +119,43 @@ def describe(self) -> str: """Returns the description of the backend.""" return f"{self.__class__.__name__}({self._config.describe()})" + @property + def recorded_batch_sizes(self) -> list[int]: + """Batch sizes with a recorded graph, in ascending order.""" + return sorted(self._recordings) + def _build(self, module: nn.Module, graph_spec: GraphSpec, data: list[Sample], cache_dir: Path) -> Backend: - """Records the XPU graph once and checks it against eager.""" - if graph_spec.has_dynamic_axis(): - raise RuntimeError("XPU graph replay needs static tensor shapes") + """Records the XPU graphs and checks them against eager.""" + tensor_specs = graph_spec.input_spec.tensor_specs + graph_spec.output_spec.tensor_specs + if any(tensor_spec.has_dynamic_axis() for tensor_spec in tensor_specs): + raise RuntimeError("XPU graph replay cannot change a dynamic axis other than the batch axis") for sample_args, sample_kwargs in data: _require_xpu_tensors(sample_args, sample_kwargs) self._config.save_to_file(cache_dir) self._orig_module = module self._data = data + if graph_spec.input_spec.detected_dynamic_axis(): + self._input_spec = graph_spec.input_spec self._record() return self def _activate(self) -> None: - """Records the graph again after a checkpoint load or a deactivate.""" - self._release_graph() + """Records the graphs again after a checkpoint load or a deactivate.""" + self._release_graphs() self._record() def _deactivate(self) -> None: - """Drops the recorded graph.""" - self._release_graph() + """Drops the recorded graphs.""" + self._release_graphs() def _deploy(self) -> None: - """Records the graph for deployment and drops the tuning samples.""" + """Records the graphs for deployment and drops the tuning samples.""" self._activate() self._data = None @annotate(message="XPUGraphBackend.infer", domain="Torch Tweak") def _infer(self, *args: Any, **kwargs: Any) -> Any: - """Replays the recorded graph. + """Replays the recorded graph that fits the call. Args: *args: inference arguments @@ -111,48 +163,194 @@ def _infer(self, *args: Any, **kwargs: Any) -> Any: Returns: The replayed outputs. Tensor outputs are clones, so the next replay can reuse the static buffers. - Inputs the recording cannot take run on the original module. + Inputs the recordings cannot take run on the original module. """ with torch.no_grad(): - if not self._load_inputs(args, kwargs): + outputs = self._replay((args, kwargs)) + if outputs is None: return self._infer_eager(args, kwargs) - self._graph.replay() - # replay and this clone are both on the current stream, in submission order. - return _clone_tree(self._captured) - - def _load_inputs(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool: - """Copies args and kwargs into the recorded inputs. Copies nothing and returns False when they do not fit.""" - assert self._static_args is not None and self._static_kwargs is not None - live_args = _live_leaves(self._static_args, args) - live_kwargs = _live_leaves(self._static_kwargs, kwargs) - if live_args is None or live_kwargs is None: - return False - _copy_into(self._static_args[0], live_args) - _copy_into(self._static_kwargs[0], live_kwargs) - return True + return outputs + + def _replay(self, sample: Sample) -> Any | None: + """Replays the recordings for sample. Returns None when no recording fits.""" + assert self._layout is not None + layout = self._layout + leaves, spec = tree_flatten(sample) + batch_size = self._batch_size_of(leaves, spec) + if batch_size is None: + return None + largest = self.recorded_batch_sizes[-1] + if batch_size <= largest: + return tree_unflatten(self._replay_rows(leaves, 0, batch_size), layout.output_spec) + if not layout.can_chunk: + return None + chunks = [ + self._replay_rows(leaves, start, min(largest, batch_size - start)) + for start in range(0, batch_size, largest) + ] + merged = [ + torch.cat([chunk[index] for chunk in chunks], dim=next(iter(axes))) if axes else chunks[0][index] + for index, axes in enumerate(layout.output_axes) + ] + return tree_unflatten(merged, layout.output_spec) + + def _batch_size_of(self, leaves: list[Any], spec: TreeSpec) -> int | None: + """Returns the batch size of a call, or None when its leaves do not fit the recordings.""" + assert self._layout is not None + layout = self._layout + if spec != layout.input_spec: + return None + batch_size = 1 + for leaf, axes in zip(leaves, layout.input_axes, strict=True): + if axes: + if not isinstance(leaf, torch.Tensor): + return None + axis, multiplier = next(iter(axes.items())) + batch_size, remainder = divmod(leaf.shape[axis], multiplier) + if batch_size == 0 or remainder: + return None + break + reference = self._recordings[self.recorded_batch_sizes[0]].static_leaves + for leaf, static, axes in zip(leaves, reference, layout.input_axes, strict=True): + if not _leaf_fits(leaf, static, axes, batch_size): + return None + return batch_size + + def _replay_rows( + self, leaves: list[Any], start: int, rows: int, fill: str = "zeros", size: int | None = None + ) -> list[Any]: + """Replays rows start to start + rows of the call. Returns cloned output leaves. + + fill: "zeros" pads with zeros. "repeat" pads with copies of the first row and is used by the build check. + size: the recorded batch size to replay. Defaults to the smallest one that holds rows. + """ + assert self._layout is not None + layout = self._layout + if size is None: + size = min(recorded for recorded in self.recorded_batch_sizes if recorded >= rows) + recording = self._recordings[size] + for live, static, axes in zip(leaves, recording.static_leaves, layout.input_axes, strict=True): + if not isinstance(static, torch.Tensor): + continue + if not axes: + static.copy_(live) + continue + source = live[_rows_index(live.ndim, axes, start, rows)] + if rows == size: + static.copy_(source) + continue + if fill == "repeat": + static.copy_(_repeat_first_rows(source, axes, size)) + elif len(axes) == 1: + # Stale rows from an earlier call or a larger graph sharing the buffer must not reach the padding. + static[_rows_index(static.ndim, axes, rows, size - rows)].zero_() + else: + static.zero_() + static[_rows_index(static.ndim, axes, 0, rows)].copy_(source) + recording.graph.replay() + # replay and these clones are both on the current stream, in submission order. + return [ + captured[_rows_index(captured.ndim, axes, 0, rows)].clone() if axes else _clone_tree(captured) + for captured, axes in zip(recording.captured_leaves, layout.output_axes, strict=True) + ] def _infer_eager(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: - """Runs the original module on inputs that differ from the recorded call.""" + """Runs the original module on inputs that differ from the recorded calls.""" assert self._orig_module is not None if not self._warned_eager: - logger.warning("XPU graph inputs do not match the recorded call. Running the original module.") + logger.warning("XPU graph inputs do not match the recorded calls. Running the original module.") self._warned_eager = True with self._autocast(): return self._orig_module(*args, **kwargs) + def _eager(self, sample: Sample) -> Any: + """Runs the original module on a copy of sample, so in-place writes do not change sample.""" + assert self._orig_module is not None + args, kwargs = _clone_tree(sample) + with self._autocast(), torch.no_grad(): + return _clone_tree(self._orig_module(*args, **kwargs)) + def _record(self) -> None: - """Warm up, record, and reject the graph when replay disagrees with eager.""" + """Records every batch size, then rejects the backend when replay disagrees with eager.""" assert self._orig_module is not None assert self._data - self._orig_module.to(self._device) - args, kwargs = self._data[0] - static_args = _clone_tree(args) - static_kwargs = _clone_tree(kwargs) + try: + self._layout = self._detect_layout() + self._pool = torch.xpu.graph_pool_handle() + sizes = self._batch_sizes_to_record() if self._layout.has_batch_axis else [1] + buffer = tree_flatten(_clone_tree(self._sample_of_size(max(sizes))))[0] + # The largest graph sizes the shared pool. Smaller graphs reuse its blocks. + for size in sorted(sizes, reverse=True): + try: + self._recordings[size] = self._record_size(size, max(sizes), buffer) + except torch.OutOfMemoryError: + if size == min(sizes): + raise + logger.warning("Out of memory while recording an XPU graph for batch size %d. Skipping it.", size) + torch.xpu.empty_cache() + self._replay_matches_eager() + except Exception: + self._release_graphs() + raise + logger.info("XPU graphs recorded for batch sizes %s.", self.recorded_batch_sizes) - logger.info("Recording XPU graph.") + def _batch_sizes_to_record(self) -> list[int]: + """1, 2, 4, then multiples of 8, up to the first size that holds the largest batch size seen. + + Above batch size 8, padding adds at most 7 rows. Every power of two profiled by + HighestThroughputStrategy is on this grid, so profiling replays without padding. + """ + assert self._input_spec is not None + _, seen = recovered_batch_dim_bounds(self._input_spec.tensor_specs) + cap = self._config.max_recorded_batch_size + grid = [size for size in [1, 2, 4, *range(8, cap + 1, 8)] if size <= cap] + if grid[-1] < cap: + grid.append(cap) + largest = min(seen, cap) + return [size for size in grid if size < largest] + [min(size for size in grid if size >= largest)] + + def _sample_of_size(self, size: int) -> Sample: + """Returns the first tuning sample with batch size size.""" + assert self._data + args, kwargs = self._data[0] + if self._input_spec is None: + return _clone_tree((args, kwargs)) + return self._input_spec.make_batch(args, kwargs, size) + + def _detect_layout(self) -> _Layout: + """Finds the batch axes of every input and output leaf from samples of batch size 1 and 2.""" + one = self._sample_of_size(1) + one_leaves, input_spec = tree_flatten(one) + outputs_one, output_spec = tree_flatten(self._eager(one)) + if self._input_spec is None: + return _Layout(input_spec, [{} for _ in one_leaves], output_spec, [{} for _ in outputs_one]) + two = self._sample_of_size(2) + two_leaves, two_spec = tree_flatten(two) + outputs_two, outputs_two_spec = tree_flatten(self._eager(two)) + if two_spec != input_spec or outputs_two_spec != output_spec: + raise RuntimeError("XPU graph inputs or outputs change structure with the batch size") + layout = _Layout( + input_spec=input_spec, + input_axes=[_batch_axes(a, b) for a, b in zip(one_leaves, two_leaves, strict=True)], + output_spec=output_spec, + output_axes=[_batch_axes(a, b) for a, b in zip(outputs_one, outputs_two, strict=True)], + ) + if not layout.has_batch_axis: + raise RuntimeError("XPU graph found no batch axis in the inputs") + outputs = zip(outputs_one, layout.output_axes, strict=True) + if any(isinstance(leaf, torch.Tensor) and not axes for leaf, axes in outputs): + raise RuntimeError("XPU graph output has no batch axis, so padded rows cannot be removed") + return layout + + def _record_size(self, size: int, largest: int, buffer: list[Any]) -> _Recording: + """Warm up and record the graph for one batch size. buffer holds the input leaves at batch size largest.""" + assert self._orig_module is not None + assert self._layout is not None + static = tree_unflatten(self._static_leaves(size, largest, buffer), self._layout.input_spec) + static_args, static_kwargs = static + graph = torch.xpu.XPUGraph() try: - self._graph = torch.xpu.XPUGraph() with self._autocast(), torch.no_grad(): # Warmup. The first oneDNN GEMM for a shape selects an algorithm and allocates its workspace. # That work stays outside the recording. @@ -161,44 +359,88 @@ def _record(self) -> None: warnings.simplefilter("always") # Records this call on default_capture_stream. # Exiting the block ends the recording. - with torch.xpu.graph(self._graph): + with torch.xpu.graph(graph, pool=self._pool): captured = self._orig_module(*static_args, **static_kwargs) if any("XPU Graph is empty" in str(item.message) for item in caught): raise RuntimeError("XPU graph capture recorded no device work") - # Recording runs on default_capture_stream. replay runs on the current stream. - # Both write captured, so the current stream waits for the recording stream. - stream = torch.xpu.graph.default_capture_stream - assert stream is not None - torch.xpu.current_stream().wait_stream(stream) - self._static_args = tree_flatten(static_args) - self._static_kwargs = tree_flatten(static_kwargs) - self._captured = captured - self._replay_matches_eager(static_args, static_kwargs) except Exception: - self._release_graph() + graph.reset() raise - logger.info("XPU graph recorded.") - - def _replay_matches_eager(self, static_args: Any, static_kwargs: Any) -> None: - """Copies every sample and a perturbed copy into the static inputs and compares replay to eager.""" - assert self._orig_module is not None + # Recording runs on default_capture_stream. replay runs on the current stream. + # Both write captured, so the current stream waits for the recording stream. + stream = torch.xpu.graph.default_capture_stream + assert stream is not None + torch.xpu.current_stream().wait_stream(stream) + static_leaves, static_spec = tree_flatten(static) + captured_leaves, captured_spec = tree_flatten(captured) + if static_spec != self._layout.input_spec or captured_spec != self._layout.output_spec: + graph.reset() + raise RuntimeError("XPU graph inputs or outputs change structure with the batch size") + for leaves, axes_list in ( + (static_leaves, self._layout.input_axes), + (captured_leaves, self._layout.output_axes), + ): + for leaf, axes in zip(leaves, axes_list, strict=True): + if isinstance(leaf, torch.Tensor) and any(leaf.shape[a] != m * size for a, m in axes.items()): + graph.reset() + raise RuntimeError("XPU graph batch axis does not scale with the batch size") + return _Recording(graph=graph, static_leaves=static_leaves, captured_leaves=captured_leaves) + + def _static_leaves(self, size: int, largest: int, buffer: list[Any]) -> list[Any]: + """Returns the input leaves recorded for size. buffer holds the input leaves at batch size largest. + + A leaf without a batch axis is the buffer leaf. A leaf batched only along axis 0 is a prefix view of + the buffer leaf, which is contiguous. Graphs share that memory because replays run one at a time. + Any other batched leaf below largest gets its own tensor. + """ + assert self._layout is not None + own: list[Any] | None = None + leaves = [] + for index, (shared, axes) in enumerate(zip(buffer, self._layout.input_axes, strict=True)): + if not isinstance(shared, torch.Tensor) or not axes or size == largest: + leaves.append(shared) + elif axes.keys() == {0}: + leaves.append(shared[: axes[0] * size]) + else: + if own is None: + own = tree_flatten(_clone_tree(self._sample_of_size(size)))[0] + leaves.append(own[index]) + return leaves + + def _replay_matches_eager(self) -> None: + """Compares replay to eager on every tuning sample, a perturbed copy, and padded batches.""" assert self._data - assert self._graph is not None - for sample_args, sample_kwargs in self._data: - for args, kwargs in ( - (sample_args, sample_kwargs), - (_perturb(sample_args), _perturb(sample_kwargs)), - ): - if not self._load_inputs(args, kwargs): - raise RuntimeError("XPU graph inputs do not match the captured call") - # replay may write static_args in place. Eager must see the earlier values. - frozen_args = _clone_tree(static_args) - frozen_kwargs = _clone_tree(static_kwargs) - self._graph.replay() - with self._autocast(), torch.no_grad(): - # The module may return an input. Compare a tensor replay did not write. - expected = _clone_tree(self._orig_module(*frozen_args, **frozen_kwargs)) - _assert_close(self._captured, expected) + samples: list[Sample] = [] + for sample in self._data: + samples += [sample, _perturb(sample)] + for sample in samples: + outputs = self._replay(sample) + if outputs is None: + raise RuntimeError("XPU graph inputs do not match the captured call") + _assert_close(outputs, self._eager(sample)) + if self._layout and self._layout.has_batch_axis: + self._prefixes_match_eager() + + def _prefixes_match_eager(self) -> None: + """Replays prefixes of one perturbed batch on every graph and compares them to one eager run. + + Each graph replays its full size, then a shorter prefix padded with zeros and with repeats of + the first row. A module whose rows depend on other rows or on the padding rows fails. + """ + assert self._layout is not None + reference = _perturb(self._sample_of_size(self.recorded_batch_sizes[-1])) + leaves, _ = tree_flatten(reference) + expected, _ = tree_flatten(self._eager(reference)) + for size in self.recorded_batch_sizes: + rows = size // 2 + 1 if size > 2 else size - 1 + checks = [(size, "zeros")] + ([(rows, "zeros"), (rows, "repeat")] if rows else []) + for prefix, fill in checks: + outputs = self._replay_rows(leaves, 0, prefix, fill, size) + prefix_of_expected = [ + ref[_rows_index(ref.ndim, axes, 0, prefix)] if axes else ref + for ref, axes in zip(expected, self._layout.output_axes, strict=True) + ] + _assert_close(outputs, prefix_of_expected) def _autocast(self) -> torch.autocast: """Returns the autocast context used for warmup, capture, and the eager check.""" @@ -210,14 +452,13 @@ def _autocast(self) -> torch.autocast: cache_enabled=False, ) - def _release_graph(self) -> None: - """Drops the graph and the static tensors it captured.""" - if self._graph is not None: - self._graph.reset() - self._graph = None - self._static_args = None - self._static_kwargs = None - self._captured = None + def _release_graphs(self) -> None: + """Drops the graphs, the static tensors they captured, and their memory pool.""" + for recording in self._recordings.values(): + recording.graph.reset() + self._recordings = {} + self._layout = None + self._pool = None def to_dict(self) -> dict[str, Any]: """Returns the state_dict of the backend.""" @@ -229,6 +470,7 @@ def to_dict(self) -> dict[str, Any]: self.STATE_ORIG_MODULE: self._orig_module.state_dict(), self.STATE_DATA: self._data, self.STATE_DEVICE: self._device, + self.STATE_INPUT_SPEC: None if self._input_spec is None else self._input_spec.to_dict(), } @classmethod @@ -243,6 +485,8 @@ def from_dict(cls, module: nn.Module | None, state_dict: dict[str, Any]) -> Self backend = cls(config=config) backend._data = state_dict[cls.STATE_DATA] backend._device = state_dict[cls.STATE_DEVICE] + input_spec = state_dict.get(cls.STATE_INPUT_SPEC) + backend._input_spec = None if input_spec is None else SampleMetadata.from_dict(input_spec) backend._orig_module = module module.load_state_dict(state_dict[cls.STATE_ORIG_MODULE], strict=False) backend.state = BackendState.CHECKPOINT_LOADED @@ -254,30 +498,55 @@ def _clone_tree(tree: Any) -> Any: return tree_map_only(torch.Tensor, lambda tensor: tensor.detach().clone(), tree) -def _live_leaves(static: tuple[list[Any], TreeSpec], live: Any) -> list[Any] | None: - """Returns the leaves of live, or None when they cannot be copied into the inputs recorded by the graph.""" - static_leaves, static_spec = static - live_leaves, live_spec = tree_flatten(live) - if static_spec != live_spec: - return None - for dst, src in zip(static_leaves, live_leaves, strict=True): - if isinstance(dst, torch.Tensor) and isinstance(src, torch.Tensor): - if (dst.shape, dst.dtype, dst.device) != (src.shape, src.dtype, src.device): - return None - elif isinstance(dst, torch.Tensor) or isinstance(src, torch.Tensor): - # != of a tensor and a non-tensor returns a tensor. - # The next branch would not reject this pair. - return None - elif dst != src: - return None - return live_leaves - - -def _copy_into(static_leaves: list[Any], live_leaves: list[Any]) -> None: - """Copies live tensors into the tensors recorded by the graph. live_leaves comes from _live_leaves.""" - for dst, src in zip(static_leaves, live_leaves, strict=True): - if isinstance(dst, torch.Tensor): - dst.copy_(src) +def _batch_axes(one: Any, two: Any) -> BatchAxes: + """Returns the batch axes of a leaf from its values at batch size 1 and 2.""" + if not isinstance(one, torch.Tensor) or not isinstance(two, torch.Tensor): + return {} + if one.ndim != two.ndim: + raise RuntimeError("XPU graph tensor rank changes with the batch size") + axes: BatchAxes = {} + for axis, (size_one, size_two) in enumerate(zip(one.shape, two.shape, strict=True)): + if size_one == size_two: + continue + if size_two != 2 * size_one: + raise RuntimeError("XPU graph tensor axis does not scale with the batch size") + axes[axis] = size_one + return axes + + +def _leaf_fits(leaf: Any, static: Any, axes: BatchAxes, batch_size: int) -> bool: + """Returns True when leaf can be copied into static with batch_size rows.""" + if isinstance(static, torch.Tensor) != isinstance(leaf, torch.Tensor): + return False + if not isinstance(static, torch.Tensor): + # != of two non-tensors returns a bool. A tensor pair never reaches this line. + return leaf == static + if (leaf.dtype, leaf.device, leaf.ndim) != (static.dtype, static.device, static.ndim): + return False + for axis, size in enumerate(leaf.shape): + expected = axes[axis] * batch_size if axis in axes else static.shape[axis] + if size != expected: + return False + return True + + +def _rows_index(ndim: int, axes: BatchAxes, start: int, rows: int) -> tuple[slice, ...]: + """Index of rows start to start + rows along every batch axis.""" + index = [slice(None)] * ndim + for axis, multiplier in axes.items(): + index[axis] = slice(start * multiplier, (start + rows) * multiplier) + return tuple(index) + + +def _repeat_first_rows(source: torch.Tensor, axes: BatchAxes, size: int) -> torch.Tensor: + """Returns source grown to size rows along every batch axis by repeating its first row.""" + grown = source + for axis, multiplier in axes.items(): + first = grown.narrow(axis, 0, multiplier) + repeats = [1] * grown.ndim + repeats[axis] = size + grown = first.repeat(repeats) + return grown def _require_xpu_tensors(*trees: Any) -> None: From f37ed69b9ac9c2654bbf6dfdcdb3b8ab3c667fe7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Fri, 23 Oct 2026 18:02:21 +0200 Subject: [PATCH 3/5] fix(profiling): judge saturation on the measured events only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first call of a new batch size can compile, for example in Torch Inductor. Saturation counted that call, so a backend looked saturated and stopped profiling early. On MobileNetV3, Inductor stopped at batch size 8 and never ran 16, 32, or 64. The stop strategy now sees the same measured window that the reported throughput uses. Signed-off-by: Łukasz Ślusarczyk --- tests/unit/torch/task/test_profiling.py | 34 +++++++++++++++++++ torch_tweak/torch/task/profiling/profiling.py | 5 ++- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/task/test_profiling.py b/tests/unit/torch/task/test_profiling.py index 36866936..33e3f4b3 100644 --- a/tests/unit/torch/task/test_profiling.py +++ b/tests/unit/torch/task/test_profiling.py @@ -260,3 +260,37 @@ def model(*_args, **_kwargs): assert order == ["sync", "model", "sync"] assert events[0].execution_time >= 0 + + +def test_compilation_of_a_new_batch_size_does_not_stop_profiling(): + """The first call of a batch size compiles. Only the measured window counts for saturation.""" + model = ToyTorchModel() + graph_spec = model.graph_spec(batch_sizes=[1, 2]) + compiled = {1} + + class CompilingBackend(TorchEagerBackend): + def __init__(self): + pass + + def infer(self, *args, **kwargs): + batch_size = args[0].shape[0] + if batch_size not in compiled: + compiled.add(batch_size) + time.sleep(0.1) + time.sleep(0.005) + return model(*args) + + def describe(self): + return "compiling_backend" + + profile_config = ProfilingConfig( + batch_sizes=[1, 2, 4, 8], + measurement_stop_strategy=StableWindowMeasuringStopStrategy(window_size=3, stability_percentage=80), + profiling_stop_strategy=ThroughputSaturatedProfilingStopStrategy(throughput_cutoff_threshold=0.05), + ) + samples = [(model.inputs(batch_sizes=[1]), {})] + + result = profile_backend(CompilingBackend(), "test_model", graph_spec, samples, profile_config) # type: ignore + + assert result.status == ProfilingStatus.Status.SUCCESS + assert {entry.batch_size for entry in result.results.entries} == {1, 2, 4, 8} diff --git a/torch_tweak/torch/task/profiling/profiling.py b/torch_tweak/torch/task/profiling/profiling.py index 3ea2595c..832bd731 100644 --- a/torch_tweak/torch/task/profiling/profiling.py +++ b/torch_tweak/torch/task/profiling/profiling.py @@ -142,7 +142,10 @@ def _profile( profiling_results.entries += new_entries - if profiling_stop_strategy.should_stop(get_inference_events(new_entries)): + # The first calls of a batch size can include compilation, so saturation is judged on + # the measured events only, the same ones get_throughput_per_batch_size uses. + measured = profiling_config.measurement_stop_strategy.get_events(get_inference_events(new_entries)) + if profiling_stop_strategy.should_stop(measured): break return ProfilingStatus(status=ProfilingStatus.Status.SUCCESS, results=profiling_results) From 4ee8f0bd46d0276be1bc6c4f6f6e0b915f411dbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Fri, 23 Oct 2026 18:05:20 +0200 Subject: [PATCH 4/5] feat(strategy): route each batch size to the backend fastest for it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HighestThroughputStrategy kept one backend for every batch size. It chose by peak throughput, which is usually at the largest batch size. A backend that was faster at small batch sizes, such as an XPU graph, lost them. The strategy now compares backends at every profiled batch size. Another backend takes a range when it is at least 10% faster there. BatchSizeRouterBackend holds the routed backends and sends each call by its batch size. It saves and loads every routed backend in a checkpoint. On MobileNetV3 it routes batch sizes 1 to 16 to the XPU graph and larger ones to Inductor. Signed-off-by: Łukasz Ślusarczyk --- .../api/backends/batch_size_router_backend.md | 13 ++ docs/guides/backends/xpu_graph_backend.md | 2 + .../guides/tune_strategies/tune_strategies.md | 17 ++ docs/learn/core_functionalities.md | 1 + properdocs.yml | 1 + .../backend/test_batch_size_router_backend.py | 151 ++++++++++++++++++ torch_tweak/torch/backend/__init__.py | 2 + .../backend/batch_size_router_backend.py | 134 ++++++++++++++++ .../highest_throughput_strategy.py | 48 +++++- 9 files changed, 367 insertions(+), 2 deletions(-) create mode 100644 docs/api/backends/batch_size_router_backend.md create mode 100644 tests/unit/torch/backend/test_batch_size_router_backend.py create mode 100644 torch_tweak/torch/backend/batch_size_router_backend.py diff --git a/docs/api/backends/batch_size_router_backend.md b/docs/api/backends/batch_size_router_backend.md new file mode 100644 index 00000000..557c4fd0 --- /dev/null +++ b/docs/api/backends/batch_size_router_backend.md @@ -0,0 +1,13 @@ + + +# Batch Size Router Backend API + +## BatchSizeRouterBackend + +::: torch_tweak.torch.backend.BatchSizeRouterBackend + options: + show_root_heading: true + heading_level: 3 diff --git a/docs/guides/backends/xpu_graph_backend.md b/docs/guides/backends/xpu_graph_backend.md index 8ba1df50..734ace0c 100644 --- a/docs/guides/backends/xpu_graph_backend.md +++ b/docs/guides/backends/xpu_graph_backend.md @@ -58,6 +58,8 @@ strategy = HighestThroughputStrategy(backends=[XPUGraphBackend()]) On XPU the backend is first in `default_backends(device)`, before Torch Inductor. `FirstWinsStrategy`, `HighestThroughputStrategy`, and `tt.default_strategy(device)` use that list when no backends are given. On CPU the list leaves this backend out. +An XPU graph often wins at small batch sizes and loses at large ones. +`HighestThroughputStrategy` then sends small batches to it and large batches to the faster backend. PyTorch records a native SYCL graph when it was built with a SYCL compiler of 2026.1 or newer, except on Intel Data Center GPU Max Series. An older build records the legacy SYCL command graph. diff --git a/docs/guides/tune_strategies/tune_strategies.md b/docs/guides/tune_strategies/tune_strategies.md index f77952b3..835cb399 100644 --- a/docs/guides/tune_strategies/tune_strategies.md +++ b/docs/guides/tune_strategies/tune_strategies.md @@ -196,6 +196,23 @@ tt.tune(model, input_data) 3. Selects backend with highest throughput 4. Returns best performing backend +### Different Winners per Batch Size + +With a batch axis, every backend is profiled on every power of two from 1 up to the max batch size. +The best backend is the one with the highest throughput at any batch size, usually the largest one. +Another backend takes a range of batch sizes when it is at least 10% faster there. +The strategy then returns a `BatchSizeRouterBackend` that sends each call to the backend tuned for its batch size. +A batch size between two profiled ones goes to the backend of the larger one. + +For example, an XPU graph launches a small batch faster and Torch Inductor computes a large batch faster: + +``` +BatchSizeRouterBackend(1-16: XPUGraphBackend(), 17+: TorchInductorBackend()) +``` + +Every routed backend stays in memory. +A checkpoint stores each of them. + ### When to Use ✅ **Good for**: diff --git a/docs/learn/core_functionalities.md b/docs/learn/core_functionalities.md index 3253fc7e..b25ea937 100644 --- a/docs/learn/core_functionalities.md +++ b/docs/learn/core_functionalities.md @@ -133,6 +133,7 @@ strategy = OneBackendStrategy(backend=OpenVINOBackend()) ## HighestThroughputStrategy Profiles all compatible backends and selects the fastest. Use this when maximum throughput matters and you can afford longer tuning time. `HighestThroughputStrategy()` with no backend list compares `default_backends(device)`: XPU graph, Torch Inductor, OpenVINO, and eager on XPU, and the same without XPU graph on CPU. +When different backends win at different batch sizes, it returns a `BatchSizeRouterBackend` that sends each call to the backend tuned for its batch size. ```python from torch_tweak.torch.tune_strategy import HighestThroughputStrategy diff --git a/properdocs.yml b/properdocs.yml index fe1998df..9f2b5621 100644 --- a/properdocs.yml +++ b/properdocs.yml @@ -117,6 +117,7 @@ nav: - Torch Inductor Backend: api/backends/torch_inductor_backend.md - XPU Graph Backend: api/backends/xpu_graph_backend.md - Torch Eager Backend: api/backends/torch_eager_backend.md + - Batch Size Router Backend: api/backends/batch_size_router_backend.md - Tuning Strategies API: - FirstWinsStrategy: api/tune_strategies/first_wins_strategy.md - OneBackendStrategy: api/tune_strategies/one_backend_strategy.md diff --git a/tests/unit/torch/backend/test_batch_size_router_backend.py b/tests/unit/torch/backend/test_batch_size_router_backend.py new file mode 100644 index 00000000..599b1b96 --- /dev/null +++ b/tests/unit/torch/backend/test_batch_size_router_backend.py @@ -0,0 +1,151 @@ +# Copyright (c) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import time +from pathlib import Path +from typing import Any + +import pytest +import torch + +from tests.toy_models.torch_models import ToyTorchModel +from torch_tweak.torch.backend import Backend, BatchSizeRouterBackend +from torch_tweak.torch.backend.backend import BackendState +from torch_tweak.torch.task.profiling import NumStepsMeasuringStopStrategy +from torch_tweak.torch.task.profiling.profiling_stop_strategy import AllSamplesProfilingStopStrategy +from torch_tweak.torch.tune_strategy.highest_throughput_strategy import HighestThroughputStrategy + + +class TimedBackend(Backend): + """Sleeps per_call plus per_row times the batch size, then runs the module.""" + + requires_original_module = True + + def __init__(self, label: str = "", per_call: float = 0.0, per_row: float = 0.0): + super().__init__() + self.label, self.per_call, self.per_row = label, per_call, per_row + self.calls: list[int] = [] + self._module: Any = None + + def key(self) -> str: + return f"TimedBackend_{self.label}" + + def describe(self) -> str: + return f"TimedBackend({self.label})" + + def _build(self, module: Any, graph_spec: Any, data: list[Any], cache_dir: Path) -> Backend: + self._module = module + return self + + def _activate(self): + pass + + def _deactivate(self): + pass + + def _deploy(self): + pass + + def _infer(self, *args, **kwargs): + batch_size = args[0].shape[0] + self.calls.append(batch_size) + time.sleep(self.per_call + self.per_row * batch_size) + return self._module(*args, **kwargs) + + def to_dict(self): + return {"type": "TimedBackend", "label": self.label} + + @classmethod + def from_dict(cls, module, state_dict: dict): + backend = cls(label=state_dict["label"]) + backend._module = module + backend.state = BackendState.CHECKPOINT_LOADED + return backend + + +def built(backend: TimedBackend, module, torch_device, tmp_path) -> TimedBackend: + sample = ((torch.rand(1, 4, device=torch_device),), {}) + backend.build(module, None, [sample], torch_device, tmp_path) + return backend + + +@pytest.fixture +def model(torch_device): + return ToyTorchModel().eval().to(torch_device) + + +@pytest.fixture +def graph_spec(model, torch_device): + return model.graph_spec(batch_sizes=[1, 16], device=torch_device) + + +def test_routes_send_each_batch_size_to_its_backend(model, graph_spec, torch_device, tmp_path): + small = built(TimedBackend("small"), model, torch_device, tmp_path) + large = built(TimedBackend("large"), model, torch_device, tmp_path) + router = BatchSizeRouterBackend([(4, small), (None, large)], graph_spec.input_spec) + + assert router.is_active + for batch_size in (1, 4, 5, 64): + router.infer(model.samples([batch_size], device=torch_device)[0][0][0]) + assert small.calls == [1, 4] + assert large.calls == [5, 64] + assert router.describe() == "BatchSizeRouterBackend(1-4: TimedBackend(small), 5+: TimedBackend(large))" + assert router.requires_original_module + + +def test_last_route_must_take_every_larger_batch_size(model, graph_spec, torch_device, tmp_path): + backend = built(TimedBackend(), model, torch_device, tmp_path) + with pytest.raises(ValueError, match="every larger batch size"): + BatchSizeRouterBackend([(4, backend)], graph_spec.input_spec) + + +def test_serialization(model, graph_spec, torch_device, tmp_path): + small = built(TimedBackend("small"), model, torch_device, tmp_path) + large = built(TimedBackend("large"), model, torch_device, tmp_path) + router = BatchSizeRouterBackend([(4, small), (None, large)], graph_spec.input_spec) + + loaded = BatchSizeRouterBackend.from_dict(model, router.to_dict()) + + assert loaded.state == BackendState.CHECKPOINT_LOADED + assert loaded.describe() == router.describe() + loaded.deploy(torch_device) + sample = model.samples([2], device=torch_device)[0][0][0] + torch.testing.assert_close(loaded.infer(sample), model(sample)) + assert loaded.backend_for(2).calls == [2] + + +def tune(backends, model, graph_spec, torch_device, tmp_path) -> Backend: + strategy = HighestThroughputStrategy( + backends=backends, + measurement_stop_strategy=NumStepsMeasuringStopStrategy(num_steps=3), + profiling_stop_strategy=AllSamplesProfilingStopStrategy(), + ).enable_find_max_batch_size(False) + strategy.enable_correctness_check(False) + sample = ((model.sample(device=torch_device).unsqueeze(0),), {}) + return strategy.tune(model, "test", graph_spec, [sample], torch_device, tmp_path) + + +def test_strategy_routes_small_batches_to_the_backend_faster_there(model, graph_spec, torch_device, tmp_path): + # Per call: per_row 1 ms per row against a fixed 6 ms. Rows are faster up to batch size 4. + rows = TimedBackend("rows", per_row=1e-3) + fixed = TimedBackend("fixed", per_call=6e-3) + + tuned = tune([rows, fixed], model, graph_spec, torch_device, tmp_path) + + assert isinstance(tuned, BatchSizeRouterBackend) + assert [(largest, backend.describe()) for largest, backend in tuned.routes] == [ + (4, "TimedBackend(rows)"), + (None, "TimedBackend(fixed)"), + ] + + +def test_strategy_keeps_one_backend_when_it_wins_every_batch_size(model, graph_spec, torch_device, tmp_path): + tuned = tune( + [TimedBackend("slow", per_call=5e-3), TimedBackend("fast", per_call=1e-3)], + model, + graph_spec, + torch_device, + tmp_path, + ) + + assert tuned.describe() == "TimedBackend(fast)" diff --git a/torch_tweak/torch/backend/__init__.py b/torch_tweak/torch/backend/__init__.py index 1786a626..7faab28e 100644 --- a/torch_tweak/torch/backend/__init__.py +++ b/torch_tweak/torch/backend/__init__.py @@ -6,6 +6,7 @@ """Torch backend module.""" from torch_tweak.torch.backend.backend import Backend +from torch_tweak.torch.backend.batch_size_router_backend import BatchSizeRouterBackend from torch_tweak.torch.backend.openvino import OpenVINOBackend, OpenVINOBackendConfig from torch_tweak.torch.backend.torch_eager import TorchEagerBackend, TorchEagerBackendConfig from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend, TorchInductorBackendConfig @@ -14,6 +15,7 @@ __all__ = [ "Backend", + "BatchSizeRouterBackend", "OpenVINOBackend", "OpenVINOBackendConfig", "TorchInductorBackend", diff --git a/torch_tweak/torch/backend/batch_size_router_backend.py b/torch_tweak/torch/backend/batch_size_router_backend.py new file mode 100644 index 00000000..49b3411c --- /dev/null +++ b/torch_tweak/torch/backend/batch_size_router_backend.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Backend that sends each call to the backend tuned for its batch size.""" + +from pathlib import Path +from typing import Any, Self + +import torch +import torch.nn as nn + +from torch_tweak.torch.backend.backend import Backend, BackendState +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.utils.trace import annotate + +# The largest batch size of a route, or None for every larger batch size. +Route = tuple[int | None, Backend] + + +class BatchSizeRouterBackend(Backend): + """Holds several built backends and sends each call to the one tuned for its batch size. + + HighestThroughputStrategy creates it when different backends win at different batch sizes. + routes are in ascending order. A call goes to the first route that holds its batch size. + The last route takes every larger batch size. + """ + + STATE_TYPE = "type" + STATE_ROUTES = "routes" + STATE_INPUT_SPEC = "input_spec" + + def __init__(self, routes: list[Route], input_spec: SampleMetadata): + """Initializes backend. + + Args: + routes: (largest batch size, backend) pairs in ascending order. The last bound is None. + input_spec: Input spec of the graph spec, used to read the batch size of a call. + """ + super().__init__() + if not routes or routes[-1][0] is not None: + raise ValueError("The last route must take every larger batch size") + self._routes = routes + self._input_spec = input_spec + self._device = routes[0][1].device + if all(backend.is_active for _, backend in routes): + self.state = BackendState.ACTIVE + + @property + def routes(self) -> list[Route]: + """(largest batch size, backend) pairs in ascending order.""" + return self._routes + + @property + def requires_original_module(self) -> bool: # type: ignore[override] + """True when any routed backend runs the original module.""" + return any(backend.requires_original_module for _, backend in self._routes) + + def key(self) -> str: + """Returns the key of the backend.""" + return "_".join([self.__class__.__name__, *(backend.key() for _, backend in self._routes)]) + + def describe(self) -> str: + """Returns the description of the backend.""" + parts = [] + smallest = 1 + for largest, backend in self._routes: + sizes = f"{smallest}+" if largest is None else f"{smallest}-{largest}" + parts.append(f"{sizes}: {backend.describe()}") + smallest = (largest or 0) + 1 + return f"{self.__class__.__name__}({', '.join(parts)})" + + def backend_for(self, batch_size: int) -> Backend: + """Returns the backend that runs batch_size.""" + for largest, backend in self._routes: + if largest is None or batch_size <= largest: + return backend + raise AssertionError("The last route takes every batch size") + + def _batch_size_of(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> int | None: + """Reads the batch size from the first input tensor with a batch axis.""" + for locator, tensor_spec in self._input_spec.tensor_data: + multipliers = tensor_spec.get_batch_axis_multipliers() + if not multipliers: + continue + value = locator.get_value(args if tensor_spec.name.startswith("args") else kwargs) + axis, multiplier = next(iter(multipliers.items())) + return value.shape[axis] // multiplier + return None + + def _build(self, module: nn.Module, graph_spec: GraphSpec, data: list[Sample], cache_dir: Path) -> Backend: + raise RuntimeError(f"{self.name} is created from backends that are already built") + + def _activate(self) -> None: + for _, backend in self._routes: + backend.activate() + + def _deactivate(self) -> None: + for _, backend in self._routes: + backend.deactivate() + + def _deploy(self) -> None: + for _, backend in self._routes: + backend.deploy(self._device) + + @annotate(message="BatchSizeRouterBackend.infer", domain="Torch Tweak") + def _infer(self, *args: Any, **kwargs: Any) -> Any: + batch_size = self._batch_size_of(args, kwargs) + backend = self._routes[-1][1] if batch_size is None else self.backend_for(batch_size) + return backend.infer(*args, **kwargs) + + def to_dict(self) -> dict[str, Any]: + """Returns the state_dict of the backend.""" + return { + self.STATE_TYPE: self.__class__.__name__, + self.STATE_ROUTES: [(largest, backend.to_dict()) for largest, backend in self._routes], + self.STATE_INPUT_SPEC: self._input_spec.to_dict(), + } + + @classmethod + def from_dict(cls, module: torch.nn.Module | None, state_dict: dict[str, Any]) -> Self: + """Creates a backend from a state_dict.""" + if state_dict.get(cls.STATE_TYPE) != cls.__name__: + raise ValueError(f"Invalid state_dict type: {state_dict.get(cls.STATE_TYPE)}") + backend_classes = {backend_class.__name__: backend_class for backend_class in Backend.__subclasses__()} + routes: list[Route] = [] + for largest, backend_state in state_dict[cls.STATE_ROUTES]: + backend_type = backend_state[cls.STATE_TYPE] + if backend_type not in backend_classes: + raise ValueError(f"Unknown backend type: {backend_type}") + routes.append((largest, backend_classes[backend_type].from_dict(module, backend_state))) + backend = cls(routes, SampleMetadata.from_dict(state_dict[cls.STATE_INPUT_SPEC])) + backend.state = BackendState.CHECKPOINT_LOADED + return backend diff --git a/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py b/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py index 8a965964..17d45d0e 100644 --- a/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py +++ b/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py @@ -9,6 +9,8 @@ 2. Runs all backends with given max batch size. 3. Finds the highest throughput for each backend. 4. Returns the backend with the highest throughput. + When another backend is faster for some batch sizes, returns a BatchSizeRouterBackend + that sends those batch sizes to it. """ from collections.abc import Iterable @@ -20,7 +22,8 @@ import torch import torch.nn as nn -from torch_tweak.torch.backend import Backend +from torch_tweak.torch.backend import Backend, BatchSizeRouterBackend +from torch_tweak.torch.backend.batch_size_router_backend import Route from torch_tweak.torch.config import ( DEFAULT_STABILITY_PERCENTAGE, DEFAULT_THROUGHPUT_BACKOFF_LIMIT, @@ -29,7 +32,10 @@ ) from torch_tweak.torch.module.graph_spec import GraphSpec from torch_tweak.torch.module.recording_module import Sample -from torch_tweak.torch.task.find_max_batch_size import calculate_highest_throughput_for_backend +from torch_tweak.torch.task.find_max_batch_size import ( + calculate_highest_throughput_for_backend, + get_throughput_per_batch_size, +) from torch_tweak.torch.task.profiling import ( MeasuringStopStrategy, ModelExecutionTimeMeasuringStrategy, @@ -39,6 +45,7 @@ StableWindowMeasuringStopStrategy, ThroughputSaturatedProfilingStopStrategy, ) +from torch_tweak.torch.task.profiling.events import get_inference_events from torch_tweak.torch.tune_strategy.default_backends import default_backends from torch_tweak.torch.tune_strategy.extension import TuneStrategyFindMaxBatchSizeExtension from torch_tweak.utils.logging import control_output, log @@ -123,6 +130,7 @@ def _tune( best_backend = None best_throughput = float("-inf") best_batch_size = 1 + throughputs: list[tuple[Backend, dict[int, float]]] = [] backends = list(self._backends or default_backends(device)) @@ -150,6 +158,10 @@ def _tune( ) ) measurements += results.entries + per_batch_size = get_throughput_per_batch_size( + get_inference_events(results.entries), self._measurement_stop_strategy + ) + throughputs.append((backend, dict(per_batch_size))) log( "✅ backend profiled - throughput: %.2f samples/s, batch size: %s", throughput, @@ -193,6 +205,9 @@ def _tune( measurements=measurements, ) ) + routes = _routes_by_batch_size(throughputs, best_backend) + if len(routes) > 1: + best_backend = BatchSizeRouterBackend(routes, graph_spec.input_spec) best_backend.activate() log("🎯 Strategy %s execution finished:", self.__class__.__name__, sink=self._sink) log( @@ -227,3 +242,32 @@ def _describe_parts(self) -> list[str]: "backends:", *[f" {backend}" for backend in backends], ] + + +# Another backend takes a batch size from the best backend only when it is this much faster. +# The margin keeps measurement noise from splitting the routes. +ROUTE_SPEEDUP = 1.1 + + +def _routes_by_batch_size(throughputs: list[tuple[Backend, dict[int, float]]], best: Backend) -> list[Route]: + """Returns (largest batch size, backend) routes. A single route means best runs every batch size. + + throughputs holds samples per second by batch size for every profiled backend. + A batch size between two profiled ones goes to the route of the larger one. + """ + best_per_batch_size = next(per_batch_size for backend, per_batch_size in throughputs if backend is best) + routes: list[Route] = [] + for batch_size in sorted(best_per_batch_size): + winner, top = best, best_per_batch_size[batch_size] + for backend, per_batch_size in throughputs: + throughput = per_batch_size.get(batch_size, 0.0) + if throughput > top and throughput > ROUTE_SPEEDUP * best_per_batch_size[batch_size]: + winner, top = backend, throughput + if routes and routes[-1][1] is winner: + routes[-1] = (batch_size, winner) + else: + routes.append((batch_size, winner)) + if not routes: + return [(None, best)] + routes[-1] = (None, routes[-1][1]) + return routes From d5d18a84b4bd07af3c5b2df29e594370061fbb72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Fri, 23 Oct 2026 18:07:33 +0200 Subject: [PATCH 5/5] feat(xpu_graph): record Torch Inductor kernels with compile_with_inductor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit XPUGraphBackendConfig(compile_with_inductor=True) compiles the module with Torch Inductor and records the compiled kernels in the XPU graphs. Replay runs the fused kernels without a Python launch per kernel. It is the XPU counterpart of torch.compile(mode="reduce-overhead") on CUDA. Compilation runs in the warmup before each recording. Inductor compiles three times for all recorded sizes. Signed-off-by: Łukasz Ślusarczyk --- docs/guides/backends/xpu_graph_backend.md | 16 ++++++++++++++++ .../torch/backend/test_xpu_graph_backend.py | 19 +++++++++++++++++++ .../torch/backend/xpu_graph_backend.py | 14 ++++++++++++-- 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/docs/guides/backends/xpu_graph_backend.md b/docs/guides/backends/xpu_graph_backend.md index 734ace0c..25a06322 100644 --- a/docs/guides/backends/xpu_graph_backend.md +++ b/docs/guides/backends/xpu_graph_backend.md @@ -46,6 +46,22 @@ from torch_tweak.torch.backend import XPUGraphBackend, XPUGraphBackendConfig backend = XPUGraphBackend(XPUGraphBackendConfig(max_recorded_batch_size=64)) ``` +## Inductor Kernels + +`compile_with_inductor=True` records the kernels Torch Inductor generates instead of the eager kernels. +Replay then runs the fused Inductor kernels without launching each one from Python. +On CUDA, `torch.compile(mode="reduce-overhead")` does the same with CUDA graphs. That mode has no effect on XPU. + +Compilation runs during the warmup before each recording. +Inductor compiles about three times for all recorded sizes: once for the largest size, once with a dynamic batch axis, and once for batch size 1. +Replay is still checked against the eager module. + +```python +from torch_tweak.torch.backend import XPUGraphBackend, XPUGraphBackendConfig + +backend = XPUGraphBackend(XPUGraphBackendConfig(compile_with_inductor=True)) +``` + ## Quick Start ```python diff --git a/tests/unit/torch/backend/test_xpu_graph_backend.py b/tests/unit/torch/backend/test_xpu_graph_backend.py index fba22c01..0e90a3f7 100644 --- a/tests/unit/torch/backend/test_xpu_graph_backend.py +++ b/tests/unit/torch/backend/test_xpu_graph_backend.py @@ -392,6 +392,25 @@ def test_dynamic_axis_is_rejected(torch_device, tmp_path): XPUGraphBackend().build(module, spec, [short, long], torch_device, tmp_path) +@requires_xpu +def test_compile_with_inductor_records_the_compiled_kernels(torch_device, tmp_path): + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + module = LinearRelu() + config = XPUGraphBackendConfig(compile_with_inductor=True) + backend = batched_build(module, torch_device, tmp_path, largest=16, config=config) + + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] > 0 + assert backend.recorded_batch_sizes == [1, 2, 4, 8, 16] + assert backend.describe() == "XPUGraphBackend(compile_with_inductor=True)" + # Replay runs the recorded kernels without calling the compiled module. + backend._recorded_module = None + with torch.no_grad(): + for batch_size in (1, 3, 16, 20): + x = torch.randn(batch_size, 16, device=torch_device) + torch.testing.assert_close(backend.infer(x), module(x), rtol=1e-4, atol=1e-4) + + @requires_xpu def test_batch_axis_serialization(torch_device, tmp_path): module = LinearRelu() diff --git a/torch_tweak/torch/backend/xpu_graph_backend.py b/torch_tweak/torch/backend/xpu_graph_backend.py index 6eb67640..0d4e60af 100644 --- a/torch_tweak/torch/backend/xpu_graph_backend.py +++ b/torch_tweak/torch/backend/xpu_graph_backend.py @@ -36,11 +36,14 @@ class XPUGraphBackendConfig(BackendConfig): max_recorded_batch_size: The largest batch size recorded for a graph spec with a batch axis. Graphs are recorded for batch sizes 1, 2, 4, and every multiple of 8 up to this size. A larger batch replays in chunks of this size. + compile_with_inductor: If True, record the kernels Torch Inductor generates for the module + instead of the eager kernels. Compilation runs during the warmup before each recording. """ autocast_enabled: bool = False autocast_dtype: torch.dtype | None = None max_recorded_batch_size: int = 256 + compile_with_inductor: bool = False @classmethod def from_dict(cls, state_dict: dict[str, Any]) -> Self: @@ -108,6 +111,8 @@ def __init__(self, config: XPUGraphBackendConfig | None = None): self._layout: _Layout | None = None # Keyed by batch size. A graph spec without a batch axis has one recording under key 1. self._recordings: dict[int, _Recording] = {} + # The callable that warmup and recording run: the original module or its Inductor compilation. + self._recorded_module: Any = None self._pool: Any = None self._warned_eager = False @@ -277,6 +282,9 @@ def _record(self) -> None: self._orig_module.to(self._device) try: self._layout = self._detect_layout() + self._recorded_module = ( + torch.compile(self._orig_module) if self._config.compile_with_inductor else self._orig_module + ) self._pool = torch.xpu.graph_pool_handle() sizes = self._batch_sizes_to_record() if self._layout.has_batch_axis else [1] buffer = tree_flatten(_clone_tree(self._sample_of_size(max(sizes))))[0] @@ -354,13 +362,14 @@ def _record_size(self, size: int, largest: int, buffer: list[Any]) -> _Recording with self._autocast(), torch.no_grad(): # Warmup. The first oneDNN GEMM for a shape selects an algorithm and allocates its workspace. # That work stays outside the recording. - self._orig_module(*_clone_tree(static_args), **_clone_tree(static_kwargs)) + # With compile_with_inductor, this call also compiles for the shape. + self._recorded_module(*_clone_tree(static_args), **_clone_tree(static_kwargs)) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") # Records this call on default_capture_stream. # Exiting the block ends the recording. with torch.xpu.graph(graph, pool=self._pool): - captured = self._orig_module(*static_args, **static_kwargs) + captured = self._recorded_module(*static_args, **static_kwargs) if any("XPU Graph is empty" in str(item.message) for item in caught): raise RuntimeError("XPU graph capture recorded no device work") except Exception: @@ -457,6 +466,7 @@ def _release_graphs(self) -> None: for recording in self._recordings.values(): recording.graph.reset() self._recordings = {} + self._recorded_module = None self._layout = None self._pool = None