Add support for per token reward/advantages with a custom_reward_post_process_path

slime/backends/megatron_utils/data.py

@@ -401,6 +401,9 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc
                         val = cp_size * sum_of_sample_mean(val) / len(loss_masks)
                     else:
                         val = val.mean() * cp_size
+                elif isinstance(val[0], list):
+                    # Adding support for per token rewards
+                    val = sum(sum(v) / len(v) for v in val) / len(val)
                 else:
                     val = sum(val) / len(val)
             elif isinstance(val, torch.Tensor):

slime/backends/megatron_utils/loss.py

@@ -268,7 +268,13 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
         ]
 
     if args.advantage_estimator in ["grpo", "gspo"]:
-        rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
+        if isinstance(rewards[0], list):
+            # Adding support for per token rewards.
+            # For efficiency (making multiple tensors vs one) when it's per sequence rewards
+            # we only create multiple tensors only when its per token reward
+            rewards = [torch.tensor(r, dtype=torch.float32, device=kl[0].device) for r in rewards]
+        else:
+            rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
         returns = get_grpo_returns(rewards, kl)
         # TODO: is the copy necessary?
         advantages = [r for r in returns]