some compile-related improvements

ghstack-source-id: 3af1989
Pull Request resolved: #443

Author: tianyu-l

test_runner.py

@@ -153,6 +153,15 @@ def build_test_list():
             "1D compile",
             "1d_compile",
         ),
+        OverrideDefinitions(
+            [
+                [
+                    "--training.compile --model.norm_type=rmsnorm --activation_checkpoint.selective_ac_option=op",
+                ],
+            ],
+            "1D compile with selective op AC",
+            "1d_compile_sac_op",
+        ),
         OverrideDefinitions(
             [
                 [

torchtitan/parallelisms/parallelize_llama.py

@@ -432,22 +432,14 @@ def apply_compile(model, job_config: JobConfig):
             "fused_rmsnorm not yet compatible with torch.compile. Please use layernorm or rmsnorm."
         )
 
+    # NOTE(anijain): enable the following flag to accelarate compilation
+    torch._dynamo.config.inline_inbuilt_nn_modules = True
+
     for layer_id, transformer_block in model.layers.named_children():
         # turn on per-transformer block compile after AC wrapping and before FSDP
-        # TODO: dynamic shape have some issues so we turn it off for now.
-        # TODO: inline inbuilt nn modules does not work yet, enable it to accelarate
-        # compile time.
-        # torch._dynamo.config.inline_inbuilt_nn_modules = True
-        transformer_block = torch.compile(transformer_block, dynamic=False)
+        transformer_block = torch.compile(transformer_block, fullgraph=True)
         model.layers.register_module(layer_id, transformer_block)
 
-    ac_config = job_config.activation_checkpoint
-    if ac_config.mode == "selective" and ac_config.selective_ac_option == "op":
-        # some temp flags for torch.compile enablement + SAC
-        torch._dynamo.config._experimental_support_context_fn_in_torch_utils_checkpoint = (
-            True
-        )
-
     logger.info("Compiled each TransformerBlock with torch.compile")
     return model