bug fix: KL estimators mismatch, vram usage optimized, feat: dapo training

Author: atlas5301

recipe/dapo/dapo_ray_trainer.py

@@ -47,6 +47,15 @@ class RayDAPOTrainer(RayPPOTrainer):
     Note that this trainer runs on the driver process on a single CPU/GPU node.
     """
 
+    def fit(self):
+        # Delegate to the namespace-aware training loop when a topology or multiple namespaces are configured.
+        if self.topology_schedule or len(self.namespace_specs) > 1 or len(self.training_order) > 1:
+            return super().fit()
+        if not self.training_order:
+            print("No training namespaces configured; skipping training.")
+            return
+        return self._fit_single_namespace()
+
     def compute_kl_related_metrics(self, batch: DataProto, metrics: dict, timing_raw: dict):
         batch.batch["response_mask"] = compute_response_mask(batch)
 
@@ -73,7 +82,7 @@ def compute_kl_related_metrics(self, batch: DataProto, metrics: dict, timing_raw
 
         return batch
 
-    def fit(self):
+    def _fit_single_namespace(self):
         """
         The training loop of PPO.
         The driver process only need to call the compute functions of the worker group through RPC

recipe/dapo/main_dapo.py

@@ -22,8 +22,10 @@
 import ray
 from omegaconf import OmegaConf
 
+from verl.trainer import main_ppo as main_ppo_mod
+from verl.trainer.namespace import build_namespace_specs, namespaced_role_key
 from verl.trainer.ppo.reward import load_reward_manager
-from verl.utils.device import is_cuda_available
+from verl.trainer.ppo.utils import Role
 
 from .dapo_ray_trainer import RayDAPOTrainer
 
@@ -34,39 +36,12 @@ def main(config):
 
 
 def run_ppo(config) -> None:
-    if not ray.is_initialized():
-        # this is for local ray cluster
-        default_runtime_env = {
-            "env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN", "VLLM_LOGGING_LEVEL": "WARN"}
-        }
-        ray_init_kwargs = config.ray_kwargs.get("ray_init", {})
-        runtime_env_kwargs = ray_init_kwargs.get("runtime_env", {})
-        runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)
-        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, "runtime_env": runtime_env})
-        print(f"ray init kwargs: {ray_init_kwargs}")
-        ray.init(**OmegaConf.to_container(ray_init_kwargs))
-
-    try:
-        if (
-            is_cuda_available
-            and config.global_profiler.tool == "nsys"
-            and OmegaConf.select(config.global_profiler, "steps") is not None
-            and len(OmegaConf.select(config.global_profiler, "steps")) > 0
-        ):
-            nsight_options = OmegaConf.to_container(
-                config.global_profiler.global_tool_config.nsys.controller_nsight_options
-            )
-            runner = TaskRunner.options(runtime_env={"nsight": nsight_options}).remote()
-        else:
-            runner = TaskRunner.remote()
-        ray.get(runner.run.remote(config))
-    finally:
-        if ray.is_initialized():
-            ray.shutdown()
+    """Entry point for running DAPO with the PPO runner."""
+    task_runner_cls = ray.remote(num_cpus=1)(TaskRunner)  # type: ignore[arg-type]
+    main_ppo_mod.run_ppo(config, task_runner_class=task_runner_cls)
 
 
-@ray.remote(num_cpus=1)  # please make sure main_task is not scheduled on head
-class TaskRunner:
+class TaskRunner(main_ppo_mod.TaskRunner):
     def run(self, config):
         # print initial config
         from pprint import pprint
@@ -80,72 +55,70 @@ def run(self, config):
         pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values
         OmegaConf.resolve(config)
 
-        # download the checkpoint from hdfs
-        local_path = copy_to_local(config.actor_rollout_ref.model.path)
-
-        # 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)
-        # used for multimodal LLM, could be none
-        processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)
-
-        from verl.single_controller.ray import RayWorkerGroup
-
-        # define worker classes
-        if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}:
-            assert config.critic.strategy in {"fsdp", "fsdp2"}
-
-            from verl.workers.fsdp_workers import AsyncActorRolloutRefWorker, CriticWorker
-
-            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.workers.megatron_workers import AsyncActorRolloutRefWorker, CriticWorker
-
-            ray_worker_group_cls = RayWorkerGroup
-
-        else:
-            raise NotImplementedError
-
-        from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role
-
-        role_worker_mapping = {
-            Role.ActorRollout: ray.remote(AsyncActorRolloutRefWorker),
-            Role.Critic: ray.remote(CriticWorker),
-        }
-
-        global_pool_id = "global_pool"
-        resource_pool_spec = {
-            global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
-        }
-        mapping = {
-            Role.ActorRollout: global_pool_id,
-            Role.Critic: global_pool_id,
-        }
-
-        # we should adopt a multi-source reward function here
-        # - for rule-based rm, we directly call a reward score
-        # - for model-based rm, we call a model
-        # - for code related prompt, we send to a sandbox if there are test cases
-        # - finally, we combine all the rewards together
-        # - The reward type depends on the tag of the data
+        namespace_specs = build_namespace_specs(config)
+
+        # instantiate tokenizer/processor per namespace
+        tokenizers = {}
+        processors = {}
+        for name, spec in namespace_specs.items():
+            local_path = copy_to_local(
+                spec.config.actor_rollout_ref.model.path,
+                use_shm=spec.config.actor_rollout_ref.model.get("use_shm", False),
+            )
+            tokenizers[name] = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)
+            processors[name] = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)
+
+        active_namespace = config.trainer.get("namespace", "default")
+        tokenizer = tokenizers[active_namespace]
+        processor = processors.get(active_namespace)
+
+        self.role_worker_mapping = {}
+        self.mapping = {}
+
+        # Register actor-like workers and collect the worker group class.
+        ray_worker_group_cls = None
+        for spec in namespace_specs.values():
+            if not spec.spawn_roles:
+                continue
+            actor_cls, rg_cls = self._select_actor_worker_impl(spec.config)
+            ray_worker_group_cls = rg_cls if ray_worker_group_cls is None else ray_worker_group_cls
+            if rg_cls is not None and ray_worker_group_cls != rg_cls:
+                raise ValueError("All namespaces must share the same RayWorkerGroup class")
+
+            critic_cls = self._select_critic_worker_impl(spec.config)
+
+            for role in spec.spawn_roles:
+                key = namespaced_role_key(spec.name, role)
+                if role == Role.Critic:
+                    self.role_worker_mapping[key] = ray.remote(critic_cls)
+                else:
+                    self.role_worker_mapping[key] = ray.remote(actor_cls)
+                self.mapping[key] = spec.resource_pool
+
+        # reward model
         if config.reward_model.enable:
-            if config.reward_model.strategy in {"fsdp", "fsdp2"}:
-                from verl.workers.fsdp_workers import RewardModelWorker
-            elif config.reward_model.strategy == "megatron":
-                from verl.workers.megatron_workers import RewardModelWorker
+            use_legacy_worker_impl = config.trainer.get("use_legacy_worker_impl", "auto")
+            if use_legacy_worker_impl in ["auto", "enable"]:
+                if config.reward_model.strategy in {"fsdp", "fsdp2"}:
+                    from verl.workers.fsdp_workers import RewardModelWorker
+                elif config.reward_model.strategy == "megatron":
+                    from verl.workers.megatron_workers import RewardModelWorker
+                else:
+                    raise NotImplementedError
+            elif use_legacy_worker_impl == "disable":
+                from verl.workers.roles import RewardModelWorker
+
+                print("Using new worker implementation")
             else:
-                raise NotImplementedError
-            role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker)
-            mapping[Role.RewardModel] = global_pool_id
+                raise ValueError(f"Invalid use_legacy_worker_impl: {use_legacy_worker_impl}")
 
-        # reference model
-        if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
-            role_worker_mapping[Role.RefPolicy] = ray.remote(AsyncActorRolloutRefWorker)
-            mapping[Role.RefPolicy] = global_pool_id
+            available_pools = [spec.resource_pool for spec in namespace_specs.values() if spec.spawn_roles]
+            reward_pool = "reward_pool" if config.reward_model.enable_resource_pool else available_pools[0]
+            self.role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker)
+            self.mapping[Role.RewardModel] = reward_pool
 
         reward_fn = load_reward_manager(
             config,
@@ -163,17 +136,72 @@ def run(self, config):
             max_resp_len=config.data.max_response_length,
             overlong_buffer_cfg=config.reward_model.overlong_buffer,
         )
-        resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
+        reward_fn_map = {spec.name: reward_fn for spec in namespace_specs.values()}
+        val_reward_fn_map = {spec.name: val_reward_fn for spec in namespace_specs.values()}
+
+        base_rm_cfg = OmegaConf.to_container(config.reward_model, resolve=True)
+        base_custom_cfg = OmegaConf.to_container(config.custom_reward_function, resolve=True)
+        for spec in namespace_specs.values():
+            rm_cfg = OmegaConf.to_container(spec.config.reward_model, resolve=True)
+            custom_cfg = OmegaConf.to_container(spec.config.custom_reward_function, resolve=True)
+            if rm_cfg != base_rm_cfg or custom_cfg != base_custom_cfg:
+                reward_fn_map[spec.name] = load_reward_manager(
+                    spec.config,
+                    tokenizer,
+                    0,
+                    max_resp_len=spec.config.data.max_response_length,
+                    overlong_buffer_cfg=spec.config.reward_model.overlong_buffer,
+                )
+                val_reward_fn_map[spec.name] = load_reward_manager(
+                    spec.config,
+                    tokenizer,
+                    1,
+                    max_resp_len=spec.config.data.max_response_length,
+                    overlong_buffer_cfg=spec.config.reward_model.overlong_buffer,
+                )
+
+        resource_pool_manager = self.init_resource_pool_mgr(config, namespace_specs=namespace_specs)
+
+        from verl.utils.dataset.rl_dataset import collate_fn
+        # Create training/validation datasets when only one namespace is present.
+        train_dataset = val_dataset = train_sampler = None
+        if len(namespace_specs) == 1:
+            train_dataset = main_ppo_mod.create_rl_dataset(
+                config.data.train_files,
+                config.data,
+                tokenizer,
+                processor,
+                is_train=True,
+                max_samples=config.data.get("train_max_samples", -1),
+            )
+            val_dataset = main_ppo_mod.create_rl_dataset(
+                config.data.val_files,
+                config.data,
+                tokenizer,
+                processor,
+                is_train=False,
+                max_samples=config.data.get("val_max_samples", -1),
+            )
+            train_sampler = main_ppo_mod.create_rl_sampler(config.data, train_dataset)
 
         trainer = RayDAPOTrainer(
             config=config,
             tokenizer=tokenizer,
             processor=processor,
-            role_worker_mapping=role_worker_mapping,
+            role_worker_mapping=self.role_worker_mapping,
             resource_pool_manager=resource_pool_manager,
             ray_worker_group_cls=ray_worker_group_cls,
             reward_fn=reward_fn,
             val_reward_fn=val_reward_fn,
+            reward_fn_map=reward_fn_map,
+            val_reward_fn_map=val_reward_fn_map,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            collate_fn=collate_fn,
+            train_sampler=train_sampler,
+            namespace_specs=namespace_specs,
+            tokenizers_by_namespace=tokenizers,
+            processors_by_namespace=processors,
         )
         trainer.init_workers()
         trainer.fit()

verl/trainer/ppo/ray_trainer.py

@@ -20,6 +20,7 @@
 
 import json
 import os
+import time
 import uuid
 from collections import defaultdict
 from copy import deepcopy
@@ -1334,43 +1335,43 @@ def _run_single_step(self, runtime: NamespaceRuntime, batch_dict, logger, progre
 
         next_step = self.global_steps + 1
         is_last_step = next_step >= runtime.total_training_steps
-        with marked_timer("step", timing_raw):
-            with marked_timer("gen", timing_raw, color="red"):
-                if not self.async_rollout_mode:
-                    gen_batch_output = self.rollout_wg.generate_sequences(gen_batch_output)
-                else:
-                    gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch_output)
+        step_start = time.perf_counter()
+        with marked_timer("gen", timing_raw, color="red"):
+            if not self.async_rollout_mode:
+                gen_batch_output = self.rollout_wg.generate_sequences(gen_batch_output)
+            else:
+                gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch_output)
 
-                timing_raw.update(gen_batch_output.meta_info.get("timing", {}))
-                gen_batch_output.meta_info.pop("timing", None)
+            timing_raw.update(gen_batch_output.meta_info.get("timing", {}))
+            gen_batch_output.meta_info.pop("timing", None)
 
-            if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
-                if self.reward_fn is None:
-                    raise ValueError("A reward_fn is required for REMAX advantage estimation.")
+        if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
+            if self.reward_fn is None:
+                raise ValueError("A reward_fn is required for REMAX advantage estimation.")
 
-                with marked_timer("gen_max", timing_raw, color="purple"):
-                    gen_baseline_batch = deepcopy(gen_batch)
-                    gen_baseline_batch.meta_info["do_sample"] = False
-                    if not self.async_rollout_mode:
-                        gen_baseline_output = self.rollout_wg.generate_sequences(gen_baseline_batch)
-                    else:
-                        gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)
-                    batch = batch.union(gen_baseline_output)
-                    rm_scores = None
-                    if self.use_rm and "rm_scores" not in batch.batch.keys():
-                        rm_scores = self.rm_wg.compute_rm_score(batch)
-                        batch = batch.union(rm_scores)
-                    reward_baseline_tensor, _ = compute_reward(batch, self.reward_fn)
-                    reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
-
-                    keys_to_pop = set(gen_baseline_output.batch.keys())
-                    if rm_scores is not None:
-                        keys_to_pop.update(rm_scores.batch.keys())
-                    batch.pop(batch_keys=list(keys_to_pop))
-
-                    batch.batch["reward_baselines"] = reward_baseline_tensor
-
-                    del rm_scores, gen_baseline_batch, gen_baseline_output
+            with marked_timer("gen_max", timing_raw, color="purple"):
+                gen_baseline_batch = deepcopy(gen_batch)
+                gen_baseline_batch.meta_info["do_sample"] = False
+                if not self.async_rollout_mode:
+                    gen_baseline_output = self.rollout_wg.generate_sequences(gen_baseline_batch)
+                else:
+                    gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)
+                batch = batch.union(gen_baseline_output)
+                rm_scores = None
+                if self.use_rm and "rm_scores" not in batch.batch.keys():
+                    rm_scores = self.rm_wg.compute_rm_score(batch)
+                    batch = batch.union(rm_scores)
+                reward_baseline_tensor, _ = compute_reward(batch, self.reward_fn)
+                reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
+
+                keys_to_pop = set(gen_baseline_output.batch.keys())
+                if rm_scores is not None:
+                    keys_to_pop.update(rm_scores.batch.keys())
+                batch.pop(batch_keys=list(keys_to_pop))
+
+                batch.batch["reward_baselines"] = reward_baseline_tensor
+
+                del rm_scores, gen_baseline_batch, gen_baseline_output
         # 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)
@@ -1586,6 +1587,8 @@ def _cached_logprob(ns: str):
         if rollout_data_dir:
             self._log_rollout_data(batch, reward_extra_infos_dict, timing_raw, rollout_data_dir)
 
+        timing_raw["step"] = time.perf_counter() - step_start
+
         next_step = self.global_steps + 1
         if (
             self.val_reward_fn is not None
@@ -1616,7 +1619,14 @@ def _cached_logprob(ns: str):
         metrics.update(compute_data_metrics(batch=metrics_source, use_critic=self.use_critic))
         metrics.update(compute_timing_metrics(batch=metrics_source, timing_raw=timing_raw))
         n_gpus = self.resource_pool_manager.get_n_gpus()
-        metrics.update(compute_throughout_metrics(batch=metrics_source, timing_raw=timing_raw, n_gpus=n_gpus))
+        perf_metrics = compute_throughout_metrics(batch=metrics_source, timing_raw=timing_raw, n_gpus=n_gpus)
+        if getattr(self, "topology_step", None):
+            topo_name = getattr(self.topology_step, "name", None)
+            if topo_name:
+                perf_metrics.update(
+                    {f"perf-{topo_name}/{k.split('/', 1)[1]}": v for k, v in perf_metrics.items() if "/" in k}
+                )
+        metrics.update(perf_metrics)
 
         if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):
             self.train_dataloader.sampler.update(batch=metrics_source)