fix varlen op saving for simple fsdp (#2396)

liangel-02 · web-flow · commit 238fa9953e4a · 2026-02-23T11:01:46.000-05:00
diff --git a/tests/unit_tests/test_activation_checkpoint.py b/tests/unit_tests/test_activation_checkpoint.py
@@ -28,7 +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,
+    torch.ops.torch_attn._varlen_attn.default,
 }
 
 
diff --git a/torchtitan/experiments/simple_fsdp/llama3/parallelize.py b/torchtitan/experiments/simple_fsdp/llama3/parallelize.py
@@ -33,7 +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,
+    torch.ops.torch_attn._varlen_attn.default,
     torch._higher_order_ops.inductor_compiled_code,
 }
 

Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@`
`28`	`28`	`# used to compute the scaling factor for quantization.`
`29`	`29`	`torch.ops.aten.max.default,`
`30`	`30`	`torch._higher_order_ops.flex_attention,`
`31`		`- torch.ops.torch_attn._varlen_attn,`
	`31`	`+ torch.ops.torch_attn._varlen_attn.default,`
`32`	`32`	`}`
`33`	`33`
`34`	`34`
Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@`
`33`	`33`	`# used to compute the scaling factor for quantization.`
`34`	`34`	`torch.ops.aten.max.default,`
`35`	`35`	`torch._higher_order_ops.flex_attention,`
`36`		`- torch.ops.torch_attn._varlen_attn,`
	`36`	`+ torch.ops.torch_attn._varlen_attn.default,`
`37`	`37`	`torch._higher_order_ops.inductor_compiled_code,`
`38`	`38`	`}`
`39`	`39`