Fix loss_parallel with BF16 logits #130550
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/130550
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit f3eb4ef with merge base 72d9135 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
oncall: distributed
| # only update grad_input to -1 if not masked | ||
| assert partial_placement.mask_buffer.data is not None | ||
| grad_update = partial_placement.mask_buffer.data.float() - 1.0 | ||
| grad_update = partial_placement.mask_buffer.data.to(grad_input.dtype) - 1.0 |
There was a problem hiding this comment.
@tianyu-l I think we should avoid using .data here since its intent can be ambiguous. If you want it to not be part of autograd graph, then you should use .detach().
There was a problem hiding this comment.
I agree I think detach is a safer option.
but I'll defer to @tianyu-l
There was a problem hiding this comment.
@tianyu-l I think we should avoid using
.datahere since its intent can be ambiguous. If you want it to not be part of autograd graph, then you should use.detach().
Hmm the .data is just a class variable (of type torch.Tensor) of the class MaskBuffer. So it shouldn't be a concern?
https://github.com/pytorch/pytorch/blob/main/torch/distributed/_tensor/ops/embedding_ops.py#L28
There was a problem hiding this comment.
yeah the MuskBuffer is not a torch.Tensor and data is just coincidentally a field name of that custom object, we can probably rename that field later to avoid confusing with tensor.data
|
@pytorchbot merge |
Merge failedReason: This PR needs a If not, please add the To add a label, you can comment to pytorchbot, for example For more information, see Details for Dev Infra teamRaised by workflow job |
|
I did not see a TP specific label, so I used the DTensor one. |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Fixes pytorch#130549 This PR uses the specific dtype for the `grad_input` buffer and fixes the error Pull Request resolved: pytorch#130550 Approved by: https://github.com/tianyu-l
6 participants
Fixes #130549
This PR uses the specific dtype for the
grad_inputbuffer and fixes the errorcc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o