Merge advantage computation logic into __call__ function

Author: yanxi-chen

trinity/algorithm/advantage_fn/grpo_advantage.py

@@ -24,13 +24,48 @@ def __call__(
         exps: DataProto,
         **kwargs,
     ) -> Tuple[DataProto, Dict]:
-        advantages, returns = compute_grpo_outcome_advantage(
-            token_level_rewards=exps.batch["token_level_rewards"],
-            eos_mask=exps.batch["response_mask"],
-            index=exps.non_tensor_batch["uid"],
-        )
-        exps.batch["advantages"] = advantages
-        exps.batch["returns"] = returns
+        """
+        Compute advantage for GRPO, operating only on Outcome reward
+        (with only one scalar reward for each response).
+
+            token_level_rewards: `(torch.Tensor)`
+                shape: (bs, response_length)
+            eos_mask: `(torch.Tensor)`
+                shape: (bs, response_length)
+            scores: `(torch.Tensor)`
+                shape: (bs, response_length)
+        """
+        token_level_rewards = exps.batch["token_level_rewards"]
+        eos_mask = exps.batch["response_mask"]
+        index = exps.non_tensor_batch["uid"]
+        epsilon = 1e-6
+
+        response_length = token_level_rewards.shape[-1]
+        scores = token_level_rewards.sum(dim=-1)
+
+        id2score = defaultdict(list)
+        id2mean = {}
+        id2std = {}
+
+        with torch.no_grad():
+            bsz = scores.shape[0]
+            for i in range(bsz):
+                id2score[index[i]].append(scores[i])
+            for idx in id2score:
+                if len(id2score[idx]) == 1:
+                    id2mean[idx] = torch.tensor(0.0)
+                    id2std[idx] = torch.tensor(1.0)
+                elif len(id2score[idx]) > 1:
+                    id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
+                    id2std[idx] = torch.std(torch.tensor([id2score[idx]]))
+                else:
+                    raise ValueError(f"no score in prompt index: {idx}")
+            for i in range(bsz):
+                scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
+            scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask
+
+        exps.batch["advantages"] = scores
+        exps.batch["returns"] = scores
 
         metrics = {
             # TODO: add meaningful metrics
@@ -41,52 +76,3 @@ def __call__(
     @classmethod
     def default_args(cls) -> Dict:
         return {}
-
-
-# NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar.
-def compute_grpo_outcome_advantage(
-    token_level_rewards: torch.Tensor,
-    eos_mask: torch.Tensor,
-    index: torch.Tensor,
-    epsilon: float = 1e-6,
-):
-    """
-    Compute advantage for GRPO, operating only on Outcome reward
-    (with only one scalar reward for each response).
-    Args:
-        token_level_rewards: `(torch.Tensor)`
-            shape: (bs, response_length)
-        eos_mask: `(torch.Tensor)`
-            shape: (bs, response_length)
-
-    Returns:
-        advantages: `(torch.Tensor)`
-            shape: (bs, response_length)
-        Returns: `(torch.Tensor)`
-            shape: (bs, response_length)
-    """
-    response_length = token_level_rewards.shape[-1]
-    scores = token_level_rewards.sum(dim=-1)
-
-    id2score = defaultdict(list)
-    id2mean = {}
-    id2std = {}
-
-    with torch.no_grad():
-        bsz = scores.shape[0]
-        for i in range(bsz):
-            id2score[index[i]].append(scores[i])
-        for idx in id2score:
-            if len(id2score[idx]) == 1:
-                id2mean[idx] = torch.tensor(0.0)
-                id2std[idx] = torch.tensor(1.0)
-            elif len(id2score[idx]) > 1:
-                id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
-                id2std[idx] = torch.std(torch.tensor([id2score[idx]]))
-            else:
-                raise ValueError(f"no score in prompt index: {idx}")
-        for i in range(bsz):
-            scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
-        scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask
-
-    return scores, scores

trinity/algorithm/advantage_fn/opmd_advantage.py

@@ -21,16 +21,59 @@ def __call__(
         exps: DataProto,
         **kwargs,
     ) -> Tuple[DataProto, Dict]:
-        advantages, returns = compute_opmd_outcome_advantage(
-            token_level_rewards=exps.batch["token_level_rewards"],
-            eos_mask=exps.batch["response_mask"],
-            # TODO (yanxi): check consistency with exps.batch["attention_mask"][:, -response_length:] in original implementation
-            index=exps.non_tensor_batch["uid"],
-            opmd_baseline="mean",
-            tau=1.0,
-        )
-        exps.batch["advantages"] = advantages
-        exps.batch["returns"] = returns
+        """Modified from compute_grpo_outcome_advantage
+
+        Compute advantage for OPMD, operating only on Outcome reward
+        (with only one scalar reward for each response).
+
+            token_level_rewards: `(torch.Tensor)`
+                shape: (bs, response_length)
+            eos_mask: `(torch.Tensor)`
+                shape: (bs, response_length)
+            scores: `(torch.Tensor)`
+                shape: (bs, response_length)
+        """
+        token_level_rewards = exps.batch["token_level_rewards"]
+        eos_mask = exps.batch["response_mask"]
+        # TODO (yanxi): confirm consistency with exps.batch["attention_mask"][:, -response_length:] in original implementation
+        index = exps.non_tensor_batch["uid"]
+        opmd_baseline = "mean"
+        tau = 1.0
+
+        response_length = token_level_rewards.shape[-1]
+        scores = token_level_rewards.sum(dim=-1)
+
+        id2score = defaultdict(list)
+        id2baseline = {}
+
+        with torch.no_grad():
+            bsz = scores.shape[0]
+            for i in range(bsz):
+                id2score[index[i]].append(scores[i])
+            for idx in id2score:
+                if len(id2score[idx]) == 1:
+                    id2baseline[idx] = torch.tensor(0.0)
+                    # TODO: consider id2baseline[idx] = id2score[idx] (so that this sample won't take effect?)
+                elif len(id2score[idx]) > 1:
+                    if opmd_baseline == "mean":
+                        id2baseline[idx] = torch.mean(torch.tensor(id2score[idx]))
+                    elif opmd_baseline == "logavgexp":
+                        rewards_tensor = torch.tensor(id2score[idx])
+                        # here we use the fact that logavgexp(x) = logsumexp(x) - log(len(x))
+                        id2baseline[idx] = tau * (
+                            torch.logsumexp(rewards_tensor / tau, dim=-1)
+                            - torch.log(torch.tensor(len(id2score[idx])))
+                        )
+                    else:
+                        raise NotImplementedError
+                else:
+                    raise ValueError(f"no score in prompt index: {idx}")
+            for i in range(bsz):
+                scores[i] = scores[i] - id2baseline[index[i]]
+            scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask
+
+        exps.batch["advantages"] = scores
+        exps.batch["returns"] = scores
 
         metrics = {
             # TODO: add meaningful metrics
@@ -41,63 +84,3 @@ def __call__(
     @classmethod
     def default_args(cls) -> Dict:
         return {}
-
-
-def compute_opmd_outcome_advantage(
-    token_level_rewards: torch.Tensor,
-    eos_mask: torch.Tensor,
-    index: torch.Tensor,
-    opmd_baseline: str = "mean",
-    tau: float = 1.0,
-):
-    """Modified from compute_grpo_outcome_advantage
-
-    Compute advantage for OPMD, operating only on Outcome reward
-    (with only one scalar reward for each response).
-    Args:
-        token_level_rewards: `(torch.Tensor)`
-            shape: (bs, response_length)
-        eos_mask: `(torch.Tensor)`
-            shape: (bs, response_length)
-
-    Returns:
-        advantages: `(torch.Tensor)`
-            shape: (bs, response_length)
-        Returns: `(torch.Tensor)`
-            shape: (bs, response_length)
-    """
-    response_length = token_level_rewards.shape[-1]
-    scores = token_level_rewards.sum(dim=-1)
-
-    id2score = defaultdict(list)
-    id2baseline = {}
-
-    with torch.no_grad():
-        bsz = scores.shape[0]
-        for i in range(bsz):
-            id2score[index[i]].append(scores[i])
-        for idx in id2score:
-            if len(id2score[idx]) == 1:
-                id2baseline[idx] = torch.tensor(0.0)
-                # TODO: consider id2baseline[idx] = id2score[idx] (so that this sample won't take effect?)
-            elif len(id2score[idx]) > 1:
-                if opmd_baseline == "mean":
-                    id2baseline[idx] = torch.mean(torch.tensor(id2score[idx]))
-                elif opmd_baseline == "logavgexp":
-                    rewards_tensor = torch.tensor(id2score[idx])
-                    # NOTE: we use the fact that logavgexp(x) = logsumexp(x) - log(len(x)).
-                    # Hopefully the logsumexp calculation is numerically stable (as claimed by PyTorch's doc)
-                    # in cases where tau is small...
-                    id2baseline[idx] = tau * (
-                        torch.logsumexp(rewards_tensor / tau, dim=-1)
-                        - torch.log(torch.tensor(len(id2score[idx])))
-                    )
-                else:
-                    raise NotImplementedError
-            else:
-                raise ValueError(f"no score in prompt index: {idx}")
-        for i in range(bsz):
-            scores[i] = scores[i] - id2baseline[index[i]]
-        scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask
-
-    return scores, scores

trinity/algorithm/advantage_fn/ppo_advantage.py

@@ -27,40 +27,7 @@ def __call__(
         exps: DataProto,
         **kwargs,
     ) -> Tuple[DataProto, Dict]:
-        advantages, returns = compute_gae_advantage_return(
-            token_level_rewards=exps.batch["token_level_rewards"],
-            values=exps.batch["values"],
-            eos_mask=exps.batch["response_mask"],
-            gamma=self.gamma,
-            lam=self.lam,
-        )
-        exps.batch["advantages"] = advantages
-        exps.batch["returns"] = returns
-
-        metrics = {
-            # TODO: add meaningful metrics
-        }
-
-        return exps, metrics
-
-    @classmethod
-    def default_args(cls) -> Dict:
-        return {
-            "gamma": 1.0,
-            "lam": 1.0,
-        }
-
-
-def compute_gae_advantage_return(
-    token_level_rewards: torch.Tensor,
-    values: torch.Tensor,
-    eos_mask: torch.Tensor,
-    gamma: float,
-    lam: float,
-):
-    """Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py
-
-    Args:
+        """
         token_level_rewards: `(torch.Tensor)`
             shape: (bs, response_length)
         values: `(torch.Tensor)`
@@ -71,31 +38,49 @@ def compute_gae_advantage_return(
             discounted factor used in RL
         lam: `(float)`
             lambda value when computing Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438)
-
-    Returns:
         advantages: `(torch.Tensor)`
             shape: (bs, response_length)
-        Returns: `(torch.Tensor)`
+        returns: `(torch.Tensor)`
             shape: (bs, response_length)
+        """
+        token_level_rewards = exps.batch["token_level_rewards"]
+        values = exps.batch["values"]
+        eos_mask = exps.batch["response_mask"]
+        gamma = self.gamma
+        lam = self.lam
+
+        with torch.no_grad():
+            lastgaelam = 0
+            advantages_reversed = []
+            gen_len = token_level_rewards.shape[-1]
+
+            # values = values * eos_mask TODO: may use in multi-turn
+            for t in reversed(range(gen_len)):
+                nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
+                delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
+
+                lastgaelam = delta + gamma * lam * lastgaelam
+                # lastgaelam = torch.where(  # TODO: may use in multi-turn
+                #     eos_mask[:, t] == 1, delta + gamma * lam * lastgaelam, lastgaelam
+                # )
+                advantages_reversed.append(lastgaelam)
+            advantages = torch.stack(advantages_reversed[::-1], dim=1)
+
+            returns = advantages + values
+            advantages = masked_whiten(advantages, eos_mask)
 
-    """
-    with torch.no_grad():
-        lastgaelam = 0
-        advantages_reversed = []
-        gen_len = token_level_rewards.shape[-1]
+        exps.batch["advantages"] = advantages
+        exps.batch["returns"] = returns
 
-        # values = values * eos_mask TODO: may use in multi-turn
-        for t in reversed(range(gen_len)):
-            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
-            delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
+        metrics = {
+            # TODO: add meaningful metrics
+        }
 
-            lastgaelam = delta + gamma * lam * lastgaelam
-            # lastgaelam = torch.where(  # TODO: may use in multi-turn
-            #     eos_mask[:, t] == 1, delta + gamma * lam * lastgaelam, lastgaelam
-            # )
-            advantages_reversed.append(lastgaelam)
-        advantages = torch.stack(advantages_reversed[::-1], dim=1)
+        return exps, metrics
 
-        returns = advantages + values
-        advantages = masked_whiten(advantages, eos_mask)
-    return advantages, returns
+    @classmethod
+    def default_args(cls) -> Dict:
+        return {
+            "gamma": 1.0,
+            "lam": 1.0,
+        }