[core] refactor step method
#76
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
- add safety checker + manual device assignment
|
The documentation is not available anymore as the PR was closed or merged. |
lvwerra
reviewed
Jan 5, 2023
trl/trainer/ppo_trainer.py
Outdated
| @@ -252,6 +255,17 @@ def _step_safety_checker( | |||
| f"Batch size ({batch_size}) does not match number of examples - but got {len(tensor_list)} for: {name}" | |||
| ) | |||
|
|
|||
| # set scores on the correct device | |||
| if name == "scores": | |||
| scores = [score.to(self.accelerator.device) for score in scores] | |||
There was a problem hiding this comment.
I think we can do the same for queries and values as well. although they should be you don't know what the user might do before passing them to step and then the behaviour is consistent for all inputs.
Do you know what happens when the tensor is already on device? Will it copy it again or do nothing?
There was a problem hiding this comment.
Makes sense, in PT it should do nothing if they are on the same device so this should be cost-free
There was a problem hiding this comment.
Should be now handled in b32bfbf
lvwerra
reviewed
Jan 5, 2023
lvwerra
reviewed
Jan 5, 2023
Co-authored-by: Leandro von Werra <lvwerra@users.noreply.github.com>
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.
This PR adds a new safety checker inside
stepmethod to make sure the rewards are set on the correct device. Regarding the queries and the responses users should retrieve the dataloader from the trainer and use that instead as the device assignment is performed directly at the dataloader level. Since the reward is not part of the dataloader the device assignment needs to be performed manually.This PR also adds inside the safety checker a new check. Before this PR if a user pass a reward tensor with a dimension different from 0 (e.g.
torch.tensor([1.0])) it would break the training loop. Therefore now we force the reward tensor to be with the desired shape. We now also throw a value error if the dimension of the reward is > 1.cc @lewtun @lvwerra @edbeeching