Make local compile backend configurable

scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py

@@ -99,6 +99,7 @@ def run_worker(args):
         training=training,
         parallelism=parallelism,
         local_compile_regions=config.local_compile_regions,
+        local_compile_backend=config.local_compile_backend,
         activation_remat=DefaultActivationRematConfig(),
         dump_folder="/tmp",
     )

tests/rl/unit_tests/gpu/test_bitwise_parity.py

@@ -143,6 +143,7 @@ def build_trainer_model(
         training=trainer_config.training,
         parallelism=parallelism,
         local_compile_regions=model_config.local_compile_regions,
+        local_compile_backend=model_config.local_compile_backend,
         activation_remat=trainer_config.activation_remat,
         dump_folder=config.dump_folder,
     )

tests/unit_tests/cpu/config/test_compile_config.py

@@ -8,7 +8,11 @@
 import torch
 
 from torchtitan.components.loss import compute_logprobs, cross_entropy_loss, mse_loss
-from torchtitan.distributed.local_compile import apply_local_compile, local_compile
+from torchtitan.distributed.local_compile import (
+    apply_local_compile,
+    compiled,
+    local_compile,
+)
 from torchtitan.models.common.moe import _microbatch_load_balance_local_stats
 
 
@@ -24,7 +28,7 @@ def test_apply_local_compile_rejects_unknown_name() -> None:
         apply_local_compile(["foo"])
 
 
-def test_local_compile_binds_once_from_existing_config(monkeypatch) -> None:
+def test_local_compile_uses_default_backend(monkeypatch) -> None:
     compiled_calls = []
 
     @local_compile("test_local_compile", batch_invariant=False)
@@ -43,19 +47,17 @@ def compiled(value: int) -> tuple[str, int]:
     apply_local_compile(["test_local_compile"])
 
     assert fn(3) == ("compiled", 3)
-    assert len(compiled_calls) == 1
-    assert "backend" not in compiled_calls[0][1]
+    assert compiled_calls[0][1]["backend"] == "inductor"
     assert compiled_calls[0][1]["fullgraph"] is True
 
 
-def test_local_compile_forwards_compile_kwargs(monkeypatch) -> None:
+def test_local_compile_forwards_backend_and_compile_kwargs(monkeypatch) -> None:
     compiled_calls = []
     options = {"max_autotune": True}
 
     @local_compile(
         "test_compile_kwargs",
         batch_invariant=False,
-        backend="eager",
         dynamic=True,
         fullgraph=True,
         options=options,
@@ -68,11 +70,11 @@ def fake_compile(reference, **kwargs):
         return reference
 
     monkeypatch.setattr(torch, "compile", fake_compile)
-    apply_local_compile(["test_compile_kwargs"])
+    apply_local_compile(["test_compile_kwargs"], backend="aot_eager")
 
     assert fn(3) == 3
     assert compiled_calls[0][1] == {
-        "backend": "eager",
+        "backend": "aot_eager",
         "fullgraph": True,
         "dynamic": True,
         "options": options,
@@ -88,6 +90,91 @@ def test_local_compile_rejects_fullgraph_false() -> None:
         )
 
 
+def test_local_compile_rejects_per_region_backend() -> None:
+    with pytest.raises(ValueError, match="apply_local_compile"):
+        local_compile(
+            "test_backend",
+            batch_invariant=False,
+            backend="aot_eager",
+        )
+
+
+def test_compiled_helper_binds_lazily_and_rebinds_backend(monkeypatch) -> None:
+    compiled_calls = []
+    inductor_options = {"max_autotune": True}
+
+    def reference(value: int) -> tuple[str, int]:
+        return "eager", value
+
+    def fake_compile(reference, **kwargs):
+        compiled_calls.append((reference, kwargs))
+        backend = kwargs["backend"]
+
+        def compiled_fn(value: int) -> tuple[str, int]:
+            return backend, value
+
+        return compiled_fn
+
+    monkeypatch.setattr(torch, "compile", fake_compile)
+    helper = compiled(reference, inductor_options=inductor_options)
+
+    assert compiled_calls == []
+    assert helper(3) == ("inductor", 3)
+    assert compiled_calls == [
+        (
+            reference,
+            {"backend": "inductor", "options": inductor_options},
+        )
+    ]
+
+    inductor_options["max_autotune"] = False
+    apply_local_compile([])
+
+    assert compiled_calls[-1] == (
+        reference,
+        {"backend": "inductor", "options": {"max_autotune": False}},
+    )
+
+    apply_local_compile([], backend="custom_device_backend")
+
+    assert helper(4) == ("custom_device_backend", 4)
+    assert compiled_calls[-1] == (
+        reference,
+        {"backend": "custom_device_backend"},
+    )
+
+
+def test_torchtitan_helpers_use_configured_backend(monkeypatch) -> None:
+    compiled_calls = []
+
+    def fake_compile(reference, **kwargs):
+        compiled_calls.append((reference, kwargs))
+        return reference
+
+    monkeypatch.setattr(torch, "compile", fake_compile)
+
+    from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+    from torchtitan.models.common import vision_encoder  # noqa: F401
+
+    from torchtitan.models.common.attention import attention
+
+    assert compiled_calls == []
+
+    apply_local_compile([], backend="custom_device_backend")
+
+    helper_calls = [
+        call
+        for call in compiled_calls
+        if call[0] in (flex_attention, create_block_mask)
+    ]
+    assert helper_calls == [
+        (flex_attention, {"backend": "custom_device_backend"}),
+        (create_block_mask, {"backend": "custom_device_backend"}),
+        (create_block_mask, {"backend": "custom_device_backend"}),
+    ]
+    assert attention.FlexInnerAttention._compiled_flex_attn is not flex_attention
+
+
 def test_local_compile_rejects_non_batch_invariant_region(monkeypatch) -> None:
     @local_compile("test_non_batch_invariant", batch_invariant=False)
     def fn(value: int) -> int:
@@ -114,7 +201,7 @@ def fake_compile(reference, **kwargs):
     monkeypatch.setattr(torch, "compile", fake_compile)
     apply_local_compile(["loss"])
 
-    assert compiled_names == [
+    assert compiled_names[:3] == [
         cross_entropy_loss.__name__,
         mse_loss.__name__,
         compute_logprobs.__name__,
@@ -131,9 +218,7 @@ def fake_compile(reference, **kwargs):
     monkeypatch.setattr(torch, "compile", fake_compile)
     apply_local_compile(["moe_aux_loss"])
 
-    assert compiled == [
-        (
-            _microbatch_load_balance_local_stats.__wrapped__,
-            {"fullgraph": True},
-        )
-    ]
+    assert compiled[0] == (
+        _microbatch_load_balance_local_stats.__wrapped__,
+        {"backend": "inductor", "fullgraph": True},
+    )

tests/unit_tests/cpu/test_trainer.py

@@ -756,14 +756,31 @@ def test_graph_training_engine_rejects_local_compile_regions() -> None:
     ):
         GraphTrainingEngine(
             config,
-            model_config=SimpleNamespace(local_compile_regions=["loss"]),
+            model_config=SimpleNamespace(
+                local_compile_regions=["loss"],
+                local_compile_backend="inductor",
+            ),
             max_num_documents=None,
             output_dir="",
         )
 
     init.assert_not_called()
 
 
+def test_parallelize_compile_kwargs_include_backend() -> None:
+    engine = SimpleNamespace(
+        model_config=SimpleNamespace(
+            local_compile_regions=["loss"],
+            local_compile_backend="custom_device_backend",
+        )
+    )
+
+    assert TrainingEngine._parallelize_compile_kwargs(engine) == {
+        "local_compile_regions": ["loss"],
+        "local_compile_backend": "custom_device_backend",
+    }
+
+
 def test_optim_update_clips_before_parameter_update() -> None:
     events = []
     optimizers = MagicMock()

torchtitan/distributed/local_compile.py

@@ -15,10 +15,13 @@
 from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode
 
 
-def apply_local_compile(regions: list[str]) -> None:
+_BACKEND: str = "inductor"
+
+
+def apply_local_compile(regions: list[str], *, backend: str = "inductor") -> None:
     """Bind registered functions to eager or compiled implementations.
 
-    Process-wide: a later call replaces this choice for every model in the process.
+    Process-wide: a later call replaces these choices for every model in the process.
     """
     unknown = [name for name in regions if name not in _LOCAL_COMPILE_CALLBACKS]
     if unknown:
@@ -27,12 +30,48 @@ def apply_local_compile(regions: list[str]) -> None:
             f"registered values are {sorted(_LOCAL_COMPILE_CALLBACKS)}"
         )
 
+    global _BACKEND
+    _BACKEND = backend
     for callbacks in _LOCAL_COMPILE_CALLBACKS.values():
         for bind_local_compile_fn in callbacks:
             bind_local_compile_fn(regions)
+    for bind_compiled_helper in _COMPILED_HELPERS:
+        bind_compiled_helper()
 
 
 _LOCAL_COMPILE_CALLBACKS: dict[str, list[Callable[[list[str]], None]]] = {}
+_COMPILED_HELPERS: list[Callable[[], None]] = []
+
+
+def compiled(
+    reference: Callable[..., Any],
+    *,
+    inductor_options: dict[str, Any] | None = None,
+) -> Callable[..., Any]:
+    """Return ``reference`` compiled with the configured backend.
+
+    The function binds on first use and rebinds whenever ``apply_local_compile``
+    changes the process-wide backend. ``inductor_options`` are passed only when
+    that backend is inductor.
+    """
+    fn: Callable[..., Any] | None = None
+
+    def bind_compiled_helper() -> None:
+        nonlocal fn
+        compile_kwargs: dict[str, Any] = {"backend": _BACKEND}
+        if _BACKEND == "inductor" and inductor_options is not None:
+            compile_kwargs["options"] = inductor_options
+        fn = torch.compile(reference, **compile_kwargs)
+
+    @functools.wraps(reference)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+        if fn is None:
+            bind_compiled_helper()
+        assert fn is not None
+        return fn(*args, **kwargs)
+
+    _COMPILED_HELPERS.append(bind_compiled_helper)
+    return wrapped
 
 
 def local_compile(
@@ -52,6 +91,8 @@ def local_compile(
     """
     if compile_kwargs.pop("fullgraph", True) is not True:
         raise ValueError("local_compile requires fullgraph=True.")
+    if "backend" in compile_kwargs:
+        raise ValueError("local_compile backend is configured by apply_local_compile.")
 
     def decorate(reference: Callable[..., Any]) -> Callable[..., Any]:
         fn = reference
@@ -66,7 +107,12 @@ def bind_local_compile(regions: list[str]) -> None:
                     "batch-invariant mode; remove it from local_compile_regions."
                 )
             if enabled:
-                fn = torch.compile(reference, fullgraph=True, **compile_kwargs)
+                fn = torch.compile(
+                    reference,
+                    fullgraph=True,
+                    backend=_BACKEND,
+                    **compile_kwargs,
+                )
             else:
                 fn = reference
 
@@ -80,4 +126,4 @@ def wrapped(*args: Any, **kwargs: Any) -> Any:
     return decorate
 
 
-__all__ = ["apply_local_compile", "local_compile"]
+__all__ = ["apply_local_compile", "compiled", "local_compile"]

torchtitan/distributed/pipeline_parallel.py

@@ -79,6 +79,7 @@ def pipeline_llm(
     training: TrainingConfig,
     parallelism: ParallelismConfig,
     local_compile_regions: list[str],
+    local_compile_backend: str = "inductor",
     activation_remat: ActivationRematConfig | None,
     dump_folder: str,
     device: torch.device,
@@ -145,6 +146,7 @@ def pipeline_llm(
             training=training,
             parallelism=parallelism,
             local_compile_regions=local_compile_regions,
+            local_compile_backend=local_compile_backend,
             activation_remat=activation_remat,
             dump_folder=dump_folder,
         )

torchtitan/distributed/utils.py

@@ -149,11 +149,11 @@ def set_determinism(
         # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html
         os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
-        from torch.nn.attention.flex_attention import flex_attention
-
         from torchtitan.models.common.attention import FlexInnerAttention
 
         if torch.version.hip is not None:
+            from torch.nn.attention.flex_attention import flex_attention
+
             # Compiled ROCm flex attention is not deterministic.
             # Falling back to eager (non-compiled) flex_attention for determinism on ROCm.
             logger.info(
@@ -167,10 +167,6 @@ def set_determinism(
             # no options) so the regional-inductor scoop configs are preserved.
             FlexInnerAttention.inductor_configs["max_autotune"] = False
             FlexInnerAttention.inductor_configs["coordinate_descent_tuning"] = False
-            # pyrefly: ignore [no-matching-overload]
-            FlexInnerAttention._compiled_flex_attn = torch.compile(
-                flex_attention, options=FlexInnerAttention.inductor_configs
-            )
 
     if debug_config.detect_anomaly:
         logger.warning(

torchtitan/experiments/transformers_modeling_backend/parallelize.py

@@ -88,6 +88,7 @@ def parallelize_hf_transformers(
     training: TrainingConfig,
     parallelism: ParallelismConfig,
     local_compile_regions: list[str],
+    local_compile_backend: str = "inductor",
     activation_remat: ActivationRematConfig | None,
     dump_folder: str,
     **kwargs: Any,
@@ -102,7 +103,10 @@ def parallelize_hf_transformers(
     5. Apply AC and FSDP
     """
     # Bind local implementations early; torch.compile traces on first use.
-    apply_local_compile(local_compile_regions)
+    apply_local_compile(
+        local_compile_regions,
+        backend=local_compile_backend,
+    )
     # Flex attention supports FSDP, TP, CP, and PP (in any combination). Under CP
     # the flex kernel's local SPMD boundary redistributes
     # k/v from seq-sharded to CP-Replicate (all-gather); see _attach_flex_kernel

torchtitan/experiments/transformers_modeling_backend/pipeline.py

@@ -297,6 +297,7 @@ def pipeline_hf_transformers(
     training: TrainingConfig,
     parallelism: ParallelismConfig,
     local_compile_regions: list[str],
+    local_compile_backend: str = "inductor",
     activation_remat: ActivationRematConfig | None,
     dump_folder: str,
     device: torch.device,
@@ -391,6 +392,7 @@ def pipeline_hf_transformers(
             training=training,
             parallelism=parallelism,
             local_compile_regions=local_compile_regions,
+            local_compile_backend=local_compile_backend,
             activation_remat=activation_remat,
             dump_folder=dump_folder,
         )