Fix DDPOptimizer fake_mode execution #92986
Conversation
When running compiled submods for the purpose of producing outputs to pass to the compilation step for the next submod, we use fake parameters and assume fake inputs, but we forgot to activate our fake_mode during execution. This caused certain edge cases where tensors other than activations or parameters got created during execution, such as scalar->tensor expansion in the case of executing torch.where(tensor, scalar, scalar). Also add a test and clarify behavior of DDPOptimizer via comments. [ghstack-poisoned]
wconstab
requested review from
H-Huang,
awgu,
kwen2501,
mrshenli,
rohan-varma,
wanchaol and
zhaojuanmao
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/92986
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit cac8faa: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
When running compiled submods for the purpose of producing outputs to pass to the compilation step for the next submod, we use fake parameters and assume fake inputs, but we forgot to activate our fake_mode during execution. This caused certain edge cases where tensors other than activations or parameters got created during execution, such as scalar->tensor expansion in the case of executing torch.where(tensor, scalar, scalar). Also add a test and clarify behavior of DDPOptimizer via comments. ghstack-source-id: ce20aa8 Pull Request resolved: #92986
wconstab
requested review from
bdhirsh,
ezyang,
ngimel and
voznesenskym
and removed request for
H-Huang,
awgu,
kwen2501,
mrshenli,
rohan-varma,
wanchaol and
zhaojuanmao
|
@voznesenskym Semi-related to this PR, is the code that infers a FakeMode from input tensors or else creates a new one (https://github.com/pytorch/pytorch/blob/master/torch/_dynamo/optimizations/distributed.py#L148) good enough going forward, or do we want to align this FakeMode better with one that lives in dynamo? I'm not sure what the current design is. |
| # Finally, we have to produce inputs for use compiling the next submodule, | ||
| # and these need to be FakeTensors, so we execute the module under fake_mode | ||
| with fake_mode: | ||
| return curr_submod(*new_args, **kwargs) |
There was a problem hiding this comment.
This code is executed at runtime, right?
There was a problem hiding this comment.
not sure what you mean.
Since torch dynamo defers its compilation until the first execution, then in a way yes, this code happens "at runtime".
But this code only happens as a part of the compilation flow, which in a simple (static model) scenario only happens once. The second time a user calls their compiled ddp model, none of this code should run, since we're not recompiling.
There was a problem hiding this comment.
maybe you confused it with 'WrapperModule.forwrad' - that's the only piece of code in the whole `ddp_optimizer' file that I'd expect to run repeatedly on every runtime. (all it does is unwrap the tuple output from the compiled subgraph)
There was a problem hiding this comment.
yeah, wrappermodule.forward is the place I was thinking of.
This looks fine to me.
When Ed and I were working on it - it was very confusing which part of this was compile time, and which was runtime.
fake_mode = fake_mode_from_tensors(example_inputs) Is the right way to do things :) I added it for this purpose. The idea is that you do a best effort to get the current fake mode, and if there isn't one, you can make it. There's a few useful comments in that definition, one should say something along the lines of maybe having it always provide a fake_mode... |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
wconstab
changed the title
Stack from ghstack (oldest at bottom):
When running compiled submods for the purpose of producing outputs to pass
to the compilation step for the next submod, we use fake parameters and
assume fake inputs, but we forgot to activate our fake_mode during execution.
This caused certain edge cases where tensors other than activations or parameters
got created during execution, such as scalar->tensor expansion in the case
of executing torch.where(tensor, scalar, scalar).
Also add a test and clarify behavior of DDPOptimizer via comments.
Fixes #92941
cc @mlazos @soumith @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire