[fix] revert reward update and evaluation (#6295)

* Revert "rewrite reward fn"

This reverts commit d06042b.

* Revert "upgrade reward math verification"

This reverts commit a6085ff.

* Revert "fix bug"

This reverts commit 01640eb.

* Revert "reuse comm-group"

This reverts commit bd61918.

* Revert "Support evaluation during training"

This reverts commit 57a8839.

Author: YeAnbang

.gitignore

@@ -165,4 +165,3 @@ applications/ColossalChat/logs
 applications/ColossalChat/tests/logs
 applications/ColossalChat/wandb
 applications/ColossalChat/model
-applications/ColossalChat/eval

applications/ColossalChat/coati/distributed/consumer.py

@@ -36,7 +36,6 @@ def __init__(
         minibatch_size: int = 1,
         save_interval: int = 100,
         save_dir: str = "./model",
-        eval_interval: int = -1,
     ):
         self.num_producers = num_producers
         self.num_episodes = num_episodes
@@ -52,7 +51,6 @@ def __init__(
         self.save_dir = save_dir
         assert batch_size % minibatch_size == 0, "batch_size should be divisible by microbatch_size"
         self.num_microbatches = batch_size // minibatch_size
-        self.eval_interval = eval_interval
 
         self.model_config = model_config
         self.plugin_config = plugin_config
@@ -95,6 +93,7 @@ def setup(self) -> None:
                 cc.init_collective_group(self.num_producers + 1, self.num_producers, group_name="sync_model")
 
         self.buffer = []
+
         self.recv_cnt = 0
 
     def state_dict(self) -> Dict[str, torch.Tensor]:
@@ -111,27 +110,6 @@ def loop(self) -> None:
             with tqdm(range(self.num_update_per_episode), desc=f"Episode {episode}", disable=self.rank != 0) as pbar:
                 for step in pbar:
                     i = 0
-                    if self.eval_interval > 0 and step % self.eval_interval == 0:
-                        eval_statistics = None
-                        eval_global_step = None
-                        for r in range(self.num_producers):
-                            print(f"[T{dist.get_rank()}] Recv eval result episode {episode} step {step} from {r}")
-                            local_eval_result = ray_broadcast_tensor_dict(
-                                None, src=0, device=self.device, group_name=f"sync_data_{r}"
-                            )
-                            assert "consumer_global_step" in local_eval_result
-                            eval_global_step = local_eval_result.pop("consumer_global_step").item()
-                            if eval_statistics is None:
-                                eval_statistics = local_eval_result
-                            else:
-                                eval_statistics = {
-                                    k: eval_statistics[k] + local_eval_result[k] for k in eval_statistics
-                                }
-                        eval_statistics = {"eval/" + k: (v[0] / v[1]).item() for k, v in eval_statistics.items()}
-                        if dist.get_rank() == 0:
-                            if hasattr(self, "wandb_run"):
-                                self.wandb_run.log(eval_statistics, step=eval_global_step)
-                            print(f"Eval statistics: {eval_statistics}")
                     for _ in range(self.num_recv_per_update):
                         # receive data from producers
                         for r in range(self.num_producers):
@@ -217,7 +195,6 @@ def __init__(
         minibatch_size=1,
         save_interval: int = 100,
         save_dir="./model",
-        eval_interval: int = -1,
     ):
         super().__init__(
             num_producers,
@@ -232,9 +209,6 @@ def __init__(
             model_config,
             plugin_config,
             minibatch_size,
-            save_interval,
-            save_dir,
-            eval_interval,
         )
         path = model_config.pop("path")
         self.model = AutoModelForCausalLM.from_pretrained(path, **model_config)

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -40,7 +40,6 @@ def __init__(
         project_name=None,
         save_interval: int = 100,
         save_dir="./model",
-        eval_interval: int = -1,
     ):
         print(f"Using GRPO config: {grpo_config}")
         if grpo_config.get("loss_variation", "sample_level") == "token_level":
@@ -73,7 +72,6 @@ def __init__(
             minibatch_size,
             save_interval=save_interval,
             save_dir=save_dir,
-            eval_interval=eval_interval,
         )
         path = model_config.pop("path")
         self.policy_model = AutoModelForCausalLM.from_pretrained(path, **model_config)
@@ -530,5 +528,4 @@ def state_dict(self):
         self.policy_model._force_wait_all_gather()
         model = self.policy_model.unwrap()
         state_dict = model.state_dict()
-        state_dict["consumer_global_step"] = torch.tensor([self.global_step], device=self.device)
         return state_dict

applications/ColossalChat/coati/distributed/inference_backend.py

@@ -205,8 +205,7 @@ def __init__(
         generate_config = generate_config.copy()
         generate_config.update(self.FORCE_GENERATE_CONFIG)
         generate_config.update({"n": num_generations})
-        self.generate_config = generate_config
-        self.sample_params = SamplingParams(**generate_config)
+        self.generate_config = SamplingParams(**generate_config)
         self.model_config = model_config
         self.tokenizer = tokenizer
         self.num_generations = num_generations
@@ -220,9 +219,8 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
         micro_batch_input_ids_no_padding = [
             micro_batch_input_ids[i][first_non_padding_token_idx[i] :] for i in range(micro_batch_size)
         ]
-        sample_params = kwargs.get("sample_params", self.sample_params)
         outputs = self.llm.generate(
-            prompt_token_ids=micro_batch_input_ids_no_padding, sampling_params=sample_params, use_tqdm=False
+            prompt_token_ids=micro_batch_input_ids_no_padding, sampling_params=self.generate_config, use_tqdm=False
         )
         out_tokens = []
         out_len = []
@@ -268,11 +266,11 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
             "response_idx": response_idx,
         }
 
-        data = {k: v.view(micro_batch_size, -1, v.size(-1)) for k, v in data.items()}
+        data = {k: v.view(micro_batch_size, self.num_generations, 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["gt_answer"] = kwargs["gt_answer"].repeat_interleave(self.num_generations, dim=1)
         data = {k: v.to(get_current_device()) for k, v in data.items()}
         return data
 

applications/ColossalChat/coati/distributed/launch.py

@@ -34,7 +34,7 @@ def launch_distributed(
     inference_microbatch_size: int,
     train_batch_size: int,
     train_minibatch_size: int,
-    train_dataset_config: Dict[str, Any],
+    dataset_config: Dict[str, Any],
     dataloaders_config: Dict[str, Any],
     inference_model_config: Dict[str, Any],
     generate_config: Dict[str, Any],
@@ -50,9 +50,6 @@ def launch_distributed(
     project_name: Optional[str] = None,
     save_interval: int = 100,
     save_dir: str = "./model",
-    eval_dataset_config: Optional[Dict[str, Any]] = None,
-    eval_interval: int = 100,
-    eval_save_dir: Optional[str] = None,
 ):
 
     if core_algo not in ALGO_MAP:
@@ -63,9 +60,9 @@ def launch_distributed(
     train_dp_size = get_dp_size_fast(num_consumer_procs, plugin_config)
     assert (inference_batch_size * num_producers) % (train_batch_size * train_dp_size) == 0
 
-    dataset_path = train_dataset_config["path"]
+    dataset_path = dataset_config["path"]
     num_samples = get_jsonl_size_fast(dataset_path)
-    global_inference_batch_size = inference_batch_size * num_producers  # TODO: this doesn't support TP on producer
+    global_inference_batch_size = inference_batch_size * num_producers
     num_update_per_episode = num_samples // global_inference_batch_size
     num_recv_per_update = inference_batch_size // inference_microbatch_size
 
@@ -77,7 +74,7 @@ def launch_distributed(
             num_consumer_procs=num_consumer_procs,
             num_episodes=num_episodes,
             batch_size=inference_batch_size,
-            train_dataset_config=train_dataset_config,
+            dataset_config=dataset_config,
             dataloaders_config=dataloaders_config,
             model_config=inference_model_config,
             generate_config=generate_config,
@@ -86,10 +83,6 @@ def launch_distributed(
             backend=inference_backend,
             num_generations=num_generations,
             consumer_plugin_config=plugin_config,
-            eval_dataset_config=eval_dataset_config,
-            eval_interval=eval_interval * num_recv_per_update,
-            evaluation_function_type=grpo_config["reward_fn_type"],
-            eval_save_dir=eval_save_dir,
         )
         procs.append(producer)
     generate_config_consumer = copy.deepcopy(generate_config)
@@ -118,7 +111,6 @@ def launch_distributed(
             project_name=project_name,
             save_interval=save_interval,
             save_dir=save_dir,
-            eval_interval=eval_interval,
         )
         procs.append(consumer)
     ray.get([p.setup.remote() for p in procs])