From 5dc5986a85a40b027a869c5caf5744a257416921 Mon Sep 17 00:00:00 2001 From: Samanvitha Sateesha Date: Wed, 7 Oct 2026 18:44:31 +0000 Subject: [PATCH] feat: make SAC matmul recompute frequency configurable --- .../gpu/test_activation_checkpoint.py | 61 +++++++++++++++++++ .../distributed/activation_checkpoint.py | 34 +++++++++-- 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/gpu/test_activation_checkpoint.py b/tests/unit_tests/gpu/test_activation_checkpoint.py index 4692ba65f4f..708d4e5fc05 100644 --- a/tests/unit_tests/gpu/test_activation_checkpoint.py +++ b/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() diff --git a/torchtitan/distributed/activation_checkpoint.py b/torchtitan/distributed/activation_checkpoint.py index 517d47101c7..5131ca135ca 100644 --- a/torchtitan/distributed/activation_checkpoint.py +++ b/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