[fix] Cast learning-to-rank loss logits to float32 for bf16/fp16 CrossEncoder training#3800
Conversation
…sEncoder training (#3793) Training a CrossEncoder reranker in bf16/fp16 crashed in the learning-to-rank losses because the internal logits matrix defaults to float32 while the model emits bf16/fp16 logits, raising "Index put requires the source and destination dtypes match". Upcasting the logits to float32 right after the activation fixes all six affected losses (PListMLELoss, LambdaLoss, ListNetLoss, ADRMSELoss, and ListMLELoss/RankNetLoss via inheritance) and keeps exp/log/cumsum/softmax in a numerically stable precision. It is a no-op for the usual float32 training. Adds a regression test covering the six losses under bfloat16.
|
Hello! I was able to reproduce the issue, and the issue indeed seems fixed with this. I made one small change: move the .float() call to before the activation (as the activation can perhaps also be fragile with bf16). It looks good to me beyond that. Thank you!
|
There was a problem hiding this comment.
Pull request overview
This PR fixes mixed-precision (bf16/fp16) training crashes for CrossEncoder learning-to-rank/listwise losses by ensuring the internal logits computations run in float32, avoiding dtype-mismatch scatter operations and improving numerical stability.
Changes:
- Cast concatenated logits to
float32in several reranking losses before constructing/scattering into padded logits matrices. - Add a regression test that runs all affected losses against a low-precision CrossEncoder and checks dtype/finite outputs and gradient flow.
Reviewed changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
sentence_transformers/cross_encoder/losses/plist_mle.py |
Upcasts concatenated logits to float32 to prevent bf16/fp16 scatter dtype mismatch. |
sentence_transformers/cross_encoder/losses/list_net.py |
Upcasts concatenated logits to float32 for stable, dtype-consistent listwise loss computation. |
sentence_transformers/cross_encoder/losses/lambda_loss.py |
Upcasts concatenated logits to float32 to avoid dtype mismatch and overflow risk in listwise ops. |
sentence_transformers/cross_encoder/losses/adr_mse.py |
Upcasts concatenated logits to float32 before building padded matrices and rank computations. |
tests/cross_encoder/losses/test_rerank_losses_dtype.py |
Adds a regression test covering bf16/fp16 execution and backprop for multiple listwise losses. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
It worked locally for me, but I'm fine with excluding this just in case Co-authored-by: Copilot Autofix powered by AI <[email protected]>
|
Thanks so much, Tom!! Great call moving the upcast before the activation |
Fixes #3793.
Training a CrossEncoder reranker in bf16/fp16 crashes in the learning-to-rank losses. The logits matrix is built with
torch.full(...), which is float32 by default, and then the model's logits get scattered into it — so as soon as the model's in bf16 you get:I just cast the logits to float32 (
.float()after the activation) rather than thedtype=logits.dtypefrom the issue. Matching the matrix to the logits dtype doesn't fully fix it — the next line that writes the labels (labels_matrix[...] = torch.cat(labels).float()) crashes the same way in reverse. On top of that these losses doexp/cumsum/log/softmaxwith no max-shift, so bf16 is shaky and fp16 overflows. Running the loss in fp32 avoids both, and it's a no-op when you're already in float32.Same pattern was in
PListMLELoss,LambdaLoss,ListNetLossandADRMSELoss.ListMLELossandRankNetLosssubclass those two, so they're covered too.Added a test that runs all six on a bf16 model — fails on main, passes here.