Allow "must recompute" in torch.compile + selective checkpointing (SAC) #114036
Conversation
This PR needs a
|
yf225
added
the
suppress-bc-linter
|
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.
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
yf225
requested review from
soulitzer,
anijain2305 and
Chillee
and removed request for
Chillee
yf225
removed
suppress-bc-linter
|
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