Add raise_last_usage memory optimization pass to Inductor

[yf225]

This PR adds a new Inductor memory optimization pass to raise the consumer nodes if moving them up can free more memory (i.e. if its output tensor is smaller than the sum memory of all its last-usage input tensors).

Particularly, for each node in the graph, we raise its consumer node to be right after the current node, if it satisfies the following conditions:

• The consumer node's all input args have their write sites scheduled before or at the current node.
• The consumer node only writes to one output tensor.
• The consumer node's out tensor is smaller than the sum memory of all its last-usage input args (i.e. we can get memory savings by raising it).

This pass is particularly important for significantly reducing peak memory usage for Compiled FSDP2 on llama-7b model (decreasing mem usage by more than 57GB).

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/125559

• 📄 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.

❌ 1 New Failure, 1 Unrelated Failure

As of commit 5c742dc with merge base da991fa ():

NEW FAILURE - The following job has failed:

• Lint / lintrunner-noclang / linux-job (gh)
  >>> Lint for torch/_inductor/memory_passes.py:

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

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

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_timm, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  sebotnet33ts_256

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

[yf225]

Do we still need this pass if we set Inductor config reorder_for_locality=False by default for Compiled Autograd graph?

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.