[Feature] RLHF Reward Model #1315
Closed
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
facebook-github-bot
added
the
CLA Signed
Co-authored-by: Alessandro Pietro Bardelli <apbard@users.noreply.github.com>
apbard
approved these changes
Jun 26, 2023
Co-authored-by: Alessandro Pietro Bardelli <apbard@users.noreply.github.com>
vmoens
reviewed
Jun 26, 2023
torchrl/modules/models/rlhf.py
Outdated
Comment on lines
32
to
33
| """ Returns a tuple (rewards, end_scores) where `rewards` contains all rewards computed at each timestep, `end_scores` contains the reward computed at the last-non-padding token | ||
| """ |
There was a problem hiding this comment.
Suggested change
| """ Returns a tuple (rewards, end_scores) where `rewards` contains all rewards computed at each timestep, `end_scores` contains the reward computed at the last-non-padding token | |
| """ | |
| """Computes the rewards associated with some encoded sequence of tokens. | |
| Returns a tuple (rewards, end_scores) where `rewards` contains all rewards computed at each timestep, `end_scores` contains the reward computed at the last-non-padding token | |
| """ |
torchrl/modules/models/rlhf.py
Outdated
Comment on lines
35
to
45
| hidden_states = outputs[0] | ||
| rewards = self.lm_head(hidden_states).squeeze(-1) | ||
| end_scores = [] | ||
| bs = input_ids.shape[0] | ||
|
|
||
| for i in range(bs): | ||
| pad_inds = (input_ids[i] == self.PAD_ID).nonzero() | ||
| first_pad_ind = ( | ||
| pad_inds[0].item() if len(pad_inds) > 0 else input_ids.shape[1] | ||
| ) | ||
| end_scores.append(rewards[i, first_pad_ind - 1]) |
There was a problem hiding this comment.
My intuition is that there is a better (more efficient way) of coding that loop.
Let's move it to a private method such that we can easily refactor this piece of code and compare the results
There was a problem hiding this comment.
For instance, assuming that padding is only on the right:
>>> import torch
>>> z = torch.arange(12).view(3, 4)
>>> z[0, 2:] = 100
>>> z[1, 3:] = 100
>>> z[2, 1:] = 100
>>> mask = z == 100
>>> mask = torch.cat([mask, torch.ones_like(mask[..., :1])], -1) # make sure that there is one True on each row
>>> first_pad = mask[..., :-1] ^ mask[..., 1:]
>>> first_pad = first_pad.nonzero()
>>> first_pad = first_pad[:, -1]There was a problem hiding this comment.
Thanks! I've factored the slow code out into a private method as you suggested. Did I understand right that we should land as is and follow up with speed improvements?
|
|
||
| for i in range(bs): | ||
| # Check if there is any padding otherwise take length of sequence | ||
| c_inds = (chosen_ids[i] == pad_token_id).nonzero() |
There was a problem hiding this comment.
in general nonzero is expensive and shouldn't be called to often
| end_ind = max(c_ind, r_ind) | ||
|
|
||
| # Retrieve first index where trajectories diverge | ||
| divergence_ind = (chosen_ids[i] != rejected_ids[i]).nonzero()[0] |
There was a problem hiding this comment.
ditto, nonzero is expensive
torchrl/modules/models/rlhf.py
Outdated
| The loss is computed as loss = -log_sigmoid(chosen_reward - rejected_reward). | ||
| This loss is small when the reward model favours the chosen data and large if | ||
| the model favours the rejected data. | ||
| Note: the loss is computed excluding the common "prefix" subsequence to effectively disregard contribution of the original prompt. |
There was a problem hiding this comment.
Perhaps an example of a call to this?
# Conflicts: # test/test_rlhf.py
Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
tcbegley
changed the title
This was referenced Jun 28, 2023
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 builds on top of #1309, adding the
GPT2RewardModelclass and tests. Changes from that PR are needed so that we can use the dataloaders in the tests.