diff --git a/.claude/skills/torch-tweak-assistant/SKILL.md b/.claude/skills/torch-tweak-assistant/SKILL.md index aeb4ddce..a013bb78 100644 --- a/.claude/skills/torch-tweak-assistant/SKILL.md +++ b/.claude/skills/torch-tweak-assistant/SKILL.md @@ -65,15 +65,16 @@ Do NOT activate this skill when: ### TorchAO Backend - **Best for**: PyTorch-native quantization on Intel GPU / XPU -- **Simple**: `TorchAOBackend()` with no config needed ### Torch Inductor Backend - **Best for**: PyTorch compiler-based optimization on Intel GPU / XPU -- **Simple**: `TorchInductorBackend()` with no config needed + +### XPU Graph Backend +- **Best for**: replaying a module that was recorded once during tuning +- **Rejection**: failed, empty, or incorrect recordings are dropped for that module and input ### OpenVINO Backend - **Best for**: an OpenVINO model on an Intel CPU or GPU -- **Simple**: `OpenVINOBackend()` with no config needed - **Strengths**: - Converts PyTorch models to OpenVINO IR for deployment - Applies hardware-specific optimizations @@ -82,7 +83,6 @@ Do NOT activate this skill when: ### Torch Eager Backend - **Best for**: the uncompiled baseline inside a throughput search -- **Simple**: `TorchEagerBackend()` with no config needed ## Tuning Strategies @@ -93,8 +93,8 @@ Do NOT activate this skill when: ## Best Practices ### Backend Selection -- **HighestThroughputStrategy**: TorchAO, Torch Inductor, OpenVINO, and eager. It measures each one and keeps the fastest. List order does not change the winner. -- **FirstWinsStrategy**: default order is Torch Inductor, OpenVINO, then eager. The first backend that builds is kept. TorchAO is not in that default, because the first backend that builds wins and a quantization that builds would be kept without a speed comparison. +- **HighestThroughputStrategy**: the default list is Torch Inductor, XPU graph, and eager. It measures each one and keeps the fastest. List order does not change the winner. TorchAO and OpenVINO are optional additions. +- **FirstWinsStrategy**: default order is XPU graph, Torch Inductor, OpenVINO, then eager. The first backend that builds is kept. TorchAO is not in that default, because the first backend that builds wins and a quantization that builds would be kept without a speed comparison. ### Common Issues - **Graph breaks**: conditional logic in a module blocks tuning of that module @@ -124,6 +124,7 @@ tt.load(model, "tuned_model.tt") from torch_tweak.torch.backend import ( TorchAOBackend, TorchInductorBackend, + XPUGraphBackend, OpenVINOBackend, TorchEagerBackend, ) @@ -133,6 +134,7 @@ strategy = HighestThroughputStrategy( backends=[ TorchAOBackend(), TorchInductorBackend(), + XPUGraphBackend(), OpenVINOBackend(), TorchEagerBackend(), ] diff --git a/README.md b/README.md index 198ffdab..2cb5e5ba 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ This page is the landing summary. The same file is the PyPI long description. | Feature | Description | |-----------------------------|---------------------------------------------------------------------------------------------------------------------------| | Ease-of-use | Single line of code to run all possible tuning paths directly from your source code | -| Wide Backend Support | Compatible with TorchAO, Torch Inductor, OpenVINO, and eager PyTorch | +| Wide Backend Support | Compatible with TorchAO, Torch Inductor, XPU graph, OpenVINO, and eager PyTorch | | Model Tuning | Enhance the performance of models such as ResNET and BERT for efficient inference deployment | | Pipeline Tuning | Streamline Python code pipelines for models such as Stable Diffusion and Flux using seamless model wrapping and tuning | | Correctness Testing | Ensures tuned models produce correct outputs by validating on provided data samples | @@ -55,6 +55,7 @@ use frameworks like vLLM or SGLang for best performance. Use Torch Tweak for gen from torch_tweak.torch.backend import ( TorchAOBackend, TorchInductorBackend, + XPUGraphBackend, OpenVINOBackend, TorchEagerBackend, ) @@ -65,13 +66,14 @@ strategy = HighestThroughputStrategy( backends=[ TorchAOBackend(), TorchInductorBackend(), + XPUGraphBackend(), OpenVINOBackend(), TorchEagerBackend(), ] ) ``` -TorchAO, Torch Inductor, and OpenVINO configuration is in the [backends guide](docs/learn/backends.md). +TorchAO, Torch Inductor, XPU graph, and OpenVINO configuration is in the [backends guide](docs/learn/backends.md). Torch Eager accepts an optional autocast dtype and otherwise runs the module unchanged. Tune strategies are in [core functionalities](docs/learn/core_functionalities.md). diff --git a/docs/api/backends/xpu_graph_backend.md b/docs/api/backends/xpu_graph_backend.md new file mode 100644 index 00000000..0761af62 --- /dev/null +++ b/docs/api/backends/xpu_graph_backend.md @@ -0,0 +1,20 @@ + + +# XPU Graph Backend API + +## XPUGraphBackend + +::: torch_tweak.torch.backend.XPUGraphBackend + options: + show_root_heading: true + heading_level: 3 + +## XPUGraphBackendConfig + +::: torch_tweak.torch.backend.XPUGraphBackendConfig + 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 new file mode 100644 index 00000000..d3b40d04 --- /dev/null +++ b/docs/guides/backends/xpu_graph_backend.md @@ -0,0 +1,37 @@ + + +# XPU Graph Backend Guide + +The XPU graph backend records the module on an Intel GPU once during tuning and replays that recording on later inference calls. +Recording is not repeated for each inference call. + +## Overview + +- **One 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 + +## Quick Start + +```python +from torch_tweak.torch.backend import XPUGraphBackend +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. + +PyTorch records a native SYCL graph when it was built with a SYCL compiler of 2026.1 or newer, except on PVC. +An older build records the legacy SYCL command graph. +This backend uses the recording that build already performs. + +## Next Steps + +- API reference: [XPU Graph Backend](../../api/backends/xpu_graph_backend.md) diff --git a/docs/guides/tune_strategies/tune_strategies.md b/docs/guides/tune_strategies/tune_strategies.md index ed41af1d..9215c13f 100644 --- a/docs/guides/tune_strategies/tune_strategies.md +++ b/docs/guides/tune_strategies/tune_strategies.md @@ -81,7 +81,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 Torch Inductor, then 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. 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. ### Usage diff --git a/docs/guides/tuning.md b/docs/guides/tuning.md index 63c246a8..6be1f9de 100644 --- a/docs/guides/tuning.md +++ b/docs/guides/tuning.md @@ -268,6 +268,7 @@ Example output from that command, with the cache directory set to `/tmp/torch-tw 2026-10-01 06:19:51,359 - INFO - name: First Wins Strategy 2026-10-01 06:19:51,359 - INFO - description: evaluate backends in order, return first working backend 2026-10-01 06:19:51,360 - INFO - backends: +2026-10-01 06:19:51,360 - INFO - XPUGraphBackend() 2026-10-01 06:19:51,360 - INFO - TorchInductorBackend() 2026-10-01 06:19:51,360 - INFO - OpenVINOBackend() 2026-10-01 06:19:51,360 - INFO - TorchEagerBackend() diff --git a/docs/learn/backends.md b/docs/learn/backends.md index 1e9f50a8..2327d5ea 100644 --- a/docs/learn/backends.md +++ b/docs/learn/backends.md @@ -33,6 +33,20 @@ backend = TorchInductorBackend() 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. +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. + +```python +from torch_tweak.torch.backend import XPUGraphBackend + +backend = XPUGraphBackend() +``` + +Guide: [XPU Graph Backend](../guides/backends/xpu_graph_backend.md). + ## OpenVINO Backend The OpenVINO backend converts the module to an OpenVINO model and runs it on an Intel CPU or GPU. diff --git a/docs/learn/core_functionalities.md b/docs/learn/core_functionalities.md index 421c24d3..7f740a1a 100644 --- a/docs/learn/core_functionalities.md +++ b/docs/learn/core_functionalities.md @@ -110,6 +110,7 @@ from torch_tweak.torch.tune_strategy import FirstWinsStrategy strategy = FirstWinsStrategy( backends=[ + XPUGraphBackend(), TorchInductorBackend(), OpenVINOBackend(), TorchEagerBackend(), @@ -131,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 and eager only. +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. ```python from torch_tweak.torch.tune_strategy import HighestThroughputStrategy diff --git a/docs/learn/overview.md b/docs/learn/overview.md index acb642f1..62f6116f 100644 --- a/docs/learn/overview.md +++ b/docs/learn/overview.md @@ -12,7 +12,7 @@ NOTE: This file has been modified by Intel Corporation. **Torch Tweak** is an inference toolkit designed for tuning and deploying Deep Learning models with a focus on Intel XPU (GPU) accelerators. It provides model tuning capabilities through compilation and conversion paths that can significantly improve inference speed and efficiency across various AI workloads including Computer Vision, Natural Language Processing, Speech Recognition, and Generative AI. -The toolkit enables seamless tuning of PyTorch models and pipelines using TorchAO, Torch Inductor, OpenVINO, and eager through a single Python API. The resulting tuned models are ready for deployment in production environments. +The toolkit enables seamless tuning of PyTorch models and pipelines using TorchAO, Torch Inductor, XPU graph, OpenVINO, and eager through a single Python API. The resulting tuned models are ready for deployment in production environments. Torch Tweak works with your environment - relying first on your software versions - and selects the best-performing backend for your software and hardware setup, guiding you to supported technologies. @@ -31,7 +31,7 @@ The distinct capabilities of Torch Tweak are summarized in the feature matrix: | Feature | Description | |-----------------------------|---------------------------------------------------------------------------------------------------------------------------| | Ease-of-use | Single line of code to run all possible tuning paths directly from your source code | -| Wide Backend Support | Compatible with TorchAO, Torch Inductor, OpenVINO, and eager | +| Wide Backend Support | Compatible with TorchAO, Torch Inductor, XPU graph, OpenVINO, and eager | | Model Tuning | Enhance the performance of models such as ResNET and BERT for efficient inference deployment | | Pipeline Tuning | Streamline Python code pipelines for models such as Stable Diffusion and Flux using seamless model wrapping and tuning | | Correctness Testing | Ensures tuned models produce correct outputs by validating on provided data samples | diff --git a/examples/E5Large/e5large/tune.py b/examples/E5Large/e5large/tune.py index aea8da8a..c8ef63b5 100644 --- a/examples/E5Large/e5large/tune.py +++ b/examples/E5Large/e5large/tune.py @@ -11,7 +11,11 @@ import torch from torch_tweak.torch import HighestThroughputStrategy, inspect, save, tune, wrap -from torch_tweak.torch.backend import TorchEagerBackend, TorchInductorBackend +from torch_tweak.torch.backend import ( + TorchEagerBackend, + TorchInductorBackend, + XPUGraphBackend, +) from torch_tweak.torch.config import config as global_config from .cmd_args import get_parser @@ -97,6 +101,7 @@ def call_wrapper(*args, **kwargs): backends=[ TorchEagerBackend(), TorchInductorBackend(), + XPUGraphBackend(), ] ).enable_find_max_batch_size(True), ) diff --git a/examples/ESM2/esm2/tune.py b/examples/ESM2/esm2/tune.py index ae4c168b..1b1a3452 100644 --- a/examples/ESM2/esm2/tune.py +++ b/examples/ESM2/esm2/tune.py @@ -11,7 +11,11 @@ from transformers import AutoTokenizer, EsmForMaskedLM import torch_tweak.torch as tt -from torch_tweak.torch.backend import TorchEagerBackend, TorchInductorBackend +from torch_tweak.torch.backend import ( + TorchEagerBackend, + TorchInductorBackend, + XPUGraphBackend, +) DEVICE = torch.device("xpu") MODEL_NAME = "facebook/esm2_t33_650M_UR50D" @@ -61,6 +65,7 @@ def tune( strategy = tt.FirstWinsStrategy( backends=[ TorchInductorBackend(), + XPUGraphBackend(), TorchEagerBackend(), ] ) diff --git a/examples/FLUX/flux/tune.py b/examples/FLUX/flux/tune.py index 935de7ad..5c4061ff 100644 --- a/examples/FLUX/flux/tune.py +++ b/examples/FLUX/flux/tune.py @@ -17,6 +17,7 @@ OpenVINOBackendConfig, TorchEagerBackend, TorchInductorBackend, + XPUGraphBackend, ) logger = getLogger(__name__) @@ -58,6 +59,7 @@ def tune_model( strategy = tt.FirstWinsStrategy( backends=[ TorchInductorBackend(), + XPUGraphBackend(), OpenVINOBackend( OpenVINOBackendConfig( use_dynamo=False, diff --git a/examples/Minimal/demo.py b/examples/Minimal/demo.py index 176fe193..1c2a8a71 100644 --- a/examples/Minimal/demo.py +++ b/examples/Minimal/demo.py @@ -9,7 +9,11 @@ import torch.nn as nn import torch_tweak.torch as tt -from torch_tweak.torch.backend import TorchEagerBackend, TorchInductorBackend +from torch_tweak.torch.backend import ( + TorchEagerBackend, + TorchInductorBackend, + XPUGraphBackend, +) def main() -> None: @@ -19,13 +23,17 @@ 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) - # Skip the max-batch search. The two batch sizes below are enough to pick a backend. + # One batch size keeps every tensor shape static, so the XPU graph can be recorded. strategy = tt.HighestThroughputStrategy( - backends=[TorchInductorBackend(), TorchEagerBackend()], + backends=[ + TorchInductorBackend(), + XPUGraphBackend(), + TorchEagerBackend(), + ], ).enable_find_max_batch_size(False) module = tt.Module(model, name="mlp", strategy=strategy) - tt.tune(module, sample, batch_sizes=[1, 4], device=device) + tt.tune(module, sample, batch_sizes=[4], device=device) results = strategy.results[-1].highest_throughput_results if strategy.results else [] if not results: diff --git a/examples/ParakeetCTC/parakeet_ctc/tune.py b/examples/ParakeetCTC/parakeet_ctc/tune.py index 55d461ba..910fb04b 100644 --- a/examples/ParakeetCTC/parakeet_ctc/tune.py +++ b/examples/ParakeetCTC/parakeet_ctc/tune.py @@ -15,7 +15,11 @@ from parakeet_ctc.cmd_args import parse_args from parakeet_ctc.model import get_model from torch_tweak.torch import FirstWinsStrategy, TuneStrategy, inspect, save, tune, wrap -from torch_tweak.torch.backend import TorchEagerBackend, TorchInductorBackend +from torch_tweak.torch.backend import ( + TorchEagerBackend, + TorchInductorBackend, + XPUGraphBackend, +) logger = getLogger(__name__) @@ -81,6 +85,7 @@ def main(): strategy = FirstWinsStrategy( backends=[ TorchInductorBackend(), + XPUGraphBackend(), TorchEagerBackend(), ] ) diff --git a/examples/ParakeetRNNT/parakeet_rnnt/tune.py b/examples/ParakeetRNNT/parakeet_rnnt/tune.py index 6622ddef..10a2b191 100644 --- a/examples/ParakeetRNNT/parakeet_rnnt/tune.py +++ b/examples/ParakeetRNNT/parakeet_rnnt/tune.py @@ -19,6 +19,7 @@ TorchEagerBackend, TorchInductorBackend, TorchInductorBackendConfig, + XPUGraphBackend, ) logger = getLogger(__name__) @@ -48,6 +49,7 @@ def tune_model( strategy = strategy or HighestThroughputStrategy( backends=[ TorchInductorBackend(TorchInductorBackendConfig(autocast_enabled=True, autocast_dtype=torch.float16)), + XPUGraphBackend(), TorchEagerBackend(), ] ).enable_find_max_batch_size(False) diff --git a/examples/ResNet/resnet/tune.py b/examples/ResNet/resnet/tune.py index 297c3d6c..87997620 100644 --- a/examples/ResNet/resnet/tune.py +++ b/examples/ResNet/resnet/tune.py @@ -20,6 +20,7 @@ TorchEagerBackend, TorchInductorBackend, TorchInductorBackendConfig, + XPUGraphBackend, ) logger = getLogger(__name__) @@ -60,6 +61,7 @@ def tune_model( TorchInductorBackend( config=TorchInductorBackendConfig(autocast_enabled=True, autocast_dtype=torch.float16) ), + XPUGraphBackend(), TorchEagerBackend(), ] ).enable_find_max_batch_size(False), diff --git a/examples/StableDiffusion/stable_diffusion/tuning_strategy.py b/examples/StableDiffusion/stable_diffusion/tuning_strategy.py index d4b391cf..bb1bb886 100644 --- a/examples/StableDiffusion/stable_diffusion/tuning_strategy.py +++ b/examples/StableDiffusion/stable_diffusion/tuning_strategy.py @@ -13,6 +13,7 @@ TorchEagerBackend, TorchInductorBackend, TorchInductorBackendConfig, + XPUGraphBackend, ) from torch_tweak.torch.tune_strategy import HighestThroughputStrategy, OneBackendStrategy @@ -35,6 +36,7 @@ def _inductor_backend() -> TorchInductorBackend: def _backend_catalog() -> dict[str, Backend]: return { "inductor": _inductor_backend(), + "xpu_graph": XPUGraphBackend(), "openvino": _openvino_backend(use_dynamo=False), "openvino_dynamo": _openvino_backend(use_dynamo=True), "eager": TorchEagerBackend(), diff --git a/properdocs.yml b/properdocs.yml index a50d0ec5..fe1998df 100644 --- a/properdocs.yml +++ b/properdocs.yml @@ -85,6 +85,7 @@ nav: - OpenVINO Backend: guides/backends/openvino_backend.md - TorchAO Backend: guides/backends/torchao_backend.md - Torch Inductor Backend: guides/backends/torch_inductor_backend.md + - XPU Graph Backend: guides/backends/xpu_graph_backend.md - Torch Eager Backend: guides/backends/torch_eager_backend.md - Tune Strategies: guides/tune_strategies/tune_strategies.md - Deployment: guides/deployment/deployment.md @@ -114,6 +115,7 @@ nav: - OpenVINO Backend: api/backends/openvino_backend.md - TorchAO Backend: api/backends/torchao_backend.md - 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 - Tuning Strategies API: - FirstWinsStrategy: api/tune_strategies/first_wins_strategy.md diff --git a/tests/unit/torch/backend/test_xpu_graph_backend.py b/tests/unit/torch/backend/test_xpu_graph_backend.py new file mode 100644 index 00000000..eda892d5 --- /dev/null +++ b/tests/unit/torch/backend/test_xpu_graph_backend.py @@ -0,0 +1,340 @@ +# Copyright (c) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for XPUGraphBackend.""" + +import copy +import inspect +from collections.abc import Iterator + +import pytest +import torch +import torch.nn as nn + +from tests.utilities.helpers import requires_xpu +from torch_tweak.torch.backend.xpu_graph_backend import XPUGraphBackend, XPUGraphBackendConfig +from torch_tweak.torch.checkpoint.storage_tasks import TorchLoadTask, TorchSaveTask +from torch_tweak.torch.module.graph_spec import GraphSpec +from torch_tweak.torch.module.recording_module import Sample +from torch_tweak.torch.module.sample_metadata import SampleMetadata +from torch_tweak.torch.module.wrapper_module import Module +from torch_tweak.torch.module_registry import MODULE_REGISTRY +from torch_tweak.torch.tune_strategy import FirstWinsStrategy, HighestThroughputStrategy, OneBackendStrategy +from torch_tweak.torch.tuning import tune + + +class LinearRelu(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 8) + + def forward(self, x): + return torch.relu(self.linear(x)) + + +class Scaled(nn.Module): + def forward(self, x, scale=1.0): + y = x * scale + return y, y.sum() + + +class NanOutput(nn.Module): + def forward(self, x): + return torch.sqrt(-x.abs() - 1) + + +class TinyConv(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(1, 4, 3) + + def forward(self, x): + return self.conv(x) + + +@pytest.fixture(autouse=True) +def eager_calls(request, monkeypatch) -> Iterator[list[tuple]]: + """Records the inputs XPUGraphBackend runs on the original module. + + The fallback output equals eager, so a test that compares to eager cannot see it. + A test that does not take eager_calls as an argument fails when the fallback runs. + """ + calls: list[tuple] = [] + infer_eager = XPUGraphBackend._infer_eager + + def recording(self, args, kwargs): + calls.append(args) + return infer_eager(self, args, kwargs) + + monkeypatch.setattr(XPUGraphBackend, "_infer_eager", recording) + yield calls + if "eager_calls" not in inspect.signature(request.function).parameters: + assert not calls, "XPU graph backend ran the original module instead of replaying" + + +@pytest.fixture +def model(torch_device): + return LinearRelu().to(torch_device).eval() + + +@pytest.fixture +def sample_data(torch_device) -> list[Sample]: + args = (torch.randn(2, 16, device=torch_device),) + return [(args, {})] + + +def move_to_dtype(sample_data, dtype): + args, kwargs = sample_data[0] + return [((args[0].to(dtype),), kwargs)] + + +def static_graph_spec(module, data) -> GraphSpec: + args, kwargs = data[0] + with torch.no_grad(): + outputs = module(*args, **kwargs) + return GraphSpec( + name="sample", + input_spec=SampleMetadata.from_inputs(args, kwargs), + output_spec=SampleMetadata.from_outputs(outputs), + ) + + +def backend_build(backend, dtype, model, sample_data, tmp_path, device="xpu"): + device = torch.device(device) + model = model.to(device, dtype=dtype).eval() + data = move_to_dtype(sample_data, dtype) + return backend.build(model, graph_spec=static_graph_spec(model, data), data=data, device=device, cache_dir=tmp_path) + + +@requires_xpu +@pytest.mark.parametrize( + "dtype", + [torch.float16, torch.bfloat16, torch.float32], + ids=["float16", "bfloat16", "float32"], +) +def test_replay_matches_eager_on_new_inputs(dtype, model, sample_data, tmp_path): + backend = backend_build(XPUGraphBackend(), dtype, model, sample_data, tmp_path) + sample_data = move_to_dtype(sample_data, dtype) + args, kwargs = sample_data[0] + new_args = (torch.randn_like(args[0]),) + + with torch.no_grad(): + torch.testing.assert_close(backend.infer(*args, **kwargs), model(*args, **kwargs)) + torch.testing.assert_close(backend.infer(*new_args, **kwargs), model(*new_args, **kwargs)) + torch.testing.assert_close(backend.infer(*args, **kwargs), model(*args, **kwargs)) + + backend.deactivate() + + +@requires_xpu +def test_deactivate_then_activate_records_again(model, sample_data, tmp_path): + backend = backend_build(XPUGraphBackend(), torch.float32, model, sample_data, tmp_path) + args, kwargs = sample_data[0] + backend.deactivate() + backend.activate() + new_args = (torch.randn_like(args[0]),) + with torch.no_grad(): + torch.testing.assert_close(backend.infer(*new_args, **kwargs), model(*new_args, **kwargs)) + + +@requires_xpu +def test_deploy_after_load_records_and_drops_samples(model, sample_data, tmp_path, torch_device): + backend = backend_build(XPUGraphBackend(), torch.float32, model, sample_data, tmp_path) + backend.deactivate() + backend.activate() + TorchSaveTask().save(tmp_path, backend.to_dict()) + loaded = XPUGraphBackend.from_dict(model, TorchLoadTask().load(tmp_path)) + loaded.deploy(device=torch_device) + assert loaded._data is None + + args, kwargs = sample_data[0] + new_args = (torch.randn_like(args[0]),) + with torch.no_grad(): + torch.testing.assert_close(loaded.infer(*new_args, **kwargs), model(*new_args, **kwargs)) + + +@requires_xpu +def test_kwargs_and_tuple_output(torch_device, tmp_path, eager_calls): + module = Scaled().to(torch_device).eval() + args = (torch.randn(3, 5, device=torch_device),) + kwargs = {"scale": 2.0} + backend = XPUGraphBackend() + data = [(args, kwargs)] + backend = backend.build(module, static_graph_spec(module, data), data, torch_device, tmp_path) + + new_args = (torch.randn_like(args[0]),) + with torch.no_grad(): + got = backend.infer(*new_args, **kwargs) + expected = module(*new_args, **kwargs) + torch.testing.assert_close(got[0], expected[0]) + torch.testing.assert_close(got[1], expected[1]) + + with torch.no_grad(): + got = backend.infer(*new_args, scale=3.0) + expected = module(*new_args, scale=3.0) + torch.testing.assert_close(got[0], expected[0]) + torch.testing.assert_close(got[1], expected[1]) + assert len(eager_calls) == 1 + + +@requires_xpu +def test_conv_replay(torch_device, tmp_path): + module = TinyConv().to(torch_device).eval() + args = (torch.randn(2, 1, 8, 8, device=torch_device),) + backend = XPUGraphBackend() + data = [(args, {})] + backend = backend.build(module, static_graph_spec(module, data), data, torch_device, tmp_path) + new_args = (torch.randn_like(args[0]),) + with torch.no_grad(): + torch.testing.assert_close(backend.infer(*new_args), module(*new_args)) + + +@requires_xpu +def test_autocast_keeps_recorded_dtype(model, sample_data, tmp_path): + config = XPUGraphBackendConfig(autocast_enabled=True, autocast_dtype=torch.float16) + backend = backend_build(XPUGraphBackend(config), torch.float32, model, sample_data, tmp_path) + args, kwargs = sample_data[0] + output = backend.infer(*args, **kwargs) + assert output.dtype == torch.float16 + with torch.autocast(device_type="xpu", dtype=torch.float16, cache_enabled=False), torch.no_grad(): + expected = model(*args, **kwargs) + torch.testing.assert_close(output, expected) + + +@requires_xpu +def test_infer_does_not_record_again(model, sample_data, tmp_path, monkeypatch): + backend = backend_build(XPUGraphBackend(), torch.float32, model, sample_data, tmp_path) + calls = {"n": 0} + real_graph = torch.xpu.graph + + class CountingGraph: + def __init__(self, *args, **kwargs): + calls["n"] += 1 + self._graph = real_graph(*args, **kwargs) + + def __enter__(self): + return self._graph.__enter__() + + def __exit__(self, *exc): + return self._graph.__exit__(*exc) + + monkeypatch.setattr(torch.xpu, "graph", CountingGraph) + args, kwargs = sample_data[0] + backend.infer(*args, **kwargs) + backend.infer(*(torch.randn_like(args[0]),), **kwargs) + assert calls["n"] == 0 + + +@requires_xpu +def test_matching_nans_are_accepted(torch_device, tmp_path): + module = NanOutput().to(torch_device).eval() + args = (torch.randn(4, device=torch_device),) + backend = XPUGraphBackend() + data = [(args, {})] + backend = backend.build(module, static_graph_spec(module, data), data, torch_device, tmp_path) + output = backend.infer(*args) + assert torch.isnan(output).all() + + +@requires_xpu +def test_empty_recording_is_rejected(torch_device, tmp_path): + module = nn.Identity().to(torch_device).eval() + args = (torch.randn(4, device=torch_device),) + backend = XPUGraphBackend() + with pytest.raises(RuntimeError, match="no device work"): + backend.build(module, static_graph_spec(module, [(args, {})]), [(args, {})], torch_device, tmp_path) + + +@requires_xpu +def test_dynamic_shapes_are_rejected(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),) + 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) + + +@requires_xpu +def test_wrong_replay_is_rejected(torch_device, tmp_path, monkeypatch): + monkeypatch.setattr(torch.xpu.XPUGraph, "replay", lambda graph: None) + module = LinearRelu().to(torch_device).eval() + args = (torch.randn(2, 16, device=torch_device),) + backend = XPUGraphBackend() + with pytest.raises(RuntimeError, match="differs from eager"): + backend.build(module, static_graph_spec(module, [(args, {})]), [(args, {})], torch_device, tmp_path) + + +def test_cpu_is_rejected(tmp_path): + module = LinearRelu() + args = (torch.randn(2, 16),) + backend = XPUGraphBackend() + with pytest.raises(ValueError, match="not supported"): + backend.build( + module, + GraphSpec( + name="cpu", + input_spec=SampleMetadata.from_inputs(args, {}), + output_spec=SampleMetadata.from_outputs(args[0]), + ), + [(args, {})], + torch.device("cpu"), + tmp_path, + ) + + +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) + TorchSaveTask().save(tmp_path, backend.to_dict()) + loaded = XPUGraphBackend.from_dict(model, TorchLoadTask().load(tmp_path)) + loaded.activate() + + sample_data = move_to_dtype(sample_data, torch.float16) + args, kwargs = sample_data[0] + new_args = (torch.randn_like(args[0]),) + torch.testing.assert_close(backend.infer(*new_args, **kwargs), loaded.infer(*new_args, **kwargs)) + + +@requires_xpu +def test_tune_module(torch_device, tmp_path, eager_calls): + torch.manual_seed(0) + module = LinearRelu().to(torch_device).eval() + # tune routes later calls of module through the selected backend. + reference = copy.deepcopy(module) + sample = torch.randn(16, device=torch_device) + wrapped = Module( + module, + "xpu-graph-mlp", + strategy=OneBackendStrategy(XPUGraphBackend()).enable_find_max_batch_size(False), + ) + try: + tune(wrapped, sample, batch_sizes=[2], device=torch_device, ignore_failing_modules=False) + larger = torch.randn(4, 16, device=torch_device) + batch = torch.randn(2, 16, device=torch_device) + with torch.no_grad(): + torch.testing.assert_close(wrapped(larger), reference(larger)) + torch.testing.assert_close(wrapped(batch), reference(batch)) + assert [call[0].shape for call in eager_calls] == [larger.shape] + finally: + MODULE_REGISTRY.clear() diff --git a/tests/unit/torch/task/test_profiling.py b/tests/unit/torch/task/test_profiling.py index 7bdc737a..36866936 100644 --- a/tests/unit/torch/task/test_profiling.py +++ b/tests/unit/torch/task/test_profiling.py @@ -16,6 +16,7 @@ NumStepsMeasuringStopStrategy, StableWindowMeasuringStopStrategy, ) +from torch_tweak.torch.task.profiling.measuring_strategy import ModelExecutionTimeMeasuringStrategy from torch_tweak.torch.task.profiling.metrics import get_throughput, is_throughput_saturated from torch_tweak.torch.task.profiling.profiling import ProfilingStatus, profile, profile_backend from torch_tweak.torch.task.profiling.profiling_stop_strategy import ThroughputSaturatedProfilingStopStrategy @@ -243,3 +244,19 @@ def describe(self): # Verify both calls produced results assert len(result.results.entries) > 0 + + +def test_measurement_waits_for_the_device(mocker): + order = [] + mocker.patch( + "torch_tweak.torch.task.profiling.measuring_strategy.synchronize", + side_effect=lambda: order.append("sync"), + ) + + def model(*_args, **_kwargs): + order.append("model") + + events = ModelExecutionTimeMeasuringStrategy().do_measurement(1, model, ([], {})) + + assert order == ["sync", "model", "sync"] + assert events[0].execution_time >= 0 diff --git a/torch_tweak/torch/backend/__init__.py b/torch_tweak/torch/backend/__init__.py index 2edf8786..1786a626 100644 --- a/torch_tweak/torch/backend/__init__.py +++ b/torch_tweak/torch/backend/__init__.py @@ -10,6 +10,7 @@ from torch_tweak.torch.backend.torch_eager import TorchEagerBackend, TorchEagerBackendConfig from torch_tweak.torch.backend.torch_inductor_backend import TorchInductorBackend, TorchInductorBackendConfig from torch_tweak.torch.backend.torchao_backend import TorchAOBackend, TorchAOBackendConfig +from torch_tweak.torch.backend.xpu_graph_backend import XPUGraphBackend, XPUGraphBackendConfig __all__ = [ "Backend", @@ -21,4 +22,6 @@ "TorchAOBackendConfig", "TorchEagerBackend", "TorchEagerBackendConfig", + "XPUGraphBackend", + "XPUGraphBackendConfig", ] diff --git a/torch_tweak/torch/backend/backend.py b/torch_tweak/torch/backend/backend.py index 54db08f3..8374e18a 100644 --- a/torch_tweak/torch/backend/backend.py +++ b/torch_tweak/torch/backend/backend.py @@ -123,13 +123,13 @@ def to_dict(self): """Returns the state_dict of the backend.""" return asdict(self) - def to_json(self, path: Path): - """Saves the backend configuration to a file.""" - config_dict = self.to_dict() - config_dict = self._to_json(config_dict) - with open(path, "w", encoding="utf-8") as f: + def save_to_file(self, path: Path) -> None: + """Saves the backend configuration to config.json in the given directory.""" + config_path = path / "config.json" + config_dict = self._to_json(self.to_dict()) + with open(config_path, "w", encoding="utf-8") as f: json.dump(config_dict, f, indent=2) - secure_file(path) + secure_file(config_path) def _to_json(self, obj: Any) -> Any: """Convert object to JSON-serializable format. @@ -225,6 +225,8 @@ def build( try: self._assert_device(device) self._set_device(device) + if not data: + raise RuntimeError(f"Backend {self.name} needs at least one sample") ready_backend = self._build(module, graph_spec, data, cache_dir) self.state = BackendState.ACTIVE return ready_backend diff --git a/torch_tweak/torch/backend/openvino/openvino_backend.py b/torch_tweak/torch/backend/openvino/openvino_backend.py index 0ab5b417..15c55cf2 100644 --- a/torch_tweak/torch/backend/openvino/openvino_backend.py +++ b/torch_tweak/torch/backend/openvino/openvino_backend.py @@ -117,7 +117,7 @@ def describe(self) -> str: def _build(self, module: nn.Module, graph_spec: GraphSpec, data: list[Sample], cache_dir: Path) -> Backend: logger.info("Starting OpenVINO backend build") self._graph_spec = graph_spec - self._save_config(cache_dir) + self._config.save_to_file(cache_dir) module.to(self._device) self._output_object = self._get_output_object(module, data[0]) @@ -228,10 +228,6 @@ def _get_output_object(self, module: nn.Module, sample: Sample) -> Any: # `_prepare_outputs_for_return` replaces every tensor with an OpenVINO output. return copy.deepcopy(pytree.tree_map_only(torch.Tensor, lambda _: torch.empty(0), output_object)) - def _save_config(self, cache_dir: Path): - config_path = cache_dir / "config.json" - self._config.to_json(config_path) - def to_dict(self): """Returns the state_dict of the backend.""" if self._ov_model_dir is None or self._graph_spec is None: diff --git a/torch_tweak/torch/backend/torch_eager.py b/torch_tweak/torch/backend/torch_eager.py index d53c82e5..aaba2304 100644 --- a/torch_tweak/torch/backend/torch_eager.py +++ b/torch_tweak/torch/backend/torch_eager.py @@ -85,7 +85,7 @@ def _get_dtype(module: nn.Module, data: list[Sample]) -> torch.dtype: def _build(self, module: nn.Module, graph_spec: GraphSpec, data: list[Sample], cache_dir: Path) -> Backend: """Builds the model.""" - self._save_config(cache_dir) + self._config.save_to_file(cache_dir) module.to(self._device) self._orig_module = module @@ -133,11 +133,6 @@ def _deploy(self): """Deploys the backend.""" pass - def _save_config(self, cache_dir: Path): - """Store the backend configuration to a file.""" - config_path = cache_dir / "config.json" - self._config.to_json(config_path) - def to_dict(self): """Returns the state_dict of the backend.""" if self._orig_module is None: diff --git a/torch_tweak/torch/backend/torch_inductor_backend.py b/torch_tweak/torch/backend/torch_inductor_backend.py index 510ace90..b163b619 100644 --- a/torch_tweak/torch/backend/torch_inductor_backend.py +++ b/torch_tweak/torch/backend/torch_inductor_backend.py @@ -158,7 +158,7 @@ def _get_dtype(module: nn.Module, data: list[Sample]) -> torch.dtype: def _build(self, module: nn.Module, graph_spec: GraphSpec, data: list[Sample], cache_dir: Path) -> Backend: """Builds the model with torch.compile.""" - self._save_config(cache_dir) + self._config.save_to_file(cache_dir) module.to(self._device) self._output_dtype = self._get_dtype(module, data) @@ -231,12 +231,6 @@ def _deploy(self): self._data = None gc.collect() - def _save_config(self, cache_dir: Path): - """Store the backend configuration to a file.""" - config_path = cache_dir / "config.json" - self._config.to_json(config_path) - logger.info("Config saved to %s", config_path) - def to_dict(self): """Returns the state_dict of the backend.""" if not self._orig_module: diff --git a/torch_tweak/torch/backend/torchao_backend.py b/torch_tweak/torch/backend/torchao_backend.py index 9aace8f3..2ae2ccb6 100644 --- a/torch_tweak/torch/backend/torchao_backend.py +++ b/torch_tweak/torch/backend/torchao_backend.py @@ -262,7 +262,7 @@ def describe(self) -> str: def _build(self, module: nn.Module, graph_spec: GraphSpec, data: list[Sample], cache_dir: Path) -> Backend: """Builds the model with torchao quantization and torch.compile.""" - self._save_config(cache_dir) + self._config.save_to_file(cache_dir) self._orig_module = module self._data = data @@ -297,10 +297,6 @@ def _deploy(self): self._data = None gc.collect() - def _save_config(self, cache_dir: Path): - """Store the backend configuration to a file.""" - self._config.to_json(cache_dir / "config.json") - def to_dict(self): """Returns the state_dict of the backend.""" if self._orig_module is None: diff --git a/torch_tweak/torch/backend/xpu_graph_backend.py b/torch_tweak/torch/backend/xpu_graph_backend.py new file mode 100644 index 00000000..ebc1f366 --- /dev/null +++ b/torch_tweak/torch/backend/xpu_graph_backend.py @@ -0,0 +1,326 @@ +# Copyright (c) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""XPU graph backend.""" + +import warnings +from dataclasses import dataclass +from logging import getLogger +from pathlib import Path +from typing import Any, ClassVar, Self + +import torch +import torch.nn as nn +from torch.utils._pytree import TreeSpec, tree_flatten, tree_map_only + +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.utils.trace import annotate + +logger = getLogger(__name__) + + +@dataclass +class XPUGraphBackendConfig(BackendConfig): + """Configuration for an XPU graph backend. + + Args: + 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. + """ + + autocast_enabled: bool = False + autocast_dtype: torch.dtype | None = None + + @classmethod + def from_dict(cls, state_dict: dict[str, Any]) -> Self: + """Convert dict to XPUGraphBackendConfig.""" + return cls(**state_dict) + + +class XPUGraphBackend(Backend): + """Backend that records one XPU graph while tuning and replays it for inference.""" + + _devices: ClassVar[list[str]] = ["xpu"] + requires_original_module: ClassVar[bool] = True + + STATE_TYPE = "type" + STATE_CONFIG = "config" + STATE_ORIG_MODULE = "orig_module" + STATE_DATA = "data" + STATE_DEVICE = "device" + + def __init__(self, config: XPUGraphBackendConfig | None = None): + """Initializes backend. + + Args: + config: Recording and autocast configuration. + """ + super().__init__() + 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._warned_eager = False + + def key(self) -> str: + """Returns the key of the backend.""" + return f"{self.__class__.__name__}_{self._config.key()}" + + def describe(self) -> str: + """Returns the description of the backend.""" + return f"{self.__class__.__name__}({self._config.describe()})" + + 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") + 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 + self._record() + return self + + def _activate(self) -> None: + """Records the graph again after a checkpoint load or a deactivate.""" + self._release_graph() + self._record() + + def _deactivate(self) -> None: + """Drops the recorded graph.""" + self._release_graph() + + def _deploy(self) -> None: + """Records the graph 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. + + Args: + *args: inference arguments + **kwargs: inference keyword arguments + + 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. + """ + with torch.no_grad(): + if not self._load_inputs(args, kwargs): + 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 + + def _infer_eager(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + """Runs the original module on inputs that differ from the recorded call.""" + 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.") + self._warned_eager = True + with self._autocast(): + return self._orig_module(*args, **kwargs) + + def _record(self) -> None: + """Warm up, record, and reject the graph 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) + + logger.info("Recording XPU graph.") + 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. + self._orig_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(self._graph): + 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() + 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 + 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) + + def _autocast(self) -> torch.autocast: + """Returns the autocast context used for warmup, capture, and the eager check.""" + # The autocast cache aliases weight casts. Those addresses are part of the recording. + return torch.autocast( + device_type=self._device.type, + dtype=self._config.autocast_dtype, + enabled=self._config.autocast_enabled, + 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 to_dict(self) -> dict[str, Any]: + """Returns the state_dict of the backend.""" + if self._orig_module is None: + raise RuntimeError("Backend has not been properly initialized. Please call build() first.") + return { + self.STATE_TYPE: self.__class__.__name__, + self.STATE_CONFIG: self._config.to_dict(), + self.STATE_ORIG_MODULE: self._orig_module.state_dict(), + self.STATE_DATA: self._data, + self.STATE_DEVICE: self._device, + } + + @classmethod + def from_dict(cls, module: 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)}") + if module is None: + raise ValueError("Module is required to create a backend from a state_dict.") + + config = XPUGraphBackendConfig.from_dict(state_dict[cls.STATE_CONFIG]) + backend = cls(config=config) + backend._data = state_dict[cls.STATE_DATA] + backend._device = state_dict[cls.STATE_DEVICE] + backend._orig_module = module + module.load_state_dict(state_dict[cls.STATE_ORIG_MODULE], strict=False) + backend.state = BackendState.CHECKPOINT_LOADED + return backend + + +def _clone_tree(tree: Any) -> Any: + """Clones tensors and leaves every other leaf as it is.""" + 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 _require_xpu_tensors(*trees: Any) -> None: + """Rejects a recording whose inputs are not XPU tensors.""" + for tree in trees: + leaves, _ = tree_flatten(tree) + for leaf in leaves: + if isinstance(leaf, torch.Tensor) and leaf.device.type != "xpu": + raise RuntimeError("XPU graph capture needs XPU tensor inputs") + + +def _perturb(tree: Any) -> Any: + """Returns a tree with different tensor values and the same other leaves.""" + return tree_map_only(torch.Tensor, _perturb_tensor, tree) + + +def _perturb_tensor(tensor: torch.Tensor) -> torch.Tensor: + """Returns a tensor with different values and the same shape.""" + cloned = tensor.detach().clone() + if cloned.is_floating_point() or cloned.is_complex(): + return cloned + torch.randn_like(cloned) + if cloned.ndim == 0: + return ~cloned if cloned.dtype == torch.bool else cloned + 1 + return cloned.roll(1, 0) + + +def _assert_close(actual: Any, expected: Any) -> None: + """Raises when replay outputs are too far from eager outputs.""" + actual_leaves, actual_spec = tree_flatten(actual) + expected_leaves, expected_spec = tree_flatten(expected) + if actual_spec != expected_spec: + raise RuntimeError("XPU graph output structure does not match eager") + for got, ref in zip(actual_leaves, expected_leaves, strict=True): + if not isinstance(ref, torch.Tensor): + if got != ref: + raise RuntimeError("XPU graph output does not match eager") + continue + if not isinstance(got, torch.Tensor) or got.shape != ref.shape or got.dtype != ref.dtype: + raise RuntimeError("XPU graph output tensor does not match eager") + rtol, atol = (1e-2, 1e-2) if ref.dtype in (torch.float16, torch.bfloat16) else (1e-3, 1e-3) + if torch.allclose(got, ref, rtol=rtol, atol=atol, equal_nan=True): + continue + if got.is_floating_point() or got.is_complex(): + difference = (got - ref).abs().max().item() + raise RuntimeError(f"XPU graph output differs from eager (max abs {difference})") + raise RuntimeError("XPU graph output differs from eager") diff --git a/torch_tweak/torch/module/graph_spec.py b/torch_tweak/torch/module/graph_spec.py index 413f6493..fe83dd37 100644 --- a/torch_tweak/torch/module/graph_spec.py +++ b/torch_tweak/torch/module/graph_spec.py @@ -29,6 +29,10 @@ class GraphSpec: input_spec: SampleMetadata output_spec: SampleMetadata + def has_dynamic_axis(self) -> bool: + """Check if any input or output tensor has a dynamic axis.""" + return self.input_spec.detected_dynamic_axis() or self.output_spec.detected_dynamic_axis() + def update_shapes_seen(self, inputs_metadata: SampleMetadata, outputs_metadata: SampleMetadata): """Update input spec with other input spec.""" self.input_spec.update_shapes_seen(inputs_metadata) diff --git a/torch_tweak/torch/task/profiling/measuring_strategy.py b/torch_tweak/torch/task/profiling/measuring_strategy.py index 08fc27fa..90cda696 100644 --- a/torch_tweak/torch/task/profiling/measuring_strategy.py +++ b/torch_tweak/torch/task/profiling/measuring_strategy.py @@ -9,6 +9,7 @@ from collections.abc import Callable from torch_tweak.torch.task.profiling.events import ProfilingResultEvent +from torch_tweak.torch.utils.xpu import synchronize class MeasuringStrategy(ABC): @@ -46,8 +47,11 @@ def do_measurement( """Do the measurement and create a ProfilingResultEvent(s) for the model.""" args, kwargs = sample + synchronize() start = time.monotonic_ns() model(*args, **kwargs) + # Replay returns after submit. The device may still be running. + synchronize() end = time.monotonic_ns() ModelExecutionTimeMeasuringStrategy.counter += 1 diff --git a/torch_tweak/torch/tune_strategy/first_wins_strategy.py b/torch_tweak/torch/tune_strategy/first_wins_strategy.py index 8b6d77f1..ac3be7e0 100644 --- a/torch_tweak/torch/tune_strategy/first_wins_strategy.py +++ b/torch_tweak/torch/tune_strategy/first_wins_strategy.py @@ -12,7 +12,12 @@ import torch import torch.nn as nn -from torch_tweak.torch.backend import OpenVINOBackend, TorchEagerBackend, TorchInductorBackend +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 @@ -24,7 +29,7 @@ class FirstWinsStrategy(TuneStrategyFindMaxBatchSizeExtension): """Runs backends in order and keeps the first one that builds and validates. - The default order is Torch Inductor, OpenVINO, then eager. + The default order is XPU graph, Torch Inductor, OpenVINO, then eager. """ def __init__(self, backends: Iterable[Backend] | None = None, **kwargs): @@ -84,6 +89,7 @@ def _tune( def _default_backends(self) -> list[Backend]: """Returns default backends.""" return [ + XPUGraphBackend(), TorchInductorBackend(), OpenVINOBackend(), TorchEagerBackend(), diff --git a/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py b/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py index 45f32710..353de525 100644 --- a/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py +++ b/torch_tweak/torch/tune_strategy/highest_throughput_strategy.py @@ -24,6 +24,7 @@ Backend, TorchEagerBackend, TorchInductorBackend, + XPUGraphBackend, ) from torch_tweak.torch.config import ( DEFAULT_STABILITY_PERCENTAGE, @@ -221,6 +222,7 @@ def _default_backends(self) -> list[Backend]: """Returns default backends.""" return [ TorchInductorBackend(), + XPUGraphBackend(), TorchEagerBackend(), ]