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 caa9a231..a3df4d24 100644 --- a/docs/guides/backends/xpu_graph_backend.md +++ b/docs/guides/backends/xpu_graph_backend.md @@ -52,6 +52,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 @@ -65,6 +81,8 @@ You do not have to add this backend yourself. On XPU, `tt.default_strategy(device)`, `FirstWinsStrategy()`, and `HighestThroughputStrategy()` include it when you pass no backends. `FirstWinsStrategy()` tries it first. On CPU none of them use it. +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 dabcfc42..5df36c33 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 49dd84bb..05f531a5 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 a6110c05..f2204e82 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..c28e8b75 --- /dev/null +++ b/tests/unit/torch/backend/test_batch_size_router_backend.py @@ -0,0 +1,161 @@ +# 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.backend.batch_size_router_backend import routes_by_throughput +from torch_tweak.torch.checkpoint.storage_tasks import TorchLoadTask, TorchSaveTask +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 ConstantBackend(Backend): + """Stub. Returns one row of value per input row, so the output names the backend that ran. + + It sleeps per_call plus per_row times the batch size, so profiling sees a cost model. + """ + + def __init__(self, value: float = 0.0, per_call: float = 0.0, per_row: float = 0.0): + super().__init__() + self.value, self.per_call, self.per_row = value, per_call, per_row + + def key(self) -> str: + return f"ConstantBackend_{self.value}" + + def describe(self) -> str: + return f"ConstantBackend({self.value})" + + def _build(self, module: Any, graph_spec: Any, data: list[Any], cache_dir: Path) -> Backend: + return self + + def _activate(self): + pass + + def _deactivate(self): + pass + + def _deploy(self): + pass + + def _infer(self, x): + time.sleep(self.per_call + self.per_row * x.shape[0]) + return torch.full((x.shape[0], 1), self.value, device=x.device) + + def to_dict(self): + return {"type": "ConstantBackend", "value": self.value} + + @classmethod + def from_dict(cls, module, state_dict: dict): + backend = cls(state_dict["value"]) + backend.state = BackendState.CHECKPOINT_LOADED + return backend + + +def built(backend: ConstantBackend) -> ConstantBackend: + backend.build(None, None, [((torch.zeros(1, 4),), {})], torch.device("cpu"), Path()) + return backend + + +@pytest.fixture +def input_spec(): + return ToyTorchModel().graph_spec(batch_sizes=[1, 16]).input_spec + + +def router(input_spec) -> BatchSizeRouterBackend: + """Batch sizes 1 to 4 return 1.0. Larger batch sizes return 2.0.""" + return BatchSizeRouterBackend([(4, built(ConstantBackend(1.0))), (None, built(ConstantBackend(2.0)))], input_spec) + + +@pytest.mark.parametrize(("batch_size", "value"), [(1, 1.0), (4, 1.0), (5, 2.0), (64, 2.0)]) +def test_each_batch_size_runs_on_its_route(input_spec, batch_size, value): + backend = router(input_spec) + + actual = backend.infer(torch.zeros(batch_size, 4)) + + torch.testing.assert_close(actual, torch.full((batch_size, 1), value)) + + +def test_describe_lists_the_batch_sizes_of_every_route(input_spec): + assert router(input_spec).describe() == ( + "BatchSizeRouterBackend(1-4: ConstantBackend(1.0), 5+: ConstantBackend(2.0))" + ) + + +def test_last_route_must_take_every_larger_batch_size(input_spec): + with pytest.raises(ValueError, match="every larger batch size"): + BatchSizeRouterBackend([(4, built(ConstantBackend()))], input_spec) + + +def test_checkpoint_round_trip_keeps_the_routes(input_spec, tmp_path): + TorchSaveTask().save(tmp_path, router(input_spec).to_dict()) + + loaded = BatchSizeRouterBackend.from_dict(None, TorchLoadTask().load(tmp_path)) + loaded.deploy(torch.device("cpu")) + actual = loaded.infer(torch.zeros(5, 4)) + + torch.testing.assert_close(actual, torch.full((5, 1), 2.0)) + + +def test_router_needs_the_original_module_when_a_route_needs_it(input_spec): + class NeedsModule(ConstantBackend): + requires_original_module = True + + backend = BatchSizeRouterBackend([(4, built(ConstantBackend())), (None, built(NeedsModule()))], input_spec) + + assert backend.requires_original_module is True + + +def test_one_backend_faster_at_every_batch_size_is_one_route(): + fast, slow = ConstantBackend(1.0), ConstantBackend(2.0) + + routes = routes_by_throughput([(fast, {1: 200.0, 8: 1600.0}), (slow, {1: 100.0, 8: 800.0})], fast) + + assert routes == [(None, fast)] + + +def test_batch_sizes_where_another_backend_is_faster_get_their_own_route(): + best, small = ConstantBackend(1.0), ConstantBackend(2.0) + throughputs = [ + (best, {1: 100.0, 2: 200.0, 4: 400.0, 8: 800.0}), + (small, {1: 500.0, 2: 600.0, 4: 420.0, 8: 100.0}), + ] + + routes = routes_by_throughput(throughputs, best) + + assert routes == [(2, small), (None, best)] + + +def test_a_backend_within_the_noise_margin_does_not_get_a_route(): + best, other = ConstantBackend(1.0), ConstantBackend(2.0) + + routes = routes_by_throughput([(best, {1: 100.0}), (other, {1: 109.0})], best) + + assert routes == [(None, best)] + + +def test_strategy_routes_small_batches_to_the_backend_faster_there(torch_device, tmp_path): + # 1 ms per row against a fixed 6 ms per call: per row is faster up to batch size 4. + per_row = ConstantBackend(1.0, per_row=1e-3) + per_call = ConstantBackend(2.0, per_call=6e-3) + model = ToyTorchModel().eval().to(torch_device) + strategy = HighestThroughputStrategy( + backends=[per_row, per_call], + 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),), {}) + graph_spec = model.graph_spec(batch_sizes=[1, 16], device=torch_device) + + tuned = strategy.tune(model, "test", graph_spec, [sample], torch_device, tmp_path) + + assert tuned.describe() == "BatchSizeRouterBackend(1-4: ConstantBackend(1.0), 5+: ConstantBackend(2.0))" diff --git a/tests/unit/torch/backend/test_xpu_graph_backend.py b/tests/unit/torch/backend/test_xpu_graph_backend.py index 3ede60cf..bf962676 100644 --- a/tests/unit/torch/backend/test_xpu_graph_backend.py +++ b/tests/unit/torch/backend/test_xpu_graph_backend.py @@ -537,6 +537,30 @@ def test_dynamic_axis_is_rejected(model, torch_device, tmp_path): XPUGraphBackend().build(model, spec, [short, long], torch_device, tmp_path) +@requires_xpu +def test_compile_with_inductor_matches_eager(model, torch_device, tmp_path): + x = rows(3, 16, device=torch_device) + expected = eager(model, x) + config = XPUGraphBackendConfig(compile_with_inductor=True) + backend = batched_build(model, torch_device, tmp_path, largest=16, config=config) + + actual = backend.infer(x) + + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + + +@requires_xpu +def test_compile_with_inductor_compiles_three_times_for_every_size(model, torch_device, tmp_path): + # The largest size, a dynamic batch axis for the other sizes, and batch size 1. + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + config = XPUGraphBackendConfig(compile_with_inductor=True) + + batched_build(model, torch_device, tmp_path, largest=16, config=config) + + assert torch._dynamo.utils.counters["stats"]["unique_graphs"] == 3 + + @requires_xpu def test_batch_axis_checkpoint_round_trip(model, torch_device, tmp_path): x = rows(6, 16, device=torch_device) diff --git a/tests/unit/torch/task/test_profiling.py b/tests/unit/torch/task/test_profiling.py index 1fb2a2c2..495e8e71 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/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..07cf193c --- /dev/null +++ b/torch_tweak/torch/backend/batch_size_router_backend.py @@ -0,0 +1,165 @@ +# 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] + +# 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_throughput(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, best included. + 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 + + +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 + states = {backend.state for _, backend in routes} + if states == {BackendState.ACTIVE}: + self.state = BackendState.ACTIVE + elif states <= {BackendState.ACTIVE, BackendState.INACTIVE}: + self.state = BackendState.INACTIVE + + @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/backend/xpu_graph_backend.py b/torch_tweak/torch/backend/xpu_graph_backend.py index 7d0ddd1c..e5083a2c 100644 --- a/torch_tweak/torch/backend/xpu_graph_backend.py +++ b/torch_tweak/torch/backend/xpu_graph_backend.py @@ -38,11 +38,14 @@ class XPUGraphBackendConfig(BackendConfig): Graphs are recorded for 1, 2, 4, every multiple of 8 below the limit, and the limit itself. The limit is this value or the largest batch size seen during tuning, whichever is smaller. A larger batch replays in chunks of the largest recorded 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 def __post_init__(self) -> None: """Rejects a max_recorded_batch_size below 1.""" @@ -115,6 +118,8 @@ def __init__(self, config: XPUGraphBackendConfig | None = None): self._layout: _Layout | None = None # Keyed by batch size in ascending order. 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 @@ -256,11 +261,15 @@ def _infer_eager(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: 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.""" + """Runs the recorded callable without a graph on a copy of sample, so in-place writes do not change sample. + + Before recording starts, that callable is the original module. + """ assert self._orig_module is not None + module = self._orig_module if self._recorded_module is None else self._recorded_module args, kwargs = _clone_tree(sample) with self._autocast(), torch.no_grad(): - return _clone_tree(self._orig_module(*args, **kwargs)) + return _clone_tree(module(*args, **kwargs)) def _record(self) -> None: """Records every batch size, then rejects the backend when replay disagrees with eager.""" @@ -269,6 +278,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] largest = sizes[-1] @@ -352,13 +364,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") captured_leaves, captured_spec = tree_flatten(captured) @@ -460,6 +473,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 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) diff --git a/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py b/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py index 4b46596f..cdb73fc6 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 routes_by_throughput 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, @@ -194,6 +206,9 @@ def _tune( measurements=measurements, ) ) + routes = routes_by_throughput(throughputs, best_backend) + if len(routes) > 1: + best_backend = BatchSizeRouterBackend(routes, graph_spec.input_spec) # A later build can move `module`. # `OpenVINOBackend._build` offloads it to CPU. # `activate` rebuilds the winner on the weights `module` holds now.