[fsdp][torch.compile] FSDP changes #115497

voznesenskym · 2023-12-10T07:39:56Z

Stack from ghstack (oldest at bottom):

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225

[ghstack-poisoned]

pytorch-bot · 2023-12-10T07:39:59Z

🔗 Helpful Links

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

📄 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! (1 Unrelated Failure)

As of commit 2c65d1a with merge base 87ea6fb ():

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

pull / linux-focal-py3.11-clang10 / test (dynamo, 2, 7, linux.2xlarge, unstable) (gh)
test_unary_ufuncs.py::TestUnaryUfuncsCPU::test_reference_numerics_large__refs_special_multigammaln_mvlgamma_p_1_cpu_bfloat16

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

[ghstack-poisoned]

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

torch/distributed/fsdp/_common_utils.py

aten/src/ATen/native/native_functions.yaml

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

albanD · 2023-12-18T14:37:43Z

torch/distributed/fsdp/_runtime_utils.py

@@ -441,7 +441,9 @@ def _pre_forward_unshard(
    if not handle._prefetched:
        _unshard(state, handle, state._unshard_stream, state._pre_unshard_stream)
    handle._needs_pre_forward_unshard = False
-    state._device_handle.current_stream().wait_stream(state._unshard_stream)
+    # Don't wait during trace


How is the trace valid if it contains stream information but not stream synchronization?
I guess we don't care for now since it's ignored by the backend but most likely will need fixing.

Yes, this is sound for now - but the longer term version of this will need to teach inductor about .wait_stream and plumb the whole thing through.

Hi, @voznesenskym @albanD

Yes, this is sound for now - but the longer term version of this will need to teach inductor about .wait_stream and plumb the whole thing through.

May i know why here bypassing the stream methods? Does it mean for the longer term plan, the stream and other runtime primitives, like event, will be compiled through the torch.compile(dynamo+AOTAutograd+inductor)?
May i know if there is a plan for supporting stream/event in whole torch.compile stack?

Thank you. :)

albanD · 2023-12-18T14:38:27Z

torch/distributed/fsdp/_runtime_utils.py

+        and hasattr(handle, "_ran_pre_backward_hook")
+        and handle._ran_pre_backward_hook
+    ):
+        log.warning("%s %s", id(state), "Not Running pre backward! Already Ran!")


Should this really be a warning?

Debug is better

albanD · 2023-12-18T14:39:13Z

torch/distributed/utils.py

@@ -327,3 +333,7 @@ def _replace_by_prefix(
        new_key = new_prefix + key[len(old_prefix) :]
        state_dict[new_key] = state_dict[key]
        del state_dict[key]
+
+
+def _data_ptr_allocated(tensor: torch.Tensor) -> bool:


Is this missing an allow_in_graph? Or is it not used?

Doesn't need it for now - It will in the future, but this makes it easier for me to track where I will eventually need to add it.

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

albanD

SGTM!

voznesenskym · 2023-12-19T16:09:42Z

@pytorchbot merge

pytorchmergebot · 2023-12-19T16:11: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

Pull Request resolved: pytorch#115497 Approved by: https://github.com/albanD

flat_param_part_view is unused in pytorch repo: https://fburl.com/ssaomd7x it became unused since refactoring in pytorch#115497 before that, the original code is below. Since flat_param is 1D, we do not need .view for reshaping ``` self.flat_param.data = padded_unsharded_flat_param[ : unsharded_size.numel() ].view( unsharded_size ) ``` unit test: pytest test/distributed/fsdp/test_fsdp_core.py

flat_param_part_view is unused in pytorch repo: https://fburl.com/ssaomd7x it became unused since refactoring in #115497 before that, the original code is below. Since flat_param is 1D, we do not need .view for reshaping ``` self.flat_param.data = padded_unsharded_flat_param[ : unsharded_size.numel() ].view( unsharded_size ) ``` unit test: pytest test/distributed/fsdp/test_fsdp_core.py Pull Request resolved: #117082 Approved by: https://github.com/awgu, https://github.com/wconstab, https://github.com/Skylion007

[fsdp][torch.compile] FSDP changes

887f38a

[ghstack-poisoned]

pytorch-bot bot added the release notes: distributed (ddp) release notes category label Dec 10, 2023

voznesenskym mentioned this pull request Dec 10, 2023

Consider storage_changed for assigning alias_of_input in aot_autograd when computing differentiable outputs that alias each other #115315

Closed

This was referenced Dec 8, 2023

[fsdp] Replace acc_grad hooking with register_post_accumulate_grad_hook on flat_param #112184

Closed

[Do not review] Top of FSDP stack - to be broken up #115410

Closed

github-actions bot added ciflow/inductor module: distributed labels Dec 10, 2023

github-actions bot requested review from albanD, antoniojkim, bdhirsh, ezyang, miladm and SherlockNoMad December 10, 2023 07:40

Update on "[fsdp][torch.compile] FSDP changes"

5c720b9

[ghstack-poisoned]

voznesenskym mentioned this pull request Dec 12, 2023

forward decl #115709

Closed

voznesenskym added 2 commits December 12, 2023 15:17

Update on "[fsdp][torch.compile] FSDP changes"

f5e799b

[ghstack-poisoned]

Update on "[fsdp][torch.compile] FSDP changes"

1219929

[ghstack-poisoned]

albanD added oncall: distributed Add this issue/PR to distributed oncall triage queue and removed module: distributed labels Dec 13, 2023

Update on "[fsdp][torch.compile] FSDP changes"

9f85a5a

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

github-actions bot added the module: distributed label Dec 13, 2023

voznesenskym added 3 commits December 13, 2023 12:48

Update on "[fsdp][torch.compile] FSDP changes"

1d7b904

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

Update on "[fsdp][torch.compile] FSDP changes"

8add212

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

Update on "[fsdp][torch.compile] FSDP changes"

03c6cf9

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

fegin added the ciflow/periodic Trigger jobs ran periodically on master (periodic.yml) on the PR label Dec 14, 2023

Update on "[fsdp][torch.compile] FSDP changes"

df34fe4

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

voznesenskym mentioned this pull request Dec 15, 2023

Support non grapharg and intermediary grad access #115898

Closed

Update on "[fsdp][torch.compile] FSDP changes"

e88b44d

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

albanD reviewed Dec 15, 2023

View reviewed changes

torch/distributed/fsdp/_common_utils.py Outdated Show resolved Hide resolved

aten/src/ATen/native/native_functions.yaml Outdated Show resolved Hide resolved

aten/src/ATen/native/native_functions.yaml Outdated Show resolved Hide resolved

voznesenskym added 3 commits December 15, 2023 14:42

Update on "[fsdp][torch.compile] FSDP changes"

973c59d

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

Update on "[fsdp][torch.compile] FSDP changes"

8507bac

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

Update on "[fsdp][torch.compile] FSDP changes"

29e4417

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

voznesenskym requested a review from albanD December 16, 2023 18:41

albanD reviewed Dec 18, 2023

View reviewed changes

voznesenskym added 2 commits December 18, 2023 14:22

Update on "[fsdp][torch.compile] FSDP changes"

5b84241

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

Update on "[fsdp][torch.compile] FSDP changes"

2c65d1a

cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]

albanD approved these changes Dec 18, 2023

View reviewed changes

albanD removed the module: distributed label Dec 19, 2023

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Dec 19, 2023

pytorchmergebot added the merging label Dec 19, 2023

pytorchmergebot added the Merged label Dec 19, 2023

pytorchmergebot removed the merging label Dec 19, 2023

pytorchmergebot closed this in 77d5f60 Dec 19, 2023

dmenig pushed a commit to dmenig/pytorch that referenced this pull request Dec 21, 2023

[fsdp][torch.compile] FSDP changes (pytorch#115497)

d95b6c5

Pull Request resolved: pytorch#115497 Approved by: https://github.com/albanD

facebook-github-bot deleted the gh/voznesenskym/294/head branch December 23, 2023 15:22

weifengpy mentioned this pull request Jan 10, 2024

[FSDP] remove unused flat_param_part_view #117082

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[fsdp][torch.compile] FSDP changes #115497

[fsdp][torch.compile] FSDP changes #115497

Uh oh!

voznesenskym commented Dec 10, 2023 •

edited

Loading

Uh oh!

pytorch-bot bot commented Dec 10, 2023 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

albanD Dec 18, 2023

Uh oh!

voznesenskym Dec 18, 2023

Uh oh!

zejun-chen Feb 29, 2024 •

edited

Loading

Uh oh!

albanD Dec 18, 2023

Uh oh!

voznesenskym Dec 18, 2023

Uh oh!

albanD Dec 18, 2023

Uh oh!

voznesenskym Dec 18, 2023

Uh oh!

albanD left a comment

Uh oh!

voznesenskym commented Dec 19, 2023

Uh oh!

pytorchmergebot commented Dec 19, 2023

Uh oh!

Uh oh!

[fsdp][torch.compile] FSDP changes #115497

[fsdp][torch.compile] FSDP changes #115497

Uh oh!

Conversation

voznesenskym commented Dec 10, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Dec 10, 2023 • 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/115497

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

albanD Dec 18, 2023

Choose a reason for hiding this comment

Uh oh!

voznesenskym Dec 18, 2023

Choose a reason for hiding this comment

Uh oh!

zejun-chen Feb 29, 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!

albanD Dec 18, 2023

Choose a reason for hiding this comment

Uh oh!

voznesenskym Dec 18, 2023

Choose a reason for hiding this comment

Uh oh!

albanD Dec 18, 2023

Choose a reason for hiding this comment

Uh oh!

voznesenskym Dec 18, 2023

Choose a reason for hiding this comment

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

voznesenskym commented Dec 19, 2023

Uh oh!

pytorchmergebot commented Dec 19, 2023

Merge started

Uh oh!

Uh oh!

voznesenskym commented Dec 10, 2023 •

edited

Loading

pytorch-bot bot commented Dec 10, 2023 •

edited

Loading

zejun-chen Feb 29, 2024 •

edited

Loading