[Bug] Scale the cross-device contrastive loss by world size in the kl_div branch

[gss10282025]

Summary

In the cross-device kl_div branch, the contrastive loss is computed on gathered embeddings that only carry gradients for the local slice, so after DDP's averaging its gradient is 1/world_size of the single-GPU gradient. The local KD loss is not affected, so distillation gets relatively more weight as GPUs are added.

This PR multiplies the contrastive term by self.world_size before it is added to the KD loss, the same compensation Tevatron applies after its identical gather. The KD term, the m3_kd_loss branch and runs without distillation are unchanged. The encoder_only/base_same_dataset.sh and decoder_only/icl_same_dataset.sh examples go through this branch.

Issue links

Fixes #1603

Validation

At fd1a2bdf, BiEncoderOnlyEmbedderModel with cross-device negatives and kl_div KD, eight queries split over 1 or 2 GPUs with the same number of KD targets per GPU, one AdamW step with clipping, FP32: compared with a reference computed on the whole global batch, the 2-GPU gradient differs by 0.406442 (relative L2) before the change and 1.50857e-6 after; the clipped gradient, Adam moments and update match as well. KD alone already agrees on 1 and 2 GPUs. Reproduced on a second machine. I tested plain DDP, not DeepSpeed.

Checks on this branch

• The one-line change applies cleanly to fd1a2bdf and the file passes a Python syntax check.
• The numbers above come from my earlier runs of this same change; I have not run a full fine-tuning job on this branch.

Prepared with AI assistance; opened as a draft for code review.