[trainer, cfg] fix: wire critic batching keys into TrainingWorkerConfig

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

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=}")
 

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}

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}

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

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}

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}
 

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

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

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.

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.