feat: make SAC matmul recompute frequency configurable

tests/unit_tests/gpu/test_activation_checkpoint.py

@@ -321,6 +321,67 @@ def get_recomputed(force_recompute_fqns):
         # Early-stop skips the terminal output once backward has all its tensors.
         self.assertEqual(get_recomputed(["output"]), {"wq"})
 
+    def test_matmul_recompute_freq(self):
+        """matmul_recompute_freq controls which mms are recomputed; 0 saves
+        every mm.
+
+        Same aten.mm counting approach as test_force_recompute_mm_fqns.
+        mm order in ToyModule: gate (1st), wq (2nd), output (3rd).
+        """
+        from torch.utils._python_dispatch import TorchDispatchMode
+
+        class MmWeightTracker(TorchDispatchMode):
+            def __init__(self, ptrs):
+                super().__init__()
+                self._ptrs = ptrs
+                self.counts = {n: 0 for n in ptrs.values()}
+
+            def __torch_dispatch__(self, func, types, args, kwargs=None):
+                if func == torch.ops.aten.mm.default:
+                    for arg in args:
+                        name = self._ptrs.get(arg.data_ptr())
+                        if name is not None:
+                            self.counts[name] += 1
+                            break
+                return func(*args, **(kwargs or {}))
+
+        def get_recomputed(matmul_recompute_freq):
+            m = ToyModule()
+            SelectiveAC.Config(
+                force_recompute_mm_shapes_by_fqns=[],
+                matmul_recompute_freq=matmul_recompute_freq,
+            ).build().apply(m)
+            ptr_to_name = {
+                mod.weight.data_ptr(): fqn.rsplit(".", 1)[-1]
+                for fqn, mod in m.named_modules()
+                if isinstance(mod, Linear)
+            }
+            x = torch.randn(64, 512, requires_grad=True)
+            out = m(x)
+            tracker = MmWeightTracker(ptr_to_name)
+            with tracker:
+                out.backward()
+            return {n for n, c in tracker.counts.items() if c == 2}
+
+        # Default freq=2 recomputes the 2nd mm (wq): backwards compatible.
+        self.assertEqual(get_recomputed(2), {"wq"})
+        # freq=0 saves every mm (recompute none).
+        self.assertEqual(get_recomputed(0), set())
+        # freq=1 recomputes every mm. As in test_force_recompute_mm_fqns,
+        # early-stop skips the terminal output mm, so only gate and wq are
+        # observably replayed.
+        self.assertEqual(get_recomputed(1), {"gate", "wq"})
+        # freq=3 targets the 3rd mm (output), so gate and wq are both saved:
+        # wq flips from recomputed at freq=2 to saved here.
+        self.assertEqual(get_recomputed(3), set())
+
+    def test_matmul_recompute_freq_must_be_non_negative(self):
+        with self.assertRaises(ValueError):
+            SelectiveAC.Config(
+                force_recompute_mm_shapes_by_fqns=[],
+                matmul_recompute_freq=-1,
+            ).build().apply(ToyModule())
+
 
 if __name__ == "__main__":
     unittest.main()

torchtitan/distributed/activation_checkpoint.py

@@ -197,8 +197,9 @@ class SelectiveAC(ActivationCheckpointing):
 
     Saves the outputs of compute/communication ops that are expensive to
     recompute (see ``get_save_ops``) while recomputing the rest, and recomputes
-    every second matmul to balance memory and compute. Override ``get_save_ops``
-    in a subclass to tune which ops are saved.
+    every ``matmul_recompute_freq``-th matmul (default: every second) to balance
+    memory and compute. Override ``get_save_ops`` in a subclass to tune which
+    ops are saved.
     """
 
     @dataclass(kw_only=True, slots=True)
@@ -216,6 +217,17 @@ class Config(ActivationCheckpointing.Config):
         ANY mm with shape matching (*, in) x (in, out) will be force recomputed.
         """
 
+        matmul_recompute_freq: int = 2
+        """
+        Recompute one matmul/linear out of every ``matmul_recompute_freq`` under
+        the default SAC policy (the rest are saved). ``2`` (default) recomputes
+        every second matmul, matching the historical policy; larger values
+        recompute fewer matmuls (save more); ``0`` saves every matmul (recompute
+        none) -- used by memory-headroom configs that need all mms saved.
+        Matmuls matched by ``force_recompute_mm_shapes_by_fqns`` are always
+        recomputed regardless of this value. Must be >= 0.
+        """
+
     def get_save_ops(self) -> set:
         """Returns the set of ops whose activations should be saved. Override
         to customize the save set."""
@@ -225,6 +237,15 @@ def _wrap_block(
         self, module: nn.Module, *, base_fqn: str | None = None
     ) -> nn.Module:
         config = cast("SelectiveAC.Config", self.config)
+        # Read out as a plain int rather than referencing config inside the
+        # policy closure below: the closure is pickled by the precompile
+        # packager, and capturing the Config would drag the nested
+        # SelectiveAC.Config class into the pickle.
+        mm_recompute_freq = config.matmul_recompute_freq
+        if mm_recompute_freq < 0:
+            raise ValueError(
+                f"matmul_recompute_freq must be >= 0, got {mm_recompute_freq}"
+            )
         save_ops = self.get_save_ops()
 
         # Collect weight shapes to force-recompute, stored as mm RHS shape
@@ -281,9 +302,14 @@ def wrapped_policy(ctx, func, *args, **kwargs) -> CheckpointPolicy:
                         return CheckpointPolicy.PREFER_RECOMPUTE
                     meta[mm_count_key] += 1
 
-                # Save all compute/comm ops, except every second mm/linear.
+                # Save all compute/comm ops, except every
+                # matmul_recompute_freq-th mm/linear (freq=0 saves every mm).
                 if func in save_ops:
-                    if func in mm_ops and meta[mm_count_key] % 2 == 0:
+                    if (
+                        func in mm_ops
+                        and mm_recompute_freq != 0
+                        and meta[mm_count_key] % mm_recompute_freq == 0
+                    ):
                         return CheckpointPolicy.PREFER_RECOMPUTE
                     return CheckpointPolicy.MUST_SAVE
                 return CheckpointPolicy.PREFER_RECOMPUTE