From 53e5f4e78383b650d113364cec92d244d1b21f46 Mon Sep 17 00:00:00 2001 From: EphraiemSarabamoun Date: Sat, 30 May 2026 06:58:53 -0700 Subject: [PATCH] Fix KTO compute_kl to average KL across all batches In KTOTrainer.compute_kl the loop over self.random_dataloader reassigned self.kl on every iteration, so only the last batch's KL survived instead of an average over all batches. Accumulate the per batch clamped KL and divide by the batch count once after the loop. Reported in issue #215. --- align_anything/trainers/text_to_text/kto.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/align_anything/trainers/text_to_text/kto.py b/align_anything/trainers/text_to_text/kto.py index d484886a..bbf0a3d6 100644 --- a/align_anything/trainers/text_to_text/kto.py +++ b/align_anything/trainers/text_to_text/kto.py @@ -67,6 +67,8 @@ def compute_kl(self): sampler=DistributedSampler(random_dataset, shuffle=True), batch_size=self.cfgs.train_cfgs.per_device_kl_batch_size, ) + kl_sum = 0 + num_batches = 0 for batch in self.random_dataloader: log_probs = self.compute_log_probs( # size = (2 * B, L - 1) self.model.module, @@ -78,7 +80,10 @@ def compute_kl(self): ) kl = (log_probs - ref_log_probs).mean() - self.kl = max(kl, 0) + kl_sum = kl_sum + max(kl, 0) + num_batches += 1 + + self.kl = kl_sum / num_batches def loss( # pylint: disable=too-many-locals self,