What changed

In compute_advantages_and_returns (slime/backends/megatron_utils/loss.py), the PPO estimator now builds its per-token reward tensor out-of-place:

for reward, per_token_kl in zip(old_rewards, kl, strict=False):
    token_level_rewards = per_token_kl * kl_coef
    if cp_rank == 0:
        token_level_rewards[-1] += reward
    rewards.append(token_level_rewards)

per_token_kl is the raw approximate KL tensor stored in rollout_data["kl"]; token_level_rewards is the separate PPO reward signal (-kl_coef * kl, plus the scalar reward at the final token). The reward math is unchanged.

A CPU unit test (tests/test_ppo_kl_metric.py) is added and registered in the cpu-unittest matrix.

Why

compute_advantages_and_returns stores approximate KL in rollout_data["kl"], which log_rollout_data later reduces and logs as rollout/kl. The previous PPO branch used k *= kl_coef and then k[-1] += reward; because k aliased rollout_data["kl"][i], it overwrote the logged raw KL with the reward-shaped tensor.

The final form is intentionally explicit rather than relying on local-name shadowing. It matches the convention used by the other KL-penalized estimators: get_reinforce_plus_plus_returns derives token_level_rewards from KL, and get_reinforce_plus_plus_baseline_advantages computes from kl_tensor without mutating the stored KL.

Validation

tests/test_ppo_kl_metric.py runs the PPO estimator with a known KL (k1, kl_coef=0.05) and asserts rollout_data["kl"] still equals the freshly computed KL afterward. It fails on the old in-place update and passes with this fix. Megatron is stubbed through megatron.core.mpu with cp_size=1, mirroring tests/test_value_temperature.py; no GPU is required.

Local run:

PYTHONPATH=$PWD uvx --python 3.10 --with torch --with packaging --with numpy --with httpx pytest tests/test_ppo_kl_metric.py