[worker] feat: New engine share actor and ref for LoRA (#4867)

HollowMan6 · web-flow · commit 145a28b5c508 · 2026-01-11T14:06:08.000+08:00
### What does this PR do? Continuation of #4673, now sharing actor and ref for LoRA is also supported in new engine ### Checklist Before Starting - [X] Search for similar PRs. Paste at least one query link here: ... - [X] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) - `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward` - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` ### Test > For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc. ### API and Usage Example > Demonstrate how the API changes if any, and provide usage example(s) if possible. ```python # Add code snippet or script demonstrating how to use this ``` ### Design & Code Changes > Demonstrate the high-level design if this PR is complex, and list the specific changes. ### Checklist Before Submitting > [!IMPORTANT] > Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. - [X] Read the [Contribute Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md). - [X] Apply [pre-commit checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` - [ ] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs). - [ ] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... - [ ] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ). (If not accessible, please try [the Feishu group (飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).) - [ ] If your PR is related to the `recipe` submodule, please also update the reference to the submodule commit via `git submodule update --remote` or `cd recipe && git pull origin main`. Signed-off-by: Hollow Man <hollowman@opensuse.org>
diff --git a/examples/split_placement/main_ppo_split.py b/examples/split_placement/main_ppo_split.py
@@ -23,6 +23,7 @@
 
 from verl import DataProto
 from verl.trainer.ppo.ray_trainer import RayPPOTrainer
+from verl.trainer.ppo.utils import need_reference_policy
 from verl.utils.reward_score import gsm8k, math_reward
 
 
@@ -171,7 +172,7 @@ def main_task(config):
     }
 
     # use reference model
-    if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
+    if need_reference_policy(config):
         role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker)
         mapping[Role.RefPolicy] = actor_rollout_ref_pool_id
 
diff --git a/verl/experimental/fully_async_policy/fully_async_main.py b/verl/experimental/fully_async_policy/fully_async_main.py
@@ -26,7 +26,7 @@
 from verl.experimental.fully_async_policy.fully_async_trainer import FullyAsyncTrainer
 from verl.experimental.fully_async_policy.message_queue import MessageQueue, MessageQueueClient
 from verl.trainer.ppo.ray_trainer import ResourcePoolManager
-from verl.trainer.ppo.utils import Role
+from verl.trainer.ppo.utils import Role, need_reference_policy
 from verl.utils.fs import copy_to_local
 
 
@@ -122,7 +122,7 @@ def create_role_worker_mapping(config):
         role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker)
 
     # Add reference policy (if KL loss or reward is required)
-    if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
+    if need_reference_policy(config):
         role_worker_mapping[Role.RefPolicy] = ray.remote(DetachActorWorker)
 
     return role_worker_mapping, ray_worker_group_cls
diff --git a/verl/experimental/fully_async_policy/fully_async_trainer.py b/verl/experimental/fully_async_policy/fully_async_trainer.py
@@ -73,7 +73,7 @@ def __init__(
 
         self.role_worker_mapping = role_worker_mapping
         self.resource_pool_manager = resource_pool_manager
-        self.use_reference_policy = need_reference_policy(self.role_worker_mapping)
+        self.use_reference_policy = need_reference_policy(self.config)
         self.use_rm = need_reward_model(self.role_worker_mapping)
         self.use_critic = need_critic(self.config)
         self.ray_worker_group_cls = ray_worker_group_cls
diff --git a/verl/experimental/one_step_off_policy/main_ppo.py b/verl/experimental/one_step_off_policy/main_ppo.py
@@ -124,7 +124,7 @@ def create_role_worker_mapping(config):
         role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker)
 
     # Add reference policy (if KL loss or reward is required)
-    if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
+    if need_reference_policy(config):
         role_worker_mapping[Role.RefPolicy] = ray.remote(DetachActorWorker)
 
     return role_worker_mapping, ray_worker_group_cls
@@ -151,7 +151,7 @@ def run(self, config):
         # validate config
         validate_config(
             config=config,
-            use_reference_policy=need_reference_policy(role_worker_mapping),
+            use_reference_policy=need_reference_policy(config),
             use_critic=need_critic(config),
         )
 
diff --git a/verl/experimental/one_step_off_policy/ray_trainer.py b/verl/experimental/one_step_off_policy/ray_trainer.py
@@ -106,7 +106,7 @@ def __init__(
 
         self.role_worker_mapping = role_worker_mapping
         self.resource_pool_manager = resource_pool_manager
-        self.use_reference_policy = need_reference_policy(self.role_worker_mapping)
+        self.use_reference_policy = need_reference_policy(self.config)
         self.use_rm = need_reward_model(self.role_worker_mapping)
         self.use_critic = need_critic(config)
         self.ray_worker_group_cls = ray_worker_group_cls
diff --git a/verl/experimental/transfer_queue/main_ppo.py b/verl/experimental/transfer_queue/main_ppo.py
@@ -23,13 +23,8 @@
 from omegaconf import OmegaConf
 
 from verl.trainer.constants_ppo import get_ppo_ray_runtime_env
-from verl.trainer.main_ppo import (
-    TaskRunner as MainTaskRunner,
-)
-from verl.trainer.main_ppo import (
-    create_rl_dataset,
-    create_rl_sampler,
-)
+from verl.trainer.main_ppo import TaskRunner as MainTaskRunner
+from verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler
 from verl.trainer.ppo.reward import load_reward_manager
 from verl.trainer.ppo.utils import need_critic, need_reference_policy
 from verl.utils.config import validate_config
@@ -148,7 +143,7 @@ def run(self, config):
         # validate config
         validate_config(
             config=config,
-            use_reference_policy=need_reference_policy(self.role_worker_mapping),
+            use_reference_policy=need_reference_policy(config),
             use_critic=need_critic(config),
         )
 
diff --git a/verl/experimental/transfer_queue/ray_trainer.py b/verl/experimental/transfer_queue/ray_trainer.py
@@ -369,7 +369,7 @@ def __init__(
 
         self.role_worker_mapping = role_worker_mapping
         self.resource_pool_manager = resource_pool_manager
-        self.use_reference_policy = need_reference_policy(self.role_worker_mapping)
+        self.use_reference_policy = need_reference_policy(self.config)
         self.use_rm = need_reward_model(self.role_worker_mapping)
         self.use_critic = need_critic(self.config)
         self.ray_worker_group_cls = ray_worker_group_cls
diff --git a/verl/trainer/main_ppo.py b/verl/trainer/main_ppo.py
@@ -133,9 +133,14 @@ def add_actor_rollout_worker(self, config):
 
             actor_rollout_cls = ActorRolloutRefWorker
             ray_worker_group_cls = RayWorkerGroup
+
+            lora_rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0)
+            if lora_rank <= 0:
+                lora_rank = config.actor_rollout_ref.model.get("lora_rank", 0)
+            ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get("lora_adapter_path") is not None
             # NOTE: In new model engine, ref policy and actor rollout are in same ActorRolloutRefWorker,
             # while in legacy model engine, ref policy is in a separate ActorRolloutRefWorker.
-            if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
+            if need_reference_policy(config) and not ref_in_actor:
                 role = Role.ActorRolloutRef
             else:
                 role = Role.ActorRollout
@@ -249,7 +254,7 @@ def add_ref_policy_worker(self, config, ref_policy_cls):
         if use_legacy_worker_impl == "disable":
             return
 
-        if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss:
+        if need_reference_policy(config):
             self.role_worker_mapping[Role.RefPolicy] = ray.remote(ref_policy_cls)
             self.mapping[Role.RefPolicy] = "global_pool"
 
@@ -291,7 +296,7 @@ def run(self, config):
         # validate config
         validate_config(
             config=config,
-            use_reference_policy=need_reference_policy(self.role_worker_mapping),
+            use_reference_policy=need_reference_policy(config),
             use_critic=need_critic(config),
         )
 
diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py
@@ -339,7 +339,7 @@ def __init__(
 
         self.role_worker_mapping = role_worker_mapping
         self.resource_pool_manager = resource_pool_manager
-        self.use_reference_policy = need_reference_policy(self.role_worker_mapping)
+        self.use_reference_policy = need_reference_policy(self.config)
         # legacy reward model implementation
         self.use_rm = need_reward_model(self.role_worker_mapping)
         self.use_reward_loop = self.config.reward_model.use_reward_loop
@@ -1244,8 +1244,14 @@ def _compute_ref_log_prob(self, batch: DataProto) -> DataProto:
             # step 2: convert from padding to nopadding
             batch_td = left_right_2_no_padding(batch_td)
             # step 3: add meta info
-            tu.assign_non_tensor(batch_td, calculate_entropy=False, compute_loss=False)
-            output = self.ref_policy_wg.compute_ref_log_prob(batch_td)
+            metadata = {"calculate_entropy": False, "compute_loss": False}
+            if self.ref_in_actor:
+                metadata["no_lora_adapter"] = True
+            tu.assign_non_tensor(batch_td, **metadata)
+            if self.ref_in_actor:
+                output = self.actor_rollout_wg.compute_log_prob(batch_td)
+            else:
+                output = self.ref_policy_wg.compute_ref_log_prob(batch_td)
             # gather output
             log_probs = tu.get(output, "log_probs")
             # step 4. No padding to padding
diff --git a/verl/trainer/ppo/utils.py b/verl/trainer/ppo/utils.py
@@ -70,10 +70,10 @@ def from_string(cls, name: str):
 
 
 def need_reference_policy(
-    role_worker_mapping: dict[Role, WorkerType],
+    config: DictConfig,
 ) -> bool:
-    """Given a role worker mapping, do we need ref policy."""
-    return Role.RefPolicy in role_worker_mapping or Role.ActorRolloutRef in role_worker_mapping
+    """Given the config, do we need ref policy."""
+    return config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss
 
 
 def need_reward_model(
diff --git a/verl/workers/engine/base.py b/verl/workers/engine/base.py
@@ -16,7 +16,8 @@
 """
 
 from abc import abstractmethod
-from typing import Any, Callable, Generator, Optional
+from contextlib import nullcontext
+from typing import Any, Callable, ContextManager, Generator, Optional
 
 import torch
 from tensordict import TensorDict
@@ -218,6 +219,12 @@ def is_mp_src_rank_with_outputs(self):
         """
         raise NotImplementedError
 
+    def disable_adapter(self) -> ContextManager:
+        """
+        Disable all adapters temporarily under the context in the model for LoRA
+        """
+        return nullcontext()
+
 
 class BaseEngineCtx:
     def __init__(self, engine: BaseEngine, mode, **kwargs):
diff --git a/verl/workers/engine/fsdp/transformer_impl.py b/verl/workers/engine/fsdp/transformer_impl.py
@@ -20,7 +20,7 @@
 import os
 import warnings
 from contextlib import nullcontext
-from typing import Callable, Optional
+from typing import Callable, ContextManager, Optional
 
 import torch
 import torch.distributed
@@ -38,10 +38,7 @@
 from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager
 from verl.utils.dataset.dataset_utils import DatasetPadMode
 from verl.utils.debug import log_gpu_memory_usage
-from verl.utils.device import (
-    get_device_id,
-    get_device_name,
-)
+from verl.utils.device import get_device_id, get_device_name
 from verl.utils.fsdp_utils import (
     CPUOffloadPolicy,
     FSDPModule,
@@ -673,6 +670,9 @@ def get_per_tensor_param(self, layered_summon=False, base_sync_done=False):
             )
         return per_tensor_param, peft_config
 
+    def disable_adapter(self) -> ContextManager:
+        return self.module.disable_adapter()
+
 
 class EngineEvalModeCtx(BaseEngineCtx):
     def __init__(self, engine: FSDPEngine, **kwargs):
diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py
@@ -15,7 +15,7 @@
 import logging
 import os
 from functools import partial
-from typing import Any, Callable, Iterator, Optional
+from typing import Any, Callable, ContextManager, Iterator, Optional
 
 import torch
 import torch.distributed
@@ -547,6 +547,9 @@ def get_per_tensor_param(self, **kwargs):
         # TODO: support megatron LoRA
         return per_tensor_param, None
 
+    def disable_adapter(self) -> ContextManager:
+        return self.peft_cls.disable_adapter(self.module)
+
     def forward_step(self, batch_iter, model, postprocess_micro_batch_func):
         raise NotImplementedError("forward_step must be implemented in subclass")
 
diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 import logging
 import os
+from contextlib import nullcontext
 from functools import partial
 from itertools import chain
 from typing import Any, Optional
@@ -310,6 +311,7 @@ def infer_batch(self, data: TensorDict) -> TensorDict:
         global_token_num = tu.get(data, key="global_token_num")
         compute_loss = tu.get(data, key="compute_loss", default=True)
         disable_auto_offload = tu.get(data, key="disable_auto_offload", default=False)
+        no_lora_adapter = tu.pop(data, key="no_lora_adapter", default=False)
 
         default_keys = dict(
             use_remove_padding=self.model_config.use_remove_padding,
@@ -330,7 +332,9 @@ def infer_batch(self, data: TensorDict) -> TensorDict:
             self.engine.eval_mode(disable_auto_offload=disable_auto_offload),
             Timer(name="eval_batch", logger=None) as timer,
         ):
-            output = self.engine.infer_batch(data, loss_function=loss_function)
+            adapter_ctx = self.engine.disable_adapter() if no_lora_adapter else nullcontext()
+            with adapter_ctx:
+                output = self.engine.infer_batch(data, loss_function=loss_function)
         delta_time = timer.last
 
         if self.engine.is_mp_src_rank_with_outputs():
