fix bug

Author: YeAnbang

applications/ColossalChat/coati/distributed/launch.py

@@ -87,7 +87,7 @@ def launch_distributed(
             num_generations=num_generations,
             consumer_plugin_config=plugin_config,
             eval_dataset_config=eval_dataset_config,
-            eval_interval=eval_interval,
+            eval_interval=eval_interval * num_recv_per_update,
             evaluation_function_type=grpo_config["reward_fn_type"],
             eval_save_dir=eval_save_dir,
         )

applications/ColossalChat/coati/distributed/producer.py

@@ -129,7 +129,7 @@ def __init__(
         else:
             raise ValueError(f"Unexpected backend {backend}")
 
-        self.consumer_pp_size = consumer_plugin_config["pp_size"]  # consumer pp size
+        self.consumer_pp_size = consumer_plugin_config.get("pp_size", 1)  # consumer pp size
 
     def setup(self) -> None:
         cc.init_collective_group(1 + self.num_consumer_procs, 0, group_name=f"sync_data_{self.producer_idx}")
@@ -250,14 +250,11 @@ def loop(self) -> None:
                 # linear annealing for 1 episode, temperature from initial to 0.9
                 if episode <= 0:
                     ratio = 1 - (len(self.train_dataloader) - i) / len(self.train_dataloader)
-                    if isinstance(self.model.generate_config.temperature, dict):
-                        self.model.generate_config["temperature"] = (1 - ratio) * self.generate_config[
-                            "temperature"
-                        ] + ratio * 0.9
-                    else:
-                        self.model.generate_config.temperature = (1 - ratio) * self.generate_config[
-                            "temperature"
-                        ] + ratio * 0.9
+                    self.model.generate_config["temperature"] = (1 - ratio) * self.generate_config[
+                        "temperature"
+                    ] + ratio * 0.9
+                    if hasattr(self.model, "sample_params"):
+                        self.model.sample_params.temperature = self.model.generate_config["temperature"]
 
 
 @ray.remote
@@ -310,8 +307,8 @@ def __init__(
     @torch.no_grad()
     def rollout(self, input_ids, attention_mask, **kwargs):
         rollouts = self.model.generate(input_ids, attention_mask, **kwargs)
-        # if self.producer_idx == 1:
-        #     print("Rollout example:\n", self.tokenizer.decode(rollouts["input_ids"][0][0], skip_special_tokens=True))
+        if self.producer_idx == 1:
+            print("Rollout example:\n", self.tokenizer.decode(rollouts["input_ids"][0][0], skip_special_tokens=True))
 
         return rollouts