Add ReplicatedParameter and ReplicatedTensor to DDP. #74787
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) [ghstack-poisoned]
pritamdamania87
requested review from
H-Huang,
mingzhe09088,
mrshenli,
rohan-varma and
zhaojuanmao
as code owners
Contributor
🔗 Helpful links
💊 CI failures summary and remediationsAs of commit b248f29 (more details on the Dr. CI page):
🕵️ 8 new failures recognized by patternsThe following CI failures do not appear to be due to upstream breakages:
|
| Job | Step | Action |
|---|---|---|
 pull / pytorch-xla-linux-bionic-py3.7-clang8 / test (xla, 1, 1, linux.2xlarge) |
Test | 🔁 rerun |
| Lint / flake8-py3 |
Fail if there were any warnings | 🔁 rerun |
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
pritamdamania87
pushed a commit
that referenced
this pull request
Mar 26, 2022
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) ghstack-source-id: 152286552 Pull Request resolved: #74787
pritamdamania87
commented
Mar 26, 2022
|
|
||
| return True | ||
|
|
||
| class ReplicatedParameter(ReplicatedTensor, torch.nn.Parameter): |
Contributor
Author
There was a problem hiding this comment.
@albanD I couldn't pass in a ReplicatedTensor to torch.nn.Parameter and have isinstance() still return ReplicatedTensor. As a result, I had to do something messy like this. Not sure if there is a better way to do this.
Contributor
There was a problem hiding this comment.
This is a bit delayed due to my medical leave, but wanted to leave an update here for posterity: #73459 is in now - please try it out if you'd like :)
mrshenli
reviewed
Mar 28, 2022
Contributor
mrshenli
left a comment
mrshenli
left a comment
There was a problem hiding this comment.
Is this a test or do we plan to land this soon? If it's the latter, shall we do some micro-benchmarks to make sure there is no perf regression?
kumpera
reviewed
Mar 28, 2022
| ) | ||
|
|
||
| # After syncing, tag them as replicated. | ||
| from torch.distributed._shard.replicated_tensor import ReplicatedTensor, ReplicatedParameter |
Contributor
There was a problem hiding this comment.
Shouldn't we check if each param or buffer is replicated before converting them?
Similarly, shouldn't we record which ones we wrapped and use that when unwrapping, instead of relying on parameters_to_ignore.
Another thing, is using parameters_to_ignore to filter buffers just a naming artifact?
Contributor
Author
There was a problem hiding this comment.
Similarly, shouldn't we record which ones we wrapped and use that when unwrapping, instead of relying on parameters_to_ignore.
So parameters_to_ignore is slightly orthogonal to tagging something as replicated or not since that option is actually controlling which parameters DDP ignores completely. That is why we always need to rely on it and ensure we don't tag parameters_to_ignore as Replicated at all.
Another thing, is using parameters_to_ignore to filter buffers just a naming artifact?
It's an unfortunate name since it is filtering buffers as well in other places. I'll probably fix this name in a follow up PR.
Contributor
There was a problem hiding this comment.
Oops, I added self.parameters_to_ignore a couple of years ago and agree that a better name is parameters_and_buffers_to_ignore.
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) [ghstack-poisoned]
pritamdamania87
changed the title
pritamdamania87
changed the title
Contributor
Author
|
@mrshenli The plan is to hopefully land this change, however I think there are some tricky cases to take care of first. I'll probably ensure I can make CI happy first (tagged the PR as WIP) and then try a benchmark as well to ensure no regressions. |
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) [ghstack-poisoned]
pritamdamania87
pushed a commit
that referenced
this pull request
Mar 30, 2022
Pull Request resolved: #74787 As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. ghstack-source-id: 152544758 Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/)
Contributor
Author
|
@pytorchbot retest this please |
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) [ghstack-poisoned]
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) [ghstack-poisoned]
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) [ghstack-poisoned]
As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/) [ghstack-poisoned]
pritamdamania87
pushed a commit
that referenced
this pull request
Mar 30, 2022
Pull Request resolved: #74787 As per the design in #72138, introduce ReplicatedParameter and replace DDP parameters with ReplicatedParameters and DDP buffers as ReplicatedTensors. ghstack-source-id: 152632043 Differential Revision: [D35162813](https://our.internmc.facebook.com/intern/diff/D35162813/)
rohan-varma
reviewed
Apr 6, 2022
Contributor
rohan-varma
left a comment
rohan-varma
left a comment
There was a problem hiding this comment.
Are we considering enabling this change with a flag such as DDP_USE_REPLICATED to phase the rollout? I would advocate for this unless we're very confident about performance implications and reliability of ReplicatedTensor, as it would also have a large blast radius in OSS, and we may not be completely aware of it until the 1.12 release happens.
| self.assertIsInstance(out, ReplicatedTensor) | ||
|
|
||
| # Validate buffers not replicated after forward. | ||
| self.assertNotIsInstance(ddp.module.foo, ReplicatedTensor) |
Contributor
There was a problem hiding this comment.
Don't we synchronize the buffers so it should be replicated?
| params = list(model.parameters()) | ||
| self.assertEqual(2, len(params)) | ||
| self.assertEqual(model.weight, params[0]) | ||
| self.assertEqual(model.bias, params[1]) |
Contributor
There was a problem hiding this comment.
should we add self.assertTrue(all(isinstance(replicatedParam, param) for param in model.parameters())?
| self.gradient_as_bucket_view = gradient_as_bucket_view | ||
| if hasattr(module, "_ddp_params_and_buffers_to_ignore"): | ||
| self.parameters_to_ignore = module._ddp_params_and_buffers_to_ignore | ||
| if hasattr(self.module, "_ddp_params_and_buffers_to_ignore"): |
Contributor
There was a problem hiding this comment.
Curious why this change is needed?
| # In debug mode, build a mapping of parameter index -> parameter. | ||
| param_to_name_mapping = self._build_debug_param_to_name_mapping(parameters) | ||
| # Sync params and buffers. Ensures all DDP models start off at the same value. | ||
| self._sync_params_and_buffers(authoritative_rank=0) |
Contributor
There was a problem hiding this comment.
Shall we only mark this after the sync completes successfully?
| if not buffers_only: | ||
| for param_name, param in module.named_parameters(recurse=False): | ||
| if (param.requires_grad and f'{module_name}.{param_name}' | ||
| not in self.parameters_to_ignore and not isinstance(param, ReplicatedParameter)): |
Contributor
There was a problem hiding this comment.
does the latter check mean that it is possible to call this function on already replicated parameters? Should we try to make it so that the application doesn't do this because it would be unnecessary?
| ) | ||
|
|
||
| # After syncing, tag them as replicated. | ||
| from torch.distributed._shard.replicated_tensor import ReplicatedTensor, ReplicatedParameter |
Contributor
There was a problem hiding this comment.
Oops, I added self.parameters_to_ignore a couple of years ago and agree that a better name is parameters_and_buffers_to_ignore.
| for param_name, param in module.named_parameters(recurse=False): | ||
| if (param.requires_grad and f'{module_name}.{param_name}' | ||
| not in self.parameters_to_ignore and not isinstance(param, ReplicatedParameter)): | ||
| module.register_parameter(param_name, ReplicatedParameter(param)) |
Contributor
There was a problem hiding this comment.
Do we not de-register the original, non-replicated parameters?
Contributor
Author
|
Closing in favor of: #75299 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
As per the design in #72138,
introduce ReplicatedParameter and replace DDP parameters with
ReplicatedParameters and DDP buffers as ReplicatedTensors.
Differential Revision: D35162813