Add a "division by zero" check in chunked loss handling in kd_losses.py #2239
Conversation
- test kd_loss when all mask is zero
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2239
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit f5befaf with merge base c152248 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
@felipemello1 , @joecummings , PTAL |
| @@ -137,4 +137,7 @@ def forward( | |||
| student_chunk, teacher_chunk, label_chunk, normalize=False | |||
| ) | |||
|
|
|||
| if torch.sum(mask.view(-1), dim=0) == 0: | |||
| return torch.tensor(0.0, device=student_logits[0].device) | |||
There was a problem hiding this comment.
In the discussion on #2094, it seems like if we hit this point, there's probably something wrong with the training data. In that case, should we error out here? Or drop a warning? Or let people shoot themselves in the foot?
cc @lindawangg and @ebsmothers
There was a problem hiding this comment.
I will wait for inputs, and will update #2094 after this PR is resolved.
There was a problem hiding this comment.
Any inputs?
@lindawangg , @ebsmothers
There was a problem hiding this comment.
Thanks for the bump here. Personally I think this is the right way to handle it. Yes, we could error/raise a warning, but honestly this behavior should be pretty obvious in a loss curve as it will clearly just drop to zero on any such iterations (plus I think zero loss would technically be the "correct" loss value if every token is being ignored). No need to overdo it on handling this edge case here imo
| @@ -17,7 +17,7 @@ def random(): | |||
|
|
|||
|
|
|||
| class TestForwardKLWithChunkedOutputLoss: | |||
| def test_forward_kl_loss(self): | |||
| def setup_forward_kl_loss(self, set_all_masks_zero: bool = False): | |||
There was a problem hiding this comment.
nit: I would do something like "ignore_all_tokens" instead of "set_all_masks_zero" as I think it's clearer
There was a problem hiding this comment.
Sounds better, and updated.
|
Thank you @ebsmothers , updated the change. |
joecummings
changed the title
Context
What is the purpose of this PR? Is it to
Please link to any issues this PR addresses.
Changelog
What are the changes made in this PR?
Test plan
Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.
pre-commit install)pytest testspytest -v tests/torchtune/modules/loss/test_kd_losses.pypytest tests -m integration_testUX
If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example