Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add BloomModel hydra support #129

Merged
merged 1 commit into from
Dec 9, 2022
Merged

Conversation

jon-tow
Copy link
Collaborator

@jon-tow jon-tow commented Dec 8, 2022

This PR adds hydra-based PPO model branching support for BloomModels.

wandb reports:

Note: The <Arch>ModelBranch implementations should be refactored in the future to share a common interface if we plan on adding more in the near future.

@jon-tow jon-tow marked this pull request as draft December 8, 2022 22:51
@jon-tow jon-tow marked this pull request as ready for review December 9, 2022 02:39
@LouisCastricato
Copy link
Contributor

@Dahoas can you give this a quick glance

@Dahoas
Copy link
Collaborator

Dahoas commented Dec 9, 2022

Bloom rewards seem a bit low relative to other models :(. But Not much we can really do I guess

@LouisCastricato
Copy link
Contributor

To he expected

# 1.0 in head_mask indicate we keep the head
# attention_probs has shape batch_size x num_heads x N x N
# head_mask has shape n_layer x batch x num_heads x N x N
head_mask = self.get_head_mask(head_mask, hf_get_num_hidden_layers(self.config))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the head mask?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The head mask is a binary mask that can be used to drop the self-attention weights (softmax(qk)) from specified heads before computing the full attention output. For example, see here. get_head_mask just expands the dimensions to line-up with proper shape.

else:
attention_mask = attention_mask.to(hidden_states.device)

alibi = modeling_bloom.build_alibi_tensor(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't know bloom used alibi

@Dahoas
Copy link
Collaborator

Dahoas commented Dec 9, 2022

Looks good to me! I just left some questions so I can better understand how things are functioning.

@Dahoas Dahoas merged commit 1a3461d into CarperAI:main Dec 9, 2022
@jon-tow jon-tow deleted the add-bloom-support branch December 9, 2022 20:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants