Make TorchTitan local compile configuration extensible

tests/rl/unit_tests/cpu/test_generator.py

@@ -576,8 +576,8 @@ def test_qwen36_27b_perf_config():
     config = rl_grpo_qwen3_6_27b_varlen_perf()
 
     assert config.hf_assets_path.endswith("Qwen3.6-27B")
-    assert "offset_rmsnorm" in config.trainer.model.local_compile_regions
-    assert "offset_rmsnorm" in config.generator.model.local_compile_regions
+    assert "offset_rmsnorm" in config.trainer.model.local_compile_regions.regions
+    assert "offset_rmsnorm" in config.generator.model.local_compile_regions.regions
     assert config.trainer.parallelism.data_parallel_shard_degree == 2
     assert config.trainer.parallelism.tensor_parallel_degree == 2
     assert config.generator.parallelism.tensor_parallel_degree == 4

tests/rl/unit_tests/cpu/test_policy_loss_configs.py

@@ -75,5 +75,5 @@ def test_batch_invariant_configs_disable_local_compile(module_name, factory_name
     factory = getattr(importlib.import_module(module_name), factory_name)
     config = factory()
 
-    assert config.trainer.model.local_compile_regions == []
-    assert config.generator.model.local_compile_regions == []
+    assert config.trainer.model.local_compile_regions.regions == []
+    assert config.generator.model.local_compile_regions.regions == []

tests/rl/unit_tests/gpu/test_gdn_cudagraph.py

@@ -301,7 +301,7 @@ def recurrent(*args, **kwargs):
     model_config = build_model_config(
         "0.8B", seq_len=256, attn_backend="varlen", parallelism=ParallelismConfig()
     )
-    model_config.local_compile_regions = []
+    model_config.local_compile_regions.regions = []
     registry.register_to_vllm(
         model_config,
         parallelism=registry.InferenceParallelismConfig(tensor_parallel_degree=1),

tests/unit_tests/cpu/config/test_compile_config.py

@@ -4,24 +4,46 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from collections.abc import Callable
+from typing import Any
+
 import pytest
 import torch
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 
 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 local_compile, LocalCompileConfig
+from torchtitan.models.common.attention import FlexInnerAttention
+from torchtitan.models.common.attention.attention import _compiled_create_block_mask
 from torchtitan.models.common.moe import _microbatch_load_balance_local_stats
+from torchtitan.models.common.vision_encoder import compiled_create_block_mask
+
+
+class CustomDeviceLocalCompileConfig(LocalCompileConfig):
+    """Example configuration supplied by an out-of-tree backend."""
+
+    def _compile(
+        self,
+        reference: Callable[..., Any],
+        *,
+        inductor_options: dict[str, Any] | None = None,
+        **compile_kwargs: Any,
+    ) -> Callable[..., Any]:
+        del inductor_options
+        compile_kwargs["backend"] = "custom_device_backend"
+        return torch.compile(reference, **compile_kwargs)
 
 
 @pytest.fixture(autouse=True)
 def reset_local_compile():
-    apply_local_compile([])
+    LocalCompileConfig().apply()
     yield
-    apply_local_compile([])
+    LocalCompileConfig().apply()
 
 
-def test_apply_local_compile_rejects_unknown_name() -> None:
+def test_local_compile_config_rejects_unknown_name() -> None:
     with pytest.raises(ValueError, match=r"foo.*registered values"):
-        apply_local_compile(["foo"])
+        LocalCompileConfig(regions=["foo"]).apply()
 
 
 def test_local_compile_binds_once_from_existing_config(monkeypatch) -> None:
@@ -40,12 +62,13 @@ def compiled(value: int) -> tuple[str, int]:
         return compiled
 
     monkeypatch.setattr(torch, "compile", fake_compile)
-    apply_local_compile(["test_local_compile"])
+    LocalCompileConfig(regions=["test_local_compile"]).apply()
 
     assert fn(3) == ("compiled", 3)
-    assert len(compiled_calls) == 1
-    assert "backend" not in compiled_calls[0][1]
-    assert compiled_calls[0][1]["fullgraph"] is True
+    named_calls = [call for call in compiled_calls if call[0] is fn.__wrapped__]
+    assert len(named_calls) == 1
+    assert "backend" not in named_calls[0][1]
+    assert named_calls[0][1]["fullgraph"] is True
 
 
 def test_local_compile_forwards_compile_kwargs(monkeypatch) -> None:
@@ -68,10 +91,11 @@ def fake_compile(reference, **kwargs):
         return reference
 
     monkeypatch.setattr(torch, "compile", fake_compile)
-    apply_local_compile(["test_compile_kwargs"])
+    LocalCompileConfig(regions=["test_compile_kwargs"]).apply()
 
     assert fn(3) == 3
-    assert compiled_calls[0][1] == {
+    named_calls = [call for call in compiled_calls if call[0] is fn.__wrapped__]
+    assert named_calls[0][1] == {
         "backend": "eager",
         "fullgraph": True,
         "dynamic": True,
@@ -100,7 +124,7 @@ def fn(value: int) -> int:
     with pytest.raises(
         ValueError, match=r"test_non_batch_invariant.*local_compile_regions"
     ):
-        apply_local_compile(["test_non_batch_invariant"])
+        LocalCompileConfig(regions=["test_non_batch_invariant"]).apply()
 
 
 def test_loss_functions_use_local_compile(monkeypatch) -> None:
@@ -112,9 +136,18 @@ def fake_compile(reference, **kwargs):
         return reference
 
     monkeypatch.setattr(torch, "compile", fake_compile)
-    apply_local_compile(["loss"])
-
-    assert compiled_names == [
+    LocalCompileConfig(regions=["loss"]).apply()
+
+    assert [
+        name
+        for name in compiled_names
+        if name
+        in {
+            cross_entropy_loss.__name__,
+            mse_loss.__name__,
+            compute_logprobs.__name__,
+        }
+    ] == [
         cross_entropy_loss.__name__,
         mse_loss.__name__,
         compute_logprobs.__name__,
@@ -129,11 +162,41 @@ def fake_compile(reference, **kwargs):
         return reference
 
     monkeypatch.setattr(torch, "compile", fake_compile)
-    apply_local_compile(["moe_aux_loss"])
+    LocalCompileConfig(regions=["moe_aux_loss"]).apply()
 
-    assert compiled == [
+    assert [
+        call
+        for call in compiled
+        if call[0] is _microbatch_load_balance_local_stats.__wrapped__
+    ] == [
         (
             _microbatch_load_balance_local_stats.__wrapped__,
             {"fullgraph": True},
         )
     ]
+
+
+def test_custom_config_controls_all_torchtitan_compile_boundaries(
+    monkeypatch,
+) -> None:
+    compiled_calls = []
+
+    @local_compile("test_custom_backend", batch_invariant=False)
+    def region() -> None:
+        pass
+
+    def fake_compile(reference, **kwargs):
+        compiled_calls.append((reference, kwargs))
+        return reference
+
+    monkeypatch.setattr(torch, "compile", fake_compile)
+    CustomDeviceLocalCompileConfig(regions=["test_custom_backend"]).apply()
+
+    assert callable(FlexInnerAttention._compiled_flex_attn)
+    assert callable(_compiled_create_block_mask)
+    assert callable(compiled_create_block_mask)
+    expected = {"backend": "custom_device_backend"}
+    assert (region.__wrapped__, {**expected, "fullgraph": True}) in compiled_calls
+    assert (flex_attention, expected) in compiled_calls
+    assert (create_block_mask, expected) in compiled_calls
+    assert sum(reference is create_block_mask for reference, _ in compiled_calls) == 2

tests/unit_tests/cpu/config/test_context_parallel_validation.py

@@ -186,6 +186,7 @@ class TestGptOssUlysses(unittest.TestCase):
     def _parallelize(self, inner_attention):
         from types import SimpleNamespace
 
+        from torchtitan.distributed.local_compile import LocalCompileConfig
         from torchtitan.models.gpt_oss.model import GptOssModel
 
         with self.assertRaisesRegex(NotImplementedError, "Ulysses CP"):
@@ -198,7 +199,7 @@ def _parallelize(self, inner_attention):
                 parallelism_context=SimpleNamespace(cp_enabled=True),
                 training=None,
                 parallelism=None,
-                local_compile_regions=[],
+                local_compile_regions=LocalCompileConfig(),
                 activation_remat=None,
                 dump_folder="",
             )

tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py

@@ -9,13 +9,13 @@
 import torch
 import torch.nn.functional as F
 
-from torchtitan.distributed.local_compile import apply_local_compile
+from torchtitan.distributed.local_compile import LocalCompileConfig
 from torchtitan.models.common import GatedRMSNorm, Sigmoid, SiLU
 
 
 class TestGatedRMSNorm(unittest.TestCase):
     def setUp(self):
-        apply_local_compile([])
+        LocalCompileConfig().apply()
 
     def test_configurable_activation_matches_reference(self):
         x = torch.randn(4, 3, 8)

tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py

@@ -215,7 +215,7 @@ def test_lm_head_converter_swaps_only_lm_head():
     # The head keeps 2 pieces: summed over the vocab, a third doesn't help.
     assert config.lm_head.backward_mode == "hi_mid"
     # It turns on the split's local_compile region.
-    assert "fp32_to_bf16_split" in config.local_compile_regions
+    assert "fp32_to_bf16_split" in config.local_compile_regions.regions
     # Converting an already converted head is a no-op.
     converted = config.lm_head
     LMHeadFP32OutputConverter.Config().build().convert(config)

tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py

@@ -21,7 +21,7 @@ def test_kimi_k3_enables_attention_residual_local_compile() -> None:
     config = kimi_k3_debugmodel()
     model_config = cast(KimiK3Model.Config, config.model)
 
-    assert model_config.local_compile_regions == [
+    assert model_config.local_compile_regions.regions == [
         "loss",
         "gated_rmsnorm",
         "fused_binary_activation",

tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py

@@ -12,7 +12,7 @@
 def test_qwen35_architecture_enables_offset_rmsnorm_local_compile() -> None:
     for config_factory in (qwen35_debugmodel, qwen36_debugmodel, qwen38_debugmodel):
         config = config_factory()
-        assert config.model.local_compile_regions == [
+        assert config.model.local_compile_regions.regions == [
             "loss",
             "fused_binary_activation",
             "gated_rmsnorm",

tests/unit_tests/cpu/quantization/test_quantization.py

@@ -479,7 +479,7 @@ def test_qwen3_recipes_resolve(monkeypatch, recipe):
         assert isinstance(dataset.source, HuggingFaceRandomAccessSource.Config)
         assert dataset.source.path == "openai/gsm8k"
         assert config.checkpointer.initial_load_in_hf
-        assert config.model.local_compile_regions == [
+        assert config.model.local_compile_regions.regions == [
             "loss",
             "fused_binary_activation",
             "cos_sin_rope",
@@ -987,7 +987,7 @@ def test_deepseek_nvfp4_recipes_preserve_quantization_and_routing(
         if bf16_tail_fraction == 0
         else len(config.model.layers) // 2
     )
-    assert config.model.local_compile_regions == ["loss"]
+    assert config.model.local_compile_regions.regions == ["loss"]
     grouped = list(config.model.traverse(GroupedLinear.Config))
     assert grouped
     quantized_cls = _get_nvfp4_grouped_linear_cls(GroupedLinear)

tests/unit_tests/cpu/test_integration_test_definitions.py

@@ -54,7 +54,7 @@ def test_spmd_typechecking_config_disables_local_compile() -> None:
     config = llama3_debugmodel_default()
 
     assert config.debug.spmd_typechecking
-    assert config.model.local_compile_regions == []
+    assert config.model.local_compile_regions.regions == []
     config.__post_init__()
 
 

tests/unit_tests/cpu/test_trainer.py

@@ -756,7 +756,9 @@ 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=SimpleNamespace(regions=["loss"])
+            ),
             max_num_documents=None,
             output_dir="",
         )

tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py

@@ -8,17 +8,17 @@
 
 import torch
 
-from torchtitan.distributed.local_compile import apply_local_compile
+from torchtitan.distributed.local_compile import LocalCompileConfig
 from torchtitan.models.common import GatedRMSNorm, Sigmoid
 
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
 class TestGatedRMSNormCompile(unittest.TestCase):
     def setUp(self):
-        apply_local_compile(["gated_rmsnorm"])
+        LocalCompileConfig(regions=["gated_rmsnorm"]).apply()
 
     def tearDown(self):
-        apply_local_compile([])
+        LocalCompileConfig().apply()
         torch._dynamo.reset()
 
     @staticmethod

tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py

@@ -18,7 +18,7 @@
 from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
 from torch.fx.experimental.proxy_tensor import make_fx
 
-from torchtitan.distributed.local_compile import apply_local_compile
+from torchtitan.distributed.local_compile import LocalCompileConfig
 from torchtitan.models.common import hi_mid_lo_linear
 from torchtitan.models.common.hi_mid_lo_linear import HiMidLoLinear
 
@@ -177,11 +177,11 @@ def linear(input, weight):
 @contextlib.contextmanager
 def _compiled_split():
     """Turn on the split's local_compile region, as the models that use HiMidLoLinear do."""
-    apply_local_compile(["fp32_to_bf16_split"])
+    LocalCompileConfig(regions=["fp32_to_bf16_split"]).apply()
     try:
         yield
     finally:
-        apply_local_compile([])
+        LocalCompileConfig().apply()
 
 
 @pytest.mark.parametrize("dim", [0, 1])

tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py

@@ -8,14 +8,14 @@
 
 import torch
 
-from torchtitan.distributed.local_compile import apply_local_compile
+from torchtitan.distributed.local_compile import LocalCompileConfig
 from torchtitan.models.common.moe import _microbatch_load_balance_local_stats
 
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
 class TestMoEAuxLossLocalCompile(unittest.TestCase):
     def tearDown(self):
-        apply_local_compile([])
+        LocalCompileConfig().apply()
         torch._dynamo.reset()
 
     def test_forward_and_backward_match_eager(self):
@@ -43,7 +43,7 @@ def test_forward_and_backward_match_eager(self):
             generator=generator,
         )
 
-        apply_local_compile([])
+        LocalCompileConfig().apply()
         eager_counts_E, eager_prob_sums_E = _microbatch_load_balance_local_stats(
             scores_TE,
             routing_map_TE,
@@ -53,7 +53,7 @@ def test_forward_and_backward_match_eager(self):
             eager_prob_sums_E, scores_TE, grad_prob_sums_E
         )[0]
 
-        apply_local_compile(["moe_aux_loss"])
+        LocalCompileConfig(regions=["moe_aux_loss"]).apply()
         compiled_counts_E, compiled_prob_sums_E = _microbatch_load_balance_local_stats(
             scores_TE,
             routing_map_TE,

tests/unit_tests/gpu/models/common/test_situglu.py

@@ -8,17 +8,17 @@
 
 import torch
 
-from torchtitan.distributed.local_compile import apply_local_compile
+from torchtitan.distributed.local_compile import LocalCompileConfig
 from torchtitan.models.common.activation import SiTUGLU
 
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
 class TestSiTUGLULocalCompile(unittest.TestCase):
     def setUp(self):
-        apply_local_compile(["fused_binary_activation"])
+        LocalCompileConfig(regions=["fused_binary_activation"]).apply()
 
     def tearDown(self):
-        apply_local_compile([])
+        LocalCompileConfig().apply()
         torch._dynamo.reset()
 
     def test_forward_and_backward_emit_triton(self):