From 9b563f21f09979c7233009e094aa3f0e879c44e4 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Wed, 17 Jun 2026 08:31:13 +0000 Subject: [PATCH] Fix double loss normalization in server policy_loss/importance_sampling Server forward/backward defers loss normalization: each micro-batch returns a raw summed loss and optim_step divides the accumulated gradient once by total valid tokens. The policy_loss and importance_sampling calls omitted loss_reducer, so the loss defaulted to a local per-micro-batch token-mean and was normalized twice -> grad_norm collapses and the model fails to learn at any batch size. Pass loss_reducer=token_sum_reducer to both calls, matching causallm_loss and opd_loss. --- src/xorl/server/runner/model_runner.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index e5b28236..48ac29ea 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -1904,6 +1904,7 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): ce_mode=self.ce_mode, compute_kl_stats=compute_kl_stats, lm_head_fp32=self.lm_head_fp32, + loss_reducer=token_sum_reducer, metric_reducer=token_sum_reducer, ) local_loss_sum = _result.loss @@ -2010,6 +2011,7 @@ def _compute_micro_batch_loss(self, micro_batch, loss_fn, loss_fn_params): compute_kl_stats=compute_kl_stats, lm_head_fp32=self.lm_head_fp32, icepop_beta=icepop_beta, + loss_reducer=token_sum_reducer, metric_reducer=token_sum_reducer, ) local_loss_sum = _result.loss