diff --git a/tests/workers/config/test_critic_engine_batching_on_cpu.py b/tests/workers/config/test_critic_engine_batching_on_cpu.py new file mode 100644 index 00000000000..06fb01181ab --- /dev/null +++ b/tests/workers/config/test_critic_engine_batching_on_cpu.py @@ -0,0 +1,70 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from verl.workers.config.critic import CriticConfig +from verl.workers.config.engine import EngineConfig +from verl.workers.config.optimizer import OptimizerConfig + + +def _make_critic(**kwargs) -> CriticConfig: + defaults = dict( + strategy="fsdp2", + use_dynamic_bsz=True, + ppo_micro_batch_size_per_gpu=2, + optim=OptimizerConfig(lr=1e-5), + ) + defaults.update(kwargs) + return CriticConfig(**defaults) + + +def test_apply_engine_batching_copies_static_and_dynamic_knobs(): + critic = _make_critic( + use_dynamic_bsz=False, + ppo_micro_batch_size_per_gpu=4, + ppo_infer_micro_batch_size_per_gpu=2, + ppo_max_token_len_per_gpu=4096, + ppo_infer_max_token_len_per_gpu=2048, + forward_max_token_len_per_gpu=1024, + ) + engine = EngineConfig() + critic.apply_engine_batching(engine) + + assert engine.use_dynamic_bsz is False + assert engine.micro_batch_size_per_gpu == 4 + assert engine.infer_micro_batch_size_per_gpu == 2 + assert engine.max_token_len_per_gpu == 4096 + assert engine.infer_max_token_len_per_gpu == 2048 + + +def test_apply_engine_batching_falls_back_infer_micro_batch_to_training(): + critic = _make_critic( + use_dynamic_bsz=False, + ppo_micro_batch_size_per_gpu=8, + ppo_infer_micro_batch_size_per_gpu=None, + ) + engine = EngineConfig() + critic.apply_engine_batching(engine) + assert engine.infer_micro_batch_size_per_gpu == 8 + + +def test_v1_train_budget_is_not_overwritten_by_infer_budget(): + """V1 previously assigned max_token_len_per_gpu = ppo_infer_max_token_len_per_gpu.""" + critic = _make_critic( + ppo_max_token_len_per_gpu=8192, + ppo_infer_max_token_len_per_gpu=1024, + ) + engine = EngineConfig() + critic.apply_engine_batching(engine) + assert engine.max_token_len_per_gpu == 8192 + assert engine.infer_max_token_len_per_gpu == 1024 diff --git a/verl/experimental/separation/ray_trainer.py b/verl/experimental/separation/ray_trainer.py index 42b6e8aae5d..6ff35952cb5 100644 --- a/verl/experimental/separation/ray_trainer.py +++ b/verl/experimental/separation/ray_trainer.py @@ -155,8 +155,7 @@ def _create_critic_class(self): self.orig_critic_cfg = critic_cfg if self.orig_critic_cfg.strategy == "fsdp": engine_config: FSDPEngineConfig = self.orig_critic_cfg.model.fsdp_config - engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu - engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu + self.orig_critic_cfg.apply_engine_batching(engine_config) else: raise NotImplementedError(f"Unknown strategy {self.orig_critic_cfg.strategy=}") diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index f8228734e72..e69f795f03a 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -647,6 +647,10 @@ critic: use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false} ppo_max_token_len_per_gpu: 32768 forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu} + + ppo_infer_max_token_len_per_gpu: ${.forward_max_token_len_per_gpu} + + ppo_infer_micro_batch_size_per_gpu: ${.ppo_micro_batch_size_per_gpu} ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1} shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false} data_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null} diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index 0247606a4d8..b4faeaf0fb7 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -541,6 +541,10 @@ critic: use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false} ppo_max_token_len_per_gpu: 32768 forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu} + + ppo_infer_max_token_len_per_gpu: ${.forward_max_token_len_per_gpu} + + ppo_infer_micro_batch_size_per_gpu: ${.ppo_micro_batch_size_per_gpu} ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1} shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false} data_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null} diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index d6b3fb72dde..9af3899e55b 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -603,6 +603,10 @@ critic: use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false} ppo_max_token_len_per_gpu: 32768 forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu} + + ppo_infer_max_token_len_per_gpu: ${.forward_max_token_len_per_gpu} + + ppo_infer_micro_batch_size_per_gpu: ${.ppo_micro_batch_size_per_gpu} ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1} shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false} data_loader_seed: 42 diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index f2a8f11dd64..2f8552e220e 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -572,6 +572,10 @@ critic: use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false} ppo_max_token_len_per_gpu: 32768 forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu} + + ppo_infer_max_token_len_per_gpu: ${.forward_max_token_len_per_gpu} + + ppo_infer_micro_batch_size_per_gpu: ${.ppo_micro_batch_size_per_gpu} ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1} shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false} data_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null} diff --git a/verl/trainer/config/critic/critic.yaml b/verl/trainer/config/critic/critic.yaml index 8dca3ee8aa7..d6e09acfe35 100644 --- a/verl/trainer/config/critic/critic.yaml +++ b/verl/trainer/config/critic/critic.yaml @@ -48,6 +48,12 @@ ppo_max_token_len_per_gpu: 32768 # Max token length per GPU in forward pass forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu} +# Infer-time max tokens per GPU. Defaults to the documented forward budget. +ppo_infer_max_token_len_per_gpu: ${.forward_max_token_len_per_gpu} + +# Infer-time local per-GPU micro batch size. Defaults to the training micro-batch. +ppo_infer_micro_batch_size_per_gpu: ${.ppo_micro_batch_size_per_gpu} + # Number of PPO epochs per batch ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1} diff --git a/verl/trainer/config/critic/dp_critic.yaml b/verl/trainer/config/critic/dp_critic.yaml index 49337419c9e..0f1bcc6b70e 100644 --- a/verl/trainer/config/critic/dp_critic.yaml +++ b/verl/trainer/config/critic/dp_critic.yaml @@ -31,6 +31,9 @@ forward_micro_batch_size: ${oc.select:.ppo_micro_batch_size,null} # Forward-only batch size during inference (per GPU) forward_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size_per_gpu,null} +# Infer-time local per-GPU micro batch size. Prefer the forward-only knob. +ppo_infer_micro_batch_size_per_gpu: ${.forward_micro_batch_size_per_gpu} + # Sequence parallelism size for Ulysses-style model parallelism # [DEPRECATED] use fsdp_config.ulysses_sequence_parallel_size instead ulysses_sequence_parallel_size: 1 diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index 908d36b2f54..b5614721322 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -807,8 +807,7 @@ def init_workers(self): orig_critic_cfg = critic_cfg engine_config: EngineConfig = orig_critic_cfg.engine - engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu - engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu + orig_critic_cfg.apply_engine_batching(engine_config) # Build the critic profiler config via the hydra path (same as the actor / ref / SFT), # so its tool_config entries are real dataclass instances the torch/nsys/npu backends can diff --git a/verl/trainer/ppo/v1/trainer_base.py b/verl/trainer/ppo/v1/trainer_base.py index f510875d150..2b1e0845477 100644 --- a/verl/trainer/ppo/v1/trainer_base.py +++ b/verl/trainer/ppo/v1/trainer_base.py @@ -248,8 +248,7 @@ def _setup(self): # 2. define critic class if self.use_critic: critic_cfg: CriticConfig = omega_conf_to_dataclass(self.config.critic) - critic_cfg.engine.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu - critic_cfg.engine.max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu + critic_cfg.apply_engine_batching(critic_cfg.engine) # Wire the critic profiler config via the hydra path (real dataclass tool_config), so the # standalone critic TrainingWorker gets a working DistProfiler instead of a silent no-op. diff --git a/verl/workers/config/critic.py b/verl/workers/config/critic.py index f056b14ac52..8934b742ed9 100644 --- a/verl/workers/config/critic.py +++ b/verl/workers/config/critic.py @@ -128,6 +128,26 @@ def validate(self, n_gpus: int, train_batch_size: int): f"critic.ppo_mini_batch_size ({self.ppo_mini_batch_size})" ) + def apply_engine_batching(self, engine_config): + """Copy critic-level batching knobs onto the engine TrainingWorker reads. + + Actor/ref already assign these fields in ``engine_workers.py``. The critic + worker is a bare ``TrainingWorker``, so trainers must do the same here. + """ + engine_config.use_dynamic_bsz = self.use_dynamic_bsz + engine_config.micro_batch_size_per_gpu = self.ppo_micro_batch_size_per_gpu + infer_mbs = self.ppo_infer_micro_batch_size_per_gpu + if infer_mbs is None: + infer_mbs = getattr(self, "forward_micro_batch_size_per_gpu", None) + if infer_mbs is None: + infer_mbs = self.ppo_micro_batch_size_per_gpu + engine_config.infer_micro_batch_size_per_gpu = infer_mbs + engine_config.max_token_len_per_gpu = self.ppo_max_token_len_per_gpu + infer_tokens = self.ppo_infer_max_token_len_per_gpu + forward_tokens = getattr(self, "forward_max_token_len_per_gpu", None) + engine_config.infer_max_token_len_per_gpu = infer_tokens if infer_tokens is not None else forward_tokens + return engine_config + @staticmethod def _check_mutually_exclusive(mbs, mbs_per_gpu, name: str): """Validate mutually exclusive micro batch size configuration options.