fix backward pass for cached losses #3114
Open
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
tomaarsen
reviewed
Dec 4, 2024
| @@ -230,14 +230,15 @@ def embed_minibatch_iter( | |||
| def calculate_loss_and_cache_gradients(self, reps: list[list[Tensor]], reps_guided: list[list[Tensor]]) -> Tensor: | |||
| """Generalized function to calculate the cross-entropy loss and cache the gradients wrt. the embeddings.""" | |||
| loss = self.calculate_loss(reps, reps_guided) | |||
There was a problem hiding this comment.
Suggested change
| loss = self.calculate_loss(reps, reps_guided) | |
| loss = self.calculate_loss(reps, reps_guided, with_backward=True) |
I believe this is missing for CGIST.
|
My initial tests with CMNRL and CMNRL + Matryoshka are promising, and I can indeed also reproduce the high memory usage on Matryoshka on top of CMNRL only seems to add about 20%-25% training time, which seems fine. CMNRL does seem a decent bit slower than pure MNRL (more than the 10-20% that I thought it was), but based on your changes in this PR and the previous, that shouldn't be related to you at all.
|
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.
Here is a draft for a fix of the backward pass in the cached losses, while still maintaining the compatibility with the matryoshka loss. The problem is a bit more difficult than I originally thought because we need to detach the tensor before doing the minibatch loss computation.
If anyone has suggestions for improvements, let me know 😀