Update yaml and doc; fix bugs

Author: yanxi-chen

examples/grpo_gsm8k_trainable_ruler/README.md

@@ -4,22 +4,31 @@
 Ref: ART's RULER; Kimi-k2.
 
 
-Simulate a scenario where only a fraction of tasks have ground-truth answers for rule-based reward.
-
+Simulate a scenario where only a fraction (`PROBABILITY_GROUND_TRUTH_AVAILABLE = 0.2`) of tasks have ground-truth answers.
+Two RL objectives are optimized jointly: one for solution generation, the other for RULER-reward generation.
 
 
 ## Configurations and Metrics
 
-The config files are located in [`gsm8k_ruler.yaml`](gsm8k_ruler.yaml) and [`train_gsm8k_ruler.yaml`](train_gsm8k_ruler.yaml).
+The config files are located in [`gsm8k_ruler.yaml`](gsm8k_ruler.yaml) and [`train_gsm8k_trainable_ruler.yaml`](train_gsm8k_trainable_ruler.yaml).
 
 Some key configs in this example are:
 
-(TODO)
+* `default_workflow_type`: set to `math_trainable_ruler_workflow`
+* `std_threshold` for GRPO advantage: set to small value, filter out group of experiences with same rewards (e.g., when RULER fails to return valid scores, they are set to all zero)
+* `sync_style`: use `dynamic_by_explorer`, due to filtering of experiences
+* `train_batch_size`: set to 960; note that one explore step can generate more than 96 * 8 = 768 experiences
+* `lr`: set to small value (2e-6) for stability, as rewards can be noisy
+
 
 
 Some important metrics to pay attention to are:
 
-(TODO)
+* `reward`: reward calculated by rule or by RULER
+* `gold_reward`: sum of `accuracy_reward` and `format_reward`, rule-based calculation with ground truth
+* `judge_success`: whether RULER successfully returns a valid score (a coarse estimation, mix up two types of experiences)
+* `reward_for_judger`: reward for the LLM working as a RULER reward model, calculated by mean absolute error (MAE) distance from gold scores
+* `eval_accuracy`: accuracy on the evaluation set (ultimate metric for success of RL)
 
 
 ## Results
@@ -32,4 +41,4 @@ Compare with baseline: previous RULER workflow with Qwen2.5-1.5B-Instruct as LLM
 
 ## Potential improvements
 
-balance number of samples / loss weights for generation vs RULER
+balance number of samples / loss weights for generation vs for RULER

examples/grpo_gsm8k_trainable_ruler/gsm8k_ruler.yaml

@@ -1,5 +1,5 @@
-project: "Trinity-RFT-gsm8k-ruler"
-name: "qwen2.5-1.5B-gsm8k-ruler"
+project: "Trinity-RFT-gsm8k-trainable-ruler"
+name: "qwen2.5-1.5B-gsm8k-trainable-ruler"
 checkpoint_root_dir: /PATH/TO/CHECKPOINT/
 algorithm:
   algorithm_type: grpo
@@ -8,14 +8,16 @@ algorithm:
   repeat_times: 8
 model:
   model_path: /PATH/TO/MODEL/
-  max_response_tokens: 1024
-  max_model_len: 1280
+  max_prompt_tokens: 12288
+  max_response_tokens: 12288
+  max_model_len: 16000  # slightly smaller than ppo_max_token_len_per_gpu (16384)
 cluster:
   node_num: 1
   gpu_per_node: 8
 buffer:
   total_epochs: 1
   batch_size: 96
+  train_batch_size: 960
   explorer_input:
     taskset:
       name: gsm8k
@@ -37,7 +39,7 @@ buffer:
       format:
         prompt_key: 'question'
         response_key: 'answer'
-    default_workflow_type: 'math_ruler_workflow'
+    default_workflow_type: 'math_trainable_ruler_workflow'
   trainer_input:
     experience_buffer:
       name: gsm8k_buffer
@@ -47,26 +49,18 @@ explorer:
   runner_num: 32
   rollout_model:
     engine_type: vllm_async
-    engine_num: 2
+    engine_num: 4
     tensor_parallel_size: 1
     enable_prefix_caching: false
     enforce_eager: true
     dtype: bfloat16
     seed: 42
-  auxiliary_models:
-    - model_path: /PATH/TO/Qwen2.5-32B-Instruct
-      engine_num: 1
-      tensor_parallel_size: 2
-      enable_thinking: false
-      max_prompt_tokens: 12288
-      max_response_tokens: 12288
-      max_model_len: 16384
 synchronizer:
   sync_style: dynamic_by_explorer
   sync_method: 'nccl'
   sync_interval: 5
   sync_timeout: 3600
 trainer:
   trainer_type: 'verl'
-  trainer_config_path: 'examples/grpo_gsm8k_ruler/train_gsm8k_ruler.yaml'
+  trainer_config_path: 'examples/grpo_gsm8k_trainable_ruler/train_gsm8k_trainable_ruler.yaml'
   save_interval: 100

…k_trainable_ruler/train_gsm8k_ruler.yaml …e_ruler/train_gsm8k_trainable_ruler.yamlexamples/grpo_gsm8k_trainable_ruler/train_gsm8k_ruler.yaml renamed to examples/grpo_gsm8k_trainable_ruler/train_gsm8k_trainable_ruler.yaml examples/grpo_gsm8k_trainable_ruler/train_gsm8k_ruler.yaml renamed to examples/grpo_gsm8k_trainable_ruler/train_gsm8k_trainable_ruler.yaml

@@ -12,6 +12,7 @@
 from .eval_workflow import MathEvalWorkflow
 from .math_rm_workflow import MathRMWorkflow
 from .math_ruler_workflow import MathRULERWorkflow
+from .math_trainable_ruler_workflow import MathTrainableRULERWorkflow
 from .simple_mm_workflow import SimpleMMWorkflow
 from .workflow import WORKFLOWS, MathWorkflow, SimpleWorkflow, Task, Workflow
 
@@ -34,5 +35,6 @@
     "AgentScopeReactV2MathWorkflow",
     "EmailSearchWorkflow",
     "MathRULERWorkflow",
+    "MathTrainableRULERWorkflow",
     "SimpleMMWorkflow",
 ]

trinity/common/workflows/__init__.py

@@ -146,6 +146,11 @@ def get_ruler_scores(
         try:
             scores = ast.literal_eval(lst_as_str)
             scores = [max(0.0, min(1.0, score)) for score in scores]  # clip to range [0, 1]
+            if len(scores) != num_responses:
+                logger.warning(
+                    "The length of list in judger response does not match num_responses."
+                )
+                return False, [0.0 for _ in range(num_responses)]
             return True, scores
         except Exception:
             logger.warning("Unable to parse the list in judger response, set scores to all zero.")

trinity/common/workflows/math_ruler_workflow.py

@@ -3,8 +3,8 @@
 import ast
 from copy import deepcopy
 from typing import Any, List, Optional, Tuple
-import numpy as np
 
+import numpy as np
 import openai
 
 from trinity.common.experience import Experience
@@ -16,7 +16,8 @@
 logger = get_logger(__name__)
 
 # the probability that the ground truth is assumed to be available for RL
-PROB_GROUND_TRUTH_AVAILABLE = 0.2
+PROBABILITY_GROUND_TRUTH_AVAILABLE = 0.2
+
 
 @WORKFLOWS.register_module("math_trainable_ruler_workflow")
 class MathTrainableRULERWorkflow(SimpleWorkflow):
@@ -71,74 +72,77 @@ def run(self) -> List[Experience]:
             gold_reward = sum(gold_reward_dict.values())
             response.metrics.update({"gold_reward": gold_reward})
 
-            response.eid.task = self.task.task_id  # task_id is set explicitly within workflow!
+            # set task_id explicitly within workflow!
+            response.eid.task = str(self.task.task_id)
             response.eid.run = i + self.run_id_base
 
             gold_rewards.append(gold_reward)
-            gold_scores_scaled.append((gold_reward + 0.1) / 1.2)  # scale from range [-0.1, 1.1] to [0, 1]
+            gold_scores_scaled.append(
+                (gold_reward + 0.1) / 1.2
+            )  # scale from range [-0.1, 1.1] to [0, 1]
 
         # Part 2: get and use RULER scores
         ruler_rollout_args = deepcopy(self.rollout_args)
-        ground_truth_is_available = np.random.rand() < PROB_GROUND_TRUTH_AVAILABLE
+        ground_truth_is_available = np.random.rand() < PROBABILITY_GROUND_TRUTH_AVAILABLE
 
         if ground_truth_is_available:
+            # Assuming that ground truth is accessible to RL:
             # - set exp's reward to gold reward
             # - generate RULER scores for repeat_times, construct ruler_responses
             # - return responses + ruler_responses
 
             judge_success_rate, ruler_responses, ruler_scores = self.get_ruler_responses(
-                responses=responses, 
+                responses=responses,
                 judger=self.model,  # use the policy model itself as judger!
                 ruler_rollout_args=ruler_rollout_args,
                 gold_scores=gold_scores_scaled,
             )
 
             for i, response in enumerate(responses):
                 response.reward = gold_rewards[i]
-                response.metrics.update({"judge_success": float(judge_success_rate)})
-            
+                response.metrics.update({"judge_success": judge_success_rate})
+
             for i, ruler_response in enumerate(ruler_responses):
-                if ruler_response.metrics is None:
-                    ruler_response.metrics = {}
-                ruler_response.metrics.update(
-                    {
-                        "judge_success": judge_success_rate,
-                        "reward_for_judger": ruler_response.reward,
-                    }
-                )
-                ruler_response.eid.task = -1 * self.task.task_id  # HACK to distinguish two types of experiences
+                # if ruler_response.metrics is None:
+                # ruler_response.metrics = {}
+                # ruler_response.metrics.update({"judge_success": judge_success_rate})
+                # ruler_response.metrics.update({"reward_for_judger": ruler_response.reward})
+
+                # set task_id explicitly, to distinguish two types of experiences!
+                ruler_response.eid.task = str(self.task.task_id) + "-ruler"
                 ruler_response.eid.run = i + self.run_id_base
 
             return responses + ruler_responses
 
         else:
+            # Assuming that ground truth is not accessible to RL:
             # - generate RULER scores only once
-            # - set exp's reward to RULER scores
+            # - set exp's reward to RULER score
             # - return responses
 
             ruler_rollout_args.n = 1
             judge_success_rate, ruler_responses, ruler_scores = self.get_ruler_responses(
-                responses=responses, 
+                responses=responses,
                 judger=self.model,  # use the policy model itself as judger!
                 ruler_rollout_args=ruler_rollout_args,
                 gold_scores=None,
             )
 
             for i, response in enumerate(responses):
                 response.reward = ruler_scores[i]
-                response.metrics.update({"judge_success": float(judge_success_rate)})
+                response.metrics.update({"judge_success": judge_success_rate})
 
             return responses
 
     def get_ruler_responses(
-        self, 
-        responses: List[Experience], 
+        self,
+        responses: List[Experience],
         judger: Any,
         ruler_rollout_args: Any,
         gold_scores: Optional[List[float]] = None,
-    ) -> Tuple[bool, List[float]]:
+    ) -> Tuple[float, List[Experience], List[float]]:
         """Get RULER scores
-        
+
         Returns:
             judge_success_rate: float
             ruler_responses: List[Experience]
@@ -194,6 +198,7 @@ def get_ruler_responses(
 
             if (idx1 == -1) or (idx2 == -1) or (idx1 > idx2):
                 logger.warning("Unable to extract a list from judger response.")
+                break
 
             lst_as_str = ruler_response_text[idx1 : (idx2 + 1)]
             try:
@@ -203,15 +208,23 @@ def get_ruler_responses(
                     judge_success_count += 1
                     ruler_scores = [ruler_scores[i] + scores[i] for i in range(len(ruler_scores))]
                     if gold_scores:
-                        mae_error = (np.array(ruler_scores) - np.array(gold_scores)).abs().mean()
+                        mae_error = (np.array(scores) - np.array(gold_scores)).abs().mean()
                         ruler_response.reward = 1.0 - mae_error
                 else:
-                    logger.warning("The length of list in judger response does not match num_responses.")
+                    logger.warning(
+                        "The length of list in judger response does not match num_responses."
+                    )
             except Exception:
                 logger.warning("Unable to parse the list in judger response.")
-            
+
         if judge_success_count > 0:
             ruler_scores = [score / judge_success_count for score in ruler_scores]
         judge_success_rate = 1.0 * judge_success_count / len(ruler_responses)
 
+        for ruler_response in ruler_responses:
+            if ruler_response.metrics is None:
+                ruler_response.metrics = {}
+            ruler_response.metrics.update({"judge_success": judge_success_rate})
+            ruler_response.metrics.update({"reward_for_judger": ruler_response.reward})
+
         return judge_success_rate, ruler_responses, ruler_scores