Skip to content

feat: make SAC matmul recompute frequency configurable - #8

Closed
samansg-aws wants to merge 1 commit into
mainfrom
samansg/sac-matmul-recompute-freq
Closed

samansg-aws wants to merge 1 commit into
mainfrom
samansg/sac-matmul-recompute-freq

Conversation

@samansg-aws

Copy link
Copy Markdown
Collaborator

Issue #, if available: n/a

Description of changes:

SelectiveAC hardcodes "recompute every second matmul" (meta[mm_count_key] % 2 == 0),
which is the only memory/compute tradeoff per-op SAC exposes.
force_recompute_mm_shapes_by_fqns can only push more matmuls into
recomputation, never fewer, so there is no way to say "save all the matmuls".

Adds one field to SelectiveAC.Config:

  • matmul_recompute_freq: int = 2 -- recompute one matmul out of every n.
  • 2 (default) reduces to the old predicate, so existing configs are bit-identical.
  • 1 recomputes every matmul; 0 saves every matmul; negative raises ValueError.
  • Matmuls matched by force_recompute_mm_shapes_by_fqns are still always
    recomputed, so the two knobs compose as before.

Tests in tests/unit_tests/gpu/test_activation_checkpoint.py count aten.mm
calls per weight during backward (same technique as test_force_recompute_mm_fqns)
for freq in {0, 1, 2, 3}, plus the ValueError case.

pre-commit run --files passes on both changed files.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

@samansg-aws
samansg-aws force-pushed the samansg/sac-matmul-recompute-freq branch from 341a33b to 5dc5986 Compare October 7, 2026 18:44
@samansg-aws
samansg-aws deleted the samansg/sac-matmul-recompute-freq branch October 19, 2026 20:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant