From d211792a2de24d55c0ab1ef5bcf827ee0916c277 Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Thu, 15 Oct 2026 18:56:01 +0000 Subject: [PATCH 1/6] Make local compile backend configurable --- .../cpu/config/test_compile_config.py | 95 +++++++++++++++++-- tests/unit_tests/cpu/test_trainer.py | 27 +++++- torchtitan/distributed/local_compile.py | 78 ++++++++++++++- torchtitan/distributed/utils.py | 4 - .../transformers_modeling_backend/model.py | 1 + .../tests/test_flex_attention.py | 21 ++++ torchtitan/models/base/model.py | 3 + .../models/common/attention/attention.py | 8 +- torchtitan/models/common/vision_encoder.py | 3 +- torchtitan/rl/generator/vllm/model_wrapper.py | 2 + torchtitan/training_engine.py | 2 + 11 files changed, 220 insertions(+), 24 deletions(-) diff --git a/tests/unit_tests/cpu/config/test_compile_config.py b/tests/unit_tests/cpu/config/test_compile_config.py index 4cbcf0ce629..79c95ec5e31 100644 --- a/tests/unit_tests/cpu/config/test_compile_config.py +++ b/tests/unit_tests/cpu/config/test_compile_config.py @@ -6,16 +6,27 @@ 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 ( + apply_local_compile, + compiled, + local_compile, + set_local_compile_backend, +) +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 @pytest.fixture(autouse=True) def reset_local_compile(): + set_local_compile_backend("inductor") apply_local_compile([]) yield + set_local_compile_backend("inductor") apply_local_compile([]) @@ -24,7 +35,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_rebinds_active_region_when_backend_changes(monkeypatch) -> None: compiled_calls = [] @local_compile("test_local_compile", batch_invariant=False) @@ -33,19 +44,26 @@ def fn(value: int) -> tuple[str, int]: def fake_compile(reference, **kwargs): compiled_calls.append((reference, kwargs)) + backend = kwargs["backend"] def compiled(value: int) -> tuple[str, int]: - return "compiled", value + return backend, value return compiled monkeypatch.setattr(torch, "compile", fake_compile) 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]["fullgraph"] is True + assert fn(3) == ("inductor", 3) + + set_local_compile_backend("custom_device_backend") + apply_local_compile(["test_local_compile"]) + + assert fn(3) == ("custom_device_backend", 3) + assert [kwargs for _, kwargs in compiled_calls] == [ + {"backend": "inductor", "fullgraph": True}, + {"backend": "custom_device_backend", "fullgraph": True}, + ] def test_local_compile_forwards_compile_kwargs(monkeypatch) -> None: @@ -134,6 +152,67 @@ def fake_compile(reference, **kwargs): assert compiled == [ ( _microbatch_load_balance_local_stats.__wrapped__, - {"fullgraph": True}, + {"backend": "inductor", "fullgraph": True}, ) ] + + +def test_compiled_helper_only_forwards_options_to_inductor(monkeypatch) -> None: + compiled_calls = [] + options = {"max_autotune": False} + + def reference(value: int) -> int: + return value + + def fake_compile(fn, **kwargs): + compiled_calls.append((fn, kwargs)) + return fn + + monkeypatch.setattr(torch, "compile", fake_compile) + helper = compiled(reference, inductor_options=options) + set_local_compile_backend("custom_device_backend") + + assert compiled_calls == [] + assert helper(3) == 3 + set_local_compile_backend("inductor") + assert helper(3) == 3 + assert compiled_calls == [ + ( + reference, + { + "backend": "custom_device_backend", + }, + ), + ( + reference, + { + "backend": "inductor", + "options": options, + }, + ), + ] + + +def test_attention_helpers_do_not_forward_inductor_options_to_custom_backend( + monkeypatch, +) -> None: + compiled_calls = [] + + def fake_compile(reference, **kwargs): + if "options" in kwargs: + raise TypeError("custom backend got unexpected keyword argument 'options'") + compiled_calls.append((reference, kwargs)) + return lambda *args, **kwargs: None + + monkeypatch.setattr(torch, "compile", fake_compile) + set_local_compile_backend("custom_device_backend") + + FlexInnerAttention._compiled_flex_attn() + _compiled_create_block_mask() + compiled_create_block_mask() + + assert compiled_calls == [ + (flex_attention, {"backend": "custom_device_backend"}), + (create_block_mask, {"backend": "custom_device_backend"}), + (create_block_mask, {"backend": "custom_device_backend"}), + ] diff --git a/tests/unit_tests/cpu/test_trainer.py b/tests/unit_tests/cpu/test_trainer.py index c79f1b0798f..deb91aaf1a2 100644 --- a/tests/unit_tests/cpu/test_trainer.py +++ b/tests/unit_tests/cpu/test_trainer.py @@ -764,6 +764,16 @@ def test_graph_training_engine_rejects_local_compile_regions() -> None: init.assert_not_called() +def test_parallelize_compile_kwargs_include_regions_only() -> None: + engine = SimpleNamespace( + model_config=SimpleNamespace(local_compile_regions=["loss"]) + ) + + assert TrainingEngine._parallelize_compile_kwargs(engine) == { + "local_compile_regions": ["loss"], + } + + def test_optim_update_clips_before_parameter_update() -> None: events = [] optimizers = MagicMock() @@ -1096,20 +1106,27 @@ def test_seed_checkpoint_initialize_skips_forward_backward(): ), _dist_moe_runtime=None, state_dict_adapter=None, + model_config=SimpleNamespace(local_compile_backend="custom_device_backend"), ), ) - TrainingEngine.initialize( - engine, - hf_assets_path="", - create_seed_checkpoint=True, - ) + with patch( + "torchtitan.training_engine.set_local_compile_backend", + side_effect=lambda backend: events.append(("backend", backend)), + ) as set_backend: + TrainingEngine.initialize( + engine, + hf_assets_path="", + create_seed_checkpoint=True, + ) assert events == [ + ("backend", "custom_device_backend"), "model", "model_memory", "optim", "checkpointer", ] + set_backend.assert_called_once_with("custom_device_backend") assert engine.model_device_mem_stats is model_mem_stats engine._initialize_model.assert_called_once_with( hf_assets_path="", diff --git a/torchtitan/distributed/local_compile.py b/torchtitan/distributed/local_compile.py index 6774c21ebf6..9ccb307c539 100644 --- a/torchtitan/distributed/local_compile.py +++ b/torchtitan/distributed/local_compile.py @@ -15,6 +15,22 @@ from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode +_BACKEND: str = "inductor" +_COMPILED_HELPER_RESETS: list[Callable[[], None]] = [] + + +def _reset_compiled_helpers() -> None: + for reset in _COMPILED_HELPER_RESETS: + reset() + + +def set_local_compile_backend(backend: str) -> None: + """Set the process-wide backend used by local compile helpers.""" + global _BACKEND + _BACKEND = backend + _reset_compiled_helpers() + + def apply_local_compile(regions: list[str]) -> None: """Bind registered functions to eager or compiled implementations. @@ -52,6 +68,7 @@ def local_compile( """ if compile_kwargs.pop("fullgraph", True) is not True: raise ValueError("local_compile requires fullgraph=True.") + explicit_backend = compile_kwargs.pop("backend", None) def decorate(reference: Callable[..., Any]) -> Callable[..., Any]: fn = reference @@ -66,7 +83,13 @@ 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) + backend = _BACKEND if explicit_backend is None else explicit_backend + fn = torch.compile( + reference, + fullgraph=True, + backend=backend, + **compile_kwargs, + ) else: fn = reference @@ -80,4 +103,55 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: return decorate -__all__ = ["apply_local_compile", "local_compile"] +def compiled( + reference: Callable[..., Any], + *, + inductor_options: dict[str, Any] | None = None, + **compile_kwargs: Any, +) -> Callable[..., Any]: + """Compile ``reference`` lazily with the configured local backend. + + ``inductor_options`` are passed as ``torch.compile(options=...)`` only when + the effective backend is Inductor. + """ + explicit_backend = compile_kwargs.pop("backend", None) + explicit_options = compile_kwargs.pop("options", None) + fn: Callable[..., Any] | None = None + + def reset() -> None: + nonlocal fn + fn = None + + @functools.wraps(reference) + def wrapped(*args: Any, **kwargs: Any) -> Any: + nonlocal fn + if fn is None: + backend = _BACKEND if explicit_backend is None else explicit_backend + options = explicit_options + if backend == "inductor" and inductor_options is not None: + options = inductor_options + if options is None: + fn = torch.compile( + reference, + backend=backend, + **compile_kwargs, + ) + else: + fn = torch.compile( + reference, + backend=backend, + options=options, + **compile_kwargs, + ) + return fn(*args, **kwargs) + + _COMPILED_HELPER_RESETS.append(reset) + return wrapped + + +__all__ = [ + "apply_local_compile", + "compiled", + "local_compile", + "set_local_compile_backend", +] diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py index 77c4dea011e..9b645e2430f 100644 --- a/torchtitan/distributed/utils.py +++ b/torchtitan/distributed/utils.py @@ -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( diff --git a/torchtitan/experiments/transformers_modeling_backend/model.py b/torchtitan/experiments/transformers_modeling_backend/model.py index 6ec6c1080d2..4d6c7e75302 100644 --- a/torchtitan/experiments/transformers_modeling_backend/model.py +++ b/torchtitan/experiments/transformers_modeling_backend/model.py @@ -286,6 +286,7 @@ def __init__( self.sharding_config = None # HF modules run eager; the loss and the swapped-in TorchTitan MoE (SwiGLU) can compile. self.local_compile_regions = ["loss", "fused_binary_activation"] + self.local_compile_backend = "inductor" assert model_config is not None, "model_config is required" diff --git a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py index 8f1c0b5b63d..fcf8d2b21f5 100644 --- a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py +++ b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py @@ -78,6 +78,27 @@ def test_packed_documents_do_not_leak(self): self.assertFalse(torch.allclose(out1[:, :len_a], out2[:, :len_a])) +class TestHFTransformerConfig(unittest.TestCase): + def test_local_compile_backend_defaults_to_inductor(self): + from torchtitan.experiments.transformers_modeling_backend import ( + TitanModelConfig, + ) + from torchtitan.experiments.transformers_modeling_backend.model import ( + HFTransformerModel, + ) + + config = HFTransformerModel.Config( + model_config=TitanModelConfig( + dim=256, + n_layers=2, + n_heads=16, + n_kv_heads=16, + ), + ) + + self.assertEqual(config.local_compile_backend, "inductor") + + class TestAttnMaskTypeValidation(unittest.TestCase): """A block_causal config builds a flex BlockMask and is supported. diff --git a/torchtitan/models/base/model.py b/torchtitan/models/base/model.py index 52760ed6a96..deb6b04de45 100644 --- a/torchtitan/models/base/model.py +++ b/torchtitan/models/base/model.py @@ -92,6 +92,9 @@ class Config(Module.Config): FlexAttention manages its own compilation and is not controlled by this list. """ + local_compile_backend: str = "inductor" + """Backend for TorchTitan-owned local and attention compile sites.""" + @abstractmethod def get_nparams_and_flops(self, model: Module, seq_len: int) -> tuple[int, int]: pass diff --git a/torchtitan/models/common/attention/attention.py b/torchtitan/models/common/attention/attention.py index 0258daedfc1..395434af279 100644 --- a/torchtitan/models/common/attention/attention.py +++ b/torchtitan/models/common/attention/attention.py @@ -39,6 +39,7 @@ ) from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode +from torchtitan.distributed.local_compile import compiled from torchtitan.distributed.parallelism_context import MeshAxisName from torchtitan.models.base.module import Module from torchtitan.models.common.linear import Linear @@ -344,10 +345,9 @@ def build_attention_metadata( "triton.cudagraphs": False, } - # pyrefly: ignore[no-matching-overload] - _compiled_flex_attn: ClassVar[Callable] = torch.compile( + _compiled_flex_attn: ClassVar[Callable] = compiled( flex_attention, - options=inductor_configs, + inductor_options=inductor_configs, ) def __init__(self, config: Config) -> None: @@ -642,7 +642,7 @@ def sliding_window_mod( return sliding_window_mod -_compiled_create_block_mask = torch.compile(create_block_mask) +_compiled_create_block_mask = compiled(create_block_mask) def create_attention_mask(*args, **kwargs): diff --git a/torchtitan/models/common/vision_encoder.py b/torchtitan/models/common/vision_encoder.py index a1d4b8c1fa3..ca3221a0077 100644 --- a/torchtitan/models/common/vision_encoder.py +++ b/torchtitan/models/common/vision_encoder.py @@ -28,6 +28,7 @@ import torch_remat as remat from torch.nn.attention.flex_attention import BlockMask, create_block_mask +from torchtitan.distributed.local_compile import compiled from torchtitan.distributed.parallelism_context import MeshAxisName from torchtitan.distributed.spmd_types import spmd_mesh_group from torchtitan.models.base.module import Module @@ -35,7 +36,7 @@ from torchtitan.models.common.linear import Linear from torchtitan.models.common.nn_modules import GELU, LayerNorm, RMSNorm -compiled_create_block_mask = torch.compile(create_block_mask) +compiled_create_block_mask = compiled(create_block_mask) # Applies rotary position embedding: (query, key, rope_cache) -> (query, key). RopeApply = Callable[ diff --git a/torchtitan/rl/generator/vllm/model_wrapper.py b/torchtitan/rl/generator/vllm/model_wrapper.py index 8438e0745dc..b2c7c5260c2 100644 --- a/torchtitan/rl/generator/vllm/model_wrapper.py +++ b/torchtitan/rl/generator/vllm/model_wrapper.py @@ -38,6 +38,7 @@ from torchtitan.config import TrainingConfig from torchtitan.distributed import maybe_apply_numa_binding from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode +from torchtitan.distributed.local_compile import set_local_compile_backend from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import current_spmd_mesh from torchtitan.models.base.module import Module @@ -339,6 +340,7 @@ def __init__( training_parallelism = parallelism.to_training() model_config = copy.deepcopy(model_config) self.config = _replace_vllm_layer_configs(model_config) + set_local_compile_backend(self.config.local_compile_backend) # Build ParallelismContext from the translated ParallelismConfig so TP/EP # sharding sees the same mesh shape as vLLM. The DP axis is represented diff --git a/torchtitan/training_engine.py b/torchtitan/training_engine.py index 721c8a2e6be..0dc9e209222 100644 --- a/torchtitan/training_engine.py +++ b/torchtitan/training_engine.py @@ -35,6 +35,7 @@ wrap_fwd_bwd_with_cuda_graph, wrap_with_cuda_graph, ) +from torchtitan.distributed.local_compile import set_local_compile_backend from torchtitan.models.base import BaseModel from torchtitan.observability import structured_logger as sl from torchtitan.observability.metrics import ( @@ -293,6 +294,7 @@ def initialize( create_seed_checkpoint: bool = False, ) -> None: """Initialize model execution and the state required to train it.""" + set_local_compile_backend(self.model_config.local_compile_backend) self._initialize_model( hf_assets_path=hf_assets_path, create_seed_checkpoint=create_seed_checkpoint, From d5be8b188c18e4aff62164cf8b239a301eb0a25d Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 21:59:50 +0000 Subject: [PATCH 2/6] Use a local compile policy for custom backends --- .../cpu/config/test_compile_config.py | 38 ++++-- tests/unit_tests/cpu/test_trainer.py | 23 ++-- torchtitan/distributed/local_compile.py | 119 +++++++++++------- .../transformers_modeling_backend/model.py | 3 +- .../tests/test_flex_attention.py | 8 +- torchtitan/models/base/model.py | 10 +- torchtitan/rl/generator/vllm/model_wrapper.py | 3 +- torchtitan/training_engine.py | 3 +- 8 files changed, 129 insertions(+), 78 deletions(-) diff --git a/tests/unit_tests/cpu/config/test_compile_config.py b/tests/unit_tests/cpu/config/test_compile_config.py index 79c95ec5e31..1c79fbb9ad3 100644 --- a/tests/unit_tests/cpu/config/test_compile_config.py +++ b/tests/unit_tests/cpu/config/test_compile_config.py @@ -4,6 +4,9 @@ # 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 @@ -13,7 +16,7 @@ apply_local_compile, compiled, local_compile, - set_local_compile_backend, + LocalCompileConfig, ) from torchtitan.models.common.attention import FlexInnerAttention from torchtitan.models.common.attention.attention import _compiled_create_block_mask @@ -21,13 +24,29 @@ from torchtitan.models.common.vision_encoder import compiled_create_block_mask +class CustomDeviceLocalCompileConfig(LocalCompileConfig): + """Example policy supplied by an out-of-tree device backend.""" + + def compile( + self, + reference: Callable[..., Any], + *, + inductor_options: dict[str, Any] | None = None, + **compile_kwargs: Any, + ) -> Callable[..., Any]: + compile_kwargs.setdefault("backend", "custom_device_backend") + return super().compile( + reference, + inductor_options=inductor_options, + **compile_kwargs, + ) + + @pytest.fixture(autouse=True) def reset_local_compile(): - set_local_compile_backend("inductor") - apply_local_compile([]) + LocalCompileConfig().apply() yield - set_local_compile_backend("inductor") - apply_local_compile([]) + LocalCompileConfig().apply() def test_apply_local_compile_rejects_unknown_name() -> None: @@ -52,11 +71,12 @@ def compiled(value: int) -> tuple[str, int]: return compiled monkeypatch.setattr(torch, "compile", fake_compile) + LocalCompileConfig().apply() apply_local_compile(["test_local_compile"]) assert fn(3) == ("inductor", 3) - set_local_compile_backend("custom_device_backend") + CustomDeviceLocalCompileConfig().apply() apply_local_compile(["test_local_compile"]) assert fn(3) == ("custom_device_backend", 3) @@ -170,11 +190,11 @@ def fake_compile(fn, **kwargs): monkeypatch.setattr(torch, "compile", fake_compile) helper = compiled(reference, inductor_options=options) - set_local_compile_backend("custom_device_backend") + CustomDeviceLocalCompileConfig().apply() assert compiled_calls == [] assert helper(3) == 3 - set_local_compile_backend("inductor") + LocalCompileConfig().apply() assert helper(3) == 3 assert compiled_calls == [ ( @@ -205,7 +225,7 @@ def fake_compile(reference, **kwargs): return lambda *args, **kwargs: None monkeypatch.setattr(torch, "compile", fake_compile) - set_local_compile_backend("custom_device_backend") + CustomDeviceLocalCompileConfig().apply() FlexInnerAttention._compiled_flex_attn() _compiled_create_block_mask() diff --git a/tests/unit_tests/cpu/test_trainer.py b/tests/unit_tests/cpu/test_trainer.py index deb91aaf1a2..1e0fa784c41 100644 --- a/tests/unit_tests/cpu/test_trainer.py +++ b/tests/unit_tests/cpu/test_trainer.py @@ -1106,27 +1106,26 @@ def test_seed_checkpoint_initialize_skips_forward_backward(): ), _dist_moe_runtime=None, state_dict_adapter=None, - model_config=SimpleNamespace(local_compile_backend="custom_device_backend"), + model_config=SimpleNamespace( + local_compile=SimpleNamespace( + apply=lambda: events.append("local_compile") + ) + ), ), ) - with patch( - "torchtitan.training_engine.set_local_compile_backend", - side_effect=lambda backend: events.append(("backend", backend)), - ) as set_backend: - TrainingEngine.initialize( - engine, - hf_assets_path="", - create_seed_checkpoint=True, - ) + TrainingEngine.initialize( + engine, + hf_assets_path="", + create_seed_checkpoint=True, + ) assert events == [ - ("backend", "custom_device_backend"), + "local_compile", "model", "model_memory", "optim", "checkpointer", ] - set_backend.assert_called_once_with("custom_device_backend") assert engine.model_device_mem_stats is model_mem_stats engine._initialize_model.assert_called_once_with( hf_assets_path="", diff --git a/torchtitan/distributed/local_compile.py b/torchtitan/distributed/local_compile.py index 9ccb307c539..a6fc2f5da6b 100644 --- a/torchtitan/distributed/local_compile.py +++ b/torchtitan/distributed/local_compile.py @@ -4,10 +4,11 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Function-scoped torch.compile registration.""" +"""Function-scoped torch.compile registration and configuration.""" import functools from collections.abc import Callable +from dataclasses import dataclass from typing import Any import torch @@ -15,8 +16,10 @@ from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode -_BACKEND: str = "inductor" _COMPILED_HELPER_RESETS: list[Callable[[], None]] = [] +_LOCAL_COMPILE_CALLBACKS: dict[ + str, list[Callable[["LocalCompileConfig", list[str]], None]] +] = {} def _reset_compiled_helpers() -> None: @@ -24,31 +27,66 @@ def _reset_compiled_helpers() -> None: reset() -def set_local_compile_backend(backend: str) -> None: - """Set the process-wide backend used by local compile helpers.""" - global _BACKEND - _BACKEND = backend - _reset_compiled_helpers() +@dataclass(kw_only=True, slots=True) +class LocalCompileConfig: + """Policy for TorchTitan-owned local ``torch.compile`` boundaries.""" + regions: list[str] | None = None + """Optional region override. -def apply_local_compile(regions: list[str]) -> None: - """Bind registered functions to eager or compiled implementations. - - Process-wide: a later call replaces this choice for every model in the process. + ``None`` preserves the model's existing ``local_compile_regions`` value. """ - unknown = [name for name in regions if name not in _LOCAL_COMPILE_CALLBACKS] - if unknown: - raise ValueError( - f"Unknown local_compile_regions entries {unknown}; " - f"registered values are {sorted(_LOCAL_COMPILE_CALLBACKS)}" - ) - for callbacks in _LOCAL_COMPILE_CALLBACKS.values(): - for bind_local_compile_fn in callbacks: - bind_local_compile_fn(regions) + def compile( + self, + reference: Callable[..., Any], + *, + inductor_options: dict[str, Any] | None = None, + **compile_kwargs: Any, + ) -> Callable[..., Any]: + """Compile one TorchTitan-owned region. + + Subclasses can override this method to select an out-of-tree backend. + ``inductor_options`` are forwarded only when the effective backend is + Inductor. + """ + backend = compile_kwargs.setdefault("backend", "inductor") + if backend == "inductor" and inductor_options is not None: + compile_kwargs.setdefault("options", inductor_options) + return torch.compile(reference, **compile_kwargs) + + def apply(self, fallback_regions: list[str] | None = None) -> None: + """Activate this policy and bind all registered named regions. + + The active policy is process-wide. A later call replaces it for lazy + helpers such as FlexAttention and ``create_block_mask``. + """ + global _ACTIVE_LOCAL_COMPILE_CONFIG + + regions = self.regions if self.regions is not None else fallback_regions + if regions is None: + regions = [] + + unknown = [name for name in regions if name not in _LOCAL_COMPILE_CALLBACKS] + if unknown: + raise ValueError( + f"Unknown local_compile_regions entries {unknown}; " + f"registered values are {sorted(_LOCAL_COMPILE_CALLBACKS)}" + ) + + _ACTIVE_LOCAL_COMPILE_CONFIG = self + _reset_compiled_helpers() + for callbacks in _LOCAL_COMPILE_CALLBACKS.values(): + for bind_local_compile_fn in callbacks: + bind_local_compile_fn(self, regions) + + +_ACTIVE_LOCAL_COMPILE_CONFIG = LocalCompileConfig() -_LOCAL_COMPILE_CALLBACKS: dict[str, list[Callable[[list[str]], None]]] = {} +def apply_local_compile(regions: list[str]) -> None: + """Bind regions using the active local compile policy.""" + _ACTIVE_LOCAL_COMPILE_CONFIG.apply(regions) def local_compile( @@ -60,7 +98,8 @@ def local_compile( """Register a function that can be compiled independently. Args: - name: Name used to enable the function in a model's ``local_compile_regions``. + name: Name used to enable the function in a model's + ``local_compile_regions`` or the policy's ``regions`` override. batch_invariant: Whether the compiled function preserves batch invariance. **compile_kwargs: Additional ``torch.compile`` keyword arguments; ``fullgraph`` is fixed to ``True`` so each function forms one complete @@ -68,12 +107,14 @@ def local_compile( """ if compile_kwargs.pop("fullgraph", True) is not True: raise ValueError("local_compile requires fullgraph=True.") - explicit_backend = compile_kwargs.pop("backend", None) def decorate(reference: Callable[..., Any]) -> Callable[..., Any]: fn = reference - def bind_local_compile(regions: list[str]) -> None: + def bind_local_compile( + config: LocalCompileConfig, + regions: list[str], + ) -> None: nonlocal fn enabled = name in regions batch_invariant_mode = is_in_batch_invariant_mode() @@ -83,11 +124,9 @@ def bind_local_compile(regions: list[str]) -> None: "batch-invariant mode; remove it from local_compile_regions." ) if enabled: - backend = _BACKEND if explicit_backend is None else explicit_backend - fn = torch.compile( + fn = config.compile( reference, fullgraph=True, - backend=backend, **compile_kwargs, ) else: @@ -114,8 +153,6 @@ def compiled( ``inductor_options`` are passed as ``torch.compile(options=...)`` only when the effective backend is Inductor. """ - explicit_backend = compile_kwargs.pop("backend", None) - explicit_options = compile_kwargs.pop("options", None) fn: Callable[..., Any] | None = None def reset() -> None: @@ -126,23 +163,11 @@ def reset() -> None: def wrapped(*args: Any, **kwargs: Any) -> Any: nonlocal fn if fn is None: - backend = _BACKEND if explicit_backend is None else explicit_backend - options = explicit_options - if backend == "inductor" and inductor_options is not None: - options = inductor_options - if options is None: - fn = torch.compile( - reference, - backend=backend, - **compile_kwargs, - ) - else: - fn = torch.compile( - reference, - backend=backend, - options=options, - **compile_kwargs, - ) + fn = _ACTIVE_LOCAL_COMPILE_CONFIG.compile( + reference, + inductor_options=inductor_options, + **compile_kwargs, + ) return fn(*args, **kwargs) _COMPILED_HELPER_RESETS.append(reset) @@ -152,6 +177,6 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: __all__ = [ "apply_local_compile", "compiled", + "LocalCompileConfig", "local_compile", - "set_local_compile_backend", ] diff --git a/torchtitan/experiments/transformers_modeling_backend/model.py b/torchtitan/experiments/transformers_modeling_backend/model.py index 4d6c7e75302..defed88bdae 100644 --- a/torchtitan/experiments/transformers_modeling_backend/model.py +++ b/torchtitan/experiments/transformers_modeling_backend/model.py @@ -26,6 +26,7 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.model import BaseModel from torchtitan.models.base.module import Module, ModuleDict @@ -286,7 +287,7 @@ def __init__( self.sharding_config = None # HF modules run eager; the loss and the swapped-in TorchTitan MoE (SwiGLU) can compile. self.local_compile_regions = ["loss", "fused_binary_activation"] - self.local_compile_backend = "inductor" + self.local_compile = LocalCompileConfig() assert model_config is not None, "model_config is required" diff --git a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py index fcf8d2b21f5..bcb17369185 100644 --- a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py +++ b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py @@ -79,7 +79,7 @@ def test_packed_documents_do_not_leak(self): class TestHFTransformerConfig(unittest.TestCase): - def test_local_compile_backend_defaults_to_inductor(self): + def test_local_compile_defaults_to_inductor_policy(self): from torchtitan.experiments.transformers_modeling_backend import ( TitanModelConfig, ) @@ -96,7 +96,11 @@ def test_local_compile_backend_defaults_to_inductor(self): ), ) - self.assertEqual(config.local_compile_backend, "inductor") + self.assertEqual( + config.local_compile_regions, + ["loss", "fused_binary_activation"], + ) + self.assertIsNone(config.local_compile.regions) class TestAttnMaskTypeValidation(unittest.TestCase): diff --git a/torchtitan/models/base/model.py b/torchtitan/models/base/model.py index deb6b04de45..4cad411061d 100644 --- a/torchtitan/models/base/model.py +++ b/torchtitan/models/base/model.py @@ -14,7 +14,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import apply_local_compile, LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from .module import Module @@ -92,8 +92,12 @@ class Config(Module.Config): FlexAttention manages its own compilation and is not controlled by this list. """ - local_compile_backend: str = "inductor" - """Backend for TorchTitan-owned local and attention compile sites.""" + local_compile: LocalCompileConfig = field(default_factory=LocalCompileConfig) + """Policy for TorchTitan-owned local and attention compile sites. + + Out-of-tree device backends can supply a subclass that overrides + ``compile``. + """ @abstractmethod def get_nparams_and_flops(self, model: Module, seq_len: int) -> tuple[int, int]: diff --git a/torchtitan/rl/generator/vllm/model_wrapper.py b/torchtitan/rl/generator/vllm/model_wrapper.py index b2c7c5260c2..cff016cfdc9 100644 --- a/torchtitan/rl/generator/vllm/model_wrapper.py +++ b/torchtitan/rl/generator/vllm/model_wrapper.py @@ -38,7 +38,6 @@ from torchtitan.config import TrainingConfig from torchtitan.distributed import maybe_apply_numa_binding from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode -from torchtitan.distributed.local_compile import set_local_compile_backend from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import current_spmd_mesh from torchtitan.models.base.module import Module @@ -340,7 +339,7 @@ def __init__( training_parallelism = parallelism.to_training() model_config = copy.deepcopy(model_config) self.config = _replace_vllm_layer_configs(model_config) - set_local_compile_backend(self.config.local_compile_backend) + self.config.local_compile.apply() # Build ParallelismContext from the translated ParallelismConfig so TP/EP # sharding sees the same mesh shape as vLLM. The DP axis is represented diff --git a/torchtitan/training_engine.py b/torchtitan/training_engine.py index 0dc9e209222..b4835eefb59 100644 --- a/torchtitan/training_engine.py +++ b/torchtitan/training_engine.py @@ -35,7 +35,6 @@ wrap_fwd_bwd_with_cuda_graph, wrap_with_cuda_graph, ) -from torchtitan.distributed.local_compile import set_local_compile_backend from torchtitan.models.base import BaseModel from torchtitan.observability import structured_logger as sl from torchtitan.observability.metrics import ( @@ -294,7 +293,7 @@ def initialize( create_seed_checkpoint: bool = False, ) -> None: """Initialize model execution and the state required to train it.""" - set_local_compile_backend(self.model_config.local_compile_backend) + self.model_config.local_compile.apply() self._initialize_model( hf_assets_path=hf_assets_path, create_seed_checkpoint=create_seed_checkpoint, From 91337ab8de11c2ba7d03475443f61df9fcaf822c Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 22:30:14 +0000 Subject: [PATCH 3/6] Unify local compile configuration --- .../numerical_tests_qwen3_5_shard.py | 2 +- tests/rl/unit_tests/cpu/test_generator.py | 4 +- .../cpu/test_policy_loss_configs.py | 4 +- .../rl/unit_tests/gpu/test_bitwise_parity.py | 2 +- tests/rl/unit_tests/gpu/test_gdn_cudagraph.py | 2 +- .../cpu/config/test_compile_config.py | 49 ++++-------- .../test_context_parallel_validation.py | 1 - .../cpu/models/common/test_gated_rmsnorm.py | 4 +- .../models/common/test_hi_mid_lo_linear.py | 2 +- .../cpu/models/kimi_k3/test_kimi_k3_config.py | 2 +- .../cpu/models/qwen3_5/test_offset_rmsnorm.py | 2 +- .../cpu/quantization/test_quantization.py | 4 +- .../cpu/test_integration_test_definitions.py | 2 +- tests/unit_tests/cpu/test_trainer.py | 24 ++---- .../gpu/models/common/test_gated_rmsnorm.py | 6 +- .../models/common/test_hi_mid_lo_linear.py | 6 +- .../common/test_moe_aux_loss_compile.py | 8 +- .../gpu/models/common/test_situglu.py | 6 +- .../gpu/models/common/test_swiglu.py | 6 +- .../test_kimi_k3_attention_residual.py | 6 +- .../gpu/models/qwen3_5/test_offset_rmsnorm.py | 10 +-- .../gpu/models/qwen3_5/test_rope_compile.py | 10 +-- .../gpu/models/qwen3_5/test_sigmoid_gate.py | 6 +- torchtitan/config/transform/lm_head_fp32.py | 4 +- torchtitan/config/validation.py | 4 +- torchtitan/distributed/local_compile.py | 78 +++++++------------ torchtitan/distributed/pipeline_parallel.py | 2 - .../experiments/graph_trainer/configs.py | 8 +- .../graph_trainer/precompile_main.py | 1 + .../experiments/graph_trainer/trainer.py | 15 ++-- .../transformers_modeling_backend/model.py | 5 +- .../parallelize.py | 4 - .../transformers_modeling_backend/pipeline.py | 2 - .../tests/test_flex_attention.py | 8 +- torchtitan/models/base/model.py | 16 +--- torchtitan/models/common/multimodal.py | 4 - torchtitan/models/deepseek_v3/model.py | 17 ++-- torchtitan/models/deepseek_v4/model.py | 15 ++-- torchtitan/models/flux/README.md | 2 +- torchtitan/models/flux/model.py | 9 +-- torchtitan/models/gpt_oss/model.py | 19 ++--- torchtitan/models/kimi_k2_7/model.py | 2 - torchtitan/models/kimi_k3/model.py | 23 +++--- torchtitan/models/llama3/model.py | 7 +- torchtitan/models/muse_glimmer/model.py | 9 ++- torchtitan/models/qwen3/model.py | 17 ++-- torchtitan/models/qwen3_5/model.py | 26 +++---- torchtitan/rl/generator/generator.py | 2 +- torchtitan/rl/generator/vllm/model_wrapper.py | 5 +- torchtitan/training_engine.py | 13 ++-- torchtitan_recipes/models/deepseek_v3.py | 2 +- torchtitan_recipes/rl/alphabet_sort.py | 8 +- torchtitan_recipes/tests/__init__.py | 2 +- .../tests/graph_trainer/numerics.py | 12 +-- .../tests/models/deepseek_v3.py | 4 +- torchtitan_recipes/tests/rl/alphabet_sort.py | 12 +-- 56 files changed, 229 insertions(+), 296 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py b/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py index 5d83ef4d73e..ce03e455076 100644 --- a/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py +++ b/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py @@ -91,6 +91,7 @@ def run_worker(args): with parallelism_context.activate_spmd(): model = config.build() + config.local_compile.apply() model.to_empty(device="cuda") model.init_weights(buffer_device=torch.device("cuda")) @@ -98,7 +99,6 @@ def run_worker(args): parallelism_context=parallelism_context, training=training, parallelism=parallelism, - local_compile_regions=config.local_compile_regions, activation_remat=DefaultActivationRematConfig(), dump_folder="/tmp", ) diff --git a/tests/rl/unit_tests/cpu/test_generator.py b/tests/rl/unit_tests/cpu/test_generator.py index 3d973912d6f..a7b2fd55f26 100644 --- a/tests/rl/unit_tests/cpu/test_generator.py +++ b/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 + assert "offset_rmsnorm" in config.generator.model.local_compile.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 diff --git a/tests/rl/unit_tests/cpu/test_policy_loss_configs.py b/tests/rl/unit_tests/cpu/test_policy_loss_configs.py index b79c8ef600b..bdbbfff8d7a 100644 --- a/tests/rl/unit_tests/cpu/test_policy_loss_configs.py +++ b/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 == [] + assert config.generator.model.local_compile.regions == [] diff --git a/tests/rl/unit_tests/gpu/test_bitwise_parity.py b/tests/rl/unit_tests/gpu/test_bitwise_parity.py index 797845b57c2..b668b904553 100644 --- a/tests/rl/unit_tests/gpu/test_bitwise_parity.py +++ b/tests/rl/unit_tests/gpu/test_bitwise_parity.py @@ -137,12 +137,12 @@ def build_trainer_model( utils.set_default_dtype(TORCH_DTYPE_MAP[trainer_config.training.dtype]), ): model = model_config.build() + model_config.local_compile.apply() model = model.parallelize( parallelism_context=parallelism_context, training=trainer_config.training, parallelism=parallelism, - local_compile_regions=model_config.local_compile_regions, activation_remat=trainer_config.activation_remat, dump_folder=config.dump_folder, ) diff --git a/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py b/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py index 176cbc6c46e..2a7f435f080 100644 --- a/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py +++ b/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 = [] registry.register_to_vllm( model_config, parallelism=registry.InferenceParallelismConfig(tensor_parallel_degree=1), diff --git a/tests/unit_tests/cpu/config/test_compile_config.py b/tests/unit_tests/cpu/config/test_compile_config.py index 1c79fbb9ad3..05b3f046b31 100644 --- a/tests/unit_tests/cpu/config/test_compile_config.py +++ b/tests/unit_tests/cpu/config/test_compile_config.py @@ -4,16 +4,12 @@ # 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, compiled, local_compile, LocalCompileConfig, @@ -24,24 +20,6 @@ from torchtitan.models.common.vision_encoder import compiled_create_block_mask -class CustomDeviceLocalCompileConfig(LocalCompileConfig): - """Example policy supplied by an out-of-tree device backend.""" - - def compile( - self, - reference: Callable[..., Any], - *, - inductor_options: dict[str, Any] | None = None, - **compile_kwargs: Any, - ) -> Callable[..., Any]: - compile_kwargs.setdefault("backend", "custom_device_backend") - return super().compile( - reference, - inductor_options=inductor_options, - **compile_kwargs, - ) - - @pytest.fixture(autouse=True) def reset_local_compile(): LocalCompileConfig().apply() @@ -49,9 +27,9 @@ def reset_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_rebinds_active_region_when_backend_changes(monkeypatch) -> None: @@ -71,13 +49,14 @@ def compiled(value: int) -> tuple[str, int]: return compiled monkeypatch.setattr(torch, "compile", fake_compile) - LocalCompileConfig().apply() - apply_local_compile(["test_local_compile"]) + LocalCompileConfig(regions=["test_local_compile"]).apply() assert fn(3) == ("inductor", 3) - CustomDeviceLocalCompileConfig().apply() - apply_local_compile(["test_local_compile"]) + LocalCompileConfig( + regions=["test_local_compile"], + backend="custom_device_backend", + ).apply() assert fn(3) == ("custom_device_backend", 3) assert [kwargs for _, kwargs in compiled_calls] == [ @@ -106,7 +85,7 @@ 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] == { @@ -136,9 +115,9 @@ def fn(value: int) -> int: lambda: True, ) with pytest.raises( - ValueError, match=r"test_non_batch_invariant.*local_compile_regions" + 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: @@ -150,7 +129,7 @@ def fake_compile(reference, **kwargs): return reference monkeypatch.setattr(torch, "compile", fake_compile) - apply_local_compile(["loss"]) + LocalCompileConfig(regions=["loss"]).apply() assert compiled_names == [ cross_entropy_loss.__name__, @@ -167,7 +146,7 @@ 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 == [ ( @@ -190,7 +169,7 @@ def fake_compile(fn, **kwargs): monkeypatch.setattr(torch, "compile", fake_compile) helper = compiled(reference, inductor_options=options) - CustomDeviceLocalCompileConfig().apply() + LocalCompileConfig(backend="custom_device_backend").apply() assert compiled_calls == [] assert helper(3) == 3 @@ -225,7 +204,7 @@ def fake_compile(reference, **kwargs): return lambda *args, **kwargs: None monkeypatch.setattr(torch, "compile", fake_compile) - CustomDeviceLocalCompileConfig().apply() + LocalCompileConfig(backend="custom_device_backend").apply() FlexInnerAttention._compiled_flex_attn() _compiled_create_block_mask() diff --git a/tests/unit_tests/cpu/config/test_context_parallel_validation.py b/tests/unit_tests/cpu/config/test_context_parallel_validation.py index 0575d338eb8..d345c412761 100644 --- a/tests/unit_tests/cpu/config/test_context_parallel_validation.py +++ b/tests/unit_tests/cpu/config/test_context_parallel_validation.py @@ -198,7 +198,6 @@ def _parallelize(self, inner_attention): parallelism_context=SimpleNamespace(cp_enabled=True), training=None, parallelism=None, - local_compile_regions=[], activation_remat=None, dump_folder="", ) diff --git a/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py b/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py index 4d1ecbae692..7fd0f24639e 100644 --- a/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py +++ b/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) diff --git a/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py b/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py index c2664043867..1a7753acde1 100644 --- a/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py +++ b/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 # Converting an already converted head is a no-op. converted = config.lm_head LMHeadFP32OutputConverter.Config().build().convert(config) diff --git a/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py b/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py index e22edb1d94b..0c93f8ab84a 100644 --- a/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py +++ b/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 == [ "loss", "gated_rmsnorm", "fused_binary_activation", diff --git a/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py b/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py index b43336cb5ad..e0e52a20e8a 100644 --- a/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py +++ b/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 == [ "loss", "fused_binary_activation", "gated_rmsnorm", diff --git a/tests/unit_tests/cpu/quantization/test_quantization.py b/tests/unit_tests/cpu/quantization/test_quantization.py index 61603ae7ced..9318bb35ecf 100644 --- a/tests/unit_tests/cpu/quantization/test_quantization.py +++ b/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 == [ "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 == ["loss"] grouped = list(config.model.traverse(GroupedLinear.Config)) assert grouped quantized_cls = _get_nvfp4_grouped_linear_cls(GroupedLinear) diff --git a/tests/unit_tests/cpu/test_integration_test_definitions.py b/tests/unit_tests/cpu/test_integration_test_definitions.py index 0eb38f0db00..b17a2beb965 100644 --- a/tests/unit_tests/cpu/test_integration_test_definitions.py +++ b/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 == [] config.__post_init__() diff --git a/tests/unit_tests/cpu/test_trainer.py b/tests/unit_tests/cpu/test_trainer.py index 1e0fa784c41..8128e8d31fc 100644 --- a/tests/unit_tests/cpu/test_trainer.py +++ b/tests/unit_tests/cpu/test_trainer.py @@ -747,16 +747,18 @@ def test_graph_training_engine_rejects_optimizer_cuda_graph() -> None: init.assert_not_called() -def test_graph_training_engine_rejects_local_compile_regions() -> None: +def test_graph_training_engine_rejects_local_compile() -> None: config = SimpleNamespace(optim=SimpleNamespace(enable_cuda_graph=False)) with ( patch.object(TrainingEngine, "__init__") as init, - pytest.raises(ValueError, match="local_compile_regions"), + pytest.raises(ValueError, match=r"local_compile\.regions"), ): GraphTrainingEngine( config, - model_config=SimpleNamespace(local_compile_regions=["loss"]), + model_config=SimpleNamespace( + local_compile=SimpleNamespace(regions=["loss"]) + ), max_num_documents=None, output_dir="", ) @@ -764,14 +766,10 @@ def test_graph_training_engine_rejects_local_compile_regions() -> None: init.assert_not_called() -def test_parallelize_compile_kwargs_include_regions_only() -> None: - engine = SimpleNamespace( - model_config=SimpleNamespace(local_compile_regions=["loss"]) - ) +def test_parallelize_compile_kwargs_are_empty_for_regular_models() -> None: + engine = SimpleNamespace() - assert TrainingEngine._parallelize_compile_kwargs(engine) == { - "local_compile_regions": ["loss"], - } + assert TrainingEngine._parallelize_compile_kwargs(engine) == {} def test_optim_update_clips_before_parameter_update() -> None: @@ -1106,11 +1104,6 @@ def test_seed_checkpoint_initialize_skips_forward_backward(): ), _dist_moe_runtime=None, state_dict_adapter=None, - model_config=SimpleNamespace( - local_compile=SimpleNamespace( - apply=lambda: events.append("local_compile") - ) - ), ), ) TrainingEngine.initialize( @@ -1120,7 +1113,6 @@ def test_seed_checkpoint_initialize_skips_forward_backward(): ) assert events == [ - "local_compile", "model", "model_memory", "optim", diff --git a/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py b/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py index 3d1604aa825..3badec1bc25 100644 --- a/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py +++ b/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 diff --git a/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py b/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py index 1cd17b19b6b..2821b03c20c 100644 --- a/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py +++ b/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]) diff --git a/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py b/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py index b5c8145489f..4482a1e710a 100644 --- a/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py +++ b/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, diff --git a/tests/unit_tests/gpu/models/common/test_situglu.py b/tests/unit_tests/gpu/models/common/test_situglu.py index 3bd6f42355a..a9b262c96f3 100644 --- a/tests/unit_tests/gpu/models/common/test_situglu.py +++ b/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): diff --git a/tests/unit_tests/gpu/models/common/test_swiglu.py b/tests/unit_tests/gpu/models/common/test_swiglu.py index dde74bd4d75..85b09654cb8 100644 --- a/tests/unit_tests/gpu/models/common/test_swiglu.py +++ b/tests/unit_tests/gpu/models/common/test_swiglu.py @@ -9,17 +9,17 @@ 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.activation import ClampedSwiGLU, SiTUGLU, SwiGLU @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestSwiGLULocalCompile(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): diff --git a/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py b/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py index ea9ef9c9b20..45456e29bdb 100644 --- a/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py +++ b/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py @@ -8,7 +8,7 @@ import torch -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.kimi_k3.model import AttentionResidual @@ -45,11 +45,11 @@ def _reference( @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestKimiK3AttentionResidualLocalCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["attention_residual"]) + LocalCompileConfig(regions=["attention_residual"]).apply() self.attention_residual = AttentionResidual.Config().build() def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py b/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py index f39c6fb5ba6..c4a0fa1e1b1 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py @@ -13,7 +13,7 @@ with_comms, ) -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import set_current_spmd_mesh, set_spmd_meshes from torchtitan.models.common.decoder_sharding import head_norm_sharding_config @@ -23,10 +23,10 @@ @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestOffsetRMSNormCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["offset_rmsnorm"]) + LocalCompileConfig(regions=["offset_rmsnorm"]).apply() def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): @@ -101,7 +101,7 @@ def world_size(self) -> int: @with_comms def test_compiled_qk_norm(self): - apply_local_compile(["offset_rmsnorm"]) + LocalCompileConfig(regions=["offset_rmsnorm"]).apply() device = self.device_type config = OffsetRMSNorm.Config( dim=128, @@ -137,7 +137,7 @@ def test_compiled_qk_norm(self): output.sum().backward() self.assertEqual(output.shape, x_local.shape) - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py b/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py index c7c91aa4ad8..d6fbe401791 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py @@ -8,7 +8,7 @@ import torch -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.rope import CosSinRoPE from torchtitan.models.qwen3_5.rope import MRoPE @@ -16,10 +16,10 @@ @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestRoPELocalCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["cos_sin_rope"]) + LocalCompileConfig(regions=["cos_sin_rope"]).apply() def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): @@ -89,7 +89,7 @@ def test_forward_and_backward_match_eager(self): (grad_query_out, grad_key_out), ) - apply_local_compile([]) + LocalCompileConfig().apply() eager_rope = ( CosSinRoPE.Config( dim=128, @@ -222,7 +222,7 @@ def test_mrope_three_axis_positions_match_eager(self): (grad_query_out, grad_key_out), ) - apply_local_compile([]) + LocalCompileConfig().apply() eager_rope = config.build().cuda() eager_query = query.detach().clone().requires_grad_() eager_key = key.detach().clone().requires_grad_() diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py b/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py index 83b3ab74fca..8448defd201 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.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.qwen3_5.moe import sigmoid_gate @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestSigmoidGateLocalCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["shared_expert_gate"]) + LocalCompileConfig(regions=["shared_expert_gate"]).apply() # A local generator leaves the global RNG stream unchanged for later tests. generator = torch.Generator(device="cuda").manual_seed(0) @@ -29,7 +29,7 @@ def randn(*shape: int) -> torch.Tensor: self.grad_output = randn(64, 128) def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_match_reference(self): diff --git a/torchtitan/config/transform/lm_head_fp32.py b/torchtitan/config/transform/lm_head_fp32.py index a7cea8aa0f4..f8f09568051 100644 --- a/torchtitan/config/transform/lm_head_fp32.py +++ b/torchtitan/config/transform/lm_head_fp32.py @@ -67,6 +67,6 @@ def convert(self, model_config): "the model config. The torchtitan decoder names its output projection " f"{self._TARGET!r} (see torchtitan/models/common/decoder.py)." ) - if "fp32_to_bf16_split" not in model_config.local_compile_regions: - model_config.local_compile_regions.append("fp32_to_bf16_split") + if "fp32_to_bf16_split" not in model_config.local_compile.regions: + model_config.local_compile.regions.append("fp32_to_bf16_split") return model_config diff --git a/torchtitan/config/validation.py b/torchtitan/config/validation.py index 5bdaab9b22b..8ff534c6863 100644 --- a/torchtitan/config/validation.py +++ b/torchtitan/config/validation.py @@ -137,11 +137,11 @@ def validate_model_training_config( f"dispatcher: {type(dispatcher_config).__qualname__}." ) - if debug.spmd_typechecking and model.local_compile_regions: + if debug.spmd_typechecking and model.local_compile.regions: # TODO: Remove this once Dynamo supports tracing SPMD typechecking. raise ValueError( "Local compilation is not supported with SPMD typechecking. " - "Set model.local_compile_regions=[] or disable debug.spmd_typechecking." + "Set model.local_compile.regions=[] or disable debug.spmd_typechecking." ) validate_context_parallel(model, parallelism) diff --git a/torchtitan/distributed/local_compile.py b/torchtitan/distributed/local_compile.py index a6fc2f5da6b..aca670e2d0e 100644 --- a/torchtitan/distributed/local_compile.py +++ b/torchtitan/distributed/local_compile.py @@ -8,7 +8,7 @@ import functools from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import torch @@ -16,26 +16,19 @@ from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode -_COMPILED_HELPER_RESETS: list[Callable[[], None]] = [] -_LOCAL_COMPILE_CALLBACKS: dict[ - str, list[Callable[["LocalCompileConfig", list[str]], None]] -] = {} - - -def _reset_compiled_helpers() -> None: - for reset in _COMPILED_HELPER_RESETS: - reset() +_LOCAL_COMPILE_CALLBACKS: dict[str, list[Callable[["LocalCompileConfig"], None]]] = {} +_LOCAL_COMPILE_GENERATION = 0 @dataclass(kw_only=True, slots=True) class LocalCompileConfig: - """Policy for TorchTitan-owned local ``torch.compile`` boundaries.""" + """Configuration for TorchTitan-owned local ``torch.compile`` boundaries.""" - regions: list[str] | None = None - """Optional region override. + regions: list[str] = field(default_factory=list) + """Named ``@local_compile`` regions to compile independently.""" - ``None`` preserves the model's existing ``local_compile_regions`` value. - """ + backend: str = "inductor" + """Backend used by named regions, FlexAttention, and block-mask helpers.""" def compile( self, @@ -46,49 +39,41 @@ def compile( ) -> Callable[..., Any]: """Compile one TorchTitan-owned region. - Subclasses can override this method to select an out-of-tree backend. ``inductor_options`` are forwarded only when the effective backend is Inductor. """ - backend = compile_kwargs.setdefault("backend", "inductor") + backend = compile_kwargs.setdefault("backend", self.backend) if backend == "inductor" and inductor_options is not None: compile_kwargs.setdefault("options", inductor_options) return torch.compile(reference, **compile_kwargs) - def apply(self, fallback_regions: list[str] | None = None) -> None: - """Activate this policy and bind all registered named regions. + def apply(self) -> None: + """Activate this config and bind all registered named regions. - The active policy is process-wide. A later call replaces it for lazy + The active config is process-wide. A later call replaces it for lazy helpers such as FlexAttention and ``create_block_mask``. """ - global _ACTIVE_LOCAL_COMPILE_CONFIG - - regions = self.regions if self.regions is not None else fallback_regions - if regions is None: - regions = [] + global _ACTIVE_LOCAL_COMPILE_CONFIG, _LOCAL_COMPILE_GENERATION - unknown = [name for name in regions if name not in _LOCAL_COMPILE_CALLBACKS] + unknown = [ + name for name in self.regions if name not in _LOCAL_COMPILE_CALLBACKS + ] if unknown: raise ValueError( - f"Unknown local_compile_regions entries {unknown}; " + f"Unknown local_compile.regions entries {unknown}; " f"registered values are {sorted(_LOCAL_COMPILE_CALLBACKS)}" ) _ACTIVE_LOCAL_COMPILE_CONFIG = self - _reset_compiled_helpers() + _LOCAL_COMPILE_GENERATION += 1 for callbacks in _LOCAL_COMPILE_CALLBACKS.values(): for bind_local_compile_fn in callbacks: - bind_local_compile_fn(self, regions) + bind_local_compile_fn(self) _ACTIVE_LOCAL_COMPILE_CONFIG = LocalCompileConfig() -def apply_local_compile(regions: list[str]) -> None: - """Bind regions using the active local compile policy.""" - _ACTIVE_LOCAL_COMPILE_CONFIG.apply(regions) - - def local_compile( name: str, *, @@ -98,8 +83,7 @@ def local_compile( """Register a function that can be compiled independently. Args: - name: Name used to enable the function in a model's - ``local_compile_regions`` or the policy's ``regions`` override. + name: Name used to enable the function in ``LocalCompileConfig.regions``. batch_invariant: Whether the compiled function preserves batch invariance. **compile_kwargs: Additional ``torch.compile`` keyword arguments; ``fullgraph`` is fixed to ``True`` so each function forms one complete @@ -111,17 +95,14 @@ def local_compile( def decorate(reference: Callable[..., Any]) -> Callable[..., Any]: fn = reference - def bind_local_compile( - config: LocalCompileConfig, - regions: list[str], - ) -> None: + def bind_local_compile(config: LocalCompileConfig) -> None: nonlocal fn - enabled = name in regions + enabled = name in config.regions batch_invariant_mode = is_in_batch_invariant_mode() if enabled and batch_invariant_mode and not batch_invariant: raise ValueError( f"Local compile region {name!r} does not support " - "batch-invariant mode; remove it from local_compile_regions." + "batch-invariant mode; remove it from local_compile.regions." ) if enabled: fn = config.compile( @@ -154,28 +135,25 @@ def compiled( the effective backend is Inductor. """ fn: Callable[..., Any] | None = None - - def reset() -> None: - nonlocal fn - fn = None + applied_generation = -1 @functools.wraps(reference) def wrapped(*args: Any, **kwargs: Any) -> Any: - nonlocal fn - if fn is None: + nonlocal applied_generation, fn + if applied_generation != _LOCAL_COMPILE_GENERATION: fn = _ACTIVE_LOCAL_COMPILE_CONFIG.compile( reference, inductor_options=inductor_options, **compile_kwargs, ) + applied_generation = _LOCAL_COMPILE_GENERATION + assert fn is not None return fn(*args, **kwargs) - _COMPILED_HELPER_RESETS.append(reset) return wrapped __all__ = [ - "apply_local_compile", "compiled", "LocalCompileConfig", "local_compile", diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py index 67959106c65..2bc8c7fe6cc 100644 --- a/torchtitan/distributed/pipeline_parallel.py +++ b/torchtitan/distributed/pipeline_parallel.py @@ -78,7 +78,6 @@ def pipeline_llm( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, @@ -144,7 +143,6 @@ def pipeline_llm( parallelism_context=parallelism_context, training=training, parallelism=parallelism, - local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 62e7e36314a..28bbb665d56 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -340,8 +340,8 @@ def to_graph_trainer_config( Copies all fields from the base config and converts its model config to the GraphTrainer model config class, without local compile regions or in-place - WGRAD accumulation because GraphTrainer traces the whole step. The ``compile`` field keeps the - GraphTrainer.Config default; callers should explicitly set it. + WGRAD accumulation because GraphTrainer traces the whole step. The ``compile`` + field keeps the GraphTrainer.Config default; callers should explicitly set it. """ from .trainer import GraphTrainer @@ -352,8 +352,8 @@ def to_graph_trainer_config( for f in fields(base_config.model) } ) - # GraphTrainer compiles the whole step (config.compile), so it drops the model's local compile regions. - graph_model.local_compile_regions = [] + # GraphTrainer compiles the whole step, so it drops the model's local regions. + graph_model.local_compile.regions = [] disable_inplace_wgrad_accum(graph_model) d["model"] = graph_model diff --git a/torchtitan/experiments/graph_trainer/precompile_main.py b/torchtitan/experiments/graph_trainer/precompile_main.py index efd18ed18f4..9691d7ee9b6 100644 --- a/torchtitan/experiments/graph_trainer/precompile_main.py +++ b/torchtitan/experiments/graph_trainer/precompile_main.py @@ -130,6 +130,7 @@ def _common_setup(config): utils.set_default_dtype(TORCH_DTYPE_MAP[config.training.dtype]), ): model = model_config.build() + model_config.local_compile.apply() # For aot_fx_trace, apply_compile inside model.parallelize is a no-op # (returns model unchanged), so we pass the real compile_config. diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py index 1b7c6c8a91e..474023c3316 100644 --- a/torchtitan/experiments/graph_trainer/trainer.py +++ b/torchtitan/experiments/graph_trainer/trainer.py @@ -66,11 +66,11 @@ def __init__( ) -> None: if config.optim.enable_cuda_graph: raise ValueError("Optim CUDA graphs are not supported with GraphTrainer.") - if model_config.local_compile_regions: + if model_config.local_compile.regions: raise ValueError( "GraphTrainer traces the whole step into one graph; set " - "model.local_compile_regions = [] " - f"(got {model_config.local_compile_regions})." + "model.local_compile.regions = [] " + f"(got {model_config.local_compile.regions})." ) if config.parallelism.fsdp_symm_mem_scope is not None: raise ValueError( @@ -89,12 +89,9 @@ def __init__( def _parallelize_compile_kwargs(self) -> dict[str, Any]: """Return the compile kwargs ``model.parallelize``/``model.pipeline`` expect for this engine's models. - Regular models expect ``local_compile_regions`` (compiled regions); GraphTrainer overrides this to - return ``compile_config`` (whole-step compile). + Regular models need no compile kwargs; GraphTrainer returns its + whole-step ``compile_config``. """ - # TODO: apply local compile outside parallelize/pipeline (#5026 review). That needs - # GraphTrainer to stop reading its compile config there first (apply_compile, - # EP-overlap chunking, enable_autoparallel, GraphPP runtime); then delete this hook. return {"compile_config": self.config.compile} def _initialize_forward_backward(self) -> None: @@ -359,7 +356,7 @@ class Config(Trainer.Config): compile: GraphTrainerCompileConfig = field( default_factory=GraphTrainerCompileConfig ) - """Whole-step compile. GraphTrainer requires ``model.local_compile_regions`` to be empty.""" + """Whole-step compile. GraphTrainer requires ``model.local_compile.regions`` to be empty.""" def __post_init__(self) -> None: Trainer.Config.__post_init__(self) diff --git a/torchtitan/experiments/transformers_modeling_backend/model.py b/torchtitan/experiments/transformers_modeling_backend/model.py index defed88bdae..deac8ae6294 100644 --- a/torchtitan/experiments/transformers_modeling_backend/model.py +++ b/torchtitan/experiments/transformers_modeling_backend/model.py @@ -286,8 +286,9 @@ def __init__( ) self.sharding_config = None # HF modules run eager; the loss and the swapped-in TorchTitan MoE (SwiGLU) can compile. - self.local_compile_regions = ["loss", "fused_binary_activation"] - self.local_compile = LocalCompileConfig() + self.local_compile = LocalCompileConfig( + regions=["loss", "fused_binary_activation"] + ) assert model_config is not None, "model_config is required" diff --git a/torchtitan/experiments/transformers_modeling_backend/parallelize.py b/torchtitan/experiments/transformers_modeling_backend/parallelize.py index 46be051bc05..c23b96b499b 100644 --- a/torchtitan/experiments/transformers_modeling_backend/parallelize.py +++ b/torchtitan/experiments/transformers_modeling_backend/parallelize.py @@ -27,7 +27,6 @@ enable_fsdp_symm_mem, get_fsdp_reshard_after_forward_policy, ) -from torchtitan.distributed.local_compile import apply_local_compile logger = logging.getLogger(__name__) @@ -87,7 +86,6 @@ def parallelize_hf_transformers( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, **kwargs: Any, @@ -101,8 +99,6 @@ def parallelize_hf_transformers( 4. Single model._parallelize(parallelism_context) call -- shards states, wraps forward 5. Apply AC and FSDP """ - # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) # 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 diff --git a/torchtitan/experiments/transformers_modeling_backend/pipeline.py b/torchtitan/experiments/transformers_modeling_backend/pipeline.py index f439c413c35..663b05f71af 100644 --- a/torchtitan/experiments/transformers_modeling_backend/pipeline.py +++ b/torchtitan/experiments/transformers_modeling_backend/pipeline.py @@ -296,7 +296,6 @@ def pipeline_hf_transformers( *, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, @@ -390,7 +389,6 @@ def pipeline_hf_transformers( parallelism_context=parallelism_context, training=training, parallelism=parallelism, - local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py index bcb17369185..4febea30e20 100644 --- a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py +++ b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py @@ -79,7 +79,8 @@ def test_packed_documents_do_not_leak(self): class TestHFTransformerConfig(unittest.TestCase): - def test_local_compile_defaults_to_inductor_policy(self): + def test_local_compile_defaults_to_inductor(self): + from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.experiments.transformers_modeling_backend import ( TitanModelConfig, ) @@ -97,10 +98,9 @@ def test_local_compile_defaults_to_inductor_policy(self): ) self.assertEqual( - config.local_compile_regions, - ["loss", "fused_binary_activation"], + config.local_compile, + LocalCompileConfig(regions=["loss", "fused_binary_activation"]), ) - self.assertIsNone(config.local_compile.regions) class TestAttnMaskTypeValidation(unittest.TestCase): diff --git a/torchtitan/models/base/model.py b/torchtitan/models/base/model.py index 4cad411061d..2792fac3692 100644 --- a/torchtitan/models/base/model.py +++ b/torchtitan/models/base/model.py @@ -14,7 +14,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig -from torchtitan.distributed.local_compile import apply_local_compile, LocalCompileConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from .module import Module @@ -86,17 +86,10 @@ class Config(Module.Config): Subclasses define model-specific hyperparameters. """ - local_compile_regions: list[str] = field(default_factory=list) - """``@local_compile`` regions this model compiles; ``[]`` runs them all eager. - - FlexAttention manages its own compilation and is not controlled by this list. - """ - local_compile: LocalCompileConfig = field(default_factory=LocalCompileConfig) - """Policy for TorchTitan-owned local and attention compile sites. + """TorchTitan-owned local compile configuration. - Out-of-tree device backends can supply a subclass that overrides - ``compile``. + Covers named local regions, FlexAttention, and ``create_block_mask``. """ @abstractmethod @@ -139,13 +132,10 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: """Apply the ordered model-level parallelization lifecycle.""" - # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) if activation_remat is not None: diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py index 26ce302f79a..7e723db799c 100644 --- a/torchtitan/models/common/multimodal.py +++ b/torchtitan/models/common/multimodal.py @@ -20,7 +20,6 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.distributed.parallelism_context import ParallelismContext from .decoder import Decoder @@ -37,12 +36,9 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: - # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) encoders = [ diff --git a/torchtitan/models/deepseek_v3/model.py b/torchtitan/models/deepseek_v3/model.py index c1ba8bc60ee..805036eacf0 100644 --- a/torchtitan/models/deepseek_v3/model.py +++ b/torchtitan/models/deepseek_v3/model.py @@ -13,6 +13,7 @@ import torch_remat as remat from torch import nn +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.base.module import Module from torchtitan.models.common.attention import ( BaseAttention, @@ -272,13 +273,15 @@ class DeepSeekV3Model(MTPDecoder): class Config(MTPDecoder.Config): dim: int = 2048 vocab_size: int = 102400 - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "fp32_to_bf16_split", - "moe_aux_loss", - ] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "fp32_to_bf16_split", + "moe_aux_loss", + ] + ) ) def get_nparams_and_flops( diff --git a/torchtitan/models/deepseek_v4/model.py b/torchtitan/models/deepseek_v4/model.py index 1da25fda7a4..eba632c669b 100644 --- a/torchtitan/models/deepseek_v4/model.py +++ b/torchtitan/models/deepseek_v4/model.py @@ -14,6 +14,7 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.module import ModuleList from torchtitan.models.common.attention import ( @@ -179,12 +180,14 @@ def _apply_fsdp( class Config(Decoder.Config): dim: int vocab_size: int - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "fp32_to_bf16_split", - ] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "fp32_to_bf16_split", + ] + ) ) hc_mult: int = 4 n_mtp_layers: int = 0 diff --git a/torchtitan/models/flux/README.md b/torchtitan/models/flux/README.md index bb331f575d2..9ff56054f75 100644 --- a/torchtitan/models/flux/README.md +++ b/torchtitan/models/flux/README.md @@ -51,7 +51,7 @@ MODULE=torchtitan_recipes.tests.models.flux CONFIG=flux_debugmodel_inference \ The loss function can be compiled independently with: ```python -config.model.local_compile_regions = ["loss"] +config.model.local_compile.regions = ["loss"] ``` The model itself is not compiled by the standard Trainer. diff --git a/torchtitan/models/flux/model.py b/torchtitan/models/flux/model.py index 39c6480c931..36e4863a3fb 100644 --- a/torchtitan/models/flux/model.py +++ b/torchtitan/models/flux/model.py @@ -18,7 +18,7 @@ from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed import context_parallel from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_replicated_parameters, @@ -103,7 +103,9 @@ class FluxModel(BaseModel): class Config(BaseModel.Config): img_in: Linear.Config txt_in: Linear.Config - local_compile_regions: list[str] = field(default_factory=lambda: ["loss"]) + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig(regions=["loss"]) + ) in_channels: int = 64 out_channels: int = 64 vec_in_dim: int = 768 @@ -235,13 +237,10 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: """Apply Flux's AC-before-SPMD parallelization lifecycle.""" - # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): if activation_remat is not None: # Flux saves no regions under any policy: each block is diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py index 35bd95b26f0..6421509398b 100644 --- a/torchtitan/models/gpt_oss/model.py +++ b/torchtitan/models/gpt_oss/model.py @@ -17,6 +17,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.module import Module from torchtitan.models.common.attention import ( @@ -216,13 +217,15 @@ def _register_optimizer_hooks( class Config(Decoder.Config): dim: int = 2880 vocab_size: int = 201088 - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "cos_sin_rope", - "fp32_to_bf16_split", - ] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "cos_sin_rope", + "fp32_to_bf16_split", + ] + ) ) def get_nparams_and_flops( @@ -250,7 +253,6 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> GptOssModel: @@ -267,7 +269,6 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, - local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/kimi_k2_7/model.py b/torchtitan/models/kimi_k2_7/model.py index 6d1dd05790b..451ef173876 100644 --- a/torchtitan/models/kimi_k2_7/model.py +++ b/torchtitan/models/kimi_k2_7/model.py @@ -117,7 +117,6 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK25Model: @@ -131,7 +130,6 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, - local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index f1313a458a1..b4875bbfc3d 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -19,7 +19,7 @@ from torchtitan.config.function import Function from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import apply_local_compile, local_compile +from torchtitan.distributed.local_compile import local_compile, LocalCompileConfig from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -391,14 +391,16 @@ class Config(Decoder.Config): attention_residual: AttentionResidual.Config = field( default_factory=AttentionResidual.Config ) - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "gated_rmsnorm", - "fused_binary_activation", - "fp32_to_bf16_split", - "attention_residual", - ] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "gated_rmsnorm", + "fused_binary_activation", + "fp32_to_bf16_split", + "attention_residual", + ] + ) ) def get_nparams_and_flops( @@ -458,12 +460,9 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK3Model: - # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): annotate_replicated_parameters(self, parallelism_context) self._parallelize(parallelism_context) diff --git a/torchtitan/models/llama3/model.py b/torchtitan/models/llama3/model.py index b0335209d90..b3eb5f37e04 100644 --- a/torchtitan/models/llama3/model.py +++ b/torchtitan/models/llama3/model.py @@ -12,6 +12,7 @@ import torch_remat as remat from torch import nn +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.attention import ( FlexAttentionMetadata, VarlenAttentionMetadata, @@ -83,8 +84,10 @@ class Llama3Model(Decoder): class Config(Decoder.Config): dim: int = 4096 vocab_size: int = 128256 - local_compile_regions: list[str] = field( - default_factory=lambda: ["loss", "fused_binary_activation"] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=["loss", "fused_binary_activation"] + ) ) def get_nparams_and_flops( diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py index 72938708b72..04f363efced 100644 --- a/torchtitan/models/muse_glimmer/model.py +++ b/torchtitan/models/muse_glimmer/model.py @@ -18,6 +18,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -289,8 +290,10 @@ class MuseGlimmerModel(MultimodalModel): class Config(Decoder.Config): dim: int = 6656 vocab_size: int = 202048 - local_compile_regions: list[str] = field( - default_factory=lambda: ["loss", "fused_binary_activation"] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=["loss", "fused_binary_activation"] + ) ) # Narrows the base Decoder.Config.tok_embeddings (Embedding.Config) to the # bundled embedding+norm unit whose .embedding/.norm children carry the sharding. @@ -369,7 +372,6 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> MuseGlimmerModel: @@ -383,7 +385,6 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, - local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/qwen3/model.py b/torchtitan/models/qwen3/model.py index 0ec1c454b65..974855616c6 100644 --- a/torchtitan/models/qwen3/model.py +++ b/torchtitan/models/qwen3/model.py @@ -12,6 +12,7 @@ import torch.nn as nn import torch_remat as remat +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.attention import ( FlexAttentionMetadata, VarlenAttentionMetadata, @@ -108,13 +109,15 @@ def _register_optimizer_hooks( class Config(Decoder.Config): dim: int = 1024 vocab_size: int = 151936 - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "cos_sin_rope", - "fp32_to_bf16_split", - ] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "cos_sin_rope", + "fp32_to_bf16_split", + ] + ) ) def get_nparams_and_flops( diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index bd68938bfb6..2a49bd9f05f 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -18,7 +18,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import local_compile +from torchtitan.distributed.local_compile import local_compile, LocalCompileConfig from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -353,16 +353,18 @@ def _register_optimizer_hooks( @dataclass(kw_only=True, slots=True) class Config(Decoder.Config): vision_encoder: Qwen35VisionEncoder.Config | None = None - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "gated_rmsnorm", - "offset_rmsnorm", - "partial_rope", - "shared_expert_gate", - "fp32_to_bf16_split", - ] + local_compile: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "gated_rmsnorm", + "offset_rmsnorm", + "partial_rope", + "shared_expert_gate", + "fp32_to_bf16_split", + ] + ) ) def get_nparams_and_flops( @@ -436,7 +438,6 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Qwen35Model: @@ -444,7 +445,6 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, - local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/rl/generator/generator.py b/torchtitan/rl/generator/generator.py index 8c464632cc1..e03a9d528c6 100644 --- a/torchtitan/rl/generator/generator.py +++ b/torchtitan/rl/generator/generator.py @@ -210,7 +210,7 @@ def _prepare_generation_request_metrics( class VLLMCudaGraphConfig: """CUDA graph capture settings for the vLLM inference engine. - Local compile regions come from the model config's ``local_compile_regions``. + Local compile regions come from the model config's ``local_compile.regions``. Only CUDA graph capture, which is vLLM-specific, is controlled here. ``mode`` selects which vLLM CUDA graph mode to capture; see that field and diff --git a/torchtitan/rl/generator/vllm/model_wrapper.py b/torchtitan/rl/generator/vllm/model_wrapper.py index cff016cfdc9..1b74ba65168 100644 --- a/torchtitan/rl/generator/vllm/model_wrapper.py +++ b/torchtitan/rl/generator/vllm/model_wrapper.py @@ -32,7 +32,7 @@ from vllm.utils import torch_utils as _torch_utils # Registers the "loss" local-compile region. The generator shares the trainer's -# model config, whose local_compile_regions may name it, but never builds a loss. +# model config, whose local_compile.regions may name it, but never builds a loss. from torchtitan.components import loss as _loss # noqa: F401 from torchtitan.components.checkpointer import CheckpointManager from torchtitan.config import TrainingConfig @@ -339,7 +339,6 @@ def __init__( training_parallelism = parallelism.to_training() model_config = copy.deepcopy(model_config) self.config = _replace_vllm_layer_configs(model_config) - self.config.local_compile.apply() # Build ParallelismContext from the translated ParallelismConfig so TP/EP # sharding sees the same mesh shape as vLLM. The DP axis is represented @@ -360,6 +359,7 @@ def __init__( # Build model on meta device to avoid allocating full model on every GPU with self.parallelism_context.activate_spmd(), torch.device("meta"): self.model = self.config.build() + self.config.local_compile.apply() self.model._skip_lm_head = True if getattr(self.model, "vision_encoder", None) is not None: self.model.vision_encoder = None @@ -368,7 +368,6 @@ def __init__( parallelism_context=self.parallelism_context, training=TrainingConfig(), parallelism=training_parallelism, - local_compile_regions=self.config.local_compile_regions, activation_remat=None, dump_folder="", ) diff --git a/torchtitan/training_engine.py b/torchtitan/training_engine.py index b4835eefb59..5fa54d097ca 100644 --- a/torchtitan/training_engine.py +++ b/torchtitan/training_engine.py @@ -293,7 +293,6 @@ def initialize( create_seed_checkpoint: bool = False, ) -> None: """Initialize model execution and the state required to train it.""" - self.model_config.local_compile.apply() self._initialize_model( hf_assets_path=hf_assets_path, create_seed_checkpoint=create_seed_checkpoint, @@ -332,6 +331,8 @@ def _initialize_model( utils.set_default_dtype(TORCH_DTYPE_MAP[self.config.training.dtype]), ): model = self.model_config.build() + # Model construction imports and registers model-specific local regions. + self.model_config.local_compile.apply() self.model_cls = type(model) adapter_cls = type(model).state_dict_adapter_cls self.state_dict_adapter = ( @@ -410,13 +411,11 @@ def _initialize_model( def _parallelize_compile_kwargs(self) -> dict[str, Any]: """Return the compile kwargs ``model.parallelize``/``model.pipeline`` expect for this engine's models. - Regular models expect ``local_compile_regions`` (compiled regions); GraphTrainer overrides this to - return ``compile_config`` (whole-step compile). + Regular models need no compile kwargs because ``model.local_compile`` is + applied after model construction. GraphTrainer overrides this to return + ``compile_config`` (whole-step compile). """ - # TODO: apply local compile outside parallelize/pipeline (#5026 review). That needs - # GraphTrainer to stop reading its compile config there first (apply_compile, - # EP-overlap chunking, enable_autoparallel, GraphPP runtime); then delete this hook. - return {"local_compile_regions": self.model_config.local_compile_regions} + return {} def _initialize_optim(self) -> None: """Construct the parameter update and its state.""" diff --git a/torchtitan_recipes/models/deepseek_v3.py b/torchtitan_recipes/models/deepseek_v3.py index 5a762cdec49..e1ca1d6689b 100644 --- a/torchtitan_recipes/models/deepseek_v3.py +++ b/torchtitan_recipes/models/deepseek_v3.py @@ -268,5 +268,5 @@ def deepseek_v3_671b_nvfp4_ffn_mxfp8_attn( ) ], ) - config.model.local_compile_regions = ["loss"] + config.model.local_compile.regions = ["loss"] return config diff --git a/torchtitan_recipes/rl/alphabet_sort.py b/torchtitan_recipes/rl/alphabet_sort.py index f7125220ca0..f7fb37a771b 100644 --- a/torchtitan_recipes/rl/alphabet_sort.py +++ b/torchtitan_recipes/rl/alphabet_sort.py @@ -174,7 +174,7 @@ def rl_grpo_qwen3_0_6b_varlen(*, seq_len: int = 2048) -> Controller.Config: def rl_grpo_qwen3_0_6b_varlen_no_compile(*, seq_len: int = 2048) -> Controller.Config: config = rl_grpo_qwen3_0_6b_varlen(seq_len=seq_len) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return config @@ -299,7 +299,7 @@ def rl_grpo_qwen3_0_6b_flex_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] block_size = config.trainer.model.layers[0].attention.inner_attention.block_size config.async_loop.batcher = dataclasses.replace( config.async_loop.batcher, per_sample_pad_multiple=block_size @@ -671,7 +671,7 @@ def rl_grpo_qwen3_0_6b_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return Controller.Config( hf_assets_path="./assets/hf/Qwen3-0.6B", num_generators=3, @@ -856,7 +856,7 @@ def rl_grpo_qwen3_5_9b_varlen_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return config diff --git a/torchtitan_recipes/tests/__init__.py b/torchtitan_recipes/tests/__init__.py index 3792be993e9..8ac49cd5624 100644 --- a/torchtitan_recipes/tests/__init__.py +++ b/torchtitan_recipes/tests/__init__.py @@ -27,4 +27,4 @@ def _set_spmd_typechecking(config: Trainer.Config, *, typechecking: bool) -> Non if typechecking: # Local compile does not support global SPMD typechecking yet. # TODO: Remove this once the SPMD typechecking/Dynamo issue is fixed. - config.model.local_compile_regions = [] + config.model.local_compile.regions = [] diff --git a/torchtitan_recipes/tests/graph_trainer/numerics.py b/torchtitan_recipes/tests/graph_trainer/numerics.py index 277344d9ce7..9b3ee3f2df2 100644 --- a/torchtitan_recipes/tests/graph_trainer/numerics.py +++ b/torchtitan_recipes/tests/graph_trainer/numerics.py @@ -55,7 +55,7 @@ def _without_inplace_wgrad_accum(config): def llama3_eager_numerics(): config = llama3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile.regions = [] config.parallelism.tensor_parallel_degree = 2 config.parallelism.data_parallel_shard_degree = 4 config.training.num_tokens_per_microbatch_per_dp_rank = 16384 @@ -73,7 +73,7 @@ def llama3_graph_numerics(): def llama3_eager_hsdp_gradient_accumulation_numerics(): config = llama3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile.regions = [] config.parallelism.data_parallel_replicate_degree = 2 config.parallelism.data_parallel_shard_degree = 2 config.training.num_tokens_per_microbatch_per_dp_rank = 8192 @@ -103,7 +103,7 @@ def _llama3_hsdp_pp_numerics(config): def llama3_eager_hsdp_pp_numerics(): config = llama3_debugmodel(seq_len=2048) - config.model.local_compile_regions = [] + config.model.local_compile.regions = [] return _without_inplace_wgrad_accum(_llama3_hsdp_pp_numerics(config)) @@ -214,7 +214,7 @@ def _deepseek_v3_dist_moe_pp_numerics( runtime_config.pp_activation_slot_policy = slot_policy for _, experts, _, _ in config.model.traverse(DistMoeRoutedExperts.Config): experts.inplace_wgrad_accum = False - config.model.local_compile_regions = [] + config.model.local_compile.regions = [] config.parallelism.pipeline_parallel_degree = 2 config.parallelism.pipeline_parallel_schedule = "Interleaved1F1B" config.parallelism.num_pp_microbatches = 8 @@ -251,7 +251,7 @@ def deepseek_v3_dist_moe_graph_pp_stage_microbatch_numerics(): def qwen3_eager_numerics(): config = qwen3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile.regions = [] config.parallelism.tensor_parallel_degree = 2 config.parallelism.data_parallel_shard_degree = 4 config.training.num_tokens_per_microbatch_per_dp_rank = 16384 @@ -278,7 +278,7 @@ def qwen3_moe_eager_numerics(): ) config.dataloader.max_num_documents = None # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile.regions = [] config.training.disable_cuda_graphs = True config.training.num_tokens_per_microbatch_per_dp_rank = 16384 return _without_inplace_wgrad_accum(config) diff --git a/torchtitan_recipes/tests/models/deepseek_v3.py b/torchtitan_recipes/tests/models/deepseek_v3.py index 65225c382bc..9fe20eafc3d 100644 --- a/torchtitan_recipes/tests/models/deepseek_v3.py +++ b/torchtitan_recipes/tests/models/deepseek_v3.py @@ -481,7 +481,7 @@ def deepseek_v3_debugmodel_nvfp4_ffn_mxfp8_attn( ), ], ) - config.model.local_compile_regions = ["loss"] + config.model.local_compile.regions = ["loss"] return config @@ -530,5 +530,5 @@ def deepseek_v3_16b_nvfp4_ffn_mxfp8_attn( ) ], ) - config.model.local_compile_regions = ["loss"] + config.model.local_compile.regions = ["loss"] return config diff --git a/torchtitan_recipes/tests/rl/alphabet_sort.py b/torchtitan_recipes/tests/rl/alphabet_sort.py index ce1b21069ed..52b18fcfb65 100644 --- a/torchtitan_recipes/tests/rl/alphabet_sort.py +++ b/torchtitan_recipes/tests/rl/alphabet_sort.py @@ -156,7 +156,7 @@ def rl_grpo_gpt_oss_debug_varlen_no_compile( ) -> Controller.Config: config = rl_grpo_gpt_oss_debug_varlen(seq_len=seq_len) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return config @@ -194,7 +194,7 @@ def rl_grpo_gpt_oss_debug_varlen_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return Controller.Config( hf_assets_path="tests/assets/tokenizer", async_loop=AsyncLoopConfig( @@ -439,7 +439,7 @@ def rl_grpo_qwen3_moe_debug_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return Controller.Config( hf_assets_path="tests/assets/tokenizer", async_loop=AsyncLoopConfig( @@ -604,7 +604,7 @@ def rl_grpo_qwen3_5_debug_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return config @@ -629,7 +629,7 @@ def _kimi_k3_rl_model_configs( converters=converters, ) for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] return trainer_model, generator_model @@ -819,7 +819,7 @@ def rl_grpo_moe_debug_tp4_ep4() -> Controller.Config: generator_parallelism=config.generator.parallelism, ) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile.regions = [] # TODO: use FULL CUDA graphs once GPT-OSS CI enables a # CUDA-graph-compatible distributed MoE token dispatcher. config.generator.cuda_graph.mode = "NONE" From 2bed445a6b9f4e19f5e0f0c31af34da8d13deede Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 22:36:25 +0000 Subject: [PATCH 4/6] Simplify local compile backend configuration --- .../numerical_tests_qwen3_5_shard.py | 2 +- tests/rl/unit_tests/cpu/test_generator.py | 4 +- .../cpu/test_policy_loss_configs.py | 4 +- .../rl/unit_tests/gpu/test_bitwise_parity.py | 2 +- tests/rl/unit_tests/gpu/test_gdn_cudagraph.py | 2 +- .../cpu/config/test_compile_config.py | 135 +++++++----------- .../test_context_parallel_validation.py | 1 + .../cpu/models/common/test_gated_rmsnorm.py | 4 +- .../models/common/test_hi_mid_lo_linear.py | 2 +- .../cpu/models/kimi_k3/test_kimi_k3_config.py | 2 +- .../cpu/models/qwen3_5/test_offset_rmsnorm.py | 2 +- .../cpu/quantization/test_quantization.py | 4 +- .../cpu/test_integration_test_definitions.py | 2 +- tests/unit_tests/cpu/test_trainer.py | 14 +- .../gpu/models/common/test_gated_rmsnorm.py | 6 +- .../models/common/test_hi_mid_lo_linear.py | 6 +- .../common/test_moe_aux_loss_compile.py | 8 +- .../gpu/models/common/test_situglu.py | 6 +- .../gpu/models/common/test_swiglu.py | 6 +- .../test_kimi_k3_attention_residual.py | 6 +- .../gpu/models/qwen3_5/test_offset_rmsnorm.py | 10 +- .../gpu/models/qwen3_5/test_rope_compile.py | 10 +- .../gpu/models/qwen3_5/test_sigmoid_gate.py | 6 +- torchtitan/config/transform/lm_head_fp32.py | 4 +- torchtitan/config/validation.py | 4 +- torchtitan/distributed/local_compile.py | 64 +++++---- torchtitan/distributed/pipeline_parallel.py | 2 + torchtitan/distributed/utils.py | 1 - .../experiments/graph_trainer/configs.py | 8 +- .../graph_trainer/precompile_main.py | 1 - .../experiments/graph_trainer/trainer.py | 15 +- .../transformers_modeling_backend/model.py | 5 +- .../parallelize.py | 4 + .../transformers_modeling_backend/pipeline.py | 2 + .../tests/test_flex_attention.py | 25 ---- torchtitan/models/base/model.py | 11 +- torchtitan/models/common/multimodal.py | 4 + torchtitan/models/deepseek_v3/model.py | 17 +-- torchtitan/models/deepseek_v4/model.py | 15 +- torchtitan/models/flux/README.md | 2 +- torchtitan/models/flux/model.py | 9 +- torchtitan/models/gpt_oss/model.py | 19 ++- torchtitan/models/kimi_k2_7/model.py | 2 + torchtitan/models/kimi_k3/model.py | 23 +-- torchtitan/models/llama3/model.py | 7 +- torchtitan/models/muse_glimmer/model.py | 9 +- torchtitan/models/qwen3/model.py | 17 +-- torchtitan/models/qwen3_5/model.py | 26 ++-- torchtitan/rl/generator/generator.py | 2 +- torchtitan/rl/generator/vllm/model_wrapper.py | 4 +- torchtitan/training_engine.py | 12 +- torchtitan_recipes/models/deepseek_v3.py | 2 +- torchtitan_recipes/rl/alphabet_sort.py | 8 +- torchtitan_recipes/tests/__init__.py | 2 +- .../tests/graph_trainer/numerics.py | 12 +- .../tests/models/deepseek_v3.py | 4 +- torchtitan_recipes/tests/rl/alphabet_sort.py | 12 +- 57 files changed, 270 insertions(+), 328 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py b/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py index ce03e455076..5d83ef4d73e 100644 --- a/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py +++ b/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py @@ -91,7 +91,6 @@ def run_worker(args): with parallelism_context.activate_spmd(): model = config.build() - config.local_compile.apply() model.to_empty(device="cuda") model.init_weights(buffer_device=torch.device("cuda")) @@ -99,6 +98,7 @@ def run_worker(args): parallelism_context=parallelism_context, training=training, parallelism=parallelism, + local_compile_regions=config.local_compile_regions, activation_remat=DefaultActivationRematConfig(), dump_folder="/tmp", ) diff --git a/tests/rl/unit_tests/cpu/test_generator.py b/tests/rl/unit_tests/cpu/test_generator.py index a7b2fd55f26..3d973912d6f 100644 --- a/tests/rl/unit_tests/cpu/test_generator.py +++ b/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 + assert "offset_rmsnorm" in config.generator.model.local_compile_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 diff --git a/tests/rl/unit_tests/cpu/test_policy_loss_configs.py b/tests/rl/unit_tests/cpu/test_policy_loss_configs.py index bdbbfff8d7a..b79c8ef600b 100644 --- a/tests/rl/unit_tests/cpu/test_policy_loss_configs.py +++ b/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 == [] + assert config.generator.model.local_compile_regions == [] diff --git a/tests/rl/unit_tests/gpu/test_bitwise_parity.py b/tests/rl/unit_tests/gpu/test_bitwise_parity.py index b668b904553..797845b57c2 100644 --- a/tests/rl/unit_tests/gpu/test_bitwise_parity.py +++ b/tests/rl/unit_tests/gpu/test_bitwise_parity.py @@ -137,12 +137,12 @@ def build_trainer_model( utils.set_default_dtype(TORCH_DTYPE_MAP[trainer_config.training.dtype]), ): model = model_config.build() - model_config.local_compile.apply() model = model.parallelize( parallelism_context=parallelism_context, training=trainer_config.training, parallelism=parallelism, + local_compile_regions=model_config.local_compile_regions, activation_remat=trainer_config.activation_remat, dump_folder=config.dump_folder, ) diff --git a/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py b/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py index 2a7f435f080..176cbc6c46e 100644 --- a/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py +++ b/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 = [] registry.register_to_vllm( model_config, parallelism=registry.InferenceParallelismConfig(tensor_parallel_degree=1), diff --git a/tests/unit_tests/cpu/config/test_compile_config.py b/tests/unit_tests/cpu/config/test_compile_config.py index 05b3f046b31..ea155eecacd 100644 --- a/tests/unit_tests/cpu/config/test_compile_config.py +++ b/tests/unit_tests/cpu/config/test_compile_config.py @@ -10,7 +10,7 @@ from torchtitan.components.loss import compute_logprobs, cross_entropy_loss, mse_loss from torchtitan.distributed.local_compile import ( - compiled, + apply_local_compile, local_compile, LocalCompileConfig, ) @@ -22,17 +22,17 @@ @pytest.fixture(autouse=True) def reset_local_compile(): - LocalCompileConfig().apply() + apply_local_compile([]) yield - LocalCompileConfig().apply() + apply_local_compile([]) -def test_local_compile_config_rejects_unknown_name() -> None: +def test_apply_local_compile_rejects_unknown_name() -> None: with pytest.raises(ValueError, match=r"foo.*registered values"): - LocalCompileConfig(regions=["foo"]).apply() + apply_local_compile(["foo"]) -def test_local_compile_rebinds_active_region_when_backend_changes(monkeypatch) -> None: +def test_local_compile_binds_once_from_existing_config(monkeypatch) -> None: compiled_calls = [] @local_compile("test_local_compile", batch_invariant=False) @@ -41,27 +41,55 @@ def fn(value: int) -> tuple[str, int]: def fake_compile(reference, **kwargs): compiled_calls.append((reference, kwargs)) - backend = kwargs["backend"] def compiled(value: int) -> tuple[str, int]: - return backend, value + return "compiled", value return compiled monkeypatch.setattr(torch, "compile", fake_compile) - LocalCompileConfig(regions=["test_local_compile"]).apply() + apply_local_compile(["test_local_compile"]) - assert fn(3) == ("inductor", 3) + assert fn(3) == ("compiled", 3) + assert len(compiled_calls) == 1 + assert compiled_calls[0][1]["backend"] == "inductor" + assert compiled_calls[0][1]["fullgraph"] is True - LocalCompileConfig( - regions=["test_local_compile"], - backend="custom_device_backend", - ).apply() - assert fn(3) == ("custom_device_backend", 3) - assert [kwargs for _, kwargs in compiled_calls] == [ - {"backend": "inductor", "fullgraph": True}, - {"backend": "custom_device_backend", "fullgraph": True}, +def test_local_compile_config_applies_backend_to_all_owned_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 lambda *args, **kwargs: None + + monkeypatch.setattr(torch, "compile", fake_compile) + apply_local_compile( + LocalCompileConfig( + regions=["test_custom_backend"], + backend="custom_device_backend", + ) + ) + + region() + FlexInnerAttention._compiled_flex_attn() + _compiled_create_block_mask() + compiled_create_block_mask() + + assert compiled_calls == [ + ( + region.__wrapped__, + {"backend": "custom_device_backend", "fullgraph": True}, + ), + (flex_attention, {"backend": "custom_device_backend"}), + (create_block_mask, {"backend": "custom_device_backend"}), + (create_block_mask, {"backend": "custom_device_backend"}), ] @@ -85,7 +113,7 @@ def fake_compile(reference, **kwargs): return reference monkeypatch.setattr(torch, "compile", fake_compile) - LocalCompileConfig(regions=["test_compile_kwargs"]).apply() + apply_local_compile(["test_compile_kwargs"]) assert fn(3) == 3 assert compiled_calls[0][1] == { @@ -115,9 +143,9 @@ def fn(value: int) -> int: lambda: True, ) with pytest.raises( - ValueError, match=r"test_non_batch_invariant.*local_compile\.regions" + ValueError, match=r"test_non_batch_invariant.*local_compile_regions" ): - LocalCompileConfig(regions=["test_non_batch_invariant"]).apply() + apply_local_compile(["test_non_batch_invariant"]) def test_loss_functions_use_local_compile(monkeypatch) -> None: @@ -129,7 +157,7 @@ def fake_compile(reference, **kwargs): return reference monkeypatch.setattr(torch, "compile", fake_compile) - LocalCompileConfig(regions=["loss"]).apply() + apply_local_compile(["loss"]) assert compiled_names == [ cross_entropy_loss.__name__, @@ -146,7 +174,7 @@ def fake_compile(reference, **kwargs): return reference monkeypatch.setattr(torch, "compile", fake_compile) - LocalCompileConfig(regions=["moe_aux_loss"]).apply() + apply_local_compile(["moe_aux_loss"]) assert compiled == [ ( @@ -154,64 +182,3 @@ def fake_compile(reference, **kwargs): {"backend": "inductor", "fullgraph": True}, ) ] - - -def test_compiled_helper_only_forwards_options_to_inductor(monkeypatch) -> None: - compiled_calls = [] - options = {"max_autotune": False} - - def reference(value: int) -> int: - return value - - def fake_compile(fn, **kwargs): - compiled_calls.append((fn, kwargs)) - return fn - - monkeypatch.setattr(torch, "compile", fake_compile) - helper = compiled(reference, inductor_options=options) - LocalCompileConfig(backend="custom_device_backend").apply() - - assert compiled_calls == [] - assert helper(3) == 3 - LocalCompileConfig().apply() - assert helper(3) == 3 - assert compiled_calls == [ - ( - reference, - { - "backend": "custom_device_backend", - }, - ), - ( - reference, - { - "backend": "inductor", - "options": options, - }, - ), - ] - - -def test_attention_helpers_do_not_forward_inductor_options_to_custom_backend( - monkeypatch, -) -> None: - compiled_calls = [] - - def fake_compile(reference, **kwargs): - if "options" in kwargs: - raise TypeError("custom backend got unexpected keyword argument 'options'") - compiled_calls.append((reference, kwargs)) - return lambda *args, **kwargs: None - - monkeypatch.setattr(torch, "compile", fake_compile) - LocalCompileConfig(backend="custom_device_backend").apply() - - FlexInnerAttention._compiled_flex_attn() - _compiled_create_block_mask() - compiled_create_block_mask() - - assert compiled_calls == [ - (flex_attention, {"backend": "custom_device_backend"}), - (create_block_mask, {"backend": "custom_device_backend"}), - (create_block_mask, {"backend": "custom_device_backend"}), - ] diff --git a/tests/unit_tests/cpu/config/test_context_parallel_validation.py b/tests/unit_tests/cpu/config/test_context_parallel_validation.py index d345c412761..0575d338eb8 100644 --- a/tests/unit_tests/cpu/config/test_context_parallel_validation.py +++ b/tests/unit_tests/cpu/config/test_context_parallel_validation.py @@ -198,6 +198,7 @@ def _parallelize(self, inner_attention): parallelism_context=SimpleNamespace(cp_enabled=True), training=None, parallelism=None, + local_compile_regions=[], activation_remat=None, dump_folder="", ) diff --git a/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py b/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py index 7fd0f24639e..4d1ecbae692 100644 --- a/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py +++ b/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 LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.models.common import GatedRMSNorm, Sigmoid, SiLU class TestGatedRMSNorm(unittest.TestCase): def setUp(self): - LocalCompileConfig().apply() + apply_local_compile([]) def test_configurable_activation_matches_reference(self): x = torch.randn(4, 3, 8) diff --git a/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py b/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py index 1a7753acde1..c2664043867 100644 --- a/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py +++ b/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 # Converting an already converted head is a no-op. converted = config.lm_head LMHeadFP32OutputConverter.Config().build().convert(config) diff --git a/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py b/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py index 0c93f8ab84a..e22edb1d94b 100644 --- a/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py +++ b/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 == [ "loss", "gated_rmsnorm", "fused_binary_activation", diff --git a/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py b/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py index e0e52a20e8a..b43336cb5ad 100644 --- a/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py +++ b/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 == [ "loss", "fused_binary_activation", "gated_rmsnorm", diff --git a/tests/unit_tests/cpu/quantization/test_quantization.py b/tests/unit_tests/cpu/quantization/test_quantization.py index 9318bb35ecf..61603ae7ced 100644 --- a/tests/unit_tests/cpu/quantization/test_quantization.py +++ b/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 == [ "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 == ["loss"] grouped = list(config.model.traverse(GroupedLinear.Config)) assert grouped quantized_cls = _get_nvfp4_grouped_linear_cls(GroupedLinear) diff --git a/tests/unit_tests/cpu/test_integration_test_definitions.py b/tests/unit_tests/cpu/test_integration_test_definitions.py index b17a2beb965..0eb38f0db00 100644 --- a/tests/unit_tests/cpu/test_integration_test_definitions.py +++ b/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 == [] config.__post_init__() diff --git a/tests/unit_tests/cpu/test_trainer.py b/tests/unit_tests/cpu/test_trainer.py index 8128e8d31fc..c79f1b0798f 100644 --- a/tests/unit_tests/cpu/test_trainer.py +++ b/tests/unit_tests/cpu/test_trainer.py @@ -747,18 +747,16 @@ def test_graph_training_engine_rejects_optimizer_cuda_graph() -> None: init.assert_not_called() -def test_graph_training_engine_rejects_local_compile() -> None: +def test_graph_training_engine_rejects_local_compile_regions() -> None: config = SimpleNamespace(optim=SimpleNamespace(enable_cuda_graph=False)) with ( patch.object(TrainingEngine, "__init__") as init, - pytest.raises(ValueError, match=r"local_compile\.regions"), + pytest.raises(ValueError, match="local_compile_regions"), ): GraphTrainingEngine( config, - model_config=SimpleNamespace( - local_compile=SimpleNamespace(regions=["loss"]) - ), + model_config=SimpleNamespace(local_compile_regions=["loss"]), max_num_documents=None, output_dir="", ) @@ -766,12 +764,6 @@ def test_graph_training_engine_rejects_local_compile() -> None: init.assert_not_called() -def test_parallelize_compile_kwargs_are_empty_for_regular_models() -> None: - engine = SimpleNamespace() - - assert TrainingEngine._parallelize_compile_kwargs(engine) == {} - - def test_optim_update_clips_before_parameter_update() -> None: events = [] optimizers = MagicMock() diff --git a/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py b/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py index 3badec1bc25..3d1604aa825 100644 --- a/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py +++ b/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py @@ -8,17 +8,17 @@ import torch -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.models.common import GatedRMSNorm, Sigmoid @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestGatedRMSNormCompile(unittest.TestCase): def setUp(self): - LocalCompileConfig(regions=["gated_rmsnorm"]).apply() + apply_local_compile(["gated_rmsnorm"]) def tearDown(self): - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() @staticmethod diff --git a/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py b/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py index 2821b03c20c..1cd17b19b6b 100644 --- a/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py +++ b/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 LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile 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.""" - LocalCompileConfig(regions=["fp32_to_bf16_split"]).apply() + apply_local_compile(["fp32_to_bf16_split"]) try: yield finally: - LocalCompileConfig().apply() + apply_local_compile([]) @pytest.mark.parametrize("dim", [0, 1]) diff --git a/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py b/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py index 4482a1e710a..b5c8145489f 100644 --- a/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py +++ b/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py @@ -8,14 +8,14 @@ import torch -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile 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): - LocalCompileConfig().apply() + apply_local_compile([]) 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, ) - LocalCompileConfig().apply() + apply_local_compile([]) 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] - LocalCompileConfig(regions=["moe_aux_loss"]).apply() + apply_local_compile(["moe_aux_loss"]) compiled_counts_E, compiled_prob_sums_E = _microbatch_load_balance_local_stats( scores_TE, routing_map_TE, diff --git a/tests/unit_tests/gpu/models/common/test_situglu.py b/tests/unit_tests/gpu/models/common/test_situglu.py index a9b262c96f3..3bd6f42355a 100644 --- a/tests/unit_tests/gpu/models/common/test_situglu.py +++ b/tests/unit_tests/gpu/models/common/test_situglu.py @@ -8,17 +8,17 @@ import torch -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.models.common.activation import SiTUGLU @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestSiTUGLULocalCompile(unittest.TestCase): def setUp(self): - LocalCompileConfig(regions=["fused_binary_activation"]).apply() + apply_local_compile(["fused_binary_activation"]) def tearDown(self): - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): diff --git a/tests/unit_tests/gpu/models/common/test_swiglu.py b/tests/unit_tests/gpu/models/common/test_swiglu.py index 85b09654cb8..dde74bd4d75 100644 --- a/tests/unit_tests/gpu/models/common/test_swiglu.py +++ b/tests/unit_tests/gpu/models/common/test_swiglu.py @@ -9,17 +9,17 @@ import torch import torch.nn.functional as F -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.models.common.activation import ClampedSwiGLU, SiTUGLU, SwiGLU @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestSwiGLULocalCompile(unittest.TestCase): def setUp(self): - LocalCompileConfig(regions=["fused_binary_activation"]).apply() + apply_local_compile(["fused_binary_activation"]) def tearDown(self): - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): diff --git a/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py b/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py index 45456e29bdb..ea9ef9c9b20 100644 --- a/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py +++ b/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py @@ -8,7 +8,7 @@ import torch -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.kimi_k3.model import AttentionResidual @@ -45,11 +45,11 @@ def _reference( @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestKimiK3AttentionResidualLocalCompile(unittest.TestCase): def setUp(self): - LocalCompileConfig(regions=["attention_residual"]).apply() + apply_local_compile(["attention_residual"]) self.attention_residual = AttentionResidual.Config().build() def tearDown(self): - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py b/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py index c4a0fa1e1b1..f39c6fb5ba6 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py @@ -13,7 +13,7 @@ with_comms, ) -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import set_current_spmd_mesh, set_spmd_meshes from torchtitan.models.common.decoder_sharding import head_norm_sharding_config @@ -23,10 +23,10 @@ @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestOffsetRMSNormCompile(unittest.TestCase): def setUp(self): - LocalCompileConfig(regions=["offset_rmsnorm"]).apply() + apply_local_compile(["offset_rmsnorm"]) def tearDown(self): - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): @@ -101,7 +101,7 @@ def world_size(self) -> int: @with_comms def test_compiled_qk_norm(self): - LocalCompileConfig(regions=["offset_rmsnorm"]).apply() + apply_local_compile(["offset_rmsnorm"]) device = self.device_type config = OffsetRMSNorm.Config( dim=128, @@ -137,7 +137,7 @@ def test_compiled_qk_norm(self): output.sum().backward() self.assertEqual(output.shape, x_local.shape) - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py b/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py index d6fbe401791..c7c91aa4ad8 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py @@ -8,7 +8,7 @@ import torch -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.models.common.rope import CosSinRoPE from torchtitan.models.qwen3_5.rope import MRoPE @@ -16,10 +16,10 @@ @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestRoPELocalCompile(unittest.TestCase): def setUp(self): - LocalCompileConfig(regions=["cos_sin_rope"]).apply() + apply_local_compile(["cos_sin_rope"]) def tearDown(self): - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): @@ -89,7 +89,7 @@ def test_forward_and_backward_match_eager(self): (grad_query_out, grad_key_out), ) - LocalCompileConfig().apply() + apply_local_compile([]) eager_rope = ( CosSinRoPE.Config( dim=128, @@ -222,7 +222,7 @@ def test_mrope_three_axis_positions_match_eager(self): (grad_query_out, grad_key_out), ) - LocalCompileConfig().apply() + apply_local_compile([]) eager_rope = config.build().cuda() eager_query = query.detach().clone().requires_grad_() eager_key = key.detach().clone().requires_grad_() diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py b/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py index 8448defd201..83b3ab74fca 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py @@ -8,14 +8,14 @@ import torch -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.models.qwen3_5.moe import sigmoid_gate @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestSigmoidGateLocalCompile(unittest.TestCase): def setUp(self): - LocalCompileConfig(regions=["shared_expert_gate"]).apply() + apply_local_compile(["shared_expert_gate"]) # A local generator leaves the global RNG stream unchanged for later tests. generator = torch.Generator(device="cuda").manual_seed(0) @@ -29,7 +29,7 @@ def randn(*shape: int) -> torch.Tensor: self.grad_output = randn(64, 128) def tearDown(self): - LocalCompileConfig().apply() + apply_local_compile([]) torch._dynamo.reset() def test_forward_and_backward_match_reference(self): diff --git a/torchtitan/config/transform/lm_head_fp32.py b/torchtitan/config/transform/lm_head_fp32.py index f8f09568051..a7cea8aa0f4 100644 --- a/torchtitan/config/transform/lm_head_fp32.py +++ b/torchtitan/config/transform/lm_head_fp32.py @@ -67,6 +67,6 @@ def convert(self, model_config): "the model config. The torchtitan decoder names its output projection " f"{self._TARGET!r} (see torchtitan/models/common/decoder.py)." ) - if "fp32_to_bf16_split" not in model_config.local_compile.regions: - model_config.local_compile.regions.append("fp32_to_bf16_split") + if "fp32_to_bf16_split" not in model_config.local_compile_regions: + model_config.local_compile_regions.append("fp32_to_bf16_split") return model_config diff --git a/torchtitan/config/validation.py b/torchtitan/config/validation.py index 8ff534c6863..5bdaab9b22b 100644 --- a/torchtitan/config/validation.py +++ b/torchtitan/config/validation.py @@ -137,11 +137,11 @@ def validate_model_training_config( f"dispatcher: {type(dispatcher_config).__qualname__}." ) - if debug.spmd_typechecking and model.local_compile.regions: + if debug.spmd_typechecking and model.local_compile_regions: # TODO: Remove this once Dynamo supports tracing SPMD typechecking. raise ValueError( "Local compilation is not supported with SPMD typechecking. " - "Set model.local_compile.regions=[] or disable debug.spmd_typechecking." + "Set model.local_compile_regions=[] or disable debug.spmd_typechecking." ) validate_context_parallel(model, parallelism) diff --git a/torchtitan/distributed/local_compile.py b/torchtitan/distributed/local_compile.py index aca670e2d0e..ee8f3cae72c 100644 --- a/torchtitan/distributed/local_compile.py +++ b/torchtitan/distributed/local_compile.py @@ -8,7 +8,7 @@ import functools from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any import torch @@ -24,11 +24,17 @@ class LocalCompileConfig: """Configuration for TorchTitan-owned local ``torch.compile`` boundaries.""" - regions: list[str] = field(default_factory=list) - """Named ``@local_compile`` regions to compile independently.""" - + regions: list[str] backend: str = "inductor" - """Backend used by named regions, FlexAttention, and block-mask helpers.""" + + def __bool__(self) -> bool: + return bool(self.regions) + + def __contains__(self, region: str) -> bool: + return region in self.regions + + def append(self, region: str) -> None: + self.regions.append(region) def compile( self, @@ -37,22 +43,13 @@ def compile( inductor_options: dict[str, Any] | None = None, **compile_kwargs: Any, ) -> Callable[..., Any]: - """Compile one TorchTitan-owned region. - - ``inductor_options`` are forwarded only when the effective backend is - Inductor. - """ backend = compile_kwargs.setdefault("backend", self.backend) if backend == "inductor" and inductor_options is not None: compile_kwargs.setdefault("options", inductor_options) return torch.compile(reference, **compile_kwargs) def apply(self) -> None: - """Activate this config and bind all registered named regions. - - The active config is process-wide. A later call replaces it for lazy - helpers such as FlexAttention and ``create_block_mask``. - """ + """Activate this config and bind all registered named regions.""" global _ACTIVE_LOCAL_COMPILE_CONFIG, _LOCAL_COMPILE_GENERATION unknown = [ @@ -60,7 +57,7 @@ def apply(self) -> None: ] if unknown: raise ValueError( - f"Unknown local_compile.regions entries {unknown}; " + f"Unknown local_compile_regions entries {unknown}; " f"registered values are {sorted(_LOCAL_COMPILE_CALLBACKS)}" ) @@ -71,7 +68,21 @@ def apply(self) -> None: bind_local_compile_fn(self) -_ACTIVE_LOCAL_COMPILE_CONFIG = LocalCompileConfig() +_ACTIVE_LOCAL_COMPILE_CONFIG = LocalCompileConfig(regions=[]) + + +def _as_local_compile_config( + value: list[str] | LocalCompileConfig, +) -> LocalCompileConfig: + """Normalize the existing region-list setting to one local compile config.""" + if isinstance(value, LocalCompileConfig): + return value + return LocalCompileConfig(regions=value) + + +def apply_local_compile(value: list[str] | LocalCompileConfig) -> None: + """Bind local compile boundaries from a region list or full config.""" + _as_local_compile_config(value).apply() def local_compile( @@ -83,7 +94,7 @@ def local_compile( """Register a function that can be compiled independently. Args: - name: Name used to enable the function in ``LocalCompileConfig.regions``. + name: Name used to enable the function in a model's ``local_compile_regions``. batch_invariant: Whether the compiled function preserves batch invariance. **compile_kwargs: Additional ``torch.compile`` keyword arguments; ``fullgraph`` is fixed to ``True`` so each function forms one complete @@ -102,14 +113,10 @@ def bind_local_compile(config: LocalCompileConfig) -> None: if enabled and batch_invariant_mode and not batch_invariant: raise ValueError( f"Local compile region {name!r} does not support " - "batch-invariant mode; remove it from local_compile.regions." + "batch-invariant mode; remove it from local_compile_regions." ) if enabled: - fn = config.compile( - reference, - fullgraph=True, - **compile_kwargs, - ) + fn = config.compile(reference, fullgraph=True, **compile_kwargs) else: fn = reference @@ -129,11 +136,7 @@ def compiled( inductor_options: dict[str, Any] | None = None, **compile_kwargs: Any, ) -> Callable[..., Any]: - """Compile ``reference`` lazily with the configured local backend. - - ``inductor_options`` are passed as ``torch.compile(options=...)`` only when - the effective backend is Inductor. - """ + """Compile lazily with the active local compile backend.""" fn: Callable[..., Any] | None = None applied_generation = -1 @@ -154,7 +157,8 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: __all__ = [ + "apply_local_compile", "compiled", - "LocalCompileConfig", "local_compile", + "LocalCompileConfig", ] diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py index 2bc8c7fe6cc..67959106c65 100644 --- a/torchtitan/distributed/pipeline_parallel.py +++ b/torchtitan/distributed/pipeline_parallel.py @@ -78,6 +78,7 @@ def pipeline_llm( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, @@ -143,6 +144,7 @@ def pipeline_llm( parallelism_context=parallelism_context, training=training, parallelism=parallelism, + local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py index 9b645e2430f..df5e77e00d2 100644 --- a/torchtitan/distributed/utils.py +++ b/torchtitan/distributed/utils.py @@ -167,7 +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 - if debug_config.detect_anomaly: logger.warning( "Anomaly detection enabled. This incurs significant overhead " diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 28bbb665d56..62e7e36314a 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -340,8 +340,8 @@ def to_graph_trainer_config( Copies all fields from the base config and converts its model config to the GraphTrainer model config class, without local compile regions or in-place - WGRAD accumulation because GraphTrainer traces the whole step. The ``compile`` - field keeps the GraphTrainer.Config default; callers should explicitly set it. + WGRAD accumulation because GraphTrainer traces the whole step. The ``compile`` field keeps the + GraphTrainer.Config default; callers should explicitly set it. """ from .trainer import GraphTrainer @@ -352,8 +352,8 @@ def to_graph_trainer_config( for f in fields(base_config.model) } ) - # GraphTrainer compiles the whole step, so it drops the model's local regions. - graph_model.local_compile.regions = [] + # GraphTrainer compiles the whole step (config.compile), so it drops the model's local compile regions. + graph_model.local_compile_regions = [] disable_inplace_wgrad_accum(graph_model) d["model"] = graph_model diff --git a/torchtitan/experiments/graph_trainer/precompile_main.py b/torchtitan/experiments/graph_trainer/precompile_main.py index 9691d7ee9b6..efd18ed18f4 100644 --- a/torchtitan/experiments/graph_trainer/precompile_main.py +++ b/torchtitan/experiments/graph_trainer/precompile_main.py @@ -130,7 +130,6 @@ def _common_setup(config): utils.set_default_dtype(TORCH_DTYPE_MAP[config.training.dtype]), ): model = model_config.build() - model_config.local_compile.apply() # For aot_fx_trace, apply_compile inside model.parallelize is a no-op # (returns model unchanged), so we pass the real compile_config. diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py index 474023c3316..1b7c6c8a91e 100644 --- a/torchtitan/experiments/graph_trainer/trainer.py +++ b/torchtitan/experiments/graph_trainer/trainer.py @@ -66,11 +66,11 @@ def __init__( ) -> None: if config.optim.enable_cuda_graph: raise ValueError("Optim CUDA graphs are not supported with GraphTrainer.") - if model_config.local_compile.regions: + if model_config.local_compile_regions: raise ValueError( "GraphTrainer traces the whole step into one graph; set " - "model.local_compile.regions = [] " - f"(got {model_config.local_compile.regions})." + "model.local_compile_regions = [] " + f"(got {model_config.local_compile_regions})." ) if config.parallelism.fsdp_symm_mem_scope is not None: raise ValueError( @@ -89,9 +89,12 @@ def __init__( def _parallelize_compile_kwargs(self) -> dict[str, Any]: """Return the compile kwargs ``model.parallelize``/``model.pipeline`` expect for this engine's models. - Regular models need no compile kwargs; GraphTrainer returns its - whole-step ``compile_config``. + Regular models expect ``local_compile_regions`` (compiled regions); GraphTrainer overrides this to + return ``compile_config`` (whole-step compile). """ + # TODO: apply local compile outside parallelize/pipeline (#5026 review). That needs + # GraphTrainer to stop reading its compile config there first (apply_compile, + # EP-overlap chunking, enable_autoparallel, GraphPP runtime); then delete this hook. return {"compile_config": self.config.compile} def _initialize_forward_backward(self) -> None: @@ -356,7 +359,7 @@ class Config(Trainer.Config): compile: GraphTrainerCompileConfig = field( default_factory=GraphTrainerCompileConfig ) - """Whole-step compile. GraphTrainer requires ``model.local_compile.regions`` to be empty.""" + """Whole-step compile. GraphTrainer requires ``model.local_compile_regions`` to be empty.""" def __post_init__(self) -> None: Trainer.Config.__post_init__(self) diff --git a/torchtitan/experiments/transformers_modeling_backend/model.py b/torchtitan/experiments/transformers_modeling_backend/model.py index deac8ae6294..6ec6c1080d2 100644 --- a/torchtitan/experiments/transformers_modeling_backend/model.py +++ b/torchtitan/experiments/transformers_modeling_backend/model.py @@ -26,7 +26,6 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode -from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.model import BaseModel from torchtitan.models.base.module import Module, ModuleDict @@ -286,9 +285,7 @@ def __init__( ) self.sharding_config = None # HF modules run eager; the loss and the swapped-in TorchTitan MoE (SwiGLU) can compile. - self.local_compile = LocalCompileConfig( - regions=["loss", "fused_binary_activation"] - ) + self.local_compile_regions = ["loss", "fused_binary_activation"] assert model_config is not None, "model_config is required" diff --git a/torchtitan/experiments/transformers_modeling_backend/parallelize.py b/torchtitan/experiments/transformers_modeling_backend/parallelize.py index c23b96b499b..46be051bc05 100644 --- a/torchtitan/experiments/transformers_modeling_backend/parallelize.py +++ b/torchtitan/experiments/transformers_modeling_backend/parallelize.py @@ -27,6 +27,7 @@ enable_fsdp_symm_mem, get_fsdp_reshard_after_forward_policy, ) +from torchtitan.distributed.local_compile import apply_local_compile logger = logging.getLogger(__name__) @@ -86,6 +87,7 @@ def parallelize_hf_transformers( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, **kwargs: Any, @@ -99,6 +101,8 @@ def parallelize_hf_transformers( 4. Single model._parallelize(parallelism_context) call -- shards states, wraps forward 5. Apply AC and FSDP """ + # Bind local implementations early; torch.compile traces on first use. + apply_local_compile(local_compile_regions) # 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 diff --git a/torchtitan/experiments/transformers_modeling_backend/pipeline.py b/torchtitan/experiments/transformers_modeling_backend/pipeline.py index 663b05f71af..f439c413c35 100644 --- a/torchtitan/experiments/transformers_modeling_backend/pipeline.py +++ b/torchtitan/experiments/transformers_modeling_backend/pipeline.py @@ -296,6 +296,7 @@ def pipeline_hf_transformers( *, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, @@ -389,6 +390,7 @@ def pipeline_hf_transformers( parallelism_context=parallelism_context, training=training, parallelism=parallelism, + local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py index 4febea30e20..8f1c0b5b63d 100644 --- a/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py +++ b/torchtitan/experiments/transformers_modeling_backend/tests/test_flex_attention.py @@ -78,31 +78,6 @@ def test_packed_documents_do_not_leak(self): self.assertFalse(torch.allclose(out1[:, :len_a], out2[:, :len_a])) -class TestHFTransformerConfig(unittest.TestCase): - def test_local_compile_defaults_to_inductor(self): - from torchtitan.distributed.local_compile import LocalCompileConfig - from torchtitan.experiments.transformers_modeling_backend import ( - TitanModelConfig, - ) - from torchtitan.experiments.transformers_modeling_backend.model import ( - HFTransformerModel, - ) - - config = HFTransformerModel.Config( - model_config=TitanModelConfig( - dim=256, - n_layers=2, - n_heads=16, - n_kv_heads=16, - ), - ) - - self.assertEqual( - config.local_compile, - LocalCompileConfig(regions=["loss", "fused_binary_activation"]), - ) - - class TestAttnMaskTypeValidation(unittest.TestCase): """A block_causal config builds a flex BlockMask and is supported. diff --git a/torchtitan/models/base/model.py b/torchtitan/models/base/model.py index 2792fac3692..52760ed6a96 100644 --- a/torchtitan/models/base/model.py +++ b/torchtitan/models/base/model.py @@ -14,7 +14,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.distributed.parallelism_context import ParallelismContext from .module import Module @@ -86,10 +86,10 @@ class Config(Module.Config): Subclasses define model-specific hyperparameters. """ - local_compile: LocalCompileConfig = field(default_factory=LocalCompileConfig) - """TorchTitan-owned local compile configuration. + local_compile_regions: list[str] = field(default_factory=list) + """``@local_compile`` regions this model compiles; ``[]`` runs them all eager. - Covers named local regions, FlexAttention, and ``create_block_mask``. + FlexAttention manages its own compilation and is not controlled by this list. """ @abstractmethod @@ -132,10 +132,13 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: """Apply the ordered model-level parallelization lifecycle.""" + # Bind local implementations early; torch.compile traces on first use. + apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) if activation_remat is not None: diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py index 7e723db799c..26ce302f79a 100644 --- a/torchtitan/models/common/multimodal.py +++ b/torchtitan/models/common/multimodal.py @@ -20,6 +20,7 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.distributed.parallelism_context import ParallelismContext from .decoder import Decoder @@ -36,9 +37,12 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: + # Bind local implementations early; torch.compile traces on first use. + apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) encoders = [ diff --git a/torchtitan/models/deepseek_v3/model.py b/torchtitan/models/deepseek_v3/model.py index 805036eacf0..c1ba8bc60ee 100644 --- a/torchtitan/models/deepseek_v3/model.py +++ b/torchtitan/models/deepseek_v3/model.py @@ -13,7 +13,6 @@ import torch_remat as remat from torch import nn -from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.base.module import Module from torchtitan.models.common.attention import ( BaseAttention, @@ -273,15 +272,13 @@ class DeepSeekV3Model(MTPDecoder): class Config(MTPDecoder.Config): dim: int = 2048 vocab_size: int = 102400 - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=[ - "loss", - "fused_binary_activation", - "fp32_to_bf16_split", - "moe_aux_loss", - ] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: [ + "loss", + "fused_binary_activation", + "fp32_to_bf16_split", + "moe_aux_loss", + ] ) def get_nparams_and_flops( diff --git a/torchtitan/models/deepseek_v4/model.py b/torchtitan/models/deepseek_v4/model.py index eba632c669b..1da25fda7a4 100644 --- a/torchtitan/models/deepseek_v4/model.py +++ b/torchtitan/models/deepseek_v4/model.py @@ -14,7 +14,6 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig -from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.module import ModuleList from torchtitan.models.common.attention import ( @@ -180,14 +179,12 @@ def _apply_fsdp( class Config(Decoder.Config): dim: int vocab_size: int - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=[ - "loss", - "fused_binary_activation", - "fp32_to_bf16_split", - ] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: [ + "loss", + "fused_binary_activation", + "fp32_to_bf16_split", + ] ) hc_mult: int = 4 n_mtp_layers: int = 0 diff --git a/torchtitan/models/flux/README.md b/torchtitan/models/flux/README.md index 9ff56054f75..bb331f575d2 100644 --- a/torchtitan/models/flux/README.md +++ b/torchtitan/models/flux/README.md @@ -51,7 +51,7 @@ MODULE=torchtitan_recipes.tests.models.flux CONFIG=flux_debugmodel_inference \ The loss function can be compiled independently with: ```python -config.model.local_compile.regions = ["loss"] +config.model.local_compile_regions = ["loss"] ``` The model itself is not compiled by the standard Trainer. diff --git a/torchtitan/models/flux/model.py b/torchtitan/models/flux/model.py index 36e4863a3fb..39c6480c931 100644 --- a/torchtitan/models/flux/model.py +++ b/torchtitan/models/flux/model.py @@ -18,7 +18,7 @@ from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed import context_parallel from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_replicated_parameters, @@ -103,9 +103,7 @@ class FluxModel(BaseModel): class Config(BaseModel.Config): img_in: Linear.Config txt_in: Linear.Config - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig(regions=["loss"]) - ) + local_compile_regions: list[str] = field(default_factory=lambda: ["loss"]) in_channels: int = 64 out_channels: int = 64 vec_in_dim: int = 768 @@ -237,10 +235,13 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: """Apply Flux's AC-before-SPMD parallelization lifecycle.""" + # Bind local implementations early; torch.compile traces on first use. + apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): if activation_remat is not None: # Flux saves no regions under any policy: each block is diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py index 6421509398b..35bd95b26f0 100644 --- a/torchtitan/models/gpt_oss/model.py +++ b/torchtitan/models/gpt_oss/model.py @@ -17,7 +17,6 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.module import Module from torchtitan.models.common.attention import ( @@ -217,15 +216,13 @@ def _register_optimizer_hooks( class Config(Decoder.Config): dim: int = 2880 vocab_size: int = 201088 - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=[ - "loss", - "fused_binary_activation", - "cos_sin_rope", - "fp32_to_bf16_split", - ] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: [ + "loss", + "fused_binary_activation", + "cos_sin_rope", + "fp32_to_bf16_split", + ] ) def get_nparams_and_flops( @@ -253,6 +250,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> GptOssModel: @@ -269,6 +267,7 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, + local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/kimi_k2_7/model.py b/torchtitan/models/kimi_k2_7/model.py index 451ef173876..6d1dd05790b 100644 --- a/torchtitan/models/kimi_k2_7/model.py +++ b/torchtitan/models/kimi_k2_7/model.py @@ -117,6 +117,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK25Model: @@ -130,6 +131,7 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, + local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index b4875bbfc3d..f1313a458a1 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -19,7 +19,7 @@ from torchtitan.config.function import Function from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import local_compile, LocalCompileConfig +from torchtitan.distributed.local_compile import apply_local_compile, local_compile from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -391,16 +391,14 @@ class Config(Decoder.Config): attention_residual: AttentionResidual.Config = field( default_factory=AttentionResidual.Config ) - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=[ - "loss", - "gated_rmsnorm", - "fused_binary_activation", - "fp32_to_bf16_split", - "attention_residual", - ] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: [ + "loss", + "gated_rmsnorm", + "fused_binary_activation", + "fp32_to_bf16_split", + "attention_residual", + ] ) def get_nparams_and_flops( @@ -460,9 +458,12 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK3Model: + # Bind local implementations early; torch.compile traces on first use. + apply_local_compile(local_compile_regions) with parallelism_context.activate_spmd(): annotate_replicated_parameters(self, parallelism_context) self._parallelize(parallelism_context) diff --git a/torchtitan/models/llama3/model.py b/torchtitan/models/llama3/model.py index b3eb5f37e04..b0335209d90 100644 --- a/torchtitan/models/llama3/model.py +++ b/torchtitan/models/llama3/model.py @@ -12,7 +12,6 @@ import torch_remat as remat from torch import nn -from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.attention import ( FlexAttentionMetadata, VarlenAttentionMetadata, @@ -84,10 +83,8 @@ class Llama3Model(Decoder): class Config(Decoder.Config): dim: int = 4096 vocab_size: int = 128256 - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=["loss", "fused_binary_activation"] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: ["loss", "fused_binary_activation"] ) def get_nparams_and_flops( diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py index 04f363efced..72938708b72 100644 --- a/torchtitan/models/muse_glimmer/model.py +++ b/torchtitan/models/muse_glimmer/model.py @@ -18,7 +18,6 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -290,10 +289,8 @@ class MuseGlimmerModel(MultimodalModel): class Config(Decoder.Config): dim: int = 6656 vocab_size: int = 202048 - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=["loss", "fused_binary_activation"] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: ["loss", "fused_binary_activation"] ) # Narrows the base Decoder.Config.tok_embeddings (Embedding.Config) to the # bundled embedding+norm unit whose .embedding/.norm children carry the sharding. @@ -372,6 +369,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> MuseGlimmerModel: @@ -385,6 +383,7 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, + local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/qwen3/model.py b/torchtitan/models/qwen3/model.py index 974855616c6..0ec1c454b65 100644 --- a/torchtitan/models/qwen3/model.py +++ b/torchtitan/models/qwen3/model.py @@ -12,7 +12,6 @@ import torch.nn as nn import torch_remat as remat -from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.attention import ( FlexAttentionMetadata, VarlenAttentionMetadata, @@ -109,15 +108,13 @@ def _register_optimizer_hooks( class Config(Decoder.Config): dim: int = 1024 vocab_size: int = 151936 - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=[ - "loss", - "fused_binary_activation", - "cos_sin_rope", - "fp32_to_bf16_split", - ] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: [ + "loss", + "fused_binary_activation", + "cos_sin_rope", + "fp32_to_bf16_split", + ] ) def get_nparams_and_flops( diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index 2a49bd9f05f..bd68938bfb6 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -18,7 +18,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import local_compile, LocalCompileConfig +from torchtitan.distributed.local_compile import local_compile from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -353,18 +353,16 @@ def _register_optimizer_hooks( @dataclass(kw_only=True, slots=True) class Config(Decoder.Config): vision_encoder: Qwen35VisionEncoder.Config | None = None - local_compile: LocalCompileConfig = field( - default_factory=lambda: LocalCompileConfig( - regions=[ - "loss", - "fused_binary_activation", - "gated_rmsnorm", - "offset_rmsnorm", - "partial_rope", - "shared_expert_gate", - "fp32_to_bf16_split", - ] - ) + local_compile_regions: list[str] = field( + default_factory=lambda: [ + "loss", + "fused_binary_activation", + "gated_rmsnorm", + "offset_rmsnorm", + "partial_rope", + "shared_expert_gate", + "fp32_to_bf16_split", + ] ) def get_nparams_and_flops( @@ -438,6 +436,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, + local_compile_regions: list[str], activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Qwen35Model: @@ -445,6 +444,7 @@ def parallelize( parallelism_context=parallelism_context, training=training, parallelism=parallelism, + local_compile_regions=local_compile_regions, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/rl/generator/generator.py b/torchtitan/rl/generator/generator.py index e03a9d528c6..8c464632cc1 100644 --- a/torchtitan/rl/generator/generator.py +++ b/torchtitan/rl/generator/generator.py @@ -210,7 +210,7 @@ def _prepare_generation_request_metrics( class VLLMCudaGraphConfig: """CUDA graph capture settings for the vLLM inference engine. - Local compile regions come from the model config's ``local_compile.regions``. + Local compile regions come from the model config's ``local_compile_regions``. Only CUDA graph capture, which is vLLM-specific, is controlled here. ``mode`` selects which vLLM CUDA graph mode to capture; see that field and diff --git a/torchtitan/rl/generator/vllm/model_wrapper.py b/torchtitan/rl/generator/vllm/model_wrapper.py index 1b74ba65168..8438e0745dc 100644 --- a/torchtitan/rl/generator/vllm/model_wrapper.py +++ b/torchtitan/rl/generator/vllm/model_wrapper.py @@ -32,7 +32,7 @@ from vllm.utils import torch_utils as _torch_utils # Registers the "loss" local-compile region. The generator shares the trainer's -# model config, whose local_compile.regions may name it, but never builds a loss. +# model config, whose local_compile_regions may name it, but never builds a loss. from torchtitan.components import loss as _loss # noqa: F401 from torchtitan.components.checkpointer import CheckpointManager from torchtitan.config import TrainingConfig @@ -359,7 +359,6 @@ def __init__( # Build model on meta device to avoid allocating full model on every GPU with self.parallelism_context.activate_spmd(), torch.device("meta"): self.model = self.config.build() - self.config.local_compile.apply() self.model._skip_lm_head = True if getattr(self.model, "vision_encoder", None) is not None: self.model.vision_encoder = None @@ -368,6 +367,7 @@ def __init__( parallelism_context=self.parallelism_context, training=TrainingConfig(), parallelism=training_parallelism, + local_compile_regions=self.config.local_compile_regions, activation_remat=None, dump_folder="", ) diff --git a/torchtitan/training_engine.py b/torchtitan/training_engine.py index 5fa54d097ca..721c8a2e6be 100644 --- a/torchtitan/training_engine.py +++ b/torchtitan/training_engine.py @@ -331,8 +331,6 @@ def _initialize_model( utils.set_default_dtype(TORCH_DTYPE_MAP[self.config.training.dtype]), ): model = self.model_config.build() - # Model construction imports and registers model-specific local regions. - self.model_config.local_compile.apply() self.model_cls = type(model) adapter_cls = type(model).state_dict_adapter_cls self.state_dict_adapter = ( @@ -411,11 +409,13 @@ def _initialize_model( def _parallelize_compile_kwargs(self) -> dict[str, Any]: """Return the compile kwargs ``model.parallelize``/``model.pipeline`` expect for this engine's models. - Regular models need no compile kwargs because ``model.local_compile`` is - applied after model construction. GraphTrainer overrides this to return - ``compile_config`` (whole-step compile). + Regular models expect ``local_compile_regions`` (compiled regions); GraphTrainer overrides this to + return ``compile_config`` (whole-step compile). """ - return {} + # TODO: apply local compile outside parallelize/pipeline (#5026 review). That needs + # GraphTrainer to stop reading its compile config there first (apply_compile, + # EP-overlap chunking, enable_autoparallel, GraphPP runtime); then delete this hook. + return {"local_compile_regions": self.model_config.local_compile_regions} def _initialize_optim(self) -> None: """Construct the parameter update and its state.""" diff --git a/torchtitan_recipes/models/deepseek_v3.py b/torchtitan_recipes/models/deepseek_v3.py index e1ca1d6689b..5a762cdec49 100644 --- a/torchtitan_recipes/models/deepseek_v3.py +++ b/torchtitan_recipes/models/deepseek_v3.py @@ -268,5 +268,5 @@ def deepseek_v3_671b_nvfp4_ffn_mxfp8_attn( ) ], ) - config.model.local_compile.regions = ["loss"] + config.model.local_compile_regions = ["loss"] return config diff --git a/torchtitan_recipes/rl/alphabet_sort.py b/torchtitan_recipes/rl/alphabet_sort.py index f7fb37a771b..f7125220ca0 100644 --- a/torchtitan_recipes/rl/alphabet_sort.py +++ b/torchtitan_recipes/rl/alphabet_sort.py @@ -174,7 +174,7 @@ def rl_grpo_qwen3_0_6b_varlen(*, seq_len: int = 2048) -> Controller.Config: def rl_grpo_qwen3_0_6b_varlen_no_compile(*, seq_len: int = 2048) -> Controller.Config: config = rl_grpo_qwen3_0_6b_varlen(seq_len=seq_len) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return config @@ -299,7 +299,7 @@ def rl_grpo_qwen3_0_6b_flex_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] block_size = config.trainer.model.layers[0].attention.inner_attention.block_size config.async_loop.batcher = dataclasses.replace( config.async_loop.batcher, per_sample_pad_multiple=block_size @@ -671,7 +671,7 @@ def rl_grpo_qwen3_0_6b_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return Controller.Config( hf_assets_path="./assets/hf/Qwen3-0.6B", num_generators=3, @@ -856,7 +856,7 @@ def rl_grpo_qwen3_5_9b_varlen_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return config diff --git a/torchtitan_recipes/tests/__init__.py b/torchtitan_recipes/tests/__init__.py index 8ac49cd5624..3792be993e9 100644 --- a/torchtitan_recipes/tests/__init__.py +++ b/torchtitan_recipes/tests/__init__.py @@ -27,4 +27,4 @@ def _set_spmd_typechecking(config: Trainer.Config, *, typechecking: bool) -> Non if typechecking: # Local compile does not support global SPMD typechecking yet. # TODO: Remove this once the SPMD typechecking/Dynamo issue is fixed. - config.model.local_compile.regions = [] + config.model.local_compile_regions = [] diff --git a/torchtitan_recipes/tests/graph_trainer/numerics.py b/torchtitan_recipes/tests/graph_trainer/numerics.py index 9b3ee3f2df2..277344d9ce7 100644 --- a/torchtitan_recipes/tests/graph_trainer/numerics.py +++ b/torchtitan_recipes/tests/graph_trainer/numerics.py @@ -55,7 +55,7 @@ def _without_inplace_wgrad_accum(config): def llama3_eager_numerics(): config = llama3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile.regions = [] + config.model.local_compile_regions = [] config.parallelism.tensor_parallel_degree = 2 config.parallelism.data_parallel_shard_degree = 4 config.training.num_tokens_per_microbatch_per_dp_rank = 16384 @@ -73,7 +73,7 @@ def llama3_graph_numerics(): def llama3_eager_hsdp_gradient_accumulation_numerics(): config = llama3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile.regions = [] + config.model.local_compile_regions = [] config.parallelism.data_parallel_replicate_degree = 2 config.parallelism.data_parallel_shard_degree = 2 config.training.num_tokens_per_microbatch_per_dp_rank = 8192 @@ -103,7 +103,7 @@ def _llama3_hsdp_pp_numerics(config): def llama3_eager_hsdp_pp_numerics(): config = llama3_debugmodel(seq_len=2048) - config.model.local_compile.regions = [] + config.model.local_compile_regions = [] return _without_inplace_wgrad_accum(_llama3_hsdp_pp_numerics(config)) @@ -214,7 +214,7 @@ def _deepseek_v3_dist_moe_pp_numerics( runtime_config.pp_activation_slot_policy = slot_policy for _, experts, _, _ in config.model.traverse(DistMoeRoutedExperts.Config): experts.inplace_wgrad_accum = False - config.model.local_compile.regions = [] + config.model.local_compile_regions = [] config.parallelism.pipeline_parallel_degree = 2 config.parallelism.pipeline_parallel_schedule = "Interleaved1F1B" config.parallelism.num_pp_microbatches = 8 @@ -251,7 +251,7 @@ def deepseek_v3_dist_moe_graph_pp_stage_microbatch_numerics(): def qwen3_eager_numerics(): config = qwen3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile.regions = [] + config.model.local_compile_regions = [] config.parallelism.tensor_parallel_degree = 2 config.parallelism.data_parallel_shard_degree = 4 config.training.num_tokens_per_microbatch_per_dp_rank = 16384 @@ -278,7 +278,7 @@ def qwen3_moe_eager_numerics(): ) config.dataloader.max_num_documents = None # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile.regions = [] + config.model.local_compile_regions = [] config.training.disable_cuda_graphs = True config.training.num_tokens_per_microbatch_per_dp_rank = 16384 return _without_inplace_wgrad_accum(config) diff --git a/torchtitan_recipes/tests/models/deepseek_v3.py b/torchtitan_recipes/tests/models/deepseek_v3.py index 9fe20eafc3d..65225c382bc 100644 --- a/torchtitan_recipes/tests/models/deepseek_v3.py +++ b/torchtitan_recipes/tests/models/deepseek_v3.py @@ -481,7 +481,7 @@ def deepseek_v3_debugmodel_nvfp4_ffn_mxfp8_attn( ), ], ) - config.model.local_compile.regions = ["loss"] + config.model.local_compile_regions = ["loss"] return config @@ -530,5 +530,5 @@ def deepseek_v3_16b_nvfp4_ffn_mxfp8_attn( ) ], ) - config.model.local_compile.regions = ["loss"] + config.model.local_compile_regions = ["loss"] return config diff --git a/torchtitan_recipes/tests/rl/alphabet_sort.py b/torchtitan_recipes/tests/rl/alphabet_sort.py index 52b18fcfb65..ce1b21069ed 100644 --- a/torchtitan_recipes/tests/rl/alphabet_sort.py +++ b/torchtitan_recipes/tests/rl/alphabet_sort.py @@ -156,7 +156,7 @@ def rl_grpo_gpt_oss_debug_varlen_no_compile( ) -> Controller.Config: config = rl_grpo_gpt_oss_debug_varlen(seq_len=seq_len) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return config @@ -194,7 +194,7 @@ def rl_grpo_gpt_oss_debug_varlen_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return Controller.Config( hf_assets_path="tests/assets/tokenizer", async_loop=AsyncLoopConfig( @@ -439,7 +439,7 @@ def rl_grpo_qwen3_moe_debug_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return Controller.Config( hf_assets_path="tests/assets/tokenizer", async_loop=AsyncLoopConfig( @@ -604,7 +604,7 @@ def rl_grpo_qwen3_5_debug_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return config @@ -629,7 +629,7 @@ def _kimi_k3_rl_model_configs( converters=converters, ) for model_config in (trainer_model, generator_model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] return trainer_model, generator_model @@ -819,7 +819,7 @@ def rl_grpo_moe_debug_tp4_ep4() -> Controller.Config: generator_parallelism=config.generator.parallelism, ) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile.regions = [] + model_config.local_compile_regions = [] # TODO: use FULL CUDA graphs once GPT-OSS CI enables a # CUDA-graph-compatible distributed MoE token dispatcher. config.generator.cuda_graph.mode = "NONE" From cc97b3a3ce009cad86447ecebfd788b8042ee62f Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 22:40:17 +0000 Subject: [PATCH 5/6] Centralize local compile defaults --- torchtitan/distributed/local_compile.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/torchtitan/distributed/local_compile.py b/torchtitan/distributed/local_compile.py index ee8f3cae72c..984ed568fb5 100644 --- a/torchtitan/distributed/local_compile.py +++ b/torchtitan/distributed/local_compile.py @@ -8,7 +8,7 @@ import functools from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import torch @@ -24,7 +24,7 @@ class LocalCompileConfig: """Configuration for TorchTitan-owned local ``torch.compile`` boundaries.""" - regions: list[str] + regions: list[str] = field(default_factory=list) backend: str = "inductor" def __bool__(self) -> bool: @@ -68,7 +68,7 @@ def apply(self) -> None: bind_local_compile_fn(self) -_ACTIVE_LOCAL_COMPILE_CONFIG = LocalCompileConfig(regions=[]) +_ACTIVE_LOCAL_COMPILE_CONFIG = LocalCompileConfig() def _as_local_compile_config( From fe598acc5b1622fe10203f34b2650453c9d8bade Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 22:50:41 +0000 Subject: [PATCH 6/6] Adopt LocalCompileConfig for local compile boundaries --- tests/rl/unit_tests/cpu/test_generator.py | 4 +- .../cpu/test_policy_loss_configs.py | 4 +- tests/rl/unit_tests/gpu/test_gdn_cudagraph.py | 2 +- .../cpu/config/test_compile_config.py | 136 ++++++++++-------- .../test_context_parallel_validation.py | 3 +- .../cpu/models/common/test_gated_rmsnorm.py | 4 +- .../models/common/test_hi_mid_lo_linear.py | 2 +- .../cpu/models/kimi_k3/test_kimi_k3_config.py | 2 +- .../cpu/models/qwen3_5/test_offset_rmsnorm.py | 2 +- .../cpu/quantization/test_quantization.py | 4 +- .../cpu/test_integration_test_definitions.py | 2 +- tests/unit_tests/cpu/test_trainer.py | 4 +- .../gpu/models/common/test_gated_rmsnorm.py | 6 +- .../models/common/test_hi_mid_lo_linear.py | 6 +- .../common/test_moe_aux_loss_compile.py | 8 +- .../gpu/models/common/test_situglu.py | 6 +- .../gpu/models/common/test_swiglu.py | 6 +- .../test_kimi_k3_attention_residual.py | 6 +- .../gpu/models/qwen3_5/test_offset_rmsnorm.py | 10 +- .../gpu/models/qwen3_5/test_rope_compile.py | 10 +- .../gpu/models/qwen3_5/test_sigmoid_gate.py | 6 +- torchtitan/config/transform/lm_head_fp32.py | 4 +- torchtitan/config/validation.py | 5 +- torchtitan/distributed/local_compile.py | 76 ++++------ torchtitan/distributed/pipeline_parallel.py | 3 +- .../experiments/graph_trainer/configs.py | 3 +- .../experiments/graph_trainer/trainer.py | 11 +- .../transformers_modeling_backend/model.py | 5 +- .../parallelize.py | 6 +- .../transformers_modeling_backend/pipeline.py | 3 +- torchtitan/models/base/model.py | 14 +- torchtitan/models/common/multimodal.py | 6 +- torchtitan/models/deepseek_v3/model.py | 17 ++- torchtitan/models/deepseek_v4/model.py | 15 +- torchtitan/models/flux/README.md | 2 +- torchtitan/models/flux/model.py | 10 +- torchtitan/models/gpt_oss/model.py | 19 +-- torchtitan/models/kimi_k2_7/model.py | 3 +- torchtitan/models/kimi_k3/model.py | 24 ++-- torchtitan/models/llama3/model.py | 7 +- torchtitan/models/muse_glimmer/model.py | 9 +- torchtitan/models/qwen3/model.py | 17 ++- torchtitan/models/qwen3_5/model.py | 26 ++-- torchtitan_recipes/models/deepseek_v3.py | 2 +- torchtitan_recipes/rl/alphabet_sort.py | 8 +- torchtitan_recipes/tests/__init__.py | 2 +- .../tests/graph_trainer/numerics.py | 12 +- .../tests/models/deepseek_v3.py | 4 +- torchtitan_recipes/tests/rl/alphabet_sort.py | 12 +- 49 files changed, 296 insertions(+), 262 deletions(-) diff --git a/tests/rl/unit_tests/cpu/test_generator.py b/tests/rl/unit_tests/cpu/test_generator.py index 3d973912d6f..3caf039fd40 100644 --- a/tests/rl/unit_tests/cpu/test_generator.py +++ b/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 diff --git a/tests/rl/unit_tests/cpu/test_policy_loss_configs.py b/tests/rl/unit_tests/cpu/test_policy_loss_configs.py index b79c8ef600b..facd59e2261 100644 --- a/tests/rl/unit_tests/cpu/test_policy_loss_configs.py +++ b/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 == [] diff --git a/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py b/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py index 176cbc6c46e..3c3b7b56dfa 100644 --- a/tests/rl/unit_tests/gpu/test_gdn_cudagraph.py +++ b/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), diff --git a/tests/unit_tests/cpu/config/test_compile_config.py b/tests/unit_tests/cpu/config/test_compile_config.py index ea155eecacd..f4c697d761f 100644 --- a/tests/unit_tests/cpu/config/test_compile_config.py +++ b/tests/unit_tests/cpu/config/test_compile_config.py @@ -4,32 +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, - LocalCompileConfig, -) +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: @@ -48,49 +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 compiled_calls[0][1]["backend"] == "inductor" - assert compiled_calls[0][1]["fullgraph"] is True - - -def test_local_compile_config_applies_backend_to_all_owned_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 lambda *args, **kwargs: None - - monkeypatch.setattr(torch, "compile", fake_compile) - apply_local_compile( - LocalCompileConfig( - regions=["test_custom_backend"], - backend="custom_device_backend", - ) - ) - - region() - FlexInnerAttention._compiled_flex_attn() - _compiled_create_block_mask() - compiled_create_block_mask() - - assert compiled_calls == [ - ( - region.__wrapped__, - {"backend": "custom_device_backend", "fullgraph": True}, - ), - (flex_attention, {"backend": "custom_device_backend"}), - (create_block_mask, {"backend": "custom_device_backend"}), - (create_block_mask, {"backend": "custom_device_backend"}), - ] + 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: @@ -113,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, @@ -145,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: @@ -157,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__, @@ -174,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__, - {"backend": "inductor", "fullgraph": True}, + {"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 diff --git a/tests/unit_tests/cpu/config/test_context_parallel_validation.py b/tests/unit_tests/cpu/config/test_context_parallel_validation.py index 0575d338eb8..0a1ced4cf88 100644 --- a/tests/unit_tests/cpu/config/test_context_parallel_validation.py +++ b/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="", ) diff --git a/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py b/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py index 4d1ecbae692..7fd0f24639e 100644 --- a/tests/unit_tests/cpu/models/common/test_gated_rmsnorm.py +++ b/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) diff --git a/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py b/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py index c2664043867..490eee0b823 100644 --- a/tests/unit_tests/cpu/models/common/test_hi_mid_lo_linear.py +++ b/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) diff --git a/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py b/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py index e22edb1d94b..f8f3d452320 100644 --- a/tests/unit_tests/cpu/models/kimi_k3/test_kimi_k3_config.py +++ b/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", diff --git a/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py b/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py index b43336cb5ad..04f095b348f 100644 --- a/tests/unit_tests/cpu/models/qwen3_5/test_offset_rmsnorm.py +++ b/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", diff --git a/tests/unit_tests/cpu/quantization/test_quantization.py b/tests/unit_tests/cpu/quantization/test_quantization.py index 61603ae7ced..f304707c13f 100644 --- a/tests/unit_tests/cpu/quantization/test_quantization.py +++ b/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) diff --git a/tests/unit_tests/cpu/test_integration_test_definitions.py b/tests/unit_tests/cpu/test_integration_test_definitions.py index 0eb38f0db00..0235821f2b6 100644 --- a/tests/unit_tests/cpu/test_integration_test_definitions.py +++ b/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__() diff --git a/tests/unit_tests/cpu/test_trainer.py b/tests/unit_tests/cpu/test_trainer.py index c79f1b0798f..26b7349b019 100644 --- a/tests/unit_tests/cpu/test_trainer.py +++ b/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="", ) diff --git a/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py b/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py index 3d1604aa825..3badec1bc25 100644 --- a/tests/unit_tests/gpu/models/common/test_gated_rmsnorm.py +++ b/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 diff --git a/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py b/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py index 1cd17b19b6b..2821b03c20c 100644 --- a/tests/unit_tests/gpu/models/common/test_hi_mid_lo_linear.py +++ b/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]) diff --git a/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py b/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py index b5c8145489f..4482a1e710a 100644 --- a/tests/unit_tests/gpu/models/common/test_moe_aux_loss_compile.py +++ b/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, diff --git a/tests/unit_tests/gpu/models/common/test_situglu.py b/tests/unit_tests/gpu/models/common/test_situglu.py index 3bd6f42355a..a9b262c96f3 100644 --- a/tests/unit_tests/gpu/models/common/test_situglu.py +++ b/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): diff --git a/tests/unit_tests/gpu/models/common/test_swiglu.py b/tests/unit_tests/gpu/models/common/test_swiglu.py index dde74bd4d75..85b09654cb8 100644 --- a/tests/unit_tests/gpu/models/common/test_swiglu.py +++ b/tests/unit_tests/gpu/models/common/test_swiglu.py @@ -9,17 +9,17 @@ 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.activation import ClampedSwiGLU, SiTUGLU, SwiGLU @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestSwiGLULocalCompile(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): diff --git a/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py b/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py index ea9ef9c9b20..45456e29bdb 100644 --- a/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py +++ b/tests/unit_tests/gpu/models/kimi_k3/test_kimi_k3_attention_residual.py @@ -8,7 +8,7 @@ import torch -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.kimi_k3.model import AttentionResidual @@ -45,11 +45,11 @@ def _reference( @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestKimiK3AttentionResidualLocalCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["attention_residual"]) + LocalCompileConfig(regions=["attention_residual"]).apply() self.attention_residual = AttentionResidual.Config().build() def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py b/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py index f39c6fb5ba6..c4a0fa1e1b1 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_offset_rmsnorm.py @@ -13,7 +13,7 @@ with_comms, ) -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import set_current_spmd_mesh, set_spmd_meshes from torchtitan.models.common.decoder_sharding import head_norm_sharding_config @@ -23,10 +23,10 @@ @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestOffsetRMSNormCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["offset_rmsnorm"]) + LocalCompileConfig(regions=["offset_rmsnorm"]).apply() def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): @@ -101,7 +101,7 @@ def world_size(self) -> int: @with_comms def test_compiled_qk_norm(self): - apply_local_compile(["offset_rmsnorm"]) + LocalCompileConfig(regions=["offset_rmsnorm"]).apply() device = self.device_type config = OffsetRMSNorm.Config( dim=128, @@ -137,7 +137,7 @@ def test_compiled_qk_norm(self): output.sum().backward() self.assertEqual(output.shape, x_local.shape) - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py b/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py index c7c91aa4ad8..d6fbe401791 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_rope_compile.py @@ -8,7 +8,7 @@ import torch -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.rope import CosSinRoPE from torchtitan.models.qwen3_5.rope import MRoPE @@ -16,10 +16,10 @@ @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestRoPELocalCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["cos_sin_rope"]) + LocalCompileConfig(regions=["cos_sin_rope"]).apply() def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_emit_triton(self): @@ -89,7 +89,7 @@ def test_forward_and_backward_match_eager(self): (grad_query_out, grad_key_out), ) - apply_local_compile([]) + LocalCompileConfig().apply() eager_rope = ( CosSinRoPE.Config( dim=128, @@ -222,7 +222,7 @@ def test_mrope_three_axis_positions_match_eager(self): (grad_query_out, grad_key_out), ) - apply_local_compile([]) + LocalCompileConfig().apply() eager_rope = config.build().cuda() eager_query = query.detach().clone().requires_grad_() eager_key = key.detach().clone().requires_grad_() diff --git a/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py b/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py index 83b3ab74fca..8448defd201 100644 --- a/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.py +++ b/tests/unit_tests/gpu/models/qwen3_5/test_sigmoid_gate.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.qwen3_5.moe import sigmoid_gate @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestSigmoidGateLocalCompile(unittest.TestCase): def setUp(self): - apply_local_compile(["shared_expert_gate"]) + LocalCompileConfig(regions=["shared_expert_gate"]).apply() # A local generator leaves the global RNG stream unchanged for later tests. generator = torch.Generator(device="cuda").manual_seed(0) @@ -29,7 +29,7 @@ def randn(*shape: int) -> torch.Tensor: self.grad_output = randn(64, 128) def tearDown(self): - apply_local_compile([]) + LocalCompileConfig().apply() torch._dynamo.reset() def test_forward_and_backward_match_reference(self): diff --git a/torchtitan/config/transform/lm_head_fp32.py b/torchtitan/config/transform/lm_head_fp32.py index a7cea8aa0f4..10f2ca1db59 100644 --- a/torchtitan/config/transform/lm_head_fp32.py +++ b/torchtitan/config/transform/lm_head_fp32.py @@ -67,6 +67,6 @@ def convert(self, model_config): "the model config. The torchtitan decoder names its output projection " f"{self._TARGET!r} (see torchtitan/models/common/decoder.py)." ) - if "fp32_to_bf16_split" not in model_config.local_compile_regions: - model_config.local_compile_regions.append("fp32_to_bf16_split") + if "fp32_to_bf16_split" not in model_config.local_compile_regions.regions: + model_config.local_compile_regions.regions.append("fp32_to_bf16_split") return model_config diff --git a/torchtitan/config/validation.py b/torchtitan/config/validation.py index 5bdaab9b22b..4e0b8352940 100644 --- a/torchtitan/config/validation.py +++ b/torchtitan/config/validation.py @@ -137,11 +137,12 @@ def validate_model_training_config( f"dispatcher: {type(dispatcher_config).__qualname__}." ) - if debug.spmd_typechecking and model.local_compile_regions: + if debug.spmd_typechecking and model.local_compile_regions.regions: # TODO: Remove this once Dynamo supports tracing SPMD typechecking. raise ValueError( "Local compilation is not supported with SPMD typechecking. " - "Set model.local_compile_regions=[] or disable debug.spmd_typechecking." + "Set model.local_compile_regions.regions=[] or disable " + "debug.spmd_typechecking." ) validate_context_parallel(model, parallelism) diff --git a/torchtitan/distributed/local_compile.py b/torchtitan/distributed/local_compile.py index 984ed568fb5..2988d4c62d6 100644 --- a/torchtitan/distributed/local_compile.py +++ b/torchtitan/distributed/local_compile.py @@ -17,7 +17,7 @@ _LOCAL_COMPILE_CALLBACKS: dict[str, list[Callable[["LocalCompileConfig"], None]]] = {} -_LOCAL_COMPILE_GENERATION = 0 +_ALWAYS_COMPILE_CALLBACKS: list[Callable[["LocalCompileConfig"], None]] = [] @dataclass(kw_only=True, slots=True) @@ -25,33 +25,24 @@ class LocalCompileConfig: """Configuration for TorchTitan-owned local ``torch.compile`` boundaries.""" regions: list[str] = field(default_factory=list) - backend: str = "inductor" - def __bool__(self) -> bool: - return bool(self.regions) - - def __contains__(self, region: str) -> bool: - return region in self.regions - - def append(self, region: str) -> None: - self.regions.append(region) - - def compile( + def _compile( self, reference: Callable[..., Any], *, inductor_options: dict[str, Any] | None = None, **compile_kwargs: Any, ) -> Callable[..., Any]: - backend = compile_kwargs.setdefault("backend", self.backend) - if backend == "inductor" and inductor_options is not None: + """Compile one boundary with Inductor. + + Out-of-tree subclasses can override this hook to select another backend. + """ + if inductor_options is not None: compile_kwargs.setdefault("options", inductor_options) return torch.compile(reference, **compile_kwargs) def apply(self) -> None: - """Activate this config and bind all registered named regions.""" - global _ACTIVE_LOCAL_COMPILE_CONFIG, _LOCAL_COMPILE_GENERATION - + """Bind named regions and always-compiled helpers using this config.""" unknown = [ name for name in self.regions if name not in _LOCAL_COMPILE_CALLBACKS ] @@ -61,28 +52,11 @@ def apply(self) -> None: f"registered values are {sorted(_LOCAL_COMPILE_CALLBACKS)}" ) - _ACTIVE_LOCAL_COMPILE_CONFIG = self - _LOCAL_COMPILE_GENERATION += 1 for callbacks in _LOCAL_COMPILE_CALLBACKS.values(): for bind_local_compile_fn in callbacks: bind_local_compile_fn(self) - - -_ACTIVE_LOCAL_COMPILE_CONFIG = LocalCompileConfig() - - -def _as_local_compile_config( - value: list[str] | LocalCompileConfig, -) -> LocalCompileConfig: - """Normalize the existing region-list setting to one local compile config.""" - if isinstance(value, LocalCompileConfig): - return value - return LocalCompileConfig(regions=value) - - -def apply_local_compile(value: list[str] | LocalCompileConfig) -> None: - """Bind local compile boundaries from a region list or full config.""" - _as_local_compile_config(value).apply() + for bind_compile_fn in _ALWAYS_COMPILE_CALLBACKS: + bind_compile_fn(self) def local_compile( @@ -116,7 +90,7 @@ def bind_local_compile(config: LocalCompileConfig) -> None: "batch-invariant mode; remove it from local_compile_regions." ) if enabled: - fn = config.compile(reference, fullgraph=True, **compile_kwargs) + fn = config._compile(reference, fullgraph=True, **compile_kwargs) else: fn = reference @@ -136,28 +110,30 @@ def compiled( inductor_options: dict[str, Any] | None = None, **compile_kwargs: Any, ) -> Callable[..., Any]: - """Compile lazily with the active local compile backend.""" - fn: Callable[..., Any] | None = None - applied_generation = -1 + """Create an always-compiled helper controlled by ``LocalCompileConfig``.""" + fn = LocalCompileConfig()._compile( + reference, + inductor_options=inductor_options, + **compile_kwargs, + ) + + def bind_compile(config: LocalCompileConfig) -> None: + nonlocal fn + fn = config._compile( + reference, + inductor_options=inductor_options, + **compile_kwargs, + ) @functools.wraps(reference) def wrapped(*args: Any, **kwargs: Any) -> Any: - nonlocal applied_generation, fn - if applied_generation != _LOCAL_COMPILE_GENERATION: - fn = _ACTIVE_LOCAL_COMPILE_CONFIG.compile( - reference, - inductor_options=inductor_options, - **compile_kwargs, - ) - applied_generation = _LOCAL_COMPILE_GENERATION - assert fn is not None return fn(*args, **kwargs) + _ALWAYS_COMPILE_CALLBACKS.append(bind_compile) return wrapped __all__ = [ - "apply_local_compile", "compiled", "local_compile", "LocalCompileConfig", diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py index 67959106c65..769d5503a62 100644 --- a/torchtitan/distributed/pipeline_parallel.py +++ b/torchtitan/distributed/pipeline_parallel.py @@ -32,6 +32,7 @@ from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed import ParallelismContext from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.base.model import BaseModel from torchtitan.models.base.module import ModuleDict, ModuleList from torchtitan.models.common.decoder import Decoder @@ -78,7 +79,7 @@ def pipeline_llm( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 62e7e36314a..a3c1fd9a3ef 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -9,6 +9,7 @@ from torchtitan.components.loss import ChunkedLossWrapper from torchtitan.distributed.activation_remat import DefaultActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.experiments.graph_trainer.chunked_loss import ( ChunkedLossWrapperWithParamGrads, ) @@ -353,7 +354,7 @@ def to_graph_trainer_config( } ) # GraphTrainer compiles the whole step (config.compile), so it drops the model's local compile regions. - graph_model.local_compile_regions = [] + graph_model.local_compile_regions = LocalCompileConfig() disable_inplace_wgrad_accum(graph_model) d["model"] = graph_model diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py index 1b7c6c8a91e..bd8133f22f3 100644 --- a/torchtitan/experiments/graph_trainer/trainer.py +++ b/torchtitan/experiments/graph_trainer/trainer.py @@ -66,11 +66,11 @@ def __init__( ) -> None: if config.optim.enable_cuda_graph: raise ValueError("Optim CUDA graphs are not supported with GraphTrainer.") - if model_config.local_compile_regions: + if model_config.local_compile_regions.regions: raise ValueError( "GraphTrainer traces the whole step into one graph; set " - "model.local_compile_regions = [] " - f"(got {model_config.local_compile_regions})." + "model.local_compile_regions.regions = [] " + f"(got {model_config.local_compile_regions.regions})." ) if config.parallelism.fsdp_symm_mem_scope is not None: raise ValueError( @@ -359,7 +359,10 @@ class Config(Trainer.Config): compile: GraphTrainerCompileConfig = field( default_factory=GraphTrainerCompileConfig ) - """Whole-step compile. GraphTrainer requires ``model.local_compile_regions`` to be empty.""" + """Whole-step compile. + + GraphTrainer requires ``model.local_compile_regions.regions`` to be empty. + """ def __post_init__(self) -> None: Trainer.Config.__post_init__(self) diff --git a/torchtitan/experiments/transformers_modeling_backend/model.py b/torchtitan/experiments/transformers_modeling_backend/model.py index 6ec6c1080d2..6f135923a46 100644 --- a/torchtitan/experiments/transformers_modeling_backend/model.py +++ b/torchtitan/experiments/transformers_modeling_backend/model.py @@ -26,6 +26,7 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.model import BaseModel from torchtitan.models.base.module import Module, ModuleDict @@ -285,7 +286,9 @@ def __init__( ) self.sharding_config = None # HF modules run eager; the loss and the swapped-in TorchTitan MoE (SwiGLU) can compile. - self.local_compile_regions = ["loss", "fused_binary_activation"] + self.local_compile_regions = LocalCompileConfig( + regions=["loss", "fused_binary_activation"] + ) assert model_config is not None, "model_config is required" diff --git a/torchtitan/experiments/transformers_modeling_backend/parallelize.py b/torchtitan/experiments/transformers_modeling_backend/parallelize.py index 46be051bc05..a21cbe2bdee 100644 --- a/torchtitan/experiments/transformers_modeling_backend/parallelize.py +++ b/torchtitan/experiments/transformers_modeling_backend/parallelize.py @@ -27,7 +27,7 @@ enable_fsdp_symm_mem, get_fsdp_reshard_after_forward_policy, ) -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig logger = logging.getLogger(__name__) @@ -87,7 +87,7 @@ def parallelize_hf_transformers( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, **kwargs: Any, @@ -102,7 +102,7 @@ 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) + local_compile_regions.apply() # 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 diff --git a/torchtitan/experiments/transformers_modeling_backend/pipeline.py b/torchtitan/experiments/transformers_modeling_backend/pipeline.py index f439c413c35..8833d566bbf 100644 --- a/torchtitan/experiments/transformers_modeling_backend/pipeline.py +++ b/torchtitan/experiments/transformers_modeling_backend/pipeline.py @@ -25,6 +25,7 @@ from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed import ParallelismContext from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.pipeline_parallel import ( _build_get_mesh_callback, _build_pipeline_schedule, @@ -296,7 +297,7 @@ def pipeline_hf_transformers( *, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, diff --git a/torchtitan/models/base/model.py b/torchtitan/models/base/model.py index 52760ed6a96..2b211ef18a8 100644 --- a/torchtitan/models/base/model.py +++ b/torchtitan/models/base/model.py @@ -14,7 +14,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from .module import Module @@ -86,10 +86,12 @@ class Config(Module.Config): Subclasses define model-specific hyperparameters. """ - local_compile_regions: list[str] = field(default_factory=list) - """``@local_compile`` regions this model compiles; ``[]`` runs them all eager. + local_compile_regions: LocalCompileConfig = field( + default_factory=LocalCompileConfig + ) + """Configuration for named regions, FlexAttention, and block-mask helpers. - FlexAttention manages its own compilation and is not controlled by this list. + An empty ``regions`` list runs named regions eagerly. """ @abstractmethod @@ -132,13 +134,13 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: """Apply the ordered model-level parallelization lifecycle.""" # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) + local_compile_regions.apply() with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) if activation_remat is not None: diff --git a/torchtitan/models/common/multimodal.py b/torchtitan/models/common/multimodal.py index 26ce302f79a..9307c9a0f7e 100644 --- a/torchtitan/models/common/multimodal.py +++ b/torchtitan/models/common/multimodal.py @@ -20,7 +20,7 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from .decoder import Decoder @@ -37,12 +37,12 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) + local_compile_regions.apply() with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) encoders = [ diff --git a/torchtitan/models/deepseek_v3/model.py b/torchtitan/models/deepseek_v3/model.py index c1ba8bc60ee..a4c1d20f7f2 100644 --- a/torchtitan/models/deepseek_v3/model.py +++ b/torchtitan/models/deepseek_v3/model.py @@ -13,6 +13,7 @@ import torch_remat as remat from torch import nn +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.base.module import Module from torchtitan.models.common.attention import ( BaseAttention, @@ -272,13 +273,15 @@ class DeepSeekV3Model(MTPDecoder): class Config(MTPDecoder.Config): dim: int = 2048 vocab_size: int = 102400 - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "fp32_to_bf16_split", - "moe_aux_loss", - ] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "fp32_to_bf16_split", + "moe_aux_loss", + ] + ) ) def get_nparams_and_flops( diff --git a/torchtitan/models/deepseek_v4/model.py b/torchtitan/models/deepseek_v4/model.py index 1da25fda7a4..06ed95b7586 100644 --- a/torchtitan/models/deepseek_v4/model.py +++ b/torchtitan/models/deepseek_v4/model.py @@ -14,6 +14,7 @@ from torchtitan.config import TORCH_DTYPE_MAP, TrainingConfig from torchtitan.config.parallelism import ParallelismConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.module import ModuleList from torchtitan.models.common.attention import ( @@ -179,12 +180,14 @@ def _apply_fsdp( class Config(Decoder.Config): dim: int vocab_size: int - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "fp32_to_bf16_split", - ] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "fp32_to_bf16_split", + ] + ) ) hc_mult: int = 4 n_mtp_layers: int = 0 diff --git a/torchtitan/models/flux/README.md b/torchtitan/models/flux/README.md index bb331f575d2..abba84856a1 100644 --- a/torchtitan/models/flux/README.md +++ b/torchtitan/models/flux/README.md @@ -51,7 +51,7 @@ MODULE=torchtitan_recipes.tests.models.flux CONFIG=flux_debugmodel_inference \ The loss function can be compiled independently with: ```python -config.model.local_compile_regions = ["loss"] +config.model.local_compile_regions.regions = ["loss"] ``` The model itself is not compiled by the standard Trainer. diff --git a/torchtitan/models/flux/model.py b/torchtitan/models/flux/model.py index 39c6480c931..46bec80d43f 100644 --- a/torchtitan/models/flux/model.py +++ b/torchtitan/models/flux/model.py @@ -18,7 +18,7 @@ from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed import context_parallel from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import apply_local_compile +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_replicated_parameters, @@ -103,7 +103,9 @@ class FluxModel(BaseModel): class Config(BaseModel.Config): img_in: Linear.Config txt_in: Linear.Config - local_compile_regions: list[str] = field(default_factory=lambda: ["loss"]) + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig(regions=["loss"]) + ) in_channels: int = 64 out_channels: int = 64 vec_in_dim: int = 768 @@ -235,13 +237,13 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: """Apply Flux's AC-before-SPMD parallelization lifecycle.""" # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) + local_compile_regions.apply() with parallelism_context.activate_spmd(): if activation_remat is not None: # Flux saves no regions under any policy: each block is diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py index 35bd95b26f0..af275f76343 100644 --- a/torchtitan/models/gpt_oss/model.py +++ b/torchtitan/models/gpt_oss/model.py @@ -17,6 +17,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.models.base.module import Module from torchtitan.models.common.attention import ( @@ -216,13 +217,15 @@ def _register_optimizer_hooks( class Config(Decoder.Config): dim: int = 2880 vocab_size: int = 201088 - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "cos_sin_rope", - "fp32_to_bf16_split", - ] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "cos_sin_rope", + "fp32_to_bf16_split", + ] + ) ) def get_nparams_and_flops( @@ -250,7 +253,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> GptOssModel: diff --git a/torchtitan/models/kimi_k2_7/model.py b/torchtitan/models/kimi_k2_7/model.py index 6d1dd05790b..a621676c1bb 100644 --- a/torchtitan/models/kimi_k2_7/model.py +++ b/torchtitan/models/kimi_k2_7/model.py @@ -20,6 +20,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -117,7 +118,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK25Model: diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index f1313a458a1..889b9276e2a 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -19,7 +19,7 @@ from torchtitan.config.function import Function from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import apply_local_compile, local_compile +from torchtitan.distributed.local_compile import local_compile, LocalCompileConfig from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -391,14 +391,16 @@ class Config(Decoder.Config): attention_residual: AttentionResidual.Config = field( default_factory=AttentionResidual.Config ) - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "gated_rmsnorm", - "fused_binary_activation", - "fp32_to_bf16_split", - "attention_residual", - ] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "gated_rmsnorm", + "fused_binary_activation", + "fp32_to_bf16_split", + "attention_residual", + ] + ) ) def get_nparams_and_flops( @@ -458,12 +460,12 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK3Model: # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) + local_compile_regions.apply() with parallelism_context.activate_spmd(): annotate_replicated_parameters(self, parallelism_context) self._parallelize(parallelism_context) diff --git a/torchtitan/models/llama3/model.py b/torchtitan/models/llama3/model.py index b0335209d90..54bd7302f80 100644 --- a/torchtitan/models/llama3/model.py +++ b/torchtitan/models/llama3/model.py @@ -12,6 +12,7 @@ import torch_remat as remat from torch import nn +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.attention import ( FlexAttentionMetadata, VarlenAttentionMetadata, @@ -83,8 +84,10 @@ class Llama3Model(Decoder): class Config(Decoder.Config): dim: int = 4096 vocab_size: int = 128256 - local_compile_regions: list[str] = field( - default_factory=lambda: ["loss", "fused_binary_activation"] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=["loss", "fused_binary_activation"] + ) ) def get_nparams_and_flops( diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py index 72938708b72..0402160c20a 100644 --- a/torchtitan/models/muse_glimmer/model.py +++ b/torchtitan/models/muse_glimmer/model.py @@ -18,6 +18,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.distributed.parallelism_context import ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -289,8 +290,10 @@ class MuseGlimmerModel(MultimodalModel): class Config(Decoder.Config): dim: int = 6656 vocab_size: int = 202048 - local_compile_regions: list[str] = field( - default_factory=lambda: ["loss", "fused_binary_activation"] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=["loss", "fused_binary_activation"] + ) ) # Narrows the base Decoder.Config.tok_embeddings (Embedding.Config) to the # bundled embedding+norm unit whose .embedding/.norm children carry the sharding. @@ -369,7 +372,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> MuseGlimmerModel: diff --git a/torchtitan/models/qwen3/model.py b/torchtitan/models/qwen3/model.py index 0ec1c454b65..bcabbf086f9 100644 --- a/torchtitan/models/qwen3/model.py +++ b/torchtitan/models/qwen3/model.py @@ -12,6 +12,7 @@ import torch.nn as nn import torch_remat as remat +from torchtitan.distributed.local_compile import LocalCompileConfig from torchtitan.models.common.attention import ( FlexAttentionMetadata, VarlenAttentionMetadata, @@ -108,13 +109,15 @@ def _register_optimizer_hooks( class Config(Decoder.Config): dim: int = 1024 vocab_size: int = 151936 - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "cos_sin_rope", - "fp32_to_bf16_split", - ] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "cos_sin_rope", + "fp32_to_bf16_split", + ] + ) ) def get_nparams_and_flops( diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index bd68938bfb6..368f27ef7b8 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -18,7 +18,7 @@ from torchtitan.config import TrainingConfig from torchtitan.config.parallelism import ParallelismConfig from torchtitan.distributed.activation_remat import ActivationRematConfig -from torchtitan.distributed.local_compile import local_compile +from torchtitan.distributed.local_compile import local_compile, LocalCompileConfig from torchtitan.distributed.parallelism_context import MeshAxisName, ParallelismContext from torchtitan.distributed.spmd_types import ( annotate_input_spmd_types, @@ -353,16 +353,18 @@ def _register_optimizer_hooks( @dataclass(kw_only=True, slots=True) class Config(Decoder.Config): vision_encoder: Qwen35VisionEncoder.Config | None = None - local_compile_regions: list[str] = field( - default_factory=lambda: [ - "loss", - "fused_binary_activation", - "gated_rmsnorm", - "offset_rmsnorm", - "partial_rope", - "shared_expert_gate", - "fp32_to_bf16_split", - ] + local_compile_regions: LocalCompileConfig = field( + default_factory=lambda: LocalCompileConfig( + regions=[ + "loss", + "fused_binary_activation", + "gated_rmsnorm", + "offset_rmsnorm", + "partial_rope", + "shared_expert_gate", + "fp32_to_bf16_split", + ] + ) ) def get_nparams_and_flops( @@ -436,7 +438,7 @@ def parallelize( parallelism_context: ParallelismContext, training: TrainingConfig, parallelism: ParallelismConfig, - local_compile_regions: list[str], + local_compile_regions: LocalCompileConfig, activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Qwen35Model: diff --git a/torchtitan_recipes/models/deepseek_v3.py b/torchtitan_recipes/models/deepseek_v3.py index 5a762cdec49..9e4e8b61636 100644 --- a/torchtitan_recipes/models/deepseek_v3.py +++ b/torchtitan_recipes/models/deepseek_v3.py @@ -268,5 +268,5 @@ def deepseek_v3_671b_nvfp4_ffn_mxfp8_attn( ) ], ) - config.model.local_compile_regions = ["loss"] + config.model.local_compile_regions.regions = ["loss"] return config diff --git a/torchtitan_recipes/rl/alphabet_sort.py b/torchtitan_recipes/rl/alphabet_sort.py index f7125220ca0..29141010e90 100644 --- a/torchtitan_recipes/rl/alphabet_sort.py +++ b/torchtitan_recipes/rl/alphabet_sort.py @@ -174,7 +174,7 @@ def rl_grpo_qwen3_0_6b_varlen(*, seq_len: int = 2048) -> Controller.Config: def rl_grpo_qwen3_0_6b_varlen_no_compile(*, seq_len: int = 2048) -> Controller.Config: config = rl_grpo_qwen3_0_6b_varlen(seq_len=seq_len) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return config @@ -299,7 +299,7 @@ def rl_grpo_qwen3_0_6b_flex_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] block_size = config.trainer.model.layers[0].attention.inner_attention.block_size config.async_loop.batcher = dataclasses.replace( config.async_loop.batcher, per_sample_pad_multiple=block_size @@ -671,7 +671,7 @@ def rl_grpo_qwen3_0_6b_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return Controller.Config( hf_assets_path="./assets/hf/Qwen3-0.6B", num_generators=3, @@ -856,7 +856,7 @@ def rl_grpo_qwen3_5_9b_varlen_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return config diff --git a/torchtitan_recipes/tests/__init__.py b/torchtitan_recipes/tests/__init__.py index 3792be993e9..fb45108df37 100644 --- a/torchtitan_recipes/tests/__init__.py +++ b/torchtitan_recipes/tests/__init__.py @@ -27,4 +27,4 @@ def _set_spmd_typechecking(config: Trainer.Config, *, typechecking: bool) -> Non if typechecking: # Local compile does not support global SPMD typechecking yet. # TODO: Remove this once the SPMD typechecking/Dynamo issue is fixed. - config.model.local_compile_regions = [] + config.model.local_compile_regions.regions = [] diff --git a/torchtitan_recipes/tests/graph_trainer/numerics.py b/torchtitan_recipes/tests/graph_trainer/numerics.py index 277344d9ce7..e8d85296cce 100644 --- a/torchtitan_recipes/tests/graph_trainer/numerics.py +++ b/torchtitan_recipes/tests/graph_trainer/numerics.py @@ -55,7 +55,7 @@ def _without_inplace_wgrad_accum(config): def llama3_eager_numerics(): config = llama3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile_regions.regions = [] config.parallelism.tensor_parallel_degree = 2 config.parallelism.data_parallel_shard_degree = 4 config.training.num_tokens_per_microbatch_per_dp_rank = 16384 @@ -73,7 +73,7 @@ def llama3_graph_numerics(): def llama3_eager_hsdp_gradient_accumulation_numerics(): config = llama3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile_regions.regions = [] config.parallelism.data_parallel_replicate_degree = 2 config.parallelism.data_parallel_shard_degree = 2 config.training.num_tokens_per_microbatch_per_dp_rank = 8192 @@ -103,7 +103,7 @@ def _llama3_hsdp_pp_numerics(config): def llama3_eager_hsdp_pp_numerics(): config = llama3_debugmodel(seq_len=2048) - config.model.local_compile_regions = [] + config.model.local_compile_regions.regions = [] return _without_inplace_wgrad_accum(_llama3_hsdp_pp_numerics(config)) @@ -214,7 +214,7 @@ def _deepseek_v3_dist_moe_pp_numerics( runtime_config.pp_activation_slot_policy = slot_policy for _, experts, _, _ in config.model.traverse(DistMoeRoutedExperts.Config): experts.inplace_wgrad_accum = False - config.model.local_compile_regions = [] + config.model.local_compile_regions.regions = [] config.parallelism.pipeline_parallel_degree = 2 config.parallelism.pipeline_parallel_schedule = "Interleaved1F1B" config.parallelism.num_pp_microbatches = 8 @@ -251,7 +251,7 @@ def deepseek_v3_dist_moe_graph_pp_stage_microbatch_numerics(): def qwen3_eager_numerics(): config = qwen3_debugmodel(seq_len=2048) # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile_regions.regions = [] config.parallelism.tensor_parallel_degree = 2 config.parallelism.data_parallel_shard_degree = 4 config.training.num_tokens_per_microbatch_per_dp_rank = 16384 @@ -278,7 +278,7 @@ def qwen3_moe_eager_numerics(): ) config.dataloader.max_num_documents = None # Match GraphTrainer, which captures model and loss in one eager FX graph. - config.model.local_compile_regions = [] + config.model.local_compile_regions.regions = [] config.training.disable_cuda_graphs = True config.training.num_tokens_per_microbatch_per_dp_rank = 16384 return _without_inplace_wgrad_accum(config) diff --git a/torchtitan_recipes/tests/models/deepseek_v3.py b/torchtitan_recipes/tests/models/deepseek_v3.py index 65225c382bc..bc779ebaeb1 100644 --- a/torchtitan_recipes/tests/models/deepseek_v3.py +++ b/torchtitan_recipes/tests/models/deepseek_v3.py @@ -481,7 +481,7 @@ def deepseek_v3_debugmodel_nvfp4_ffn_mxfp8_attn( ), ], ) - config.model.local_compile_regions = ["loss"] + config.model.local_compile_regions.regions = ["loss"] return config @@ -530,5 +530,5 @@ def deepseek_v3_16b_nvfp4_ffn_mxfp8_attn( ) ], ) - config.model.local_compile_regions = ["loss"] + config.model.local_compile_regions.regions = ["loss"] return config diff --git a/torchtitan_recipes/tests/rl/alphabet_sort.py b/torchtitan_recipes/tests/rl/alphabet_sort.py index ce1b21069ed..d86fbadb948 100644 --- a/torchtitan_recipes/tests/rl/alphabet_sort.py +++ b/torchtitan_recipes/tests/rl/alphabet_sort.py @@ -156,7 +156,7 @@ def rl_grpo_gpt_oss_debug_varlen_no_compile( ) -> Controller.Config: config = rl_grpo_gpt_oss_debug_varlen(seq_len=seq_len) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return config @@ -194,7 +194,7 @@ def rl_grpo_gpt_oss_debug_varlen_batch_invariant() -> Controller.Config: ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return Controller.Config( hf_assets_path="tests/assets/tokenizer", async_loop=AsyncLoopConfig( @@ -439,7 +439,7 @@ def rl_grpo_qwen3_moe_debug_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return Controller.Config( hf_assets_path="tests/assets/tokenizer", async_loop=AsyncLoopConfig( @@ -604,7 +604,7 @@ def rl_grpo_qwen3_5_debug_varlen_batch_invariant( ) # Local compile regions do not support batch-invariant mode. for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return config @@ -629,7 +629,7 @@ def _kimi_k3_rl_model_configs( converters=converters, ) for model_config in (trainer_model, generator_model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] return trainer_model, generator_model @@ -819,7 +819,7 @@ def rl_grpo_moe_debug_tp4_ep4() -> Controller.Config: generator_parallelism=config.generator.parallelism, ) for model_config in (config.trainer.model, config.generator.model): - model_config.local_compile_regions = [] + model_config.local_compile_regions.regions = [] # TODO: use FULL CUDA graphs once GPT-OSS CI enables a # CUDA-graph-compatible distributed MoE token dispatcher. config.generator.cuda_graph.mode = "NONE"