apply reviews and fix trainer_test

Author: chenyushuo

.github/workflows/docker/docker-compose.yaml

@@ -1,6 +1,6 @@
 services:
   trinity-node-1:
-    image: trinity-rft-unittest:20251225
+    image: trinity-rft-unittest:20260115
     pull_policy: never
     command: bash -c "source /opt/venv/bin/activate && uv pip install -e .[dev] && ray start --head --dashboard-host 0.0.0.0 --include-dashboard true --block"
     environment:
@@ -30,7 +30,7 @@ services:
             capabilities: [gpu]
 
   trinity-node-2:
-    image: trinity-rft-unittest:20251225
+    image: trinity-rft-unittest:20260115
     pull_policy: never
     command: bash -c "source /opt/venv/bin/activate && uv pip install -e .[dev] && ray start --address=trinity-node-1:6379 --block"
     environment:

trinity/common/models/vllm_worker.py

@@ -13,7 +13,7 @@
 class WorkerExtension:
     def apply_patches(self):
         """Apply necessary patches to vLLM."""
-        from verl.utils.vllm_utils import patch_vllm_moe_model_weight_loader
+        from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
 
         patch_vllm_moe_model_weight_loader(self.model_runner.model)
         patch_vllm_prompt_logprobs(self.model_runner)

trinity/trainer/verl/fsdp_checkpoint_manager.py

@@ -419,6 +419,7 @@ def save_state_dict(
                 self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg
             ):
                 self._save_model(local_path, global_step)
+        self._save_tokenizer(local_path, global_step)
         ray.get(
             self.checkpoint_monitor.register_thread_count.remote(
                 global_step, state_dict_thread_count=1

trinity/trainer/verl/fsdp_workers.py

@@ -53,7 +53,6 @@
     get_device_name,
     get_nccl_backend,
     get_torch_device,
-    set_expandable_segments,
 )
 from verl.utils.flops_counter import FlopsCounter
 from verl.utils.fs import copy_to_local
@@ -75,7 +74,6 @@
 )
 from verl.utils.import_utils import import_external_libs
 from verl.utils.logger import log_with_rank
-from verl.utils.memory_utils import aggressive_empty_cache
 from verl.utils.profiler import (
     DistProfiler,
     DistProfilerExtension,
@@ -640,24 +638,6 @@ def _build_model_optimizer(  # noqa: C901
 
         return actor_module_fsdp, actor_optimizer, actor_lr_scheduler, actor_model_config
 
-    async def trainer_mode(self):  # TODO: check this
-        """Context switch hybridengine to trainer mode."""
-        # if self.config.rollout.free_cache_engine:
-        #     log_gpu_memory_usage("Before rollout offload", logger=logger)
-        #     await self.rollout.release()
-        #     log_gpu_memory_usage("After rollout offload", logger=logger)
-
-        self.actor_module_fsdp.train()
-
-        # add empty cache after each compute
-        aggressive_empty_cache(force_sync=True)
-
-        set_expandable_segments(True)
-
-        # restore random states
-        self.gen_random_states = get_torch_device().get_rng_state()
-        get_torch_device().set_rng_state(self.torch_random_states)
-
     @register(dispatch_mode=Dispatch.ONE_TO_ALL)
     def init_model(self):
         from trinity.trainer.verl.dp_actor import DataParallelPPOActor
@@ -1604,7 +1584,7 @@ def update_critic(self, data: DataProto):
             )
 
             lr = self.critic_lr_scheduler.get_last_lr()[0]
-            metrics["critic/lr"] = lr
+            metrics["critic/lr"] = lr.item() if torch.is_tensor(lr) else lr
             self.critic_lr_scheduler.step()
 
             output = DataProto(batch=None, meta_info={"metrics": metrics})

trinity/trainer/verl/megatron_actor.py

@@ -363,7 +363,7 @@ def logits_processor(logits, label, label_mask):
                     logits.div_(temperature)
                     ret = {}
                     if calculate_entropy:
-                        # The veRL fix consumes more GPU memory than our implementation 
+                        # The veRL fix consumes more GPU memory than our implementation
                         # (.clone() v.s. monkey patch on megatron function);
                         # therefore, we have temporarily commented out the veRL fix.
                         # logits_bak = logits.clone()

trinity/trainer/verl/megatron_checkpoint_manager.py

@@ -396,6 +396,7 @@ def save_state_dict(  # noqa: C901
 
         local_path = local_mkdir_safe(local_path)
         self._save_state_dict(local_path, global_step)
+        self._save_tokenizer(local_path, global_step)
         ray.get(
             self.checkpoint_monitor.register_thread_count.remote(
                 global_step, state_dict_thread_count=1

trinity/trainer/verl/megatron_workers.py

@@ -745,28 +745,6 @@ def upload_state_dict(self, trainer_step: int):
     def set_algorithm(self, algo_config: AlgorithmConfig):
         self.actor.set_algorithm(algo_config)
 
-    async def trainer_mode(self):
-        """Context switch hybridengine to trainer mode."""
-        # if self.config.rollout.free_cache_engine:
-        #     log_gpu_memory_usage("Before rollout offload", logger=logger)
-        #     await self.rollout.release()
-        #     log_gpu_memory_usage("After rollout offload", logger=logger)
-
-        for model in self.actor.actor_module:
-            model.train()
-        # add empty cache after each compute
-        aggressive_empty_cache(force_sync=True)
-
-        # FIXME(@wuxibin): megatron+sglang failed with `expandable_segments:True` in ci,
-        # can't reproduce it in dev environment, temporary disable it.
-        # https://github.com/volcengine/verl/actions/runs/17382936845/job/49344264323?pr=3285
-        if os.environ.get("MEGATRON_CI_DISABLE_EXPANDABLE_SEGMENTS", "0") == "0":
-            set_expandable_segments(True)
-
-        # restore random states
-        self.gen_random_states = get_torch_device().get_rng_state()
-        get_torch_device().set_rng_state(self.torch_random_states)
-
     @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor"))
     @GPUMemoryLogger(role="update_actor", logger=logger)
     @DistProfiler.annotate(color="red", role="actor_update")

trinity/trainer/verl/utils.py

@@ -100,7 +100,7 @@ def to_data_proto(
             )
     else:
         raise ValueError("Custom fields are not consistent across experiences.")
-    meta_info = {"model_versions": np.array([exp.info["model_version"] for exp in experiences])}
+    meta_info = {"model_versions": np.array([exp.info.get("model_version", 0) for exp in experiences])}
     return DataProto.from_single_dict(batch_dict, meta_info=meta_info)
 
 

trinity/trainer/verl_trainer.py

@@ -344,40 +344,6 @@ def init_workers(self):  # noqa: C901
             )
             self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls
 
-        # create a reward model if reward_fn is None
-        # for legacy discriminative reward model, we create a reward model worker here
-        # for reward loop discriminative reward model, we create a reward loop manager here
-        if not self.use_reward_loop:
-            # legacy reward model only handle reward-model based scenario
-            if self.use_rm:
-                # we create a RM here
-                resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
-                rm_cls = RayClassWithInitArgs(
-                    self.role_worker_mapping[Role.RewardModel], config=self.config.reward_model
-                )
-                self.resource_pool_to_cls[resource_pool][str(Role.RewardModel)] = rm_cls
-        else:
-            # reward loop handle hybrid reward scenario (rule, disrm, genrm, ...)
-            # Note: mode is always "async" since sync mode is deprecated
-            can_reward_loop_parallelize = (
-                not self.use_rm or self.config.reward_model.enable_resource_pool
-            )
-            # judge if we can asynchronously parallelize reward model with actor rollout
-            # two condition that we can parallelize reward model with actor rollout:
-            # 1. reward model is not enabled (rule-based reward can parallelize)
-            # 2. reward model is enabled but extra resource pool is enabled
-            # If we cannot parallelize, we should enable synchronous mode here, and launch a reward loop manager here
-            # else for parallelize mode, we launch a reward worker for each rollout worker (in agent loop, not here)
-            if not can_reward_loop_parallelize:
-                from verl.experimental.reward_loop import RewardLoopManager
-
-                self.config.reward_model.n_gpus_per_node = self.config.trainer.n_gpus_per_node
-                resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
-                self.reward_loop_manager = RewardLoopManager(
-                    config=self.config,
-                    rm_resource_pool=resource_pool,
-                )
-
         # initialize WorkerGroup
         # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,
         # you should not use `create_colocated_worker_cls`.
@@ -439,12 +405,6 @@ def init_workers(self):  # noqa: C901
                 assert str(Role.ActorRolloutRef) in all_wg, f"{all_wg.keys()=}"
                 self.ref_policy_wg = all_wg[str(Role.ActorRolloutRef)]
 
-        self.rm_wg = None
-        # initalization of rm_wg will be deprecated in the future
-        if self.use_rm and not self.use_reward_loop:
-            self.rm_wg = all_wg[str(Role.RewardModel)]
-            self.rm_wg.init_model()
-
         # we should create rollout at the end so that vllm can have a better estimation of kv cache memory
         self.actor_rollout_wg = all_wg[str(actor_role)]
         self.actor_rollout_wg.init_model()
@@ -515,13 +475,14 @@ async def train_step(self, batch_exps: List[Experience]) -> Dict:  # noqa C901
                 "bypass_mode", False
             )
             if bypass_recomputing_logprobs:  # Use `rollout_log_probs`
-                from verl.trainer.ppo.rollout_corr_helper import apply_bypass_mode
+                if "rollout_log_probs" in batch.batch:
+                    from verl.trainer.ppo.rollout_corr_helper import apply_bypass_mode
 
-                apply_bypass_mode(
-                    batch=batch,
-                    rollout_corr_config=rollout_corr_config,
-                    policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss,
-                )
+                    apply_bypass_mode(
+                        batch=batch,
+                        rollout_corr_config=rollout_corr_config,
+                        policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss,
+                    )
             else:  # Recompute old_log_probs  TODO: to be check
                 if (batch.meta_info["model_versions"] != self.global_steps - 1).any():
                     self.logger.warning(
@@ -551,8 +512,6 @@ async def train_step(self, batch_exps: List[Experience]) -> Dict:  # noqa C901
 
                         metrics.update(calculate_debug_metrics(batch))
 
-            assert "old_log_probs" in batch.batch, f'"old_log_prob" not in {batch.batch.keys()=}'
-
             if self.algorithm.use_reference:  # ref_logprob may not be used
                 # compute reference log_prob
                 with marked_timer(str(Role.RefPolicy), timing_raw, color="olive"):