Merge pull request #6309 from hpcaitech/grpo-eval-dev

[feat] Support evaluation during training

Author: YeAnbang

.gitignore

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

applications/ColossalChat/coati/distributed/consumer.py

@@ -93,7 +93,6 @@ 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]:
@@ -209,6 +208,8 @@ def __init__(
             model_config,
             plugin_config,
             minibatch_size,
+            save_interval,
+            save_dir,
         )
         path = model_config.pop("path")
         self.model = AutoModelForCausalLM.from_pretrained(path, **model_config)

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -33,12 +33,13 @@ def __init__(
         plugin_config,
         minibatch_size=1,
         num_generations=8,
-        use_wandb=True,
         generate_config=None,
         grpo_config={},
-        project_name=None,
         save_interval: int = 100,
         save_dir="./model",
+        project_name: str = None,
+        run_name: str = None,
+        wandb_group_name: str = None,
     ):
         print(f"Using GRPO config: {grpo_config}")
         if (
@@ -84,6 +85,9 @@ def __init__(
         self.effective_sample_count = 0
         self.effective_prompt_count = 0
         self.total_sample_count = 0
+        self.project_name = project_name
+        self.run_name = run_name
+        self.wandb_group_name = wandb_group_name
 
         self.policy_loss_fn = PolicyLoss(
             clip_eps_low=grpo_config.get("clip_eps_low", 0.2),
@@ -134,7 +138,6 @@ def __init__(
             **reward_model_kwargs,
         )
         self.global_step = 0
-        self.use_wandb = use_wandb
 
         self.lr_scheduler = CosineAnnealingWarmupLR(
             optimizer=self.optimizer,
@@ -145,13 +148,16 @@ def __init__(
 
     def setup(self):
         super().setup()
-        if self.use_wandb and (
-            (not self.plugin.pp_size > 1 and self.rank == 0)
-            or (self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0)
+        if (not self.plugin.pp_size > 1 and self.rank == 0) or (
+            self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0
         ):
-            # Initialize wandb.
-            name = f"{self.generate_config['backend']}_bs_{self.batch_size*self.dp_size}_temp_{self.generate_config['temperature']:.01f}_top_p_{self.generate_config['top_p']:.02f}"
-            self.wandb_run = wandb.init(project=self.project_name, sync_tensorboard=True, dir="./wandb", name=name)
+            self.wandb_run = wandb.init(
+                project=self.project_name,
+                sync_tensorboard=False,
+                dir="./wandb",
+                name=self.run_name,
+                group=self.wandb_group_name,
+            )
 
         self.policy_model, self.optimizer, _, _, self.lr_scheduler = self.booster.boost(
             self.policy_model, self.optimizer, lr_scheduler=self.lr_scheduler
@@ -265,7 +271,6 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
         total_samples = all_reduce_sum(torch.sum(torch.ones_like(loss_mask, device=loss_mask.device)), self.plugin)
         self.effective_sample_count += effective_samples.item()
         self.total_sample_count += total_samples.item()
-
         pbar.set_postfix(
             {
                 "Global Step": self.global_step,
@@ -506,8 +511,8 @@ def _criterion(outputs, inputs):
                     }
                     if self.policy_loss_fn.beta > 0:
                         metrics["train/kl"] = self.accum_kl.item() / self.accum_count
-
-                    self.wandb_run.log(metrics)
+                    if self.wandb_run is not None:
+                        self.wandb_run.log(metrics)
                 self.accum_loss.zero_()
                 self.accum_reward.zero_()
                 self.accum_ans_acc.zero_()
@@ -521,7 +526,6 @@ def _criterion(outputs, inputs):
                 # All gather excessive prompts index across DP ranks.
                 excessive_prompts_idx = [idx + self.dp_rank * self.minibatch_size for idx in excessive_prompts_idx]
                 excessive_prompts_idx = all_gather_tensors(excessive_prompts_idx, self.plugin)
-
             return loss_scalar, excessive_prompts_idx
         else:
             return None, excessive_prompts_idx
@@ -530,4 +534,5 @@ 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,7 +205,8 @@ def __init__(
         generate_config = generate_config.copy()
         generate_config.update(self.FORCE_GENERATE_CONFIG)
         generate_config.update({"n": num_generations})
-        self.generate_config = SamplingParams(**generate_config)
+        self.generate_config = generate_config
+        self.sample_params = SamplingParams(**generate_config)
         self.model_config = model_config
         self.tokenizer = tokenizer
         self.num_generations = num_generations
@@ -219,8 +220,9 @@ 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=self.generate_config, use_tqdm=False
+            prompt_token_ids=micro_batch_input_ids_no_padding, sampling_params=sample_params, use_tqdm=False
         )
         out_tokens = []
         out_len = []
@@ -236,7 +238,7 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
                 log_probs.append(p)
 
         # pad them
-        max_len = self.generate_config.max_tokens
+        max_len = self.sample_params.max_tokens
         action_mask = torch.ones(len(out_tokens), max_len, dtype=attention_mask.dtype)
 
         for i, new_token_ids in enumerate(out_tokens):
@@ -266,11 +268,11 @@ def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwar
             "response_idx": response_idx,
         }
 
-        data = {k: v.view(micro_batch_size, self.num_generations, v.size(-1)) for k, v in data.items()}
+        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(self.num_generations, dim=1)
+            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()}
         return data
 

applications/ColossalChat/coati/distributed/launch.py

@@ -1,4 +1,5 @@
 import copy
+import uuid
 from typing import Any, Dict, Optional
 
 import ray
@@ -34,7 +35,7 @@ def launch_distributed(
     inference_microbatch_size: int,
     train_batch_size: int,
     train_minibatch_size: int,
-    dataset_config: Dict[str, Any],
+    train_dataset_config: Dict[str, Any],
     dataloaders_config: Dict[str, Any],
     inference_model_config: Dict[str, Any],
     generate_config: Dict[str, Any],
@@ -50,8 +51,11 @@ 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,
+    eval_generation_config: Optional[Dict[str, Any]] = None,
 ):
-
     if core_algo not in ALGO_MAP:
         raise NotImplementedError(f"{core_algo} is not supported yet.")
     else:
@@ -60,12 +64,15 @@ 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 = dataset_config["path"]
+    dataset_path = train_dataset_config["path"]
     num_samples = get_jsonl_size_fast(dataset_path)
     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
 
+    run_name = f"{inference_backend}_bs_{train_batch_size * train_dp_size}_temp_{generate_config['temperature']:.01f}_top_p_{generate_config['top_p']:.02f}"
+    wandb_group_name = str(uuid.uuid4())
+
     procs = []
     for i in range(num_producers):
         producer = SimpleProducer.options(num_gpus=num_proc_per_producer).remote(
@@ -74,7 +81,7 @@ def launch_distributed(
             num_consumer_procs=num_consumer_procs,
             num_episodes=num_episodes,
             batch_size=inference_batch_size,
-            dataset_config=dataset_config,
+            train_dataset_config=train_dataset_config,
             dataloaders_config=dataloaders_config,
             model_config=inference_model_config,
             generate_config=generate_config,
@@ -83,6 +90,14 @@ 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,
+            evaluation_function_type=grpo_config["reward_fn_type"],
+            eval_save_dir=eval_save_dir,
+            eval_generation_config=eval_generation_config,
+            project_name=project_name,
+            run_name=run_name,
+            wandb_group_name=wandb_group_name,
         )
         procs.append(producer)
     generate_config_consumer = copy.deepcopy(generate_config)
@@ -108,9 +123,11 @@ def launch_distributed(
             generate_config=generate_config_consumer,
             grpo_config=grpo_config,
             num_generations=num_generations,
-            project_name=project_name,
             save_interval=save_interval,
             save_dir=save_dir,
+            project_name=project_name,
+            run_name=run_name,
+            wandb_group_name=wandb_group_name,
         )
         procs.append(consumer)
     ray.get([p.setup.remote() for p in procs])