Separate SAC Wrapping of MoE and Attention Modules to Enable Flex Attention Compilation (#1683)

Flex Attention requires compilation via torch.compile to achieve optimal
performance. Therefore, torch.compile is always applied to Flex
Attention, regardless of the compile.enable flag. However, when
Selective Activation Checkpointing (SAC) is enabled, torch.compile may
be bypassed or invalidated under certain conditions:

1. If compile.enable is set to False, SAC will ignore any torch.compile
calls within the SAC region.
2. If compile.enable is True but the transformer block includes a
Mixture of Experts (MoE) module.

To address this limitation, this PR separates the SAC wrapping of
Attention from MoE and FeedForward modules. This separation ensures that
Flex Attention can be compiled successfully even when SAC is enabled.
Attention module is wrapped with full AC if compile.enable is False.

FIX (workaround) #1631

Author: fegin

tests/unit_tests/test_activation_checkpoint.py

@@ -15,7 +15,7 @@
 
 
 # for selective op activation checkpointing
-_save_list = {
+_op_sac_save_list = {
     torch.ops.aten.mm.default,
     torch.ops.aten._scaled_dot_product_efficient_attention.default,
     torch.ops.aten._scaled_dot_product_flash_attention.default,
@@ -85,7 +85,7 @@ def get_bw_flops(model_fn):
             ac_config_no_force,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         flops_selective_ac = get_bw_flops(model_selective_ac)
 
@@ -103,7 +103,7 @@ def get_bw_flops(model_fn):
             ac_config_with_force_first,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         flops_with_force_first = get_bw_flops(model_with_force_first)
 
@@ -120,7 +120,7 @@ def get_bw_flops(model_fn):
             ac_config_with_force_last,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         flops_with_force_last = get_bw_flops(model_with_force_last)
 
@@ -135,7 +135,7 @@ def get_bw_flops(model_fn):
             ac_config_full_ac,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         flops_full_ac = get_bw_flops(model_with_full_ac)
 
@@ -178,7 +178,7 @@ def get_act_mem(model_fn):
             ac_config_no_force,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         mem_selective_ac = get_act_mem(model_selective_ac)
 
@@ -195,7 +195,7 @@ def get_act_mem(model_fn):
             ac_config_with_force_first,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         mem_with_force_first = get_act_mem(model_with_force_first)
 
@@ -211,7 +211,7 @@ def get_act_mem(model_fn):
             ac_config_with_force_last,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         mem_with_force_last = get_act_mem(model_with_force_last)
 
@@ -225,7 +225,7 @@ def get_act_mem(model_fn):
             ac_config_full_ac,
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         mem_full_ac = get_act_mem(model_with_full_ac)
 
@@ -252,7 +252,7 @@ def test_correctness(self):
             ),
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
         model_force_first = ToyModule()
         model_force_first.load_state_dict(model_no_ac.state_dict())
@@ -265,7 +265,7 @@ def test_correctness(self):
             ),
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
 
         model_force_last = ToyModule()
@@ -279,7 +279,7 @@ def test_correctness(self):
             ),
             model_compile_enabled=False,
             use_flex_attn=False,
-            save_list=_save_list,
+            op_sac_save_list=_op_sac_save_list,
         )
 
         def run_fwd_bwd(model, batch):