[FSDP][8/N] Replace `_FSDPPolicy.policy` with `_Policy._run_policy` #104969

awgu · 2023-07-11T13:58:39Z

Stack from ghstack (oldest at bottom):

This does some code organization improvement.

It renames _FSDPPolicy to _Policy to show that it is not only for FSDP but for any module-level API.
It formalizes the contract that such a policy should return something like target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]] that maps each module to wrap to its kwargs. It does so by requiring a _run_policy abstract method (this time private since users do not need to care about it). Then, our auto wrapping can just call _run_policy() to generate the dict and do any validation or post-processing.

This PR is technically BC-breaking because it removes the public ModuleWrapPolicy.policy. However, I do not think anyone was using that anyway, so this is a pretty safe breakage.

[ghstack-poisoned]

pytorch-bot · 2023-07-11T13:58:42Z

🔗 Helpful Links

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

📄 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 Unrelated Failure

As of commit ad2e8ce:

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

linux-focal-rocm5.6-py3.8 / test (default, 1, 3, linux.rocm.gpu, unstable) (gh)

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

ghstack-source-id: c41660b85e42e6de90e78d12b859a9961ddd703e Pull Request resolved: #104969

[ghstack-poisoned]

ghstack-source-id: 53a5622bc1a3603308b3d266e5840a730b6247ae Pull Request resolved: #104969

[ghstack-poisoned]

ghstack-source-id: 783d09ccc43c84621f91e801159db877b508b541 Pull Request resolved: #104969

[ghstack-poisoned]

torch/distributed/fsdp/wrap.py

ghstack-source-id: a5fcf40d6e6d5f7db475bf0b60d3cfa3a42ee3d3 Pull Request resolved: pytorch#104969

This is a code organization improvement. It formalizes the contract that an FSDP policy should return something like `target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]]` that maps each module to wrap to its FSDP kwargs. It does so by requiring a `_run_policy` abstract method (this time private since users do not need to care about it). Then, our auto wrapping can just call `_run_policy()` to generate the dict and do any validation or post-processing. This PR is technically BC-breaking because it removes the public `ModuleWrapPolicy.policy`. However, I do not think anyone was using that anyway, so this is a pretty safe breakage. [ghstack-poisoned]

ghstack-source-id: 86c9da70764481465d83b24d69224561d1f83fc3 Pull Request resolved: pytorch#104969

This is a code organization improvement. - It renames `_FSDPPolicy` to `_WrapPolicy` to show that it is not only for FSDP but for any module-level API. - It formalizes the contract that such a policy should return something like `target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]]` that maps each module to wrap to its kwargs. It does so by requiring a `_run_policy` abstract method (this time private since users do not need to care about it). Then, our auto wrapping can just call `_run_policy()` to generate the dict and do any validation or post-processing. This PR is technically BC-breaking because it removes the public `ModuleWrapPolicy.policy`. However, I do not think anyone was using that anyway, so this is a pretty safe breakage. [ghstack-poisoned]

ghstack-source-id: 8e8b698c212c8c5e8269641ec5c19c0498488878 Pull Request resolved: pytorch#104969

This does some code organization improvement. - It renames `_FSDPPolicy` to `_WrapPolicy` to show that it is not only for FSDP but for any module-level API. - It formalizes the contract that such a policy should return something like `target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]]` that maps each module to wrap to its kwargs. It does so by requiring a `_run_policy` abstract method (this time private since users do not need to care about it). Then, our auto wrapping can just call `_run_policy()` to generate the dict and do any validation or post-processing. This PR is technically BC-breaking because it removes the public `ModuleWrapPolicy.policy`. However, I do not think anyone was using that anyway, so this is a pretty safe breakage. [ghstack-poisoned]

rohan-varma · 2023-07-13T00:22:01Z

torch/distributed/fsdp/wrap.py

+                continue
+            elif isinstance(module, module_classes):
+                # Shallow copy to avoid coupling changes across modules
+                target_module_to_kwargs[module] = copy.copy(root_kwargs)


why do we need to copy the args? Do we expect root_kwargs to change?

We must copy the kwargs or else if a function overrides it (e.g. the mixed precision overrides), then it will override it for every target module. We have to avoid them sharing a reference to the same kwarg dict.

ghstack-source-id: 3c5d71e48d07274d6ed3e3f791d300427849d2da Pull Request resolved: pytorch#104969

…n_policy`" This does some code organization improvement. - It renames `_FSDPPolicy` to `_Policy` to show that it is not only for FSDP but for any module-level API. - It formalizes the contract that such a policy should return something like `target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]]` that maps each module to wrap to its kwargs. It does so by requiring a `_run_policy` abstract method (this time private since users do not need to care about it). Then, our auto wrapping can just call `_run_policy()` to generate the dict and do any validation or post-processing. This PR is technically BC-breaking because it removes the public `ModuleWrapPolicy.policy`. However, I do not think anyone was using that anyway, so this is a pretty safe breakage. [ghstack-poisoned]

ghstack-source-id: f0778e5f19edfb0a16e282f9cb133829badb591d Pull Request resolved: pytorch#104969

…n_policy`" This does some code organization improvement. - It renames `_FSDPPolicy` to `_Policy` to show that it is not only for FSDP but for any module-level API. - It formalizes the contract that such a policy should return something like `target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]]` that maps each module to wrap to its kwargs. It does so by requiring a `_run_policy` abstract method (this time private since users do not need to care about it). Then, our auto wrapping can just call `_run_policy()` to generate the dict and do any validation or post-processing. This PR is technically BC-breaking because it removes the public `ModuleWrapPolicy.policy`. However, I do not think anyone was using that anyway, so this is a pretty safe breakage. [ghstack-poisoned]

awgu · 2023-08-03T12:40:05Z

Unstable test and unrelated:
linux-focal-rocm5.6-py3.8 / test (default, 1, 3, linux.rocm.gpu, unstable)

awgu · 2023-08-03T12:40:16Z

@pytorchbot merge -f "unrelated failure"

pytorchmergebot · 2023-08-03T12:42:07Z

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

This PR adds a new `CustomPolicy` that acts like the existing `lambda_auto_wrap_policy` except it (1) leverages the new auto wrapping infrastructure and (2) allows overriding FSDP kwargs for particular instances. (1) gives it access to the validation checks (like for frozen parameters), and (2) makes it as expressive as manual wrapping. This should allow us to effectively deprecate manual wrapping if desired. The API is as follows: ``` def lambda_fn(module: nn.Module) -> Union[bool, Dict[str, Any]]: ... policy = CustomPolicy(lambda_fn) ``` The `lambda_fn` can return: - `False` or `{}` to indicate no wrapping - `True` to indicate wrapping while inheriting the root's FSDP kwargs - Non-empty `dict` to indicate wrapping while overriding the specified FSDP kwargs and inheriting the rest from the root --- After this PR, the follow-up work items for auto wrapping are: 1. Add shared parameter validation 2. (Longer-term / exploratory) Add a policy that provides a reasonable auto wrapping with "minimal" user input Pull Request resolved: #104986 Approved by: https://github.com/ezyang ghstack dependencies: #104427, #104967, #104999, #104969

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

726aba0

[ghstack-poisoned]

awgu mentioned this pull request Jul 11, 2023

[FSDP][Easy] Rename streams; add back stream sharing test #104966

Closed

pytorch-bot bot added the release notes: distributed (fsdp) release notes category label Jul 11, 2023

This was referenced Jul 11, 2023

[FSDP][6/N] Check valid param freezing for ModuleWrapPolicy #104427

Closed

[FSDP][7/N] Add warning about frozen params #104967

Closed

awgu added a commit that referenced this pull request Jul 11, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

0e02a85

ghstack-source-id: c41660b85e42e6de90e78d12b859a9961ddd703e Pull Request resolved: #104969

Update on "[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy"

4b9867c

[ghstack-poisoned]

awgu added a commit that referenced this pull request Jul 11, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

7f7c9be

ghstack-source-id: 53a5622bc1a3603308b3d266e5840a730b6247ae Pull Request resolved: #104969

awgu added the suppress-bc-linter Suppresses the failures of API backward-compatibility linter (Lint/bc_linter) label Jul 11, 2023

Update on "[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy"

c37de05

[ghstack-poisoned]

awgu added a commit that referenced this pull request Jul 11, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

b5392ff

ghstack-source-id: 783d09ccc43c84621f91e801159db877b508b541 Pull Request resolved: #104969

Update on "[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy"

d0dba7c

[ghstack-poisoned]

awgu mentioned this pull request Jul 11, 2023

[FSDP][9/N] Introduce CustomPolicy #104986

Closed

awgu changed the title ~~[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy~~ [FSDP][8/N] Replace _FSDPPolicy.policy with _run_policy Jul 11, 2023

awgu added the topic: not user facing topic category label Jul 11, 2023

awgu commented Jul 11, 2023

View reviewed changes

torch/distributed/fsdp/wrap.py Show resolved Hide resolved

awgu marked this pull request as ready for review July 11, 2023 16:04

awgu requested review from mrshenli, zhaojuanmao, rohan-varma, H-Huang, kwen2501, wanchaol, fegin, fduwjj, kiukchung and d4l3k as code owners July 11, 2023 16:04

awgu added a commit to awgu/pytorch that referenced this pull request Jul 11, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

7177b25

ghstack-source-id: a5fcf40d6e6d5f7db475bf0b60d3cfa3a42ee3d3 Pull Request resolved: pytorch#104969

awgu requested a review from yhcharles as a code owner July 11, 2023 19:52

awgu added a commit to awgu/pytorch that referenced this pull request Jul 11, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

028076b

ghstack-source-id: 86c9da70764481465d83b24d69224561d1f83fc3 Pull Request resolved: pytorch#104969

awgu mentioned this pull request Jul 11, 2023

[FSDP][Easy] Allow ModuleWrapPolicy to take Iterable #104999

Closed

awgu added a commit to awgu/pytorch that referenced this pull request Jul 11, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

227cb2c

ghstack-source-id: 8e8b698c212c8c5e8269641ec5c19c0498488878 Pull Request resolved: pytorch#104969

awgu added a commit to awgu/pytorch that referenced this pull request Jul 11, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

9bc2a0a

ghstack-source-id: 8e8b698c212c8c5e8269641ec5c19c0498488878 Pull Request resolved: pytorch#104969

rohan-varma approved these changes Jul 13, 2023

View reviewed changes

awgu added a commit to awgu/pytorch that referenced this pull request Jul 19, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

2df42b0

ghstack-source-id: 3c5d71e48d07274d6ed3e3f791d300427849d2da Pull Request resolved: pytorch#104969

awgu added a commit to awgu/pytorch that referenced this pull request Jul 21, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

7c6819d

ghstack-source-id: 3c5d71e48d07274d6ed3e3f791d300427849d2da Pull Request resolved: pytorch#104969

awgu added a commit to awgu/pytorch that referenced this pull request Jul 26, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

9c164c9

ghstack-source-id: 3c5d71e48d07274d6ed3e3f791d300427849d2da Pull Request resolved: pytorch#104969

awgu added a commit to awgu/pytorch that referenced this pull request Jul 31, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

deee7ed

ghstack-source-id: 3c5d71e48d07274d6ed3e3f791d300427849d2da Pull Request resolved: pytorch#104969

awgu changed the title ~~[FSDP][8/N] Replace _FSDPPolicy.policy with _run_policy~~ [FSDP][8/N] Replace _FSDPPolicy.policy with _Policy._run_policy Aug 2, 2023

awgu added 2 commits August 2, 2023 15:26

awgu added a commit to awgu/pytorch that referenced this pull request Aug 2, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with run_policy

c47a52f

ghstack-source-id: f0778e5f19edfb0a16e282f9cb133829badb591d Pull Request resolved: pytorch#104969

awgu added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 3, 2023

pytorchmergebot added the merging label Aug 3, 2023

pytorchmergebot added Merged and removed merging labels Aug 3, 2023

pytorchmergebot closed this in 15953fd Aug 3, 2023

facebook-github-bot deleted the gh/awgu/419/head branch August 6, 2023 14:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[FSDP][8/N] Replace `_FSDPPolicy.policy` with `_Policy._run_policy` #104969

[FSDP][8/N] Replace `_FSDPPolicy.policy` with `_Policy._run_policy` #104969

awgu commented Jul 11, 2023 •

edited

pytorch-bot bot commented Jul 11, 2023 •

edited

rohan-varma Jul 13, 2023

awgu Jul 13, 2023

awgu commented Aug 3, 2023

awgu commented Aug 3, 2023

pytorchmergebot commented Aug 3, 2023

[FSDP][8/N] Replace _FSDPPolicy.policy with _Policy._run_policy #104969

[FSDP][8/N] Replace _FSDPPolicy.policy with _Policy._run_policy #104969

Conversation

awgu commented Jul 11, 2023 • edited

pytorch-bot bot commented Jul 11, 2023 • edited

🔗 Helpful Links

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

✅ 1 Unrelated Failure

rohan-varma Jul 13, 2023

Choose a reason for hiding this comment

awgu Jul 13, 2023

Choose a reason for hiding this comment

awgu commented Aug 3, 2023

awgu commented Aug 3, 2023

pytorchmergebot commented Aug 3, 2023

Merge started

[FSDP][8/N] Replace `_FSDPPolicy.policy` with `_Policy._run_policy` #104969

[FSDP][8/N] Replace `_FSDPPolicy.policy` with `_Policy._run_policy` #104969

awgu commented Jul 11, 2023 •

edited

pytorch-bot bot commented Jul 11, 2023 •

edited