modify save list

Author: liangel-02

tests/integration_tests/features.py

@@ -350,12 +350,13 @@ def build_features_test_list() -> list[OverrideDefinitions]:
             [
                 [
                     "--parallelism.data_parallel_shard_degree=4",
-                    "--activation_checkpoint.mode='full'",
+                    "--activation_checkpoint.mode=selective",
+                    "--activation_checkpoint.selective_ac_option=op",
                     "--model.flavor=debugmodel_varlen_attn",
                 ]
             ],
-            "FSDP+VARLEN_ATTN",
-            "fsdp+varlen_attn",
+            "FSDP+VARLEN_ATTN + per op SAC",
+            "fsdp+varlen_attn+per_op_sac",
             ngpu=4,
             skip_rocm_test=True,
         ),

tests/unit_tests/test_activation_checkpoint.py

@@ -25,6 +25,7 @@
     # used to compute the scaling factor for quantization.
     torch.ops.aten.max.default,
     torch._higher_order_ops.flex_attention,
+    torch.ops.torch_attn._varlen_attn,
 }
 
 

torchtitan/experiments/simple_fsdp/llama3/parallelize.py

@@ -30,6 +30,7 @@
     # used to compute the scaling factor for quantization.
     torch.ops.aten.max.default,
     torch._higher_order_ops.flex_attention,
+    torch.ops.torch_attn._varlen_attn,
 }
 
 

torchtitan/models/llama3/infra/parallelize.py

@@ -41,6 +41,7 @@
     # used to compute the scaling factor for quantization.
     torch.ops.aten.max.default,
     torch._higher_order_ops.flex_attention,
+    torch.ops.torch_attn._varlen_attn,
 }