modify save list for varlen attn (#2082)

adding varlen attention ops to ac save list

**testing**

used DebugMode() to print out op list. verified that forward is not
being recomputed in the backward step.

```
[rank0]:forward ops
[rank0]:varlen_attn in forward: True
...
[rank0]:varlen_attn recomputed in backward: False
[rank0]:saved correctly
```

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

@@ -28,6 +28,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

@@ -33,6 +33,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

@@ -44,6 +44,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.default,
 }
 
 

torchtitan/models/qwen3/infra/parallelize.py

@@ -46,7 +46,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,
+    torch.ops.torch_attn._varlen_attn.default,
 }