Handle nested outputs in NoParallel - #4288
Closed
sylvesterkaczmarek wants to merge 1 commit into
Closed
Conversation
sylvesterkaczmarek
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 22, 2026 14:47
meta-cla
Bot
added
the
CLA Signed
tianyu-l
reviewed
Aug 22, 2026
| from torch.utils._pytree import tree_map | ||
|
|
||
|
|
||
| class NoParallel(ParallelStyle): |
Contributor
There was a problem hiding this comment.
this class will be deprecated soon
Contributor
Author
There was a problem hiding this comment.
Thanks, understood. Given that NoParallel is going to be deprecated and #3374 is now marked wontfix, I don't want to add maintenance work to a path that is being removed. I'm happy to close this PR if that's the preferred direction; otherwise I can keep the narrow compatibility fix.
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Fixes #3374.
Summary
NoParallel._prepare_output_fncurrently assumes a module returns a singleDTensorand accessesoutputs.placementsdirectly. Modules that return tuples or other nested output structures therefore fail.This change applies the existing layout redistribution and optional
to_local()conversion independently to everyDTensorleaf using pytree traversal. Non-DTensor leaves are preserved unchanged.Regression coverage
Added CPU-only coverage for nested tuple/list/dict outputs, redistribution, non-DTensor metadata, and
use_local_output=True.Validation
The branch is based directly on current upstream
mainatb5bcd76c66d0acb1a08b176d9639425bc55e0b2aand contains one focused commit. Full repository tests were not run locally; upstream CI should run the added unit coverage.