text: Title: Auto-selective activation checkpointing is not optimal for speed (issue with min_cut_rematerialization_partition)
  Body: ### 🐛 Describe the bug
  
  I try the new api described in [pytorch blog: selective activation checkpointing](https://pytorch.org/blog/activation-checkpointing-techniques/#compile-only-memory-budget-api-new)
  .
  Then I find that selective activation checkpointing is not optimal for speed. 
  
  A minimal reproducer:
  ```python
  import os
  os.environ["CUDA_VISIBLE_DEVICES"] = "0"
  os.environ["TORCH_COMPILE_DEBUG"] = "1"
  
  import torch
  from torch import nn 
  import torch.nn.functional as F
  
  import torch._functorch.config
  torch._functorch.config.activation_memory_budget = 0.99
  
  class Test1(nn.Module):
  
      def __init__(self):
          super().__init__()
          self.layer0 = nn.Linear(100, 100, bias=False)
          self.layer1 = nn.Linear(100, 100, bias=False)
          self.norm = nn.RMSNorm(100)
  
      def forward(self, x):
          x = self.norm(x)
          return self.layer0(self.layer1(x))
  
  class Test(nn.Module):
  
      def __init__(self):
          super().__init__()
  
          self.embs = nn.Embedding(1000, 100)
          self.layers = nn.ModuleList([Test1() for _ in range(32)])
  
      def forward(self, x):
          x = self.embs(x)
          for layer in self.layers:
              x = layer(x) + x
          return x.sum()
  
  x = torch.randint(0, 1000, (20,), device="cuda")
  model = Test().cuda().bfloat16()
  compiled_model = torch.compile(model)
  y = compiled_model(x)
  y.backward()
  ```
  
  In the `torch_compile_debug` backward folder, the `fx_graph_readable.py` file shows an unusual series of additions.  
  
  ```python
  class GraphModule(torch.nn.Module):
      def forward(self, primals_2: "i64[20]", primals_3: "bf16[100]", primals_6: "bf16[100]", primals_9: "bf16[100]", primals_12: "bf16[100]", primals_15: "bf16[100]", primals_18: "bf16[100]", primals_21: "bf16[100]", primals_24: "bf16[100]", primals_27: "bf16[100]", primals_30: "bf16[100]", primals_33: "bf16[100]", primals_36: "bf16[100]", primals_39: "bf16[100]", primals_42: "bf16[100]", primals_45: "bf16[100]", primals_48: "bf16[100]", primals_51: "bf16[100]", primals_54: "bf16[100]", primals_57: "bf16[100]", primals_60: "bf16[100]", primals_63: "bf16[100]", primals_66: "bf16[100]", primals_69: "bf16[100]", primals_72: "bf16[100]", primals_75: "bf16[100]", primals_78: "bf16[100]", primals_81: "bf16[100]", primals_84: "bf16[100]", primals_87: "bf16[100]", primals_90: "bf16[100]", primals_93: "bf16[100]", primals_96: "bf16[100]", embedding: "bf16[20, 100]", rsqrt: "bf16[20, 1]", mm: "bf16[20, 100]", mm_1: "bf16[20, 100]", rsqrt_1: "bf16[20, 1]", mm_2: "bf16[20, 100]", mm_3: "bf16[20, 100]", rsqrt_2: "bf16[20, 1]", mm_4: "bf16[20, 100]", mm_5: "bf16[20, 100]", rsqrt_3: "bf16[20, 1]", mm_6: "bf16[20, 100]", mm_7: "bf16[20, 100]", rsqrt_4: "bf16[20, 1]", mm_8: "bf16[20, 100]", mm_9: "bf16[20, 100]", rsqrt_5: "bf16[20, 1]", mm_10: "bf16[20, 100]", mm_11: "bf16[20, 100]", rsqrt_6: "bf16[20, 1]", mm_12: "bf16[20, 100]", mm_13: "bf16[20, 100]", rsqrt_7: "bf16[20, 1]", mm_14: "bf16[20, 100]", mm_15: "bf16[20, 100]", rsqrt_8: "bf16[20, 1]", mm_16: "bf16[20, 100]", mm_17: "bf16[20, 100]", rsqrt_9: "bf16[20, 1]", mm_18: "bf16[20, 100]", mm_19: "bf16[20, 100]", rsqrt_10: "bf16[20, 1]", mm_20: "bf16[20, 100]", mm_21: "bf16[20, 100]", rsqrt_11: "bf16[20, 1]", mm_22: "bf16[20, 100]", mm_23: "bf16[20, 100]", rsqrt_12: "bf16[20, 1]", mm_24: "bf16[20, 100]", mm_25: "bf16[20, 100]", rsqrt_13: "bf16[20, 1]", mm_26: "bf16[20, 100]", mm_27: "bf16[20, 100]", rsqrt_14: "bf16[20, 1]", mm_28: "bf16[20, 100]", mm_29: "bf16[20, 100]", rsqrt_15: "bf16[20, 1]", mm_30: "bf16[20, 100]", mm_31: "bf16[20, 100]", rsqrt_16: "bf16[20, 1]", mm_32: "bf16[20, 100]", mm_33: "bf16[20, 100]", rsqrt_17: "bf16[20, 1]", mm_34: "bf16[20, 100]", mm_35: "bf16[20, 100]", rsqrt_18: "bf16[20, 1]", mm_36: "bf16[20, 100]", mm_37: "bf16[20, 100]", rsqrt_19: "bf16[20, 1]", mm_38: "bf16[20, 100]", mm_39: "bf16[20, 100]", rsqrt_20: "bf16[20, 1]", mm_40: "bf16[20, 100]", mm_41: "bf16[20, 100]", rsqrt_21: "bf16[20, 1]", mm_42: "bf16[20, 100]", mm_43: "bf16[20, 100]", rsqrt_22: "bf16[20, 1]", mm_44: "bf16[20, 100]", mm_45: "bf16[20, 100]", rsqrt_23: "bf16[20, 1]", mm_46: "bf16[20, 100]", mm_47: "bf16[20, 100]", rsqrt_24: "bf16[20, 1]", mm_48: "bf16[20, 100]", mm_49: "bf16[20, 100]", rsqrt_25: "bf16[20, 1]", mm_50: "bf16[20, 100]", mm_51: "bf16[20, 100]", rsqrt_26: "bf16[20, 1]", mm_52: "bf16[20, 100]", mm_53: "bf16[20, 100]", rsqrt_27: "bf16[20, 1]", mm_54: "bf16[20, 100]", mm_55: "bf16[20, 100]", rsqrt_28: "bf16[20, 1]", mm_56: "bf16[20, 100]", mm_57: "bf16[20, 100]", rsqrt_29: "bf16[20, 1]", mm_58: "bf16[20, 100]", mm_59: "bf16[20, 100]", rsqrt_30: "bf16[20, 1]", mm_60: "bf16[20, 100]", mm_61: "bf16[20, 100]", rsqrt_31: "bf16[20, 1]", mm_62: "bf16[20, 100]", permute_66: "bf16[100, 100]", permute_70: "bf16[100, 100]", permute_74: "bf16[100, 100]", permute_78: "bf16[100, 100]", permute_82: "bf16[100, 100]", permute_86: "bf16[100, 100]", permute_90: "bf16[100, 100]", permute_94: "bf16[100, 100]", permute_98: "bf16[100, 100]", permute_102: "bf16[100, 100]", permute_106: "bf16[100, 100]", permute_110: "bf16[100, 100]", permute_114: "bf16[100, 100]", permute_118: "bf16[100, 100]", permute_122: "bf16[100, 100]", permute_126: "bf16[100, 100]", permute_130: "bf16[100, 100]", permute_134: "bf16[100, 100]", permute_138: "bf16[100, 100]", permute_142: "bf16[100, 100]", permute_146: "bf16[100, 100]", permute_150: "bf16[100, 100]", permute_154: "bf16[100, 100]", permute_158: "bf16[100, 100]", permute_162: "bf16[100, 100]", permute_166: "bf16[100, 100]", permute_170: "bf16[100, 100]", permute_174: "bf16[100, 100]", permute_178: "bf16[100, 100]", permute_182: "bf16[100, 100]", permute_186: "bf16[100, 100]", permute_190: "bf16[100, 100]", permute_194: "bf16[100, 100]", permute_198: "bf16[100, 100]", permute_202: "bf16[100, 100]", permute_206: "bf16[100, 100]", permute_210: "bf16[100, 100]", permute_214: "bf16[100, 100]", permute_218: "bf16[100, 100]", permute_222: "bf16[100, 100]", permute_226: "bf16[100, 100]", permute_230: "bf16[100, 100]", permute_234: "bf16[100, 100]", permute_238: "bf16[100, 100]", permute_242: "bf16[100, 100]", permute_246: "bf16[100, 100]", permute_250: "bf16[100, 100]", permute_254: "bf16[100, 100]", permute_258: "bf16[100, 100]", permute_262: "bf16[100, 100]", permute_266: "bf16[100, 100]", permute_270: "bf16[100, 100]", permute_274: "bf16[100, 100]", permute_278: "bf16[100, 100]", permute_282: "bf16[100, 100]", permute_286: "bf16[100, 100]", permute_290: "bf16[100, 100]", permute_294: "bf16[100, 100]", permute_298: "bf16[100, 100]", permute_302: "bf16[100, 100]", permute_306: "bf16[100, 100]", permute_310: "bf16[100, 100]", permute_314: "bf16[100, 100]", permute_318: "bf16[100, 100]", tangents_1: "bf16[]"):
           # File: /tmp/ipykernel_1043308/3460069279.py:38 in forward, code: return x.sum()
          expand: "bf16[20, 100]" = torch.ops.aten.expand.default(tangents_1, [20, 100]);  tangents_1 = None
          
           # File: /tmp/ipykernel_1043308/3460069279.py:24 in forward, code: return self.layer0(self.layer1(x))
          permute_64: "bf16[100, 20]" = torch.ops.aten.permute.default(expand, [1, 0])
          mm_64: "bf16[100, 100]" = torch.ops.aten.mm.default(permute_64, mm_62);  permute_64 = mm_62 = None
          permute_65: "bf16[100, 100]" = torch.ops.aten.permute.default(mm_64, [1, 0]);  mm_64 = None
          mm_65: "bf16[20, 100]" = torch.ops.aten.mm.default(expand, permute_66);  permute_66 = None
          permute_67: "bf16[100, 100]" = torch.ops.aten.permute.default(permute_65, [1, 0]);  permute_65 = None
          permute_68: "bf16[100, 20]" = torch.ops.aten.permute.default(mm_65, [1, 0])
          
           # File: /tmp/ipykernel_1043308/3460069279.py:37 in forward, code: x = layer(x) + x
          add_1: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_1, embedding);  mm_1 = None
          add_3: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_3, add_1);  mm_3 = None
          add_5: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_5, add_3);  mm_5 = None
          add_7: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_7, add_5);  mm_7 = None
          add_9: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_9, add_7);  mm_9 = None
          add_11: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_11, add_9);  mm_11 = None
          add_13: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_13, add_11);  mm_13 = None
          add_15: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_15, add_13);  mm_15 = None
          add_17: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_17, add_15);  mm_17 = None
          add_19: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_19, add_17);  mm_19 = None
          add_21: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_21, add_19);  mm_21 = None
          add_23: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_23, add_21);  mm_23 = None
          add_25: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_25, add_23);  mm_25 = None
          add_27: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_27, add_25);  mm_27 = None
          add_29: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_29, add_27);  mm_29 = None
          add_31: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_31, add_29);  mm_31 = None
          add_33: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_33, add_31);  mm_33 = None
          add_35: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_35, add_33);  mm_35 = None
          add_37: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_37, add_35);  mm_37 = None
          add_39: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_39, add_37);  mm_39 = None
          add_41: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_41, add_39);  mm_41 = None
          add_43: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_43, add_41);  mm_43 = None
          add_45: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_45, add_43);  mm_45 = None
          add_47: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_47, add_45);  mm_47 = None
          add_49: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_49, add_47);  mm_49 = None
          add_51: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_51, add_49);  mm_51 = None
          add_53: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_53, add_51);  mm_53 = None
          add_55: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_55, add_53);  mm_55 = None
          add_57: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_57, add_55);  mm_57 = None
          add_59: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_59, add_57);  mm_59 = None
          add_61: "bf16[20, 100]" = torch.ops.aten.add.Tensor(mm_61, add_59);  mm_61 = None
  ```
  
  A simple observation reveals that it has transformed into the following pattern during the forward pass:  
  x1 = x0 + y0,  
  x2 = x1 + y1,  
  x3 = x2 + y2.  
  
  Here, x0, x1, x2, and x3 are all stored for backward computation.  
  The optimal approach would be to store x0, x1, x2, and x3.  
  
  However, due to an issue in the `min cut` implementation of `torch.compile`, which supports recomputation for non-compute-intensive operations, it instead stores x0, y0, y1, and y2, while x1, x2, and x3 are recomputed.  
  
  Although both methods use the same amount of memory, the latter introduces unnecessary computations.
  
  ### Error logs
  
  _No response_
  
  ### Versions
  
  torch 2.5.1+cu124
  
  cc @chauhang @penguinwu @zou3519 @bdhirsh