[megatron, sglang, rollout, doc] feat: Add Validation StateMachine to async-rl, async ref_logp and support nccl-sync

recipe/dapo/dapo_ray_trainer.py

@@ -16,6 +16,7 @@
 This trainer supports model-agonistic model initialization with huggingface
 """
 
+import os
 import uuid
 from collections import defaultdict
 from copy import deepcopy
@@ -24,7 +25,6 @@
 import numpy as np
 import torch
 from tqdm import tqdm
-import os
 
 from verl import DataProto
 from verl.trainer.ppo.core_algos import agg_loss

recipe/grpo/grpo_ray_trainer.py

@@ -15,14 +15,13 @@
 FSDP PPO Trainer with Ray-based single controller.
 This trainer supports model-agonistic model initialization with huggingface
 """
-import ray
-import time
+
 import uuid
-from collections import defaultdict
 from copy import deepcopy
 from pprint import pprint
 
 import numpy as np
+import ray
 import torch
 from tqdm import tqdm
 
@@ -34,7 +33,14 @@
     compute_timing_metrics,
     reduce_metrics,
 )
-from verl.trainer.ppo.ray_trainer import AdvantageEstimator, RayPPOTrainer, _timer, apply_kl_penalty, compute_advantage, compute_response_mask
+from verl.trainer.ppo.ray_trainer import (
+    AdvantageEstimator,
+    RayPPOTrainer,
+    _timer,
+    apply_kl_penalty,
+    compute_advantage,
+    compute_response_mask,
+)
 from verl.trainer.ppo.reward import compute_reward, compute_reward_async
 
 
@@ -127,7 +133,9 @@ def fit(self):
 
                             del gen_baseline_batch, gen_baseline_output
 
-                    batch.non_tensor_batch["uid"] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object)
+                    batch.non_tensor_batch["uid"] = np.array(
+                        [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object
+                    )
                     # repeat to align with repeated responses in rollout
                     batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
                     batch = batch.union(gen_batch_output)
@@ -161,7 +169,9 @@ def fit(self):
                         entropys = old_log_prob.batch["entropys"]
                         response_masks = batch.batch["response_mask"]
                         loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode
-                        entropy_loss = agg_loss(loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode)
+                        entropy_loss = agg_loss(
+                            loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode
+                        )
                         old_log_prob_metrics = {"actor/entropy_loss": entropy_loss.detach().item()}
                         metrics.update(old_log_prob_metrics)
                         old_log_prob.batch.pop("entropys")
@@ -216,21 +226,39 @@ def fit(self):
                         print(f"{list(reward_extra_infos_dict.keys())=}")
                         if reward_extra_infos_dict:
                             batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})
-                            metrics.update({
-                                **{f"critic/rewards/{k}/mean": np.mean(v) for k, v in reward_extra_infos_dict.items() if '_sub' in k},
-                                **{f"critic/rewards/{k}/max": np.max(v) for k, v in reward_extra_infos_dict.items() if '_sub' in k},
-                                **{f"critic/rewards/{k}/min": np.min(v) for k, v in reward_extra_infos_dict.items() if '_sub' in k},
-                            })
+                            metrics.update(
+                                {
+                                    **{
+                                        f"critic/rewards/{k}/mean": np.mean(v)
+                                        for k, v in reward_extra_infos_dict.items()
+                                        if "_sub" in k
+                                    },
+                                    **{
+                                        f"critic/rewards/{k}/max": np.max(v)
+                                        for k, v in reward_extra_infos_dict.items()
+                                        if "_sub" in k
+                                    },
+                                    **{
+                                        f"critic/rewards/{k}/min": np.min(v)
+                                        for k, v in reward_extra_infos_dict.items()
+                                        if "_sub" in k
+                                    },
+                                }
+                            )
                         # compute rewards. apply_kl_penalty if available
                         if self.config.algorithm.use_kl_in_reward:
-                            batch, kl_metrics = apply_kl_penalty(batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty)
+                            batch, kl_metrics = apply_kl_penalty(
+                                batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
+                            )
                             metrics.update(kl_metrics)
                         else:
                             batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
 
                         # compute advantages, executed on the driver process
 
-                        norm_adv_by_std_in_grpo = self.config.algorithm.get("norm_adv_by_std_in_grpo", True)  # GRPO adv normalization factor
+                        norm_adv_by_std_in_grpo = self.config.algorithm.get(
+                            "norm_adv_by_std_in_grpo", True
+                        )  # GRPO adv normalization factor
 
                         batch = compute_advantage(
                             batch,
@@ -278,14 +306,20 @@ def fit(self):
                             )
 
                     # validate
-                    if self.val_reward_fn is not None and self.config.trainer.test_freq > 0 and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0):
+                    if (
+                        self.val_reward_fn is not None
+                        and self.config.trainer.test_freq > 0
+                        and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
+                    ):
                         with _timer("testing", timing_raw):
                             val_metrics: dict = self._validate()
                             if is_last_step:
                                 last_val_metrics = val_metrics
                         metrics.update(val_metrics)
 
-                    if self.config.trainer.save_freq > 0 and (is_last_step or self.global_steps % self.config.trainer.save_freq == 0):
+                    if self.config.trainer.save_freq > 0 and (
+                        is_last_step or self.global_steps % self.config.trainer.save_freq == 0
+                    ):
                         with _timer("save_checkpoint", timing_raw):
                             self._save_checkpoint()
 
@@ -311,4 +345,4 @@ def fit(self):
                 if is_last_step:
                     pprint(f"Final validation metrics: {last_val_metrics}")
                     progress_bar.close()
-                    return
+                    return

recipe/grpo/main_grpo.py

@@ -16,12 +16,12 @@
 """
 
 import hydra
+import pandas as pd
 import ray
+from torch.utils.data import Dataset
 
-import pandas as pd
-from .grpo_ray_trainer import RayGRPOTrainer
+from verl.trainer.ppo.ray_trainer import RayPPOTrainer
 from verl.trainer.ppo.reward import load_reward_manager
-from torch.utils.data import Dataset
 from verl.utils.dataset.rl_dataset import RLHFDataset as OriginalRLHFDataset
 
 
@@ -30,11 +30,11 @@ def _read_files_and_tokenize(self):
         dataframes = []
         for parquet_file in self.data_files:
             # read parquet files and cache
-            if parquet_file.endswith('parquet'):
+            if parquet_file.endswith("parquet"):
                 dataframe = pd.read_parquet(parquet_file)
-            elif parquet_file.endswith('json'):
+            elif parquet_file.endswith("json"):
                 dataframe = pd.read_json(parquet_file)
-            elif parquet_file.endswith('jsonl'):
+            elif parquet_file.endswith("jsonl"):
                 chunks = []
                 for chunk in pd.read_json(
                     parquet_file,
@@ -51,23 +51,25 @@ def _read_files_and_tokenize(self):
 
         print(f"dataset len: {len(self.dataframe)}")
 
-        if self.config.data.get('system_prompt', None) is not None:
+        if self.config.data.get("system_prompt", None) is not None:
             system_prompt = self.config.data.system_prompt
             self.dataframe[self.prompt_key] = self.dataframe[self.prompt_key].apply(
-                lambda x: [{'role': 'system', 'content': system_prompt}]+x
+                lambda x: [{"role": "system", "content": system_prompt}] + x
             )
         # filter out too long prompts
         if self.filter_overlong_prompts:
             tokenizer = self.tokenizer
             prompt_key = self.prompt_key
             self.dataframe = self.dataframe.filter(
-                lambda doc: len(tokenizer.apply_chat_template(doc[prompt_key], add_generation_prompt=True)) <= self.max_prompt_length,
+                lambda doc: len(tokenizer.apply_chat_template(doc[prompt_key], add_generation_prompt=True))
+                <= self.max_prompt_length,
                 num_proc=self.num_workers,
                 desc=f"Filtering prompts longer than {self.max_prompt_length} tokens",
             )
 
             print(f"filter dataset len: {len(self.dataframe)}")
 
+
 @hydra.main(config_path="config", config_name="ppo_trainer", version_base=None)
 def main(config):
     run_grpo(config)
@@ -77,7 +79,14 @@ def run_grpo(config) -> None:
     if not ray.is_initialized():
         # this is for local ray cluster
         ray.init(
-            runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN", "VLLM_LOGGING_LEVEL": "WARN", "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "true"}},
+            runtime_env={
+                "env_vars": {
+                    "TOKENIZERS_PARALLELISM": "true",
+                    "NCCL_DEBUG": "WARN",
+                    "VLLM_LOGGING_LEVEL": "WARN",
+                    "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "true",
+                }
+            },
             num_cpus=config.ray_init.num_cpus,
         )
 
@@ -103,14 +112,18 @@ def run(self, config):
         OmegaConf.resolve(config)
 
         # download the checkpoint from hdfs
-        local_path = copy_to_local(config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get("use_shm", False))
+        local_path = copy_to_local(
+            config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get("use_shm", False)
+        )
 
         # instantiate tokenizer
         from verl.utils import hf_processor, hf_tokenizer
 
         trust_remote_code = config.data.get("trust_remote_code", False)
         tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)
-        processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)  # used for multimodal LLM, could be none
+        processor = hf_processor(
+            local_path, trust_remote_code=trust_remote_code, use_fast=True
+        )  # used for multimodal LLM, could be none
 
         # vllm early verify
         if config.actor_rollout_ref.rollout.name in ["vllm"]:
@@ -126,15 +139,23 @@ def run(self, config):
             from verl.single_controller.ray import RayWorkerGroup
             from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker, CriticWorker
 
-            actor_rollout_cls = AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker
+            actor_rollout_cls = (
+                AsyncActorRolloutRefWorker
+                if config.actor_rollout_ref.rollout.mode == "async"
+                else ActorRolloutRefWorker
+            )
             ray_worker_group_cls = RayWorkerGroup
 
         elif config.actor_rollout_ref.actor.strategy == "megatron":
             assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
             from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup
             from verl.workers.megatron_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker, CriticWorker
 
-            actor_rollout_cls = AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker
+            actor_rollout_cls = (
+                AsyncActorRolloutRefWorker
+                if config.actor_rollout_ref.rollout.mode == "async"
+                else ActorRolloutRefWorker
+            )
             ray_worker_group_cls = NVMegatronRayWorkerGroup
 
         else:
@@ -177,8 +198,12 @@ def run(self, config):
             role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker)
             mapping[Role.RefPolicy] = global_pool_id
 
-        reward_fn = load_reward_manager(config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {}))
-        val_reward_fn = load_reward_manager(config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {}))
+        reward_fn = load_reward_manager(
+            config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})
+        )
+        val_reward_fn = load_reward_manager(
+            config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {})
+        )
         resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
 
         from verl.utils.dataset.rl_dataset import collate_fn
@@ -222,7 +247,10 @@ def create_rl_dataset(data_paths, data_config, tokenizer, processor):
 
         dataset_cls = load_extern_type(data_config.custom_cls.path, data_config.custom_cls.name)
         if not issubclass(dataset_cls, Dataset):
-            raise TypeError(f"The custom dataset class '{data_config.custom_cls.name}' from '{data_config.custom_cls.path}' must inherit from torch.utils.data.Dataset")
+            raise TypeError(
+                f"The custom dataset class '{data_config.custom_cls.name}' from "
+                f"'{data_config.custom_cls.path}' must inherit from torch.utils.data.Dataset"
+            )
     else:
         dataset_cls = RLHFDataset
     print(f"Using dataset class: {dataset_cls.__name__}")

recipe/langgraph_agent/chat_model.py

@@ -173,9 +173,10 @@ async def _preprocess(self, messages: list[BaseMessage], **kwargs: Any) -> tuple
             tuple[str, list[int], list[int]]: Request id, prompt ids, response mask.
         """
         # messages: [system], human, ai, human|tool, ai, human|tool, ...
-        assert messages[-1].type in ["human", "tool"], (
-            f"Last message must be human or tool, but got {messages[-1].type}"
-        )
+        assert messages[-1].type in [
+            "human",
+            "tool",
+        ], f"Last message must be human or tool, but got {messages[-1].type}"
         loop = asyncio.get_running_loop()
 
         # Case 1: initial chat completion: [system], human