diff --git a/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py b/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py index 5d83ef4d73e..4bae4eb01ba 100644 --- a/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py +++ b/scripts/checkpoint_conversion/numerical_tests_qwen3_5_shard.py @@ -99,6 +99,7 @@ def run_worker(args): training=training, parallelism=parallelism, local_compile_regions=config.local_compile_regions, + local_compile_backend=config.local_compile_backend, activation_remat=DefaultActivationRematConfig(), dump_folder="/tmp", ) diff --git a/tests/rl/unit_tests/gpu/test_bitwise_parity.py b/tests/rl/unit_tests/gpu/test_bitwise_parity.py index 797845b57c2..92093c53649 100644 --- a/tests/rl/unit_tests/gpu/test_bitwise_parity.py +++ b/tests/rl/unit_tests/gpu/test_bitwise_parity.py @@ -143,6 +143,7 @@ def build_trainer_model( training=trainer_config.training, parallelism=parallelism, local_compile_regions=model_config.local_compile_regions, + local_compile_backend=model_config.local_compile_backend, activation_remat=trainer_config.activation_remat, dump_folder=config.dump_folder, ) diff --git a/tests/unit_tests/cpu/config/test_compile_config.py b/tests/unit_tests/cpu/config/test_compile_config.py index 4cbcf0ce629..9a1e9cdbf64 100644 --- a/tests/unit_tests/cpu/config/test_compile_config.py +++ b/tests/unit_tests/cpu/config/test_compile_config.py @@ -8,7 +8,11 @@ import torch from torchtitan.components.loss import compute_logprobs, cross_entropy_loss, mse_loss -from torchtitan.distributed.local_compile import apply_local_compile, local_compile +from torchtitan.distributed.local_compile import ( + apply_local_compile, + compiled, + local_compile, +) from torchtitan.models.common.moe import _microbatch_load_balance_local_stats @@ -24,7 +28,7 @@ def test_apply_local_compile_rejects_unknown_name() -> None: apply_local_compile(["foo"]) -def test_local_compile_binds_once_from_existing_config(monkeypatch) -> None: +def test_local_compile_uses_default_backend(monkeypatch) -> None: compiled_calls = [] @local_compile("test_local_compile", batch_invariant=False) @@ -43,19 +47,17 @@ def compiled(value: int) -> tuple[str, int]: apply_local_compile(["test_local_compile"]) assert fn(3) == ("compiled", 3) - assert len(compiled_calls) == 1 - assert "backend" not in compiled_calls[0][1] + assert compiled_calls[0][1]["backend"] == "inductor" assert compiled_calls[0][1]["fullgraph"] is True -def test_local_compile_forwards_compile_kwargs(monkeypatch) -> None: +def test_local_compile_forwards_backend_and_compile_kwargs(monkeypatch) -> None: compiled_calls = [] options = {"max_autotune": True} @local_compile( "test_compile_kwargs", batch_invariant=False, - backend="eager", dynamic=True, fullgraph=True, options=options, @@ -68,11 +70,11 @@ def fake_compile(reference, **kwargs): return reference monkeypatch.setattr(torch, "compile", fake_compile) - apply_local_compile(["test_compile_kwargs"]) + apply_local_compile(["test_compile_kwargs"], backend="aot_eager") assert fn(3) == 3 assert compiled_calls[0][1] == { - "backend": "eager", + "backend": "aot_eager", "fullgraph": True, "dynamic": True, "options": options, @@ -88,6 +90,91 @@ def test_local_compile_rejects_fullgraph_false() -> None: ) +def test_local_compile_rejects_per_region_backend() -> None: + with pytest.raises(ValueError, match="apply_local_compile"): + local_compile( + "test_backend", + batch_invariant=False, + backend="aot_eager", + ) + + +def test_compiled_helper_binds_lazily_and_rebinds_backend(monkeypatch) -> None: + compiled_calls = [] + inductor_options = {"max_autotune": True} + + def reference(value: int) -> tuple[str, int]: + return "eager", value + + def fake_compile(reference, **kwargs): + compiled_calls.append((reference, kwargs)) + backend = kwargs["backend"] + + def compiled_fn(value: int) -> tuple[str, int]: + return backend, value + + return compiled_fn + + monkeypatch.setattr(torch, "compile", fake_compile) + helper = compiled(reference, inductor_options=inductor_options) + + assert compiled_calls == [] + assert helper(3) == ("inductor", 3) + assert compiled_calls == [ + ( + reference, + {"backend": "inductor", "options": inductor_options}, + ) + ] + + inductor_options["max_autotune"] = False + apply_local_compile([]) + + assert compiled_calls[-1] == ( + reference, + {"backend": "inductor", "options": {"max_autotune": False}}, + ) + + apply_local_compile([], backend="custom_device_backend") + + assert helper(4) == ("custom_device_backend", 4) + assert compiled_calls[-1] == ( + reference, + {"backend": "custom_device_backend"}, + ) + + +def test_torchtitan_helpers_use_configured_backend(monkeypatch) -> None: + compiled_calls = [] + + def fake_compile(reference, **kwargs): + compiled_calls.append((reference, kwargs)) + return reference + + monkeypatch.setattr(torch, "compile", fake_compile) + + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + from torchtitan.models.common import vision_encoder # noqa: F401 + + from torchtitan.models.common.attention import attention + + assert compiled_calls == [] + + apply_local_compile([], backend="custom_device_backend") + + helper_calls = [ + call + for call in compiled_calls + if call[0] in (flex_attention, create_block_mask) + ] + assert helper_calls == [ + (flex_attention, {"backend": "custom_device_backend"}), + (create_block_mask, {"backend": "custom_device_backend"}), + (create_block_mask, {"backend": "custom_device_backend"}), + ] + assert attention.FlexInnerAttention._compiled_flex_attn is not flex_attention + + def test_local_compile_rejects_non_batch_invariant_region(monkeypatch) -> None: @local_compile("test_non_batch_invariant", batch_invariant=False) def fn(value: int) -> int: @@ -114,7 +201,7 @@ def fake_compile(reference, **kwargs): monkeypatch.setattr(torch, "compile", fake_compile) apply_local_compile(["loss"]) - assert compiled_names == [ + assert compiled_names[:3] == [ cross_entropy_loss.__name__, mse_loss.__name__, compute_logprobs.__name__, @@ -131,9 +218,7 @@ def fake_compile(reference, **kwargs): monkeypatch.setattr(torch, "compile", fake_compile) apply_local_compile(["moe_aux_loss"]) - assert compiled == [ - ( - _microbatch_load_balance_local_stats.__wrapped__, - {"fullgraph": True}, - ) - ] + assert compiled[0] == ( + _microbatch_load_balance_local_stats.__wrapped__, + {"backend": "inductor", "fullgraph": True}, + ) diff --git a/tests/unit_tests/cpu/test_trainer.py b/tests/unit_tests/cpu/test_trainer.py index c79f1b0798f..61c9710fb08 100644 --- a/tests/unit_tests/cpu/test_trainer.py +++ b/tests/unit_tests/cpu/test_trainer.py @@ -756,7 +756,10 @@ def test_graph_training_engine_rejects_local_compile_regions() -> None: ): GraphTrainingEngine( config, - model_config=SimpleNamespace(local_compile_regions=["loss"]), + model_config=SimpleNamespace( + local_compile_regions=["loss"], + local_compile_backend="inductor", + ), max_num_documents=None, output_dir="", ) @@ -764,6 +767,20 @@ def test_graph_training_engine_rejects_local_compile_regions() -> None: init.assert_not_called() +def test_parallelize_compile_kwargs_include_backend() -> None: + engine = SimpleNamespace( + model_config=SimpleNamespace( + local_compile_regions=["loss"], + local_compile_backend="custom_device_backend", + ) + ) + + assert TrainingEngine._parallelize_compile_kwargs(engine) == { + "local_compile_regions": ["loss"], + "local_compile_backend": "custom_device_backend", + } + + def test_optim_update_clips_before_parameter_update() -> None: events = [] optimizers = MagicMock() diff --git a/torchtitan/distributed/local_compile.py b/torchtitan/distributed/local_compile.py index 6774c21ebf6..20d2978ba30 100644 --- a/torchtitan/distributed/local_compile.py +++ b/torchtitan/distributed/local_compile.py @@ -15,10 +15,13 @@ from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode -def apply_local_compile(regions: list[str]) -> None: +_BACKEND: str = "inductor" + + +def apply_local_compile(regions: list[str], *, backend: str = "inductor") -> None: """Bind registered functions to eager or compiled implementations. - Process-wide: a later call replaces this choice for every model in the process. + Process-wide: a later call replaces these choices for every model in the process. """ unknown = [name for name in regions if name not in _LOCAL_COMPILE_CALLBACKS] if unknown: @@ -27,12 +30,48 @@ def apply_local_compile(regions: list[str]) -> None: f"registered values are {sorted(_LOCAL_COMPILE_CALLBACKS)}" ) + global _BACKEND + _BACKEND = backend for callbacks in _LOCAL_COMPILE_CALLBACKS.values(): for bind_local_compile_fn in callbacks: bind_local_compile_fn(regions) + for bind_compiled_helper in _COMPILED_HELPERS: + bind_compiled_helper() _LOCAL_COMPILE_CALLBACKS: dict[str, list[Callable[[list[str]], None]]] = {} +_COMPILED_HELPERS: list[Callable[[], None]] = [] + + +def compiled( + reference: Callable[..., Any], + *, + inductor_options: dict[str, Any] | None = None, +) -> Callable[..., Any]: + """Return ``reference`` compiled with the configured backend. + + The function binds on first use and rebinds whenever ``apply_local_compile`` + changes the process-wide backend. ``inductor_options`` are passed only when + that backend is inductor. + """ + fn: Callable[..., Any] | None = None + + def bind_compiled_helper() -> None: + nonlocal fn + compile_kwargs: dict[str, Any] = {"backend": _BACKEND} + if _BACKEND == "inductor" and inductor_options is not None: + compile_kwargs["options"] = inductor_options + fn = torch.compile(reference, **compile_kwargs) + + @functools.wraps(reference) + def wrapped(*args: Any, **kwargs: Any) -> Any: + if fn is None: + bind_compiled_helper() + assert fn is not None + return fn(*args, **kwargs) + + _COMPILED_HELPERS.append(bind_compiled_helper) + return wrapped def local_compile( @@ -52,6 +91,8 @@ def local_compile( """ if compile_kwargs.pop("fullgraph", True) is not True: raise ValueError("local_compile requires fullgraph=True.") + if "backend" in compile_kwargs: + raise ValueError("local_compile backend is configured by apply_local_compile.") def decorate(reference: Callable[..., Any]) -> Callable[..., Any]: fn = reference @@ -66,7 +107,12 @@ def bind_local_compile(regions: list[str]) -> None: "batch-invariant mode; remove it from local_compile_regions." ) if enabled: - fn = torch.compile(reference, fullgraph=True, **compile_kwargs) + fn = torch.compile( + reference, + fullgraph=True, + backend=_BACKEND, + **compile_kwargs, + ) else: fn = reference @@ -80,4 +126,4 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: return decorate -__all__ = ["apply_local_compile", "local_compile"] +__all__ = ["apply_local_compile", "compiled", "local_compile"] diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py index 67959106c65..3c7e37f4813 100644 --- a/torchtitan/distributed/pipeline_parallel.py +++ b/torchtitan/distributed/pipeline_parallel.py @@ -79,6 +79,7 @@ def pipeline_llm( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, @@ -145,6 +146,7 @@ def pipeline_llm( training=training, parallelism=parallelism, local_compile_regions=local_compile_regions, + local_compile_backend=local_compile_backend, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py index 77c4dea011e..0edd1130a9d 100644 --- a/torchtitan/distributed/utils.py +++ b/torchtitan/distributed/utils.py @@ -149,11 +149,11 @@ def set_determinism( # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - from torch.nn.attention.flex_attention import flex_attention - from torchtitan.models.common.attention import FlexInnerAttention if torch.version.hip is not None: + from torch.nn.attention.flex_attention import flex_attention + # Compiled ROCm flex attention is not deterministic. # Falling back to eager (non-compiled) flex_attention for determinism on ROCm. logger.info( @@ -167,10 +167,6 @@ def set_determinism( # no options) so the regional-inductor scoop configs are preserved. FlexInnerAttention.inductor_configs["max_autotune"] = False FlexInnerAttention.inductor_configs["coordinate_descent_tuning"] = False - # pyrefly: ignore [no-matching-overload] - FlexInnerAttention._compiled_flex_attn = torch.compile( - flex_attention, options=FlexInnerAttention.inductor_configs - ) if debug_config.detect_anomaly: logger.warning( diff --git a/torchtitan/experiments/transformers_modeling_backend/parallelize.py b/torchtitan/experiments/transformers_modeling_backend/parallelize.py index 46be051bc05..9a187f41838 100644 --- a/torchtitan/experiments/transformers_modeling_backend/parallelize.py +++ b/torchtitan/experiments/transformers_modeling_backend/parallelize.py @@ -88,6 +88,7 @@ def parallelize_hf_transformers( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, **kwargs: Any, @@ -102,7 +103,10 @@ def parallelize_hf_transformers( 5. Apply AC and FSDP """ # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) + apply_local_compile( + local_compile_regions, + backend=local_compile_backend, + ) # Flex attention supports FSDP, TP, CP, and PP (in any combination). Under CP # the flex kernel's local SPMD boundary redistributes # k/v from seq-sharded to CP-Replicate (all-gather); see _attach_flex_kernel diff --git a/torchtitan/experiments/transformers_modeling_backend/pipeline.py b/torchtitan/experiments/transformers_modeling_backend/pipeline.py index f439c413c35..855c7bbc783 100644 --- a/torchtitan/experiments/transformers_modeling_backend/pipeline.py +++ b/torchtitan/experiments/transformers_modeling_backend/pipeline.py @@ -297,6 +297,7 @@ def pipeline_hf_transformers( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, device: torch.device, @@ -391,6 +392,7 @@ def pipeline_hf_transformers( training=training, parallelism=parallelism, local_compile_regions=local_compile_regions, + local_compile_backend=local_compile_backend, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/base/model.py b/torchtitan/models/base/model.py index 52760ed6a96..ccc1ac44c2f 100644 --- a/torchtitan/models/base/model.py +++ b/torchtitan/models/base/model.py @@ -89,9 +89,13 @@ class Config(Module.Config): 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. + Always-compiled helpers use ``local_compile_backend`` but are not controlled + by this list. """ + local_compile_backend: str = "inductor" + """Backend for TorchTitan-owned local compile regions and helpers.""" + @abstractmethod def get_nparams_and_flops(self, model: Module, seq_len: int) -> tuple[int, int]: pass @@ -133,12 +137,16 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", 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) + apply_local_compile( + local_compile_regions, + backend=local_compile_backend, + ) with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) if activation_remat is not None: 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/multimodal.py b/torchtitan/models/common/multimodal.py index 26ce302f79a..5289ce328f9 100644 --- a/torchtitan/models/common/multimodal.py +++ b/torchtitan/models/common/multimodal.py @@ -38,11 +38,15 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Self: # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) + apply_local_compile( + local_compile_regions, + backend=local_compile_backend, + ) with parallelism_context.activate_spmd(): self._parallelize(parallelism_context) encoders = [ 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/models/flux/model.py b/torchtitan/models/flux/model.py index 39c6480c931..473826c33d3 100644 --- a/torchtitan/models/flux/model.py +++ b/torchtitan/models/flux/model.py @@ -236,12 +236,16 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", 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) + apply_local_compile( + local_compile_regions, + backend=local_compile_backend, + ) 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..80396866b78 100644 --- a/torchtitan/models/gpt_oss/model.py +++ b/torchtitan/models/gpt_oss/model.py @@ -251,6 +251,7 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> GptOssModel: @@ -268,6 +269,7 @@ def parallelize( training=training, parallelism=parallelism, local_compile_regions=local_compile_regions, + local_compile_backend=local_compile_backend, 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..2a786b52c41 100644 --- a/torchtitan/models/kimi_k2_7/model.py +++ b/torchtitan/models/kimi_k2_7/model.py @@ -118,6 +118,7 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK25Model: @@ -132,6 +133,7 @@ def parallelize( training=training, parallelism=parallelism, local_compile_regions=local_compile_regions, + local_compile_backend=local_compile_backend, 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..0de8952f265 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -459,11 +459,15 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> KimiK3Model: # Bind local implementations early; torch.compile traces on first use. - apply_local_compile(local_compile_regions) + apply_local_compile( + local_compile_regions, + backend=local_compile_backend, + ) with parallelism_context.activate_spmd(): annotate_replicated_parameters(self, parallelism_context) self._parallelize(parallelism_context) diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py index 72938708b72..f54fde2b57a 100644 --- a/torchtitan/models/muse_glimmer/model.py +++ b/torchtitan/models/muse_glimmer/model.py @@ -370,6 +370,7 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> MuseGlimmerModel: @@ -384,6 +385,7 @@ def parallelize( training=training, parallelism=parallelism, local_compile_regions=local_compile_regions, + local_compile_backend=local_compile_backend, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py index bd68938bfb6..bcfd0a2bc9d 100644 --- a/torchtitan/models/qwen3_5/model.py +++ b/torchtitan/models/qwen3_5/model.py @@ -437,6 +437,7 @@ def parallelize( training: TrainingConfig, parallelism: ParallelismConfig, local_compile_regions: list[str], + local_compile_backend: str = "inductor", activation_remat: ActivationRematConfig | None, dump_folder: str, ) -> Qwen35Model: @@ -445,6 +446,7 @@ def parallelize( training=training, parallelism=parallelism, local_compile_regions=local_compile_regions, + local_compile_backend=local_compile_backend, activation_remat=activation_remat, dump_folder=dump_folder, ) diff --git a/torchtitan/rl/generator/vllm/model_wrapper.py b/torchtitan/rl/generator/vllm/model_wrapper.py index 8438e0745dc..2ae43f1702f 100644 --- a/torchtitan/rl/generator/vllm/model_wrapper.py +++ b/torchtitan/rl/generator/vllm/model_wrapper.py @@ -368,6 +368,7 @@ def __init__( training=TrainingConfig(), parallelism=training_parallelism, local_compile_regions=self.config.local_compile_regions, + local_compile_backend=self.config.local_compile_backend, activation_remat=None, dump_folder="", ) diff --git a/torchtitan/training_engine.py b/torchtitan/training_engine.py index 721c8a2e6be..a86160135a1 100644 --- a/torchtitan/training_engine.py +++ b/torchtitan/training_engine.py @@ -415,7 +415,10 @@ def _parallelize_compile_kwargs(self) -> dict[str, Any]: # 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 { + "local_compile_regions": self.model_config.local_compile_regions, + "local_compile_backend": self.model_config.local_compile_backend, + } def _initialize_optim(self) -> None: """Construct the parameter update and its state."""