[FSDP] Override named_parameters() for clean names in summon_full_params()
#74333
Conversation
[ghstack-poisoned]
CI Flow Status⚛️ CI FlowRuleset - Version:
|
🔗 Helpful links
💊 CI failures summary and remediationsAs of commit acd6186 (more details on the Dr. CI page): 💚 💚 Looks good so far! There are no failures yet. 💚 💚 This comment was automatically generated by Dr. CI (expand for details).Please report bugs/suggestions to the (internal) Dr. CI Users group. |
facebook-github-bot
added
the
oncall: distributed
[ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
awgu
requested review from
H-Huang,
mingzhe09088,
mrshenli,
pritamdamania87,
rohan-varma and
zhaojuanmao
as code owners
| """ | ||
| # Monkey patch `named_parameters()` | ||
| torch_named_parameters = torch.nn.Module.named_parameters | ||
| self.named_parameters = self._fsdp_named_parameters # type: ignore[assignment] |
There was a problem hiding this comment.
Adding the type: ignore[assignment] seems to be the best solution. See python/mypy#2427.
Fixes #73890 by monkey patching `torch.nn.Module.named_paramteters()`. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
| List, | ||
| Optional, | ||
| Generator, | ||
| Iterator, |
|
To-Do: look into overriding |
Fixes #73890 by monkey patching `torch.nn.Module.named_paramteters()`. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
| context manager. | ||
| """ | ||
| # Determine which logic to use based on the context at call time | ||
| if not hasattr(self, "training_state") or \ |
There was a problem hiding this comment.
nn.Modules contained in an FSDP instance may not have training_state as an attribute but still return True for isinstance(module, FullyShardedDataParallel). If you guys think this hasattr() check is too hacky, I will see if I can find a more direct solution.
There was a problem hiding this comment.
nit: this can be simplified to if getattr(self, "training_state", None) != TrainingState_.SUMMON_FULL_PARAMS:
Fixes #73890 by monkey patching `torch.nn.Module.named_paramteters()`. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
|
|
||
| def named_parameters( | ||
| self, | ||
| prefix: str = "", |
There was a problem hiding this comment.
Let's take in *args, **kwargs so we don't have to change if the API changes?
There was a problem hiding this comment.
Would it be correct to say that PyTorch Core cannot add positional arguments before prefix without breaking backward compatibility? In that case, could we only add **kwargs?
There was a problem hiding this comment.
Using (*args, **kwargs) is kind of the convention in Python to pass arguments to the parent method without any modification.
There was a problem hiding this comment.
Oh, I misunderstood. I thought we were trying to additionally pass in *args, **kwargs instead of replacing prefix and recurse.
Fixed this now.
…ummon_full_params()`" Fixes #73890. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
…ummon_full_params()`" Fixes #73890. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
…ummon_full_params()`" Fixes #73890. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
…ummon_full_params()`" Fixes #73890. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
rohan-varma
left a comment
rohan-varma
left a comment
There was a problem hiding this comment.
Looks good, just 2 minor q's for your consideration, will stamp after that
| # Remove any instances of the FSDP-specific prefix; there can | ||
| # be multiple in the case of nested FSDP modules | ||
| param_name = param_name.replace(FSDP_PREFIX, "") | ||
| yield (param_name, param) |
There was a problem hiding this comment.
Can we do the following to avoid the duplicated for loop?
in_summon = (training_state == summon_full_params)
for n, p in named_parameters():
name = name.replace(...) if in_summon else name
yield (name, p)
| """ | ||
| # Determine which logic to use based on the context at call time | ||
| if getattr(self, "training_state", None) != TrainingState_.SUMMON_FULL_PARAMS: | ||
| for param_name, param in torch.nn.Module.named_parameters( |
There was a problem hiding this comment.
Wonder if we can use super().named_parameters rather than torch.nn.Module.named_parameters call? If user writes MyModule that inherits from nn.Module with a custom named_parameters will this work?
…ummon_full_params()`" Fixes #73890. Differential Revision: [D34937201](https://our.internmc.facebook.com/intern/diff/D34937201) [ghstack-poisoned]
|
@awgu has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
|
Lgtm |
|
Hey @awgu. |
Stack from ghstack:
named_parameters()for clean names insummon_full_params()#74333 [FSDP] Overridenamed_parameters()for clean names insummon_full_params()Fixes #73890.
Differential Revision: D34937201