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

awgu · 2023-07-11T16:00:22Z

Stack from ghstack (oldest at bottom):

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:

Add shared parameter validation
(Longer-term / exploratory) Add a policy that provides a reasonable auto wrapping with "minimal" user input

[ghstack-poisoned]

pytorch-bot · 2023-07-11T16:00:25Z

🔗 Helpful Links

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

📄 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 Pending, 1 Unrelated Failure

As of commit 3fdf50f:

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: 2a97189df60fc83b9310f08928cfe3df0198297e Pull Request resolved: #104986

This PR adds a new `LambdaWrapPolicy` 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 = LambdaWrapPolicy(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 [ghstack-poisoned]

ghstack-source-id: aaeb5aa27683ba2e349bd1926345c2a6bdf8ea9b Pull Request resolved: #104986

ghstack-source-id: aaeb5aa27683ba2e349bd1926345c2a6bdf8ea9b Pull Request resolved: pytorch#104986

This PR adds a new `LambdaWrapPolicy` 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 = LambdaWrapPolicy(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 [ghstack-poisoned]

ghstack-source-id: 316176d9ee4b5903c8d0bf860e2dcf3417fb5f0c Pull Request resolved: #104986

ghstack-source-id: 316176d9ee4b5903c8d0bf860e2dcf3417fb5f0c Pull Request resolved: pytorch#104986

This PR adds a new `LambdaWrapPolicy` 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 = LambdaWrapPolicy(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 [ghstack-poisoned]

ghstack-source-id: 6dc615032a8544d4e70666c7587dad2f88adb489 Pull Request resolved: #104986

ghstack-source-id: 6dc615032a8544d4e70666c7587dad2f88adb489 Pull Request resolved: pytorch#104986

This PR adds a new `LambdaWrapPolicy` 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 = LambdaWrapPolicy(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 [ghstack-poisoned]

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: #104986

rohan-varma

Looks great overall, thanks for adding this! Just have a few questions.

rohan-varma · 2023-07-13T00:23:27Z

test/distributed/_composable/fully_shard/test_fully_shard_init.py

+            if isinstance(module, nn.Sequential):
+                return True
+            elif isinstance(module, FakeSequential):
+                return {"backward_prefetch": BackwardPrefetch.BACKWARD_POST}


should we add a test to ensure we raise a good error when we return an invalid key? like {"foo": 1}

I think one option here is to use inspect to get the arg signature of the module-level API passed to the policy. However, I think the default error message from passing in an unexpected arg/kwarg is not unreasonable, so I will not worry about this change in this PR:

TypeError: __init__() got an unexpected keyword argument 'foo'

rohan-varma · 2023-07-13T00:25:08Z

test/distributed/fsdp/test_wrap.py

+
+        if use_uniform_kwargs:
+
+            def lambda_fn(module: nn.Module):


does LambdaPolicy still get overriden and individually wrap BatchNorms if mixed precision is configured?

In general, what happens if my LambdaPolicy tells to wrap a BN, but mixed precision overrides this? What should be the expected behavior?

Currently, our mixed precision overrides run last. I am not sure if we should special case the LambdaWrapPolicy to be highest priority. What do you think?

I was leaning toward just keeping our overrides as absolute for now.

test/distributed/fsdp/test_wrap.py

rohan-varma · 2023-07-13T00:27:21Z

test/distributed/fsdp/test_wrap.py


-    def _test_frozen_params(self, use_orig_params: bool):
+    def _test_frozen_params(self, use_orig_params: bool, policy: _WrapPolicy):


maybe add a test for the mixed precision overriding as well?

rohan-varma · 2023-07-13T00:27:46Z

torch/distributed/fsdp/wrap.py

@@ -209,6 +211,48 @@ def __repr__(self) -> str:
        return super().__repr__() + f"({self._module_classes_str})"


+class LambdaWrapPolicy(_WrapPolicy):


nit: add an example?

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: pytorch#104986

ezyang · 2023-07-28T01:37:51Z

torch/distributed/fsdp/wrap.py

@@ -209,6 +211,48 @@ def __repr__(self) -> str:
        return super().__repr__() + f"({self._module_classes_str})"


+class LambdaWrapPolicy(_WrapPolicy):


You solicited some name bikeshedding over WC, so I'll try my best here. Assuming we are keeping WrapPolicy suffix, CallbackWrapPolicy or CustomWrapPolicy both convey intent here reasonably well.

You could also try PredicateWrapPolicy. This is not a perfect match: I would call lambda_fn a predicate if it only returned a bool, but you return a little bit extra information sometimes.

ezyang · 2023-07-28T01:39:40Z

torch/distributed/fsdp/wrap.py

+
+    def _run_policy(
+        self,
+        root_module: nn.Module,


OOC, is there a way to easily determine the FQN of the module at this point? Without the root module it seems like it might be difficult, whereas, the code that calls _run_policy would be able to generate it without much difficulty.

That is a good point. I do think we need access to the root module to be able to generate FQNs ("we" encompasses both in FSDP and for users).

In a downstream package I maintain, I've been allowing users to provide module FQN-based complements to a given auto_wrap_policy using a custom _FSDPPolicy (NameDrivenPolicy). I should be able to continue offering that functionality with some minor refactoring after this PR, but I thought it might be worth mentioning in the context of this PR because:

To call attention to this pattern since I've received positive user feedback regarding the utility FQN-based wrapping (to accommodate more flexible fine-tuning requirements in this case)

In case the team considers the functionality sufficiently useful to offer it natively at some point in the future (I'm not necessarily advocating that, just thinking aloud... either way this PR should help make such extensions cleaner)

Anyway, nice enhancement @awgu! A pleasure as always working with this thoughtfully designed API and constantly improving package.

Thanks for the pointer! This is super helpful, especially learning more about the customer feedback.

I am thinking that CustomPolicy can stay generic, and we can push specific logic to the lambda_fn construction. If the user wants to use FQNs, then the user can use the same trick as your implementation where you translate FQNs to Python id()s and then use the id()s in the lambda_fn, which will see the module as argument. Like you said, after some refactoring, you should still be able to make use of CustomPolicy. The question is whether we should offer some API that permits FQNs directly. I think that is something that we can wait and see. For now, we can just add CustomPolicy that can serve as a base for other extensions.

What do you think?

The question is whether we should offer some API that permits FQNs directly. I think that is something that we can wait and see. For now, we can just add CustomPolicy that can serve as a base for other extensions.

Agree completely, I think prioritizing API stability and simplicity at this juncture is the prudent approach. Adding an API that permits FQNs directly might be worth considering at some point in the future but that feature isn't sufficiently crucial to justify its inclusion at present (especially since implementation approaches like my id() translation one we discussed aren't too onerous).

Thanks for the quick and thoughtful response!

ezyang

Seems reasonable.

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: pytorch#104986

This PR adds a new `LambdaWrapPolicy` 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 = LambdaWrapPolicy(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 [ghstack-poisoned]

ghstack-source-id: 16789fae18bb459d11a17988761b6fb754c2a91c Pull Request resolved: #104986

This PR adds a new `LambdaWrapPolicy` 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 = LambdaWrapPolicy(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 [ghstack-poisoned]

ghstack-source-id: f2047acf56ad947009524531f997efad6fe5914a Pull Request resolved: #104986

ghstack-source-id: f2047acf56ad947009524531f997efad6fe5914a Pull Request resolved: pytorch#104986

This PR adds a new `LambdaWrapPolicy` 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 = LambdaWrapPolicy(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 [ghstack-poisoned]

ghstack-source-id: 3d3413d4209481f1c7248eead8a9ad92115a39a5 Pull Request resolved: #104986

awgu · 2023-08-03T12:43:30Z

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

awgu · 2023-08-03T12:43:39Z

@pytorchbot merge -f "unrelated failures"

pytorchmergebot · 2023-08-03T12:46:28Z

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

[FSDP][9/N] Introduce LambdaWrapPolicy

f671536

[ghstack-poisoned]

awgu mentioned this pull request Jul 11, 2023

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

Closed

awgu mentioned this pull request Jul 11, 2023

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

Closed

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

This was referenced Jul 11, 2023

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

Closed

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

Closed

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

[FSDP][9/N] Introduce LambdaWrapPolicy

9b4d944

ghstack-source-id: 2a97189df60fc83b9310f08928cfe3df0198297e Pull Request resolved: #104986

awgu added the topic: improvements topic category label Jul 11, 2023

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

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

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

[FSDP][9/N] Introduce LambdaWrapPolicy

676fb45

ghstack-source-id: aaeb5aa27683ba2e349bd1926345c2a6bdf8ea9b Pull Request resolved: #104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

48b8487

ghstack-source-id: aaeb5aa27683ba2e349bd1926345c2a6bdf8ea9b Pull Request resolved: pytorch#104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

60b6cea

ghstack-source-id: 316176d9ee4b5903c8d0bf860e2dcf3417fb5f0c Pull Request resolved: #104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

d70470c

ghstack-source-id: 316176d9ee4b5903c8d0bf860e2dcf3417fb5f0c Pull Request resolved: pytorch#104986

awgu mentioned this pull request Jul 11, 2023

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

Closed

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

[FSDP][9/N] Introduce LambdaWrapPolicy

2594969

ghstack-source-id: 6dc615032a8544d4e70666c7587dad2f88adb489 Pull Request resolved: #104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

059cacb

ghstack-source-id: 6dc615032a8544d4e70666c7587dad2f88adb489 Pull Request resolved: pytorch#104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

fb9e57e

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: #104986

rohan-varma reviewed Jul 13, 2023

View reviewed changes

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

[FSDP][9/N] Introduce LambdaWrapPolicy

c6207d5

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: pytorch#104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

f4d599c

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: pytorch#104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

104111e

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: pytorch#104986

ezyang reviewed Jul 28, 2023

View reviewed changes

ezyang approved these changes Jul 28, 2023

View reviewed changes

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

[FSDP][9/N] Introduce LambdaWrapPolicy

fc1f810

ghstack-source-id: 4ac5da42fcdcc1b44ed8d9fb6677295a3fb9d432 Pull Request resolved: pytorch#104986

awgu changed the title ~~[FSDP][9/N] Introduce LambdaWrapPolicy~~ [FSDP][9/N] Introduce CustomPolicy Aug 2, 2023

awgu requested review from penguinwu and yhcharles as code owners August 2, 2023 15:26

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

[FSDP][9/N] Introduce LambdaWrapPolicy

ccd3e42

ghstack-source-id: 16789fae18bb459d11a17988761b6fb754c2a91c Pull Request resolved: #104986

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

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

[FSDP][9/N] Introduce LambdaWrapPolicy

adb5157

ghstack-source-id: f2047acf56ad947009524531f997efad6fe5914a Pull Request resolved: #104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

bbb598d

ghstack-source-id: f2047acf56ad947009524531f997efad6fe5914a Pull Request resolved: pytorch#104986

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

[FSDP][9/N] Introduce LambdaWrapPolicy

bd9c35b

ghstack-source-id: 3d3413d4209481f1c7248eead8a9ad92115a39a5 Pull Request resolved: #104986

pytorchmergebot added the merging label Aug 3, 2023

pytorchmergebot added Merged and removed merging labels Aug 3, 2023

pytorchmergebot closed this in 57fba6f Aug 3, 2023

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

awgu commented Jul 11, 2023 •

edited

pytorch-bot bot commented Jul 11, 2023 •

edited

rohan-varma left a comment

rohan-varma Jul 13, 2023

awgu Aug 2, 2023

rohan-varma Jul 13, 2023

awgu Jul 13, 2023

rohan-varma Jul 13, 2023

awgu Aug 2, 2023

rohan-varma Jul 13, 2023

awgu Aug 2, 2023

ezyang Jul 28, 2023

ezyang Jul 28, 2023

ezyang Jul 28, 2023

awgu Aug 2, 2023

speediedan Aug 2, 2023

awgu Aug 2, 2023

speediedan Aug 3, 2023

ezyang left a comment

awgu commented Aug 3, 2023

awgu commented Aug 3, 2023

pytorchmergebot commented Aug 3, 2023


		def _test_frozen_params(self, use_orig_params: bool):
		def _test_frozen_params(self, use_orig_params: bool, policy: _WrapPolicy):

		@@ -209,6 +211,48 @@ def __repr__(self) -> str:
		return super().__repr__() + f"({self._module_classes_str})"


		class LambdaWrapPolicy(_WrapPolicy):

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

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

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

⏳ 1 Pending, 1 Unrelated Failure

rohan-varma left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ezyang left a comment

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][9/N] Introduce `CustomPolicy` #104986

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

awgu commented Jul 11, 2023 •

edited

pytorch-bot bot commented Jul 11, 2023 •

edited