config compile backend (#1768)

Stacked PRs:
 * __->__#1768


--- --- ---

### config compile backend

Author: IvanKobzarev

torchtitan/components/loss.py

@@ -27,7 +27,7 @@ def build_cross_entropy_loss(job_config: JobConfig):
     loss_fn = cross_entropy_loss
     if job_config.compile.enable and "loss" in job_config.compile.components:
         logger.info("Compiling the loss function with torch.compile")
-        loss_fn = torch.compile(loss_fn)
+        loss_fn = torch.compile(loss_fn, backend=job_config.compile.backend)
     return loss_fn
 
 

torchtitan/config/job_config.py

@@ -566,6 +566,7 @@ class Compile:
         default_factory=lambda: ["model", "loss"]
     )
     """Which components to compile"""
+    backend: str = "inductor"
 
 
 @dataclass

torchtitan/experiments/flux/loss.py

@@ -23,5 +23,5 @@ def build_mse_loss(job_config: JobConfig):
     loss_fn = mse_loss
     if job_config.compile.enable and "loss" in job_config.compile.components:
         logger.info("Compiling the loss function with torch.compile")
-        loss_fn = torch.compile(loss_fn)
+        loss_fn = torch.compile(loss_fn, backend=job_config.compile.backend)
     return loss_fn

torchtitan/experiments/llama4/infra/parallelize.py

@@ -19,6 +19,7 @@
     SequenceParallel,
 )
 from torchtitan.config import JobConfig, TORCH_DTYPE_MAP
+from torchtitan.config.job_config import Compile as CompileConfig
 from torchtitan.distributed import NoParallel, ParallelDims
 from torchtitan.distributed.activation_checkpoint import apply_ac
 from torchtitan.distributed.expert_parallel import (
@@ -123,7 +124,7 @@ def parallelize_llama(
 
     # turn on per-TransformerBlock compile after AC wrapping and before FSDP
     if model_compile_enabled:
-        apply_compile(model)
+        apply_compile(model, job_config.compile)
 
     dp_mesh: DeviceMesh | None = None
     if parallel_dims.fsdp_enabled or parallel_dims.ep_enabled:
@@ -502,7 +503,7 @@ def apply_moe_ep_tp(
         )
 
 
-def apply_compile(model: nn.Module):
+def apply_compile(model: nn.Module, compile_config: CompileConfig):
     """
     Apply torch.compile to each TransformerBlock, which makes compilation efficient due to
     repeated structure. Alternatively one can compile the whole model (after applying DP).
@@ -515,7 +516,11 @@ def apply_compile(model: nn.Module):
         fullgraph = True
         if transformer_block.moe_enabled:
             fullgraph = False
-        transformer_block = torch.compile(transformer_block, fullgraph=fullgraph)
+        transformer_block = torch.compile(
+            transformer_block,
+            backend=compile_config.backend,
+            fullgraph=fullgraph,
+        )
         model.layers.register_module(layer_id, transformer_block)
 
     logger.info("Compiling each TransformerBlock with torch.compile")

torchtitan/experiments/qwen3/infra/parallelize.py

@@ -118,7 +118,7 @@ def parallelize_qwen3(
 
     # turn on per-TransformerBlock compile after AC wrapping and before FSDP
     if model_compile_enabled:
-        apply_compile(model)
+        apply_compile(model, job_config.compile)
 
     if parallel_dims.fsdp_enabled:
         # apply FSDP or HSDP, potentially with Context Parallel

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

@@ -153,6 +153,6 @@ def parallelize_deepseekv3(
     if job_config.compile.enable:
         torch._inductor.config.reorder_for_peak_memory = False
         torch._dynamo.config.capture_scalar_outputs = True
-        model = torch.compile(model, fullgraph=True)
+        model = torch.compile(model, backend=job_config.compile.backend, fullgraph=True)
 
     return model

torchtitan/experiments/simple_fsdp/llama3/parallelize.py

@@ -122,6 +122,6 @@ def parallelize_llama(
 
     if job_config.compile.enable and "model" in job_config.compile.components:
         torch._inductor.config.reorder_for_peak_memory = False
-        model = torch.compile(model, fullgraph=True)
+        model = torch.compile(model, backend=job_config.compile.backend, fullgraph=True)
 
     return model

torchtitan/experiments/vlm/infra/parallelize.py

@@ -70,8 +70,8 @@ def parallelize_vlm(
 
     # turn on per-TransformerBlock compile after AC wrapping and before FSDP
     if job_config.compile.enable:
-        apply_compile(model)
-        apply_compile(model.encoder)
+        apply_compile(model, job_config.compile)
+        apply_compile(model.encoder, job_config.compile)
 
     if parallel_dims.fsdp_enabled:
         # apply FSDP or HSDP, potentially with Context Parallel

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -116,7 +116,7 @@ def parallelize_deepseekv3(
         )
 
     if model_compile_enabled:
-        apply_compile(model)
+        apply_compile(model, job_config.compile)
 
     dp_mesh: DeviceMesh | None = None
     if parallel_dims.fsdp_enabled or parallel_dims.ep_enabled:

torchtitan/models/llama3/infra/parallelize.py

@@ -23,6 +23,7 @@
 )
 
 from torchtitan.config import JobConfig, TORCH_DTYPE_MAP
+from torchtitan.config.job_config import Compile as CompileConfig
 from torchtitan.distributed import ParallelDims
 from torchtitan.distributed.activation_checkpoint import apply_ac
 from torchtitan.distributed.tensor_parallel import maybe_enable_async_tp
@@ -105,7 +106,7 @@ def parallelize_llama(
 
     # turn on per-TransformerBlock compile after AC wrapping and before FSDP
     if model_compile_enabled:
-        apply_compile(model)
+        apply_compile(model, job_config.compile)
 
     if parallel_dims.fsdp_enabled:
         # apply FSDP or HSDP, potentially with Context Parallel
@@ -234,13 +235,15 @@ def apply_tp(
     )
 
 
-def apply_compile(model: nn.Module):
+def apply_compile(model: nn.Module, compile_config: CompileConfig):
     """
     Apply torch.compile to each TransformerBlock, which makes compilation efficient due to
     repeated structure. Alternatively one can compile the whole model (after applying DP).
     """
     for layer_id, transformer_block in model.layers.named_children():
-        transformer_block = torch.compile(transformer_block, fullgraph=True)
+        transformer_block = torch.compile(
+            transformer_block, backend=compile_config.backend, fullgraph=True
+        )
         model.layers.register_module(layer_id, transformer_block)
 
     logger.info("Compiling each TransformerBlock with torch.compile")