-
Notifications
You must be signed in to change notification settings - Fork 21.9k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Allow "must recompute" in torch.compile + selective checkpointing (SAC) #114036
Conversation
This PR needs a
|
I don't understand the motivation here - when is this desired? |
@Chillee Some internal GenAI teams have a AC autotuner where they make a few trial runs and sort through what modules have low compute/high memory and then make a customized AC plan. They compressed 7 days of manual tuning and got equal or better with the autotuning in a couple hours. (cc. @lessw2020) I believe allowing user to specify both "what to recompute" (this PR) and "what not to recompute" (already exists) will allow them to have maximum flexibility. Alternatively we can expose the partitioning heuristics to them and allow them to influence / overwrite it. But it won't be as straightforward to them as asking them to just specify what ops to recompute, and let them reuse their own autotuning infra. |
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
I think that this is useful: pytorch-labs/float8_experimental#185 |
torch/_functorch/partitioners.py
Outdated
if must_recompute(node): | ||
# If user explicitly say they want to recompute a node, we honor it | ||
# by adding an 0-capacity edge from the source and an inf edge to the sink. | ||
# This way, X_in node is guaranteed to be part of the subgraph that contains "sink" | ||
# after the cut, thus guaranteeing that X op will be recomputed. | ||
add_edge("source", node.name + "_in", capacity=0) | ||
add_edge(node.name + "_in", "sink", capacity=math.inf) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
mic-cut partitioner is supposed to give better cut on top of user-provided annotation. So, we should maybe not do this.
But if you want to do this deliberately (because FSDP does something funky), a better way might be to just have another config flag that respects the user AC annotation. In that case, you can just check if flag_enabled and must_recompile()
and add the two edges, without any other changes anywhere.
Conceptually, its better than having 3 levels - must, ban and prefer. The new experimental flag, I am suggesting, says that partitioner will respect user AC annotation, so its either must or ban recompute.
…st_recompute_op
…st_recompute_op
I believe there is appetite for adding this API now, so resurrecting this PR. |
…st_recompute_op
…checkpointing (SAC) ghstack-source-id: dfc5d8f7d50844b2fa49726ed114a55c014ab89e Pull Request resolved: #129295
Closing in favor of #129295, to leverage ghstack. |
…checkpointing (SAC) ghstack-source-id: 7a9f0491de59fc37fb98c8b9e6848b6f7524a624 Pull Request resolved: #129295
…checkpointing (SAC) (#129295) Pull Request resolved: #129295 Approved by: https://github.com/Chillee
As titled. Done by setting up special edge from X_in to sink, so that min-cut partitioner will always put X_in in the backward graph (thus always recomputing X).
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov