fix racing condition

Author: YeAnbang

applications/ColossalChat/coati/distributed/comm.py

@@ -1,5 +1,6 @@
-from typing import Any, Dict
 import copy
+from typing import Any, Dict
+
 import ray
 import ray.util.collective as cc
 import torch
@@ -31,6 +32,7 @@ def ray_broadcast_object(obj: Any, src: int = 0, device=None, group_name: str =
             obj = c10d._tensor_to_object(obj, size_tensor.item())
     return obj
 
+
 def ray_broadcast_tensor_dict(
     tensor_dict: Dict[str, torch.Tensor],
     src: int = 0,
@@ -98,7 +100,7 @@ def pickup_rollout_task(self, num_tasks: int):
         queue length as data may still be generating
         """
         ret = False
-        if self.queue_size < self.buffer_size_limit:
+        if self.queue_size < (self.buffer_size_limit / max(0.1, self.signals.get("sample_utilization", 1.0))):
             ret = True
             self.queue_size += num_tasks
         return ret

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -250,6 +250,9 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
                 input_ids_forward_micro_batch = data["input_ids"][
                     forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size
                 ]
+                old_action_log_probs_micro_batch = old_action_log_probs[
+                    forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size
+                ]
                 attention_mask_forward_micro_batch = data["attention_mask"][
                     forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size
                 ]
@@ -306,17 +309,22 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
                         "action_mask": action_mask_forward_micro_batch,
                         "advantages": advantages_forward_micro_batch,
                         "loss_mask": loss_mask_forward_micro_batch,
+                        "old_action_log_probs": old_action_log_probs_micro_batch,
                         "source": self.rank,
                     }
                     if reference_action_log_probs is not None:
                         data_policy_forward["reference_action_log_probs"] = reference_action_log_probs
 
                     kl = []
-                    policy_model_logits = torch.empty_like(input_ids_forward_micro_batch, device=self.device)
 
                     def _criterion(outputs, inputs):
                         action_logits = outputs.logits
-                        policy_model_logits.copy_(action_logits)
+                        mini_batch_entropies.append(
+                            (
+                                ((entropy_from_logits(action_logits[:, -num_action:]) * inputs["action_mask"]).sum(-1))
+                                / inputs["action_mask"].sum(-1)
+                            ).detach()
+                        )
                         action_log_probs = memory_efficient_logprob(
                             action_logits / self.generate_config["temperature"],
                             inputs["input_ids"],
@@ -339,7 +347,7 @@ def _criterion(outputs, inputs):
 
                         loss, _ = self.policy_loss_fn(
                             action_log_probs,
-                            action_log_probs,
+                            inputs["old_action_log_probs"],
                             inputs["advantages"].repeat_interleave(action_log_probs.size(-1), dim=-1),
                             per_token_kl,
                             inputs["action_mask"],
@@ -415,7 +423,7 @@ def _criterion(outputs, inputs):
 
                     loss, _ = self.policy_loss_fn(
                         action_log_probs,
-                        old_action_log_probs,
+                        old_action_log_probs_micro_batch,
                         advantages_forward_micro_batch.repeat_interleave(action_log_probs.size(-1), dim=-1),
                         per_token_kl,
                         action_mask_forward_micro_batch,
@@ -455,7 +463,7 @@ def _criterion(outputs, inputs):
                 ans_acc = all_reduce_mean(ans_acc.mean(), self.plugin)
                 advantages = all_reduce_mean(advantages.mean(), self.plugin)
                 response_length = all_reduce_mean(response_length.mean(), self.plugin)
-                entropy = torch.cat(mini_batch_entropies, dim=0).mean()
+                entropy = all_reduce_mean(torch.cat(mini_batch_entropies, dim=0).mean(), self.plugin)
                 self.accum_loss.add_(sum(mean_loss) / len(mean_loss))
                 self.accum_entropy.add_(entropy.data)
                 if self.policy_loss_fn.beta > 0:

applications/ColossalChat/coati/distributed/inference_backend.py

@@ -59,6 +59,7 @@ def __init__(
         generate_config: Dict[str, Any],
         tokenizer: PreTrainedTokenizer,
         num_generations: int = 8,
+        tokenizer_config: Dict[str, Any] = None,
     ):
         model_config = update_by_default(model_config, self.DEFAULT_MODEL_CONFIG)
         model_config.update(self.FORCE_MODEL_CONFIG)
@@ -132,6 +133,7 @@ def __init__(
         generate_config: Dict[str, Any],
         tokenizer: PreTrainedTokenizer,
         num_generations: int = 8,
+        tokenizer_config: Dict[str, Any] = None,
     ):
         if sgl is None:
             raise ImportError("sglang is not installed")
@@ -196,12 +198,14 @@ def __init__(
         generate_config: Dict[str, Any],
         tokenizer: PreTrainedTokenizer,
         num_generations: int = 8,
+        tokenizer_config: Dict[str, Any] = None,
     ):
         if LLM is None:
             raise ImportError("vllm is not installed")
         model_config = update_by_default(model_config, self.DEFAULT_MODEL_CONFIG)
         path = model_config.pop("path")
-        self.llm = LLM(model=path, **model_config)
+        tokenizer_path = tokenizer_config.get("path", None) if tokenizer_config is not None else None
+        self.llm = LLM(model=path, tokenizer=tokenizer_path, **model_config)
         generate_config = generate_config.copy()
         generate_config.update(self.FORCE_GENERATE_CONFIG)
         generate_config.update({"n": num_generations})

applications/ColossalChat/coati/distributed/launch_zero_bubble.py

@@ -130,7 +130,7 @@ def launch_distributed(
             train_dataset_config=train_dataset_config,
             model_config=inference_model_config,
             generate_config=generate_config,
-            tokenizer_config=tokenizer_config,
+            tokenizer_config=copy.deepcopy(tokenizer_config),
             microbatch_size=inference_microbatch_size,
             backend=inference_backend,
             num_generations=num_generations,
@@ -158,8 +158,6 @@ def launch_distributed(
     consumer_master_ip_address = gpu_to_ip_address[0]
     print(f"Use {consumer_master_ip_address} as master address for torch DDP.")
     consumer_procs = []
-    if num_consumer_procs <= 1:
-        raise ValueError("Number of consumer processes should be greater than 1 for async rl training.")
     for i in range(num_consumer_procs):
         node_id = gpu_to_node_id[0]
         consumer_ip_address = gpu_to_ip_address[0]
@@ -180,6 +178,7 @@ def launch_distributed(
             model_config=train_model_config,
             plugin_config=plugin_config,
             minibatch_size=train_minibatch_size,
+            tokenizer_config=copy.deepcopy(tokenizer_config),
             generate_config=generate_config_consumer,
             grpo_config=grpo_config,
             num_generations=num_generations,

applications/ColossalChat/coati/distributed/loss.py

@@ -35,9 +35,9 @@ def forward(
         total_effective_tokens_in_batch: torch.Tensor = None,
     ) -> torch.Tensor:
         if action_mask is None:
-            ratio = (log_probs - log_probs.detach()).exp()
+            ratio = (log_probs - old_log_probs.detach()).exp()
         else:
-            ratio = ((log_probs - log_probs.detach()) * action_mask).exp()
+            ratio = ((log_probs - old_log_probs.detach()) * action_mask).exp()
 
         surr1 = ratio * advantages
         surr2 = ratio.clamp(1 - self.clip_eps_low, 1 + self.clip_eps_high) * advantages

applications/ColossalChat/coati/distributed/zero_bubble/consumer.py

@@ -7,17 +7,16 @@
 import ray.util.collective as cc
 import torch
 import torch.distributed as dist
+from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
 from coati.distributed.profiling_utils import CustomProfiler
+from coati.distributed.utils import bind_batch, post_recv, unbind_batch
 from tqdm import tqdm
 
 from colossalai.booster import Booster
 from colossalai.booster.plugin import HybridParallelPlugin
 from colossalai.initialize import launch
 from colossalai.utils import get_current_device
 
-from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
-from coati.distributed.utils import bind_batch, post_recv, unbind_batch
-
 
 class BaseConsumer:
     def __init__(
@@ -175,14 +174,15 @@ def loop(self) -> None:
                         raw_batch = ray.get(
                             self.shared_sync_data_actor.get_data.remote(self.data_uid)
                         )  # get the first queued data
+                        self.profiler.log(f"enter sleep")
                         while raw_batch is None:
-                            self.profiler.log(f"No data received by consumer {self.rank}, skipping")
                             print(
                                 f"[T{dist.get_rank()}] No data received by consumer {self.rank}, skipping. Consider increasing the data actor buffer limit"
                             )
                             time.sleep(1)
                             raw_batch = ray.get(self.shared_sync_data_actor.get_data.remote(self.data_uid))
                             continue
+                        self.profiler.log(f"exit sleep")
                         self.data_uid += 1
                         raw_batch = {k: v.to(self.device) for k, v in raw_batch.items()}
                         # calculate group reward et al. filtering. As only the filtered group will be used for training (which is incomplete),

applications/ColossalChat/coati/distributed/zero_bubble/distributor.py

@@ -3,12 +3,11 @@
 import ray
 import ray.util.collective as cc
 import torch
+from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
 from coati.distributed.profiling_utils import CustomProfiler
 
 from colossalai.utils import get_current_device
 
-from coati.distributed.comm import SharedVariableActor, ray_broadcast_tensor_dict
-
 
 @ray.remote
 class Distributor:
@@ -21,6 +20,7 @@ def __init__(
         enable_profiling: bool = True,
     ):
         self.distributor_id = distributor_id
+        self.weight_version = [0] * consumer_pp_size
         self.consumer_pp_size = consumer_pp_size
         self.state_dict_cpu = {}
         self.num_producers = num_producers
@@ -42,14 +42,17 @@ def init_collective_group(
         print(f"[D] Initialized {group_name} collective group", flush=True)
 
     def loop(self):
+        last_weight_version = self.get_weight_version()
         while True:
             time.sleep(1)
             signal = ray.get(self.shared_signal_actor.get_signal.remote())
             if self.consumer_pp_size > 1:
-                for i in range(self.consumer_pp_size):
-                    if signal.get(f"consumer_pp_{i}", None) == "ready_sync_model":
+                if all(
+                    [signal.get(f"consumer_pp_{i}", None) == "ready_sync_model" for i in range(self.consumer_pp_size)]
+                ):
+                    cc.barrier(group_name="distributor_pg")
+                    for i in range(self.consumer_pp_size):
                         self.profiler.enter(f"sync_model_consumer_pp_{i}")
-                        cc.barrier(group_name="distributor_pg")
                         ray.get(self.shared_signal_actor.set_signal.remote(f"consumer_pp_{i}", "not_ready_sync_model"))
                         # Broadcast the model state dict from consumer to shared variable actor
                         self.state_dict_cpu[i] = ray_broadcast_tensor_dict(
@@ -60,6 +63,7 @@ def loop(self):
                             backend="gloo",
                         )
                         self.profiler.exit(f"sync_model_consumer_pp_{i}")
+                        self.weight_version[i] += 1
                 for i in range(self.consumer_pp_size):
                     if signal.get(f"producer_{self.distributor_id}_pp_{i}", None) == "ready_sync_model":
                         self.profiler.enter(f"sync_model_producer_{self.distributor_id}_pp_{i}")
@@ -87,6 +91,7 @@ def loop(self):
                         None, 0, device=torch.device("cpu"), group_name="sync_model_consumer", backend="gloo"
                     )
                     self.profiler.exit("sync_model_consumer")
+                    self.weight_version[0] += 1
                 if signal.get(f"producer_{self.distributor_id}", None) == "ready_sync_model":
                     self.profiler.enter(f"sync_model_producer_{self.distributor_id}")
                     # Broadcast the model state dict to all producers
@@ -106,3 +111,9 @@ def loop(self):
             if signal.get("consumer", None) == "terminate":
                 self.profiler.log("terminate sync model worker")
                 break
+            if last_weight_version != self.get_weight_version():
+                last_weight_version = self.get_weight_version()
+                ray.get(self.shared_signal_actor.set_signal.remote("distributor_weight_version", last_weight_version))
+
+    def get_weight_version(self):
+        return min(self.weight_version)