fix(kto): sum chosen/rejected rewards per sample when chunk_size > 1 (#1495)

## Summary

Root cause: in `LigerFusedLinearKTOFunction.preference_loss_fn`,
`rewards` has shape `(chunk_size,)`, but it is multiplied by
`preference_labels_chunk.unsqueeze(1)`, which has shape `(chunk_size,
1)`. Broadcasting turns the product into a `(chunk_size, chunk_size)`
matrix, so `chosen_rewards_sum` becomes `num_chosen_in_chunk *
sum(rewards)` and `rejected_rewards_sum` becomes `num_rejected_in_chunk
* sum(rewards)`. With the default `chunk_size=1` the matrix is 1x1 and
the result is right, which is why the existing tests did not catch it.
With `chunk_size > 1` and mixed labels in a chunk, both reward metrics
are wrong. The loss and the gradients are not affected.

Example (B=4, labels `[True, False, True, False]`, with a reference
model): `chunk_size=1` returns chosen 0.193 and rejected -1.435, while
`chunk_size=4` returns -2.485 for both (2 x the sum of all four
rewards).

Fix: multiply by the 1D mask. This is the same change #1441 made for
`chosen_logps_sum` / `rejected_logps_sum` in
`fused_linear_unpaired_preference.py`; the two reward lines in
`kto_loss.py` still had the old `unsqueeze(1)`.

## Testing Done

Added `test_kto_rewards_sum_with_chunk_size`, which checks that the
chosen and rejected reward sums with `chunk_size` 2 and 4 match
`chunk_size=1`. On main it fails (chosen sum -85.45 with `chunk_size=2`
and -170.90 with `chunk_size=4`, vs -38.35 with `chunk_size=1`) and it
passes with the fix.

- Hardware Type: NVIDIA GeForce GTX 1060 6GB, PyTorch 2.5.1 (the new
test was also run CPU only)
- [x] `python -m pytest test/chunked_loss/test_kto_loss.py` (54 passed)
- [x] `make checkstyle` checks: `ruff check .` and `ruff format --check
.` pass (ruff 0.14.11, as pinned in `.pre-commit-config.yaml`)
- [ ] `make test` (not run; Triton is not available on this Windows
host)
- [ ] `make test-convergence` (not run; KTO is not covered by the
convergence tests)

src/liger_kernel/chunked_loss/kto_loss.py

@@ -63,8 +63,8 @@ class LigerFusedLinearKTOFunction(LigerFusedLinearUnpairedPreferenceBase):
             losses = 1 - F.sigmoid(beta * logratios_chunk * multiplier_chunk)
 
         rewards = beta * logratios_chunk
-        chosen_rewards_sum = (rewards * preference_labels_chunk.unsqueeze(1)).sum()
-        rejected_rewards_sum = (rewards * (~preference_labels_chunk).unsqueeze(1)).sum()
+        chosen_rewards_sum = (rewards * preference_labels_chunk).sum()
+        rejected_rewards_sum = (rewards * (~preference_labels_chunk)).sum()
 
         return losses.sum() / (full_target.shape[0]), chosen_rewards_sum, rejected_rewards_sum
 

test/chunked_loss/test_kto_loss.py

@@ -478,3 +478,25 @@ def test_kto_frozen_weight(compiled, chunk_size):
     assert weight_lora.grad is None
     assert_verbose_allclose(loss_full, loss_lora, atol=1e-5, rtol=1e-5)
     assert_verbose_allclose(input_full.grad, input_lora.grad, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.parametrize("chunk_size", [2, 4])
+def test_kto_rewards_sum_with_chunk_size(chunk_size):
+    B, T, H, V = 4, 16, 32, 64
+
+    _input = torch.randn(B, T, H, device=device, dtype=torch.float32)
+    _target = torch.randint(0, V, (B, T), device=device, dtype=torch.long)
+    _pref_labels = torch.tensor([True, False, True, False], device=device)
+    _weight = torch.randn(V, H, device=device, dtype=torch.float32)
+
+    # The summed rewards must not depend on how the batch is split into chunks
+    _, aux_ref = LigerFusedLinearKTOLoss(use_ref_model=False, chunk_size=1, compiled=False)(
+        _input=_input, lin_weight=_weight, target=_target, preference_labels=_pref_labels
+    )
+    _, aux = LigerFusedLinearKTOLoss(use_ref_model=False, chunk_size=chunk_size, compiled=False)(
+        _input=_input, lin_weight=_weight, target=_target, preference_labels=_pref_labels
+    )
+
+    # chosen_rewards_sum and rejected_rewards_sum
+    assert_verbose_allclose(aux[4], aux_ref[4], atol=1e-5, rtol=1e-5)
+    assert_verbose_allclose(aux[5], aux_ref[5], atol=1e-5, rtol=1e-5)