[feat] Support Code Generation RFT, Move Reward Calculation to Producer

.gitignore

@@ -167,3 +167,7 @@ applications/ColossalChat/wandb
 applications/ColossalChat/model
 applications/ColossalChat/eval
 applications/ColossalChat/rollouts
+applications/ColossalChat/*.txt
+applications/ColossalChat/*.db
+applications/ColossalChat/stdin
+applications/ColossalChat/*.zip

applications/ColossalChat/coati/dataset/loader.py

@@ -367,9 +367,9 @@ def apply_chat_template_and_mask(
     }
 
     # Format for RL.
-    gt_answer = None
-    if "messages" in chat and "gt_answer" in chat:
-        gt_answer = chat["gt_answer"]
+    if "messages" in chat:
+        gt_answer = chat.get("gt_answer", None)
+        test_cases = chat.get("test_cases", None)
         chat = [chat["messages"]]
 
     tokens = []
@@ -402,12 +402,14 @@ def apply_chat_template_and_mask(
     labels[~torch.tensor(assistant_mask, dtype=torch.bool)] = ignore_idx
 
     if gt_answer is not None:
-        gt_answer = tokenizer.encode(
-            gt_answer, padding="max_length", truncation=True, max_length=128, return_tensors="pt"
-        )
-        gt_answer = gt_answer.squeeze(1)
         return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "gt_answer": gt_answer}
-
+    elif test_cases is not None:
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "test_cases": test_cases,
+        }
     return {
         "input_ids": input_ids,
         "attention_mask": attention_mask,
@@ -440,3 +442,20 @@ def __getitem__(self, index: int):
             tokens = apply_chat_template_and_mask(self.tokenizer, message, self.max_length, self.system_prompt)
             self.tokenized_texts[index] = dict(tokens)
         return self.tokenized_texts[index]
+
+
+def collate_fn_grpo(batch):
+    input_ids = [item["input_ids"] for item in batch]
+    attention_mask = [item["attention_mask"] for item in batch]
+    labels = [item["labels"] for item in batch]
+    # Assume input_ids, attention_mask, labels are already of the same length,
+    # otherwise use pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
+    input_ids = torch.stack(input_ids)
+    attention_mask = torch.stack(attention_mask)
+    labels = torch.stack(labels)
+    ret = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
+    if "test_cases" in batch[0]:
+        ret["test_cases"] = [item["test_cases"] for item in batch]
+    if "gt_answer" in batch[0]:
+        ret["gt_answer"] = [item["gt_answer"] for item in batch]
+    return ret

applications/ColossalChat/coati/distributed/consumer.py

@@ -123,21 +123,16 @@ def loop(self) -> None:
                             # calculate group reward et al. filtering. As only the filtered group will be used for training (which is incomplete),
                             # we need to calculate the metrics before filtering here for logging
                             # [batch_size, num_generations, ...] -> [batch_size * num_generations, ...]
-                            raw_batch_with_reward = self.calculate_reward(
-                                {k: v.view(-1, v.size(-1)) if k != "temperature" else v for k, v in raw_batch.items()}
-                            )
-                            raw_batch_with_reward = {
+                            raw_batch = {
                                 k: v.view(-1, self.num_generations, v.size(-1)) if k != "temperature" else v
-                                for k, v in raw_batch_with_reward.items()
+                                for k, v in raw_batch.items()
                             }
                             # [batch_size, num_generations] -> [batch_size]
-                            reward = raw_batch_with_reward["reward"][:, :, 0]
-                            format_acc = raw_batch_with_reward["format_acc"][:, :, 0]
-                            ans_acc = raw_batch_with_reward["ans_acc"][:, :, 0]
+                            reward = raw_batch["reward"][:, :, 0]
+                            format_acc = raw_batch["format_acc"][:, :, 0]
+                            ans_acc = raw_batch["ans_acc"][:, :, 0]
                             response_len = (
-                                raw_batch_with_reward["response_idx"][:, :, 1]
-                                - raw_batch_with_reward["response_idx"][:, :, 0]
-                                + 1
+                                raw_batch["response_idx"][:, :, 1] - raw_batch["response_idx"][:, :, 0] + 1
                             ).type(torch.float32)
                             effective_group_mask = None
                             if self.filter_range is not None and self.grpo_config.get("dynamic_batching", True):
@@ -146,8 +141,8 @@ def loop(self) -> None:
                                 effective_group_mask = torch.logical_and(
                                     group_ans_acc_mean > self.filter_range[0], group_ans_acc_mean < self.filter_range[1]
                                 )
-                            raw_batch_with_reward = unbind_batch(raw_batch_with_reward)  # List[Dict[str, torch.Tensor]]
-                            for group_idx, group_with_reward in enumerate(raw_batch_with_reward):
+                            raw_batch = unbind_batch(raw_batch)  # List[Dict[str, torch.Tensor]]
+                            for group_idx, group_with_reward in enumerate(raw_batch):
                                 self.buffer.append(
                                     [
                                         (
@@ -163,7 +158,7 @@ def loop(self) -> None:
                                 )
                             if effective_group_mask is not None:
                                 print(
-                                    f"[T{dist.get_rank()}] Filter recv data: {len(raw_batch_with_reward)} -> {torch.sum(effective_group_mask).cpu().item()} effective groups"
+                                    f"[T{dist.get_rank()}] Filter recv data: {len(raw_batch)} -> {torch.sum(effective_group_mask).cpu().item()} effective groups"
                                 )
                         # mapping the effective group to the raw group for indexing
                         effective_group_to_raw_group_mapping = {}

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -1,13 +1,11 @@
 from contextlib import nullcontext
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import ray
 import torch
 import wandb
 from coati.distributed.consumer import BaseConsumer
 from coati.distributed.loss import PolicyLoss
-from coati.distributed.reward.reward_fn import boxed_math_reward_fn, math_reward_fn
-from coati.distributed.reward.verifiable_reward import VerifiableReward
 from coati.distributed.utils import calc_action_log_probs
 from coati.trainer.utils import all_reduce_mean, all_reduce_sum
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -119,20 +117,7 @@ def __init__(
                     "either max_tokens (vllm) or max_new_tokens (transformers) must be set in generate_config."
                 )
         # Initialize verifiable reward.
-        response_format_tags = grpo_config.get("response_format_tags", None)
-        reward_model_kwargs = {
-            k: v
-            for k, v in grpo_config.items()
-            if k in ["soft_over_length_punishment", "max_new_tokens", "cache_length"]
-        }
-        self.reward_model = VerifiableReward(
-            reward_fns=[
-                math_reward_fn if grpo_config.get("reward_fn_type") == "think_answer_tags" else boxed_math_reward_fn
-            ],
-            tokenizer=self.tokenizer,
-            tags=response_format_tags,
-            **reward_model_kwargs,
-        )
+        grpo_config.get("response_format_tags", None)
         self.global_step = 0
 
         self.lr_scheduler = CosineAnnealingWarmupLR(
@@ -498,40 +483,6 @@ def _criterion(outputs, inputs):
         else:
             return None
 
-    def calculate_reward(self, rollout: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Calculate the group reward for the given rollout group.
-
-        Args:
-            rollout_group (Dict[str, Any]):
-                a group of samples generated by the model from the same prompt
-                contain the following keys:
-                    "input_ids": torch.Tensor, [num_of_generation, prompt_length + response_length]
-                    "attention_mask": torch.Tensor, [num_of_generation, prompt_length + response_length]
-                    "action_mask": torch.Tensor, [num_of_generation, response_length]
-                    "action_log_probs": torch.Tensor, [num_of_generation, response_length]
-                    "response_idx": int, torch.Tensor, [num_of_generation, 2]
-                    "gt_answer": torch.Tensor, [num_of_generation, 128]
-                    "temperature": torch.Tensor, [] (scalar)
-
-        Returns:
-            Dict[str, Any]: The new group data with calculated reward.
-        """
-        reward_model_output = self.reward_model(
-            rollout["input_ids"],
-            gt_answer=rollout["gt_answer"],
-            response_idx=rollout["response_idx"],
-        )
-        # [num_of_generation]
-        reward = torch.tensor([value[0] for value in reward_model_output]).to(rollout["input_ids"].device)
-        format_acc = torch.tensor([value[1] for value in reward_model_output]).to(rollout["input_ids"].device)
-        ans_acc = torch.tensor([value[2] for value in reward_model_output]).to(rollout["input_ids"].device)
-
-        rollout["reward"] = reward.view((-1, 1))
-        rollout["format_acc"] = format_acc.view((-1, 1))
-        rollout["ans_acc"] = ans_acc.view((-1, 1))
-        return rollout
-
     def state_dict(self):
         self.policy_model._force_wait_all_gather()
         model = self.policy_model.unwrap()

applications/ColossalChat/coati/distributed/inference_backend.py

@@ -74,9 +74,8 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
         micro_batch_size = input_ids.size(0)
         input_ids = input_ids.to(get_current_device())
         attention_mask = attention_mask.to(get_current_device())
-        gt_answer = None
-        if "gt_answer" in kwargs:
-            gt_answer = kwargs.pop("gt_answer")
+        gt_answer = kwargs.pop("gt_answer", None)
+        test_cases = kwargs.pop("test_cases", None)
         if self.num_generations > 1:
             input_ids = input_ids.repeat_interleave(self.num_generations, dim=0)
             attention_mask = attention_mask.repeat_interleave(self.num_generations, dim=0)
@@ -116,8 +115,9 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
         data = {k: v.view(micro_batch_size, self.num_generations, v.size(-1)) for k, v in data.items()}
 
         if gt_answer is not None:
-            # repeat gt_answer for each prompt.
-            data["gt_answer"] = gt_answer.repeat_interleave(self.num_generations, dim=1)
+            data["gt_answer"] = gt_answer
+        if test_cases is not None:
+            data["test_cases"] = test_cases
         data = {k: v.to(get_current_device()) for k, v in data.items()}
         return data
 
@@ -269,11 +269,11 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
         }
 
         data = {k: v.view(micro_batch_size, -1, v.size(-1)) for k, v in data.items()}
-
-        if "gt_answer" in kwargs:
-            # repeat gt_answer for each prompt.
-            data["gt_answer"] = kwargs["gt_answer"].repeat_interleave(data["input_ids"].size(1), dim=1)
         data = {k: v.to(get_current_device()) for k, v in data.items()}
+        if "gt_answer" in kwargs:
+            data["gt_answer"] = kwargs["gt_answer"]
+        if "test_cases" in kwargs:
+            data["test_cases"] = kwargs["test_cases"]
         return data
 
     def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None:

applications/ColossalChat/coati/distributed/launch.py

@@ -37,7 +37,6 @@ def launch_distributed(
     train_batch_size: int,
     train_minibatch_size: int,
     train_dataset_config: Dict[str, Any],
-    dataloaders_config: Dict[str, Any],
     inference_model_config: Dict[str, Any],
     generate_config: Dict[str, Any],
     train_model_config: Dict[str, Any],
@@ -89,7 +88,6 @@ def launch_distributed(
             num_episodes=num_episodes,
             batch_size=inference_batch_size,
             train_dataset_config=train_dataset_config,
-            dataloaders_config=dataloaders_config,
             model_config=inference_model_config,
             generate_config=generate_config,
             tokenizer_config=tokenizer_config,
@@ -99,8 +97,7 @@ def launch_distributed(
             consumer_plugin_config=plugin_config,
             eval_dataset_config=eval_dataset_config,
             eval_interval=eval_interval,
-            evaluation_function_type=grpo_config["reward_fn_type"],
-            response_format_tags=grpo_config["response_format_tags"],
+            grpo_config=grpo_config,
             eval_save_dir=eval_save_dir,
             eval_generation_config=eval_generation_config,
             project_name=project_name,