fix: restore PR 607 KL loss removal (#639)

PR #619 reverted PR #607 which modified kl divergence, we fix this here.

Author: FurtherAI

src/art/loss.py

@@ -15,7 +15,6 @@ class Loss(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
     reduction: Literal["mean", "sum"]
     policy_loss: torch.Tensor
-    kl: torch.Tensor
     entropy: torch.Tensor | None
     policy_loss_sum: torch.Tensor
     probs_corr: torch.Tensor
@@ -126,17 +125,9 @@ def loss_fn(
             logprob_diff = old_logprobs - original_logprobs
             prob_ratio = torch.exp(logprob_diff)
         policy_loss *= torch.clamp(prob_ratio, max=upper_bound).detach()
-    if ref_logprobs is not None:
-        kl_div = (
-            torch.exp(ref_logprobs - new_logprobs) - (ref_logprobs - new_logprobs) - 1.0
-        )
-    else:
-        kl_div = torch.zeros_like(policy_loss)
     policy_loss = policy_loss * weights * assistant_mask
-    kl_div = kl_div * weights * assistant_mask
     denominator = assistant_mask.sum() + 1e-6 if reduction == "mean" else 1.0
     reduced_policy_loss = policy_loss.sum() / denominator
-    kl = kl_div.sum() / denominator
     # Compute reduced entropy for the current step.
     if entropies is not None:
         shifted_entropies = shift_tensor(entropies, 0.0)
@@ -146,7 +137,6 @@ def loss_fn(
     return Loss(
         reduction=reduction,
         policy_loss=reduced_policy_loss,
-        kl=kl,
         entropy=entropy,
         policy_loss_sum=policy_loss.sum(),
         probs_corr=probs_corr,

src/art/test/test_kl_advantage.py

@@ -46,6 +46,8 @@ def test_kl_advantage_no_effect_when_disabled():
 
     assert loss_no_kl.kl_policy_ref is None
     assert loss_without_ref.kl_policy_ref is None
+    assert loss_no_kl.reduction == "mean"
+    assert not hasattr(loss_no_kl, "kl")
 
 
 def test_kl_advantage_enabled():

tests/integration/megatron_oracle_worker.py

@@ -682,7 +682,6 @@ def _scaled_loss_fn(*args: Any, **kwargs: Any):
             return loss.model_copy(
                 update={
                     "policy_loss": loss.policy_loss * effective_loss_scale,
-                    "kl": loss.kl * effective_loss_scale,
                     "policy_loss_sum": loss.policy_loss_sum * effective_loss_scale,
                 }
             )