[Compiled Autograd][Traceable FSDP2][UX] Add compiled_autograd_kwargs_override Dynamo config #136967

yf225 · 2024-09-29T23:45:04Z

For Traceable FSDP2, the most common use case is to have fullgraph=False for forward pass (to allow user-level graph breaks), and fullgraph=True for compiled autograd backward pass (required for queue_callback support).

With torch._dynamo.compiled_autograd=True, previously we are not able to set different fullgraph config value for forward vs. backward pass, since rebuild_ctx just reuses the forward compile config as-is. This PR adds torch._dynamo.config.compiled_autograd_kwargs_override config to allow forcing fullgraph=True for CA Dynamo tracing.

With this PR, we can remove standalone compiled autograd ctx manager usage in Traceable FSDP2 unit tests, and consolidate on using torch._dynamo.compiled_autograd=True.

Test commands:

pytest -rA test/distributed/_composable/fsdp/test_fully_shard_compile.py::TestFullyShardCompile::test_transformer_backend_inductor_fullgraph_True

Stack from ghstack (oldest at bottom):

-> [Compiled Autograd][Traceable FSDP2][UX] Add compiled_autograd_kwargs_override Dynamo config #136967

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @rec

[ghstack-poisoned]

pytorch-bot · 2024-09-29T23:45:08Z

🔗 Helpful Links

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

📄 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 d88b5d3 with merge base f0fa460 ():

NEW FAILURE - The following job has failed:

inductor-periodic / cuda12.1-py3.10-gcc9-sm80 / test (inductor_torchbench_smoketest_perf, 1, 1, linux.gcp.a100) (gh)
running benchmark: 100% 30/30 [00:03<00:00, 7.76it/s]

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, lf.linux.2xlarge) (gh) (disabled by #134602)
test_transformers.py::TestSDPAPrivateUse1Only::test_scaled_dot_product_fused_attention_overrideable_backward

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

cc XilunWu H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

For Traceable FSDP2, the most common use case is to have `fullgraph=False` for forward pass (to allow user-level graph breaks), and `fullgraph=True` for compiled autograd backward pass (required for queue_callback support). With `torch._dynamo.compiled_autograd=True`, previously we are not able to set different `fullgraph` config value for forward vs. backward pass, since `rebuild_ctx` just reuses the forward compile config as-is. This PR adds `torch._dynamo.compiled_autograd_fullgraph` config to allow forcing `fullgraph=True` for CA Dynamo tracing. With this PR, we can remove standalone compiled autograd ctx manager usage in Traceable FSDP2 unit tests, and consolidate on using `torch._dynamo.compiled_autograd=True`. cc XilunWu H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

yf225 · 2024-09-30T00:08:22Z

torch/_dynamo/config.py

+# - If None, uses forward Dynamo `fullgraph` config value.
+#
+# This config only takes effect if `compiled_autograd=True`.
+compiled_autograd_fullgraph: Union[bool, None] = None


Alternatively we could do:

compiled_autograd_config_override = { "fullgraph": True }

which is probably more future-proof if we expect to need to override other configs in the future.

what about just passing in an override OptimizeContext, won't have to keep in sync with possible options e.g.

config.compiled_autograd_override_compiler_fn: Optional[...] = torch.compile(fullgraph=True)

actually, it's annoying to do so because fullgraph exists outside of OptimizeContext, yeah we can do this

just realized: if we were to use compiled_autograd_config_override = {...}, we need to use"nopython" as key instead of "fullgraph" due to the signature difference between torch._dynamo.optimize() and torch.compile()...

alternatively is there appetite to move rebuild_ctx() higher level to rebuild the torch.compile() call instead, so that "fullgraph" can be used as key here?

or we can stick with the compiled_autograd_fullgraph config and worry about this later

is there appetite to move rebuild_ctx() higher level to rebuild the torch.compile() call instead, so that "fullgraph" can be used as key here?

updated PR to explore this direction

For Traceable FSDP2, the most common use case is to have `fullgraph=False` for forward pass (to allow user-level graph breaks), and `fullgraph=True` for compiled autograd backward pass (required for queue_callback support). With `torch._dynamo.compiled_autograd=True`, previously we are not able to set different `fullgraph` config value for forward vs. backward pass, since `rebuild_ctx` just reuses the forward compile config as-is. This PR adds `torch._dynamo.compiled_autograd_fullgraph` config to allow forcing `fullgraph=True` for CA Dynamo tracing. With this PR, we can remove standalone compiled autograd ctx manager usage in Traceable FSDP2 unit tests, and consolidate on using `torch._dynamo.compiled_autograd=True`. cc XilunWu H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

ghstack-source-id: 21b3b0c Pull Request resolved: #136967

For Traceable FSDP2, the most common use case is to have `fullgraph=False` for forward pass (to allow user-level graph breaks), and `fullgraph=True` for compiled autograd backward pass (required for queue_callback support). With `torch._dynamo.compiled_autograd=True`, previously we are not able to set different `fullgraph` config value for forward vs. backward pass, since `rebuild_ctx` just reuses the forward compile config as-is. This PR adds `torch._dynamo.compiled_autograd_fullgraph` config to allow forcing `fullgraph=True` for CA Dynamo tracing. With this PR, we can remove standalone compiled autograd ctx manager usage in Traceable FSDP2 unit tests, and consolidate on using `torch._dynamo.compiled_autograd=True`. Test commands: - `pytest -rA test/distributed/_composable/fsdp/test_fully_shard_compile.py::TestFullyShardCompile::test_transformer_backend_inductor_fullgraph_True` cc XilunWu H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

ghstack-source-id: 5e09c58 Pull Request resolved: #136967

xmfan · 2024-10-01T18:52:19Z

having rebuild_ctx operate at the torch.compile level requires a lot of logic to prevent infinite recursion. if nopython-fullgraph is the only issue, i feel it's simpler to just map it inside of rebuild_ctx.

i'm thinking that you could do

def rebuild_ctx():
  if config.compiled_autograd_config_override:
    # override args/kwargs
  return optimize(rebuild_ctx, *args, **kwargs)

For Traceable FSDP2, the most common use case is to have `fullgraph=False` for forward pass (to allow user-level graph breaks), and `fullgraph=True` for compiled autograd backward pass (required for queue_callback support). With `torch._dynamo.compiled_autograd=True`, previously we are not able to set different `fullgraph` config value for forward vs. backward pass, since `rebuild_ctx` just reuses the forward compile config as-is. This PR adds `torch._dynamo.compiled_autograd_fullgraph` config to allow forcing `fullgraph=True` for CA Dynamo tracing. With this PR, we can remove standalone compiled autograd ctx manager usage in Traceable FSDP2 unit tests, and consolidate on using `torch._dynamo.compiled_autograd=True`. Test commands: - `pytest -rA test/distributed/_composable/fsdp/test_fully_shard_compile.py::TestFullyShardCompile::test_transformer_backend_inductor_fullgraph_True` cc XilunWu H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

pytorchmergebot · 2024-10-02T00:59:49Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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

pytorchmergebot · 2024-10-02T01:30:41Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, lf.linux.2xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

yf225 · 2024-10-02T02:54:28Z

@pytorchbot rebase

pytorchmergebot · 2024-10-02T02:56:00Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[ghstack-poisoned]

pytorchmergebot · 2024-10-02T02:56:12Z

Successfully rebased gh/yf225/128/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/136967)

ghstack-source-id: 1d1ad3b Pull Request resolved: #136967

yf225 · 2024-10-02T04:08:04Z

@pytorchbot merge

pytorchmergebot · 2024-10-02T04:09:44Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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

pytorchmergebot · 2024-10-02T04:35:43Z

Merge failed

Reason: 1 jobs have failed, first few of them are: inductor-periodic / cuda12.1-py3.10-gcc9-sm80 / test (inductor_torchbench_smoketest_perf, 1, 1, linux.gcp.a100)

Details for Dev Infra team

Raised by workflow job

yf225 · 2024-10-02T06:22:13Z

@pytorchbot merge -f "unrelated failures"

pytorchmergebot · 2024-10-02T06:23:52Z

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

For Traceable FSDP2, the most common use case is to have `fullgraph=False` for forward pass (to allow user-level graph breaks), and `fullgraph=True` for compiled autograd backward pass (required for queue_callback support). With `torch._dynamo.compiled_autograd=True`, previously we are not able to set different `fullgraph` config value for forward vs. backward pass, since `rebuild_ctx` just reuses the forward compile config as-is. This PR adds `torch._dynamo.config.compiled_autograd_kwargs_override` config to allow forcing `fullgraph=True` for CA Dynamo tracing. With this PR, we can remove standalone compiled autograd ctx manager usage in Traceable FSDP2 unit tests, and consolidate on using `torch._dynamo.compiled_autograd=True`. Test commands: - `pytest -rA test/distributed/_composable/fsdp/test_fully_shard_compile.py::TestFullyShardCompile::test_transformer_backend_inductor_fullgraph_True` Pull Request resolved: pytorch#136967 Approved by: https://github.com/xmfan

Add compiled_autograd_fullgraph Dynamo config

e98f0b1

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: distributed (fsdp) release notes category labels Sep 29, 2024

yf225 requested a review from xmfan September 30, 2024 00:05

yf225 commented Sep 30, 2024

View reviewed changes

yf225 added a commit that referenced this pull request Sep 30, 2024

Add compiled_autograd_fullgraph Dynamo config

d366e71

ghstack-source-id: 21b3b0c Pull Request resolved: #136967

yf225 added a commit that referenced this pull request Oct 1, 2024

Add compiled_autograd_fullgraph Dynamo config

b443744

ghstack-source-id: 5e09c58 Pull Request resolved: #136967

yf225 changed the title ~~Add compiled_autograd_fullgraph Dynamo config~~ Add compiled_autograd_config_override Dynamo config Oct 1, 2024

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 2, 2024

yf225 added release notes: dynamo and removed release notes: distributed (fsdp) release notes category labels Oct 2, 2024

pytorchmergebot added the merging label Oct 2, 2024

pytorchmergebot removed the merging label Oct 2, 2024

yf225 added the keep-going Don't stop on first failure, keep running tests until the end label Oct 2, 2024

yf225 closed this Oct 2, 2024

yf225 reopened this Oct 2, 2024

Update

d88b5d3

[ghstack-poisoned]

pytorchmergebot pushed a commit that referenced this pull request Oct 2, 2024

Add compiled_autograd_fullgraph Dynamo config

b9b952d

ghstack-source-id: 1d1ad3b Pull Request resolved: #136967

pytorchmergebot added the merging label Oct 2, 2024

pytorchmergebot removed the merging label Oct 2, 2024

pytorchmergebot added the merging label Oct 2, 2024

pytorchmergebot closed this in a89e3c2 Oct 2, 2024

pytorchmergebot added Merged and removed merging labels Oct 2, 2024

yf225 changed the title ~~Add compiled_autograd_kwargs_override Dynamo config~~ [Compiled Autograd][Traceable FSDP2][UX] Add compiled_autograd_kwargs_override Dynamo config Oct 14, 2024

github-actions bot deleted the gh/yf225/128/head branch November 14, 2024 02:06

[Compiled Autograd][Traceable FSDP2][UX] Add compiled_autograd_kwargs_override Dynamo config #136967

[Compiled Autograd][Traceable FSDP2][UX] Add compiled_autograd_kwargs_override Dynamo config #136967

Uh oh!

Conversation

yf225 commented Sep 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

yf225 Sep 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xmfan Sep 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xmfan Sep 30, 2024

Choose a reason for hiding this comment

Uh oh!

yf225 Sep 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yf225 Oct 1, 2024

Choose a reason for hiding this comment

Uh oh!

xmfan commented Oct 1, 2024

Uh oh!

pytorchmergebot commented Oct 2, 2024

Merge started

Uh oh!

pytorchmergebot commented Oct 2, 2024

Merge failed

Uh oh!

yf225 commented Oct 2, 2024

Uh oh!

pytorchmergebot commented Oct 2, 2024

Uh oh!

pytorchmergebot commented Oct 2, 2024

Uh oh!

yf225 commented Oct 2, 2024

Uh oh!

pytorchmergebot commented Oct 2, 2024

Merge started

Uh oh!

pytorchmergebot commented Oct 2, 2024

Merge failed

Uh oh!

yf225 commented Oct 2, 2024

Uh oh!

pytorchmergebot commented Oct 2, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yf225 commented Sep 29, 2024 •

edited

Loading

pytorch-bot bot commented Sep 29, 2024 •

edited

Loading

yf225 Sep 30, 2024 •

edited

Loading

xmfan Sep 30, 2024 •

edited

Loading

yf225 Sep 30, 2024 •

edited

Loading