From 7ce0deefd3e3189e92601aa397b1db3564c2c020 Mon Sep 17 00:00:00 2001 From: gss10282025 <326855337+gss10282025@users.noreply.github.com> Date: Sat, 26 Sep 2026 09:39:25 +0000 Subject: [PATCH] [Bug] Scale the cross-device contrastive loss by world size in the kl_div branch --- FlagEmbedding/abc/finetune/embedder/AbsModeling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/FlagEmbedding/abc/finetune/embedder/AbsModeling.py b/FlagEmbedding/abc/finetune/embedder/AbsModeling.py index 3e9e32cb3..38f87222b 100644 --- a/FlagEmbedding/abc/finetune/embedder/AbsModeling.py +++ b/FlagEmbedding/abc/finetune/embedder/AbsModeling.py @@ -226,7 +226,7 @@ def _compute_cross_device_neg_loss(self, q_reps, p_reps, teacher_targets=None, c cross_idxs = torch.arange(cross_q_reps.size(0), device=cross_q_reps.device, dtype=torch.long) cross_targets = cross_idxs * group_size # (world_size * batch_size) - loss += self.compute_loss(cross_scores, cross_targets) + loss += self.world_size * self.compute_loss(cross_scores, cross_targets) elif self.kd_loss_type == "m3_kd_loss": cross_teacher_targets = self._dist_gather_tensor(teacher_targets) # (world_size * batch_size, group_size)