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

@@ -346,7 +346,6 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
                         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

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

@@ -37,9 +37,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)

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

@@ -5,9 +5,9 @@
 import torch
 import wandb
 from coati.distributed.comm import SharedVariableActor
-from coati.distributed.zero_bubble.consumer import BaseConsumer
 from coati.distributed.loss import PolicyLoss
-from coati.distributed.utils import memory_efficient_logprob
+from coati.distributed.utils import entropy_from_logits, memory_efficient_logprob
+from coati.distributed.zero_bubble.consumer import BaseConsumer
 from coati.trainer.utils import all_reduce_mean, all_reduce_sum
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -33,6 +33,7 @@ def __init__(
         plugin_config,
         minibatch_size=1,
         num_generations=8,
+        tokenizer_config=None,
         generate_config=None,
         grpo_config={},
         save_interval: int = 100,
@@ -73,9 +74,11 @@ def __init__(
         self.policy_model = AutoModelForCausalLM.from_pretrained(path, **model_config)
         self.policy_model.train()
         self.policy_model.gradient_checkpointing_enable()
+        self.vocab_size = self.policy_model.config.vocab_size
         self.optimizer = HybridAdam(self.policy_model.parameters(), lr=grpo_config.get("lr", 1e-6))
         self.accum_loss = torch.zeros(1, device=self.device)
         self.accum_kl = torch.zeros(1, device=self.device)
+        self.accum_entropy = torch.zeros(1, device=self.device)
         self.accum_advantages = torch.zeros(1, device=self.device)
         self.raw_train_batch_reward = []
         self.raw_train_batch_format_acc = []
@@ -102,8 +105,11 @@ def __init__(
         if self.policy_loss_fn.beta > 0:
             self.reference_model = AutoModelForCausalLM.from_pretrained(path, **model_config)
             self.reference_model.eval()
-
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        if tokenizer_config is not None:
+            path = tokenizer_config.pop("path", None)
+            self.tokenizer = AutoTokenizer.from_pretrained(path, **tokenizer_config)
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.pad_token_id = self.tokenizer.pad_token_id
         self.num_generations = num_generations
         self.filter_range = grpo_config.get("filter_range", None)
@@ -243,10 +249,14 @@ def step(self, pbar: Any, **kwargs) -> Optional[float]:
             else self.booster.no_sync(self.policy_model, self.optimizer)
         )
         with ctx:
+            mini_batch_entropies = []
             for forward_micro_batch_start in range(0, data["input_ids"].size(0), train_microbatch_size):
                 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
                 ]
@@ -303,6 +313,7 @@ def step(self, 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:
@@ -312,6 +323,12 @@ def step(self, pbar: Any, **kwargs) -> Optional[float]:
 
                     def _criterion(outputs, inputs):
                         action_logits = outputs.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"],
@@ -334,7 +351,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"],
@@ -396,7 +413,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,
@@ -411,6 +428,20 @@ def _criterion(outputs, inputs):
                         kl = all_reduce_mean(kl.mean(), self.plugin)
                         mean_kl.append(kl.data)
                     mean_loss.append(loss.data)
+                    mini_batch_entropies.append(
+                        all_reduce_mean(
+                            (
+                                (
+                                    (
+                                        entropy_from_logits(policy_model_logits[:, -num_action:])
+                                        * action_mask_forward_micro_batch
+                                    ).sum(-1)
+                                )
+                                / action_mask_forward_micro_batch.sum(-1)
+                            ).detach(),
+                            self.plugin,
+                        )
+                    )
             if not self.plugin.pp_size > 1 or (
                 self.plugin.pp_size > 1
                 and self.booster.plugin.stage_manager.is_last_stage()
@@ -422,7 +453,9 @@ 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 = 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:
                     self.accum_kl.add_(sum(mean_kl) / len(mean_kl))
                 self.accum_advantages.add_(advantages.data)
@@ -465,6 +498,7 @@ def _criterion(outputs, inputs):
                         f"Response Length: {raw_batch_response_len_mean:.4f}",
                         f"Sample_utilization: {sample_utilization:.4f}",
                         f"Overlength samples ratio: {overlength_samples_ratio:.4f}",
+                        f"Entropy: {self.accum_entropy.item() / self.accum_count:.4f}",
                     ] + ([f"KL: {self.accum_kl.item() / self.accum_count:.4f}"] if self.policy_loss_fn.beta > 0 else [])
                     print("\n".join(to_log_msg))
                     metrics = {
@@ -476,15 +510,18 @@ def _criterion(outputs, inputs):
                         "train/advantages": self.accum_advantages.item() / self.accum_count,
                         "train/learning_rate": self.lr_scheduler.get_last_lr()[0],
                         "train/sample_utilization": sample_utilization,
+                        "train/entropy": self.accum_entropy.item() / self.accum_count,
                         "train/overlength_samples_ratio": overlength_samples_ratio,
                         "rollout/temperature": data["temperature"].cpu().numpy()[0][0],
                     }
                     if self.policy_loss_fn.beta > 0:
                         metrics["train/kl"] = self.accum_kl.item() / self.accum_count
                     if self.wandb_run is not None:
                         self.wandb_run.log(metrics)
+                    ray.get(self.shared_signal_actor.set_signal.remote("sample_utilization", sample_utilization))
                 self.accum_loss.zero_()
                 self.accum_kl.zero_()
+                self.accum_entropy.zero_()
                 self.accum_advantages.zero_()
                 self.accum_count = 0
             return loss_scalar