From a16634188c2d3c2e0c513b6fa89908083ad285c0 Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Thu, 15 Oct 2026 18:26:44 +0000 Subject: [PATCH 1/2] Add custom device compile backend plugins --- README.md | 1 + docs/custom_device_plugins.md | 57 +++++++++++++++++++ .../unit_tests/cpu/tools/test_tools_utils.py | 32 ++++++++++- torchtitan/tools/utils.py | 11 ++++ 4 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 docs/custom_device_plugins.md diff --git a/README.md b/README.md index 48ef93d47c3..2b724794932 100644 --- a/README.md +++ b/README.md @@ -265,6 +265,7 @@ If a node does not have 8 GPUs, also adjust `--nproc_per_node` and | Datasets and data loading | [torchtitan/components/data](torchtitan/components/data/README.md) | | Checkpointing and HF conversion | [torchtitan/components/checkpointer](torchtitan/components/checkpointer/README.md) | | Reinforcement learning | [torchtitan/rl](torchtitan/rl/README.md) | +| Custom device plugins | [docs/custom_device_plugins.md](docs/custom_device_plugins.md) | | Debugging, reproducibility, and validating numerics | [docs/debugging.md](docs/debugging.md) | | Validation and evaluation | [docs/evaluation.md](docs/evaluation.md) | | Benchmarks | [docs/benchmarks](docs/benchmarks) | diff --git a/docs/custom_device_plugins.md b/docs/custom_device_plugins.md new file mode 100644 index 00000000000..f26f8c5891e --- /dev/null +++ b/docs/custom_device_plugins.md @@ -0,0 +1,57 @@ +# Custom device plugins + +TorchTitan loads an installed compile-backend plugin before model modules bind +backend-less `torch.compile` calls. A custom device package can use this hook +to select its registered backend without patching TorchTitan. + +## Plugin package + +Register an entry point whose name matches the device type returned by PyTorch: + +```text +torchtitan-custom-device/ + pyproject.toml + torchtitan_custom_device/ + __init__.py + compile_backend.py +``` + +```toml +# pyproject.toml +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "torchtitan-custom-device" +version = "0.1.0" +requires-python = ">=3.11" +dependencies = ["torch"] + +[project.entry-points."torchtitan.compile_backends"] +custom_device = "torchtitan_custom_device.compile_backend:BACKEND" +``` + +The plugin provides the device-specific backend name: + +```python +# torchtitan_custom_device/compile_backend.py +BACKEND = "custom_device_backend" +``` + +Replace `custom_device` with the device type registered with PyTorch and +`custom_device_backend` with the registered Dynamo backend name. + +## Install and run + +Install the plugin into the same environment as TorchTitan: + +```bash +python -m pip install -e /path/to/torchtitan-custom-device +``` + +Then launch TorchTitan normally. When PyTorch reports `custom_device` as the +active device, TorchTitan discovers the matching entry point and calls +`torch.compiler.set_default_backend(BACKEND)` before importing FlexAttention. +Backend-less `torch.compile` calls therefore use `custom_device_backend`; calls +that pass `backend=` explicitly keep their requested backend. diff --git a/tests/unit_tests/cpu/tools/test_tools_utils.py b/tests/unit_tests/cpu/tools/test_tools_utils.py index 678a751f742..71aa955c5c8 100644 --- a/tests/unit_tests/cpu/tools/test_tools_utils.py +++ b/tests/unit_tests/cpu/tools/test_tools_utils.py @@ -10,7 +10,11 @@ import torch from torchtitan.tools.garbage_collector import GarbageCollector -from torchtitan.tools.utils import get_cuda_flash_attention_impl, get_local_device +from torchtitan.tools.utils import ( + _load_device_compile_backend, + get_cuda_flash_attention_impl, + get_local_device, +) class _FakeDeviceModule: @@ -21,6 +25,32 @@ def device_count(self) -> int: return self.num_devices +def test_load_device_compile_backend(monkeypatch: pytest.MonkeyPatch) -> None: + set_default_backend = Mock() + monkeypatch.setattr( + torch.compiler, + "set_default_backend", + set_default_backend, + ) + + backend = Mock() + backend.load.return_value = "custom_device_backend" + entry_points = Mock(return_value=[backend]) + monkeypatch.setattr( + "torchtitan.tools.utils.importlib.metadata.entry_points", + entry_points, + ) + + _load_device_compile_backend("custom_device") + + entry_points.assert_called_once_with( + group="torchtitan.compile_backends", + name="custom_device", + ) + backend.load.assert_called_once_with() + set_default_backend.assert_called_once_with("custom_device_backend") + + def test_gc_debug_collects_once_per_step(monkeypatch: pytest.MonkeyPatch) -> None: gc_collect = Mock() monkeypatch.setattr("torchtitan.tools.garbage_collector.gc.collect", gc_collect) diff --git a/torchtitan/tools/utils.py b/torchtitan/tools/utils.py index 5e1810eddc9..a060660ae86 100644 --- a/torchtitan/tools/utils.py +++ b/torchtitan/tools/utils.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. import contextlib +import importlib.metadata import logging import os import subprocess @@ -66,7 +67,17 @@ def get_device_info() -> tuple[str, ModuleType]: return device_type, device_module +def _load_device_compile_backend(device_type: str) -> None: + """Load an installed compile backend registered for ``device_type``.""" + for backend in importlib.metadata.entry_points( + group="torchtitan.compile_backends", + name=device_type, + ): + torch.compiler.set_default_backend(backend.load()) + + device_type, device_module = get_device_info() +_load_device_compile_backend(device_type) def get_local_device() -> torch.device: From 131150ff3d156c538a1dd501f984415c8f2d9ed0 Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Thu, 15 Oct 2026 18:34:58 +0000 Subject: [PATCH 2/2] Simplify compile backend configuration --- README.md | 1 - docs/custom_device_plugins.md | 57 ------------------- .../unit_tests/cpu/tools/test_tools_utils.py | 20 ++----- torchtitan/tools/utils.py | 14 ++--- 4 files changed, 9 insertions(+), 83 deletions(-) delete mode 100644 docs/custom_device_plugins.md diff --git a/README.md b/README.md index 2b724794932..48ef93d47c3 100644 --- a/README.md +++ b/README.md @@ -265,7 +265,6 @@ If a node does not have 8 GPUs, also adjust `--nproc_per_node` and | Datasets and data loading | [torchtitan/components/data](torchtitan/components/data/README.md) | | Checkpointing and HF conversion | [torchtitan/components/checkpointer](torchtitan/components/checkpointer/README.md) | | Reinforcement learning | [torchtitan/rl](torchtitan/rl/README.md) | -| Custom device plugins | [docs/custom_device_plugins.md](docs/custom_device_plugins.md) | | Debugging, reproducibility, and validating numerics | [docs/debugging.md](docs/debugging.md) | | Validation and evaluation | [docs/evaluation.md](docs/evaluation.md) | | Benchmarks | [docs/benchmarks](docs/benchmarks) | diff --git a/docs/custom_device_plugins.md b/docs/custom_device_plugins.md deleted file mode 100644 index f26f8c5891e..00000000000 --- a/docs/custom_device_plugins.md +++ /dev/null @@ -1,57 +0,0 @@ -# Custom device plugins - -TorchTitan loads an installed compile-backend plugin before model modules bind -backend-less `torch.compile` calls. A custom device package can use this hook -to select its registered backend without patching TorchTitan. - -## Plugin package - -Register an entry point whose name matches the device type returned by PyTorch: - -```text -torchtitan-custom-device/ - pyproject.toml - torchtitan_custom_device/ - __init__.py - compile_backend.py -``` - -```toml -# pyproject.toml -[build-system] -requires = ["setuptools>=61.0"] -build-backend = "setuptools.build_meta" - -[project] -name = "torchtitan-custom-device" -version = "0.1.0" -requires-python = ">=3.11" -dependencies = ["torch"] - -[project.entry-points."torchtitan.compile_backends"] -custom_device = "torchtitan_custom_device.compile_backend:BACKEND" -``` - -The plugin provides the device-specific backend name: - -```python -# torchtitan_custom_device/compile_backend.py -BACKEND = "custom_device_backend" -``` - -Replace `custom_device` with the device type registered with PyTorch and -`custom_device_backend` with the registered Dynamo backend name. - -## Install and run - -Install the plugin into the same environment as TorchTitan: - -```bash -python -m pip install -e /path/to/torchtitan-custom-device -``` - -Then launch TorchTitan normally. When PyTorch reports `custom_device` as the -active device, TorchTitan discovers the matching entry point and calls -`torch.compiler.set_default_backend(BACKEND)` before importing FlexAttention. -Backend-less `torch.compile` calls therefore use `custom_device_backend`; calls -that pass `backend=` explicitly keep their requested backend. diff --git a/tests/unit_tests/cpu/tools/test_tools_utils.py b/tests/unit_tests/cpu/tools/test_tools_utils.py index 71aa955c5c8..09a4b8ef6ae 100644 --- a/tests/unit_tests/cpu/tools/test_tools_utils.py +++ b/tests/unit_tests/cpu/tools/test_tools_utils.py @@ -11,7 +11,7 @@ from torchtitan.tools.garbage_collector import GarbageCollector from torchtitan.tools.utils import ( - _load_device_compile_backend, + _configure_compile_backend, get_cuda_flash_attention_impl, get_local_device, ) @@ -25,29 +25,17 @@ def device_count(self) -> int: return self.num_devices -def test_load_device_compile_backend(monkeypatch: pytest.MonkeyPatch) -> None: +def test_configure_compile_backend(monkeypatch: pytest.MonkeyPatch) -> None: set_default_backend = Mock() + monkeypatch.setenv("TORCHTITAN_COMPILE_BACKEND", "custom_device_backend") monkeypatch.setattr( torch.compiler, "set_default_backend", set_default_backend, ) - backend = Mock() - backend.load.return_value = "custom_device_backend" - entry_points = Mock(return_value=[backend]) - monkeypatch.setattr( - "torchtitan.tools.utils.importlib.metadata.entry_points", - entry_points, - ) + _configure_compile_backend() - _load_device_compile_backend("custom_device") - - entry_points.assert_called_once_with( - group="torchtitan.compile_backends", - name="custom_device", - ) - backend.load.assert_called_once_with() set_default_backend.assert_called_once_with("custom_device_backend") diff --git a/torchtitan/tools/utils.py b/torchtitan/tools/utils.py index a060660ae86..eb59202076d 100644 --- a/torchtitan/tools/utils.py +++ b/torchtitan/tools/utils.py @@ -5,7 +5,6 @@ # LICENSE file in the root directory of this source tree. import contextlib -import importlib.metadata import logging import os import subprocess @@ -67,17 +66,14 @@ def get_device_info() -> tuple[str, ModuleType]: return device_type, device_module -def _load_device_compile_backend(device_type: str) -> None: - """Load an installed compile backend registered for ``device_type``.""" - for backend in importlib.metadata.entry_points( - group="torchtitan.compile_backends", - name=device_type, - ): - torch.compiler.set_default_backend(backend.load()) +def _configure_compile_backend() -> None: + """Configure the default ``torch.compile`` backend from the environment.""" + if backend := os.environ.get("TORCHTITAN_COMPILE_BACKEND"): + torch.compiler.set_default_backend(backend) device_type, device_module = get_device_info() -_load_device_compile_backend(device_type) +_configure_compile_backend() def get_local_device() -> torch.device: