[Reopen #114036] Allow "must recompute" in torch.compile + selective checkpointing (SAC)

[yf225]

This replaces #114036.

Stack from ghstack (oldest at bottom):

• -> [Reopen #114036] Allow "must recompute" in torch.compile + selective checkpointing (SAC) #129295

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/129295

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (2 Unrelated Failures)

As of commit 5ca28a5 with merge base aa4ee2c ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 1, 3, linux.8xlarge.nvidia.gpu) (gh) (trunk failure)
  distributed/_composable/test_replicate_with_compiler.py::ReplicateTest::test_bucketing_concat_op

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• inductor / rocm6.1-py3.8-inductor / test (inductor, 1, 1, linux.rocm.gpu.2, unstable) (gh) (#128871)
  'test/inductor/test_max_autotune.py::TestTuningProcess::test_tuning_pool_multiple_devices'

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[bdhirsh]

is this an independent bugfix? Or is this part of the "partitioner should respect prefer_recompute" change.

Mostly a curiosity question: my understanding from the comment above is that if you run code like:

checkpoint_f = checkpoint(f)
checkpoint(g) = checkpoint(g)
out = f(g(inp))

Then AC requires us to save the inputs to f (outputs of g), but in our tag-based system every node would have the recompute tag: so you need some notion of "which AC subgraph does a node belong to" (ac_graph_id) to tell the partitioner that it should save the output of the first subgraph. Is that right?

[yf225]

this is part of the "recompute" tag cleanup to disentangle its two meanings: 1) recompute policy, 2) AC subgraph ID, by splitting them into two tags.

yes your understanding is exactly right :)

[bdhirsh]

It looks like right above this condition, if the op is a view or a lift_fresh_copy then we will not ban recomputation, even if the user marked as MUST_SAVE.

Do you think we should either error here, or prefer the user annotation instead?

Maybe a more general question: for ops that the partitioner is already strongly opinionated about whether they should be saved (e.g. randomness or view ops), should we error when a user tries to change the partitioners behavior for them? Or promise to always respect the user intent?

[soulitzer]

This is a good point, I think we should error when user MUST_SAVE a view or lift_fresh, etc. rather than ignore the annotation, or silently do something less optimal.

However, generally it is tricky to absolutely respect the "MUST_SAVE" condition, e.g. in the case where some code get's DCE'd.

Or another case is if the user choose to MUST_SAVE some tensor, but let's say the backward formula reduces this value before using it, so maybe we'd rather save the value post reduction.

But these two are really just the consequence of the fact that MUST_NOT_RECOMPUTE became MUST_SAVE... maybe its worth renaming it back lol.

[yf225]

yeah I feel that we should bias more on respecting user intent - if user wants "MUST_SAVE", we should try to give them MUST_SAVE as much as possible, since the user explicitly expresses it.

[yf225]

@pytorchbot merge -f "unrelated failures"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here