[feat] FSDP: add auto_wrap_bn #531
Conversation
- add an utility function to handle wrapping of BN
facebook-github-bot
added
the
CLA Signed
| return not isinstance(module, tuple(default_auto_wrap_policy.FORCE_LEAF_MODULES)) # type: ignore | ||
| else: | ||
| return is_bn and not isinstance(module, tuple(default_auto_wrap_policy.EXCLUDE_WRAP_MODULES)) # type: ignore | ||
|
|
There was a problem hiding this comment.
nice, I find the next couple of lines (config, single group, then guided auto wrap) very elegant
There was a problem hiding this comment.
Quick question: could dist.new_group(ranks=[my_rank]) impacts performance in any ways ?
There was a problem hiding this comment.
no, should not really, AFAIK the overhead is minimal
There was a problem hiding this comment.
I don't think there should be any perf impact since FSDP has special casing for world_size == 1. But perhaps @myleott can think of something else?
There was a problem hiding this comment.
I understood the question being about the perf cost of having many groups in pytorch distributed basically, vs. few, not specific to FSDP. I might be wrong, but that was the reasoning behind my reply
| @@ -54,7 +45,16 @@ def forward(self, x): | |||
| # TODO (Min): check DDP equivalency. | |||
There was a problem hiding this comment.
having been burnt a little by that, I would recommend not waiting too long for that part
There was a problem hiding this comment.
definitely. see my plan below.
| return not isinstance(module, tuple(default_auto_wrap_policy.FORCE_LEAF_MODULES)) # type: ignore | ||
| else: | ||
| return is_bn and not isinstance(module, tuple(default_auto_wrap_policy.EXCLUDE_WRAP_MODULES)) # type: ignore | ||
|
|
There was a problem hiding this comment.
Quick question: could dist.new_group(ranks=[my_rank]) impacts performance in any ways ?
|
Thank you guys @blefaudeux @myleott @tchaton for quick and high quality reviews. To forecast a bit:
|
Before submitting
What does this PR do?
Fixes # (issue).
PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.
Did you have fun?
Make sure you had fun coding 🙃