Fix MPS with sequence loss #2834
Merged
Conversation
This file contains 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
|
Upon further inspection, |
JAEarly
force-pushed
the
eval_loop_mps
branch
2 times, most recently
from
31d92e8 to
b945c5a
Compare
JAEarly
force-pushed
the
eval_loop_mps
branch
from
b945c5a to
1dd8bc0
Compare
|
@mvpatel2000 this is a relatively small change but is fairly annoying for running stuff on MPS, could it be merged in please? Thank you! |
dakinggg
reviewed
Jan 19, 2024
|
@JAEarly apologies for the delay in review! Next time please feel free to tag me as a reviewer.... I will see if we can setup auto-tagging as well. We didn't see this earlier :( |
mvpatel2000
approved these changes
Jan 20, 2024
ShashankMosaicML
pushed a commit
to ShashankMosaicML/composer
that referenced
this pull request
Feb 3, 2024
* Add MPS support for list outputs in training eval loop * Add error for invalid state outputs type in trainer * Remove raise ValueError in trainer eval loop --------- Co-authored-by: Daniel King <[email protected]>
ShashankMosaicML
pushed a commit
to ShashankMosaicML/composer
that referenced
this pull request
Feb 3, 2024
* Add MPS support for list outputs in training eval loop * Add error for invalid state outputs type in trainer * Remove raise ValueError in trainer eval loop --------- Co-authored-by: Daniel King <[email protected]>
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.
What does this PR do?
TLDR
Fixes the trainer eval loop for MPS devices when the model output is a list of tensors rather than a
torch.TensororMapping.Related
Similar issue but with dict: #2632
PR for above issue: #2706
Details
In the trainer evaluation loop, model outputs (
self.state.outputs) are moved to the CPU if they are on an MPS device (to avoid torchvision numerical errors on MPS devices). Currently this works fine ifself.state.outputsis a Mapping or a torch.Tensor. However, it fails for Sequence types (e.g. a list of torch.Tensors).According to the State class,
outputsis expected to be of typetorch.Tensor | Sequence[torch.Tensor]. So Sequence types should be supported in this process in the eval loop. Strictly,outputsshould not be aMapping, despite the eval operation supporting this. I have left theMappingsupport in place for now but happy to revisit.As it was a little difficult to debug this issue, I have added an error message which should make it clearer if an invalid output type is used (such that it cannot correctly be mapped from MPS to CPU).
Before submitting
pre-commiton your change? (see thepre-commitsection of prerequisites)