Changes to enable compilation

Author: Lucaskabela

torchtitan/experiments/rl/unified/actors/trainer.py

@@ -38,6 +38,16 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass(kw_only=True, slots=True)
+class TrainerCompileConfig:
+    """Compilation settings for the PolicyTrainer."""
+
+    enable: bool = False
+    """Enable per-layer torch.compile on the training model."""
+    backend: str = "eager"
+    """torch.compile backend (e.g. 'eager', 'aot_eager', 'inductor')."""
+
+
 class PolicyTrainer(Actor, Configurable):
     """
     Updates policy based on collected Episode using TorchTitan components.
@@ -64,6 +74,7 @@ class Config(Configurable.Config):
         parallelism: ParallelismConfig = field(default_factory=ParallelismConfig)
         comm: CommConfig = field(default_factory=CommConfig)
         """Communication configuration for distributed initialization."""
+        compile: TrainerCompileConfig = field(default_factory=TrainerCompileConfig)
 
     def __init__(
         self,
@@ -109,6 +120,8 @@ def __init__(
             model_spec, config, device_type, batch_invariant_mode, hf_assets_path
         )
         model.train()
+        if config.compile.enable:
+            model = self._compile_model(model, config.compile.backend)
         self.model = model
         self.model_parts = [model]
 
@@ -223,6 +236,20 @@ def _build_model(
 
         return model
 
+    def _compile_model(self, model: torch.nn.Module, backend: str) -> torch.nn.Module:
+        """Compile each transformer layer with torch.compile.
+
+        Args:
+            model: The model whose layers will be compiled.
+            backend: torch.compile backend (e.g. 'eager', 'aot_eager', 'inductor').
+        """
+        for layer_id in model.layers:
+            model.layers[layer_id].compile(backend=backend, fullgraph=True)
+        logger.info(
+            f"Compiled {len(model.layers)} transformer layers with {backend} backend"
+        )
+        return model
+
     @endpoint
     async def get_weights(self) -> dict:
         """Get model weights for generator.

torchtitan/experiments/rl/unified/config_registry.py

@@ -19,7 +19,10 @@
     SamplingConfig,
     VLLMGenerator,
 )
-from torchtitan.experiments.rl.unified.actors.trainer import PolicyTrainer
+from torchtitan.experiments.rl.unified.actors.trainer import (
+    PolicyTrainer,
+    TrainerCompileConfig,
+)
 from torchtitan.experiments.rl.unified.simple_grpo_sum_digits import RLTrainer
 from torchtitan.models.qwen3 import model_registry
 
@@ -41,6 +44,7 @@ def rl_grpo_qwen3_0_6b() -> RLTrainer.Config:
             parallelism=ParallelismConfig(
                 tensor_parallel_degree=2,
             ),
+            compile=TrainerCompileConfig(enable=True, backend="aot_eager"),
         ),
         generator=VLLMGenerator.Config(
             model_dtype="bfloat16",
@@ -80,6 +84,7 @@ def rl_grpo_qwen3_1_7b() -> RLTrainer.Config:
             parallelism=ParallelismConfig(
                 tensor_parallel_degree=2,
             ),
+            compile=TrainerCompileConfig(enable=True, backend="aot_eager"),
         ),
         generator=VLLMGenerator.Config(
             model_dtype="bfloat16",
@@ -119,6 +124,7 @@ def rl_grpo_qwen3_debug() -> RLTrainer.Config:
                 tensor_parallel_degree=1,
                 data_parallel_replicate_degree=1,
             ),
+            compile=TrainerCompileConfig(enable=True, backend="aot_eager"),
         ),
         generator=VLLMGenerator.Config(
             compile=GeneratorCompileConfig(