From 4760e786f69522d691123bc091446beb35dc6775 Mon Sep 17 00:00:00 2001 From: Yuge Zhang Date: Fri, 12 Dec 2025 09:53:11 +0800 Subject: [PATCH 1/4] support --- agentlightning/algorithm/verl/interface.py | 26 ++++++++++- agentlightning/verl/entrypoint.py | 54 +++++++++++++++------- agentlightning/verl/trainer.py | 10 +++- 3 files changed, 70 insertions(+), 20 deletions(-) diff --git a/agentlightning/algorithm/verl/interface.py b/agentlightning/algorithm/verl/interface.py index 47393eb11..0cea87d61 100644 --- a/agentlightning/algorithm/verl/interface.py +++ b/agentlightning/algorithm/verl/interface.py @@ -1,6 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional, Type from hydra import compose, initialize from omegaconf import OmegaConf @@ -10,6 +10,10 @@ from agentlightning.types import Dataset from agentlightning.verl.entrypoint import run_ppo # type: ignore +if TYPE_CHECKING: + from agentlightning.verl.daemon import AgentModeDaemon + from agentlightning.verl.trainer import AgentLightningTrainer + class VERL(Algorithm): """VERL-powered algorithm that delegates training to the VERL PPO runner. @@ -23,6 +27,8 @@ class VERL(Algorithm): config: Dictionary mirroring the overrides passed to the VERL CLI. The overrides are merged with VERL's packaged defaults via Hydra before launching training. + trainer_cls: Optional override for the trainer class. Experimental. + daemon_cls: Optional override for the daemon class. Experimental. Examples: ```python @@ -90,7 +96,12 @@ class VERL(Algorithm): ``` """ - def __init__(self, config: dict[str, Any]): + def __init__( + self, + config: dict[str, Any], + trainer_cls: Optional[Type[AgentLightningTrainer]] = None, + daemon_cls: Optional[Type[AgentModeDaemon]] = None, + ): super().__init__() # Compose the base config exactly like your decorator: @@ -102,6 +113,8 @@ def __init__(self, config: dict[str, Any]): # Allow adding new fields OmegaConf.set_struct(base_cfg, False) self.config = OmegaConf.merge(base_cfg, override_conf) + self.trainer_cls = trainer_cls + self.daemon_cls = daemon_cls def run( self, @@ -119,6 +132,11 @@ def run( adapter have been garbage-collected when using the V1 execution mode. """ + from agentlightning.verl.daemon import AgentModeDaemon + from agentlightning.verl.trainer import AgentLightningTrainer + + trainer_cls = self.trainer_cls or AgentLightningTrainer + daemon_cls = self.daemon_cls or AgentModeDaemon try: store = self.get_store() except Exception: @@ -130,6 +148,8 @@ def run( store=None, llm_proxy=None, adapter=None, + trainer_cls=trainer_cls, + daemon_cls=daemon_cls, ) else: print("Store is set. Assuming v1 execution mode.") @@ -142,6 +162,8 @@ def run( store=store, llm_proxy=llm_proxy, adapter=adapter, + trainer_cls=trainer_cls, + daemon_cls=daemon_cls, ) def get_client(self) -> AgentLightningClient: diff --git a/agentlightning/verl/entrypoint.py b/agentlightning/verl/entrypoint.py index 0a6644908..8bcb3b4de 100644 --- a/agentlightning/verl/entrypoint.py +++ b/agentlightning/verl/entrypoint.py @@ -1,13 +1,14 @@ # Copyright (c) Microsoft. All rights reserved. -# type: ignore +# pyright: reportUnknownVariableType=false +# pyright: reportUnknownMemberType=false +# pyright: reportUnknownArgumentType=false -from importlib.metadata import version -from typing import Any +from typing import TYPE_CHECKING, Any, Type import hydra import ray -from packaging import version as packaging_version +from ray.actor import ActorClass from verl.trainer.main_ppo import create_rl_sampler from verl.trainer.ppo.reward import load_reward_manager @@ -17,7 +18,10 @@ from agentlightning.types import Dataset from .dataset import AgentDataset, LoadedDataset -from .trainer import AgentLightningTrainer + +if TYPE_CHECKING: + from .daemon import AgentModeDaemon + from .trainer import AgentLightningTrainer __all__ = [ "main", @@ -27,8 +31,17 @@ @hydra.main(config_path="pkg://agentlightning/verl", config_name="config", version_base=None) -def main(config): - run_ppo(config, train_dataset=None, val_dataset=None, store=None, llm_proxy=None, adapter=None) +def main(config: Any): + run_ppo( + config, + train_dataset=None, + val_dataset=None, + store=None, + llm_proxy=None, + adapter=None, + trainer_cls=AgentLightningTrainer, + daemon_cls=AgentModeDaemon, + ) def run_ppo( @@ -38,6 +51,8 @@ def run_ppo( store: LightningStore | None, llm_proxy: LLMProxy | None, adapter: TraceAdapter[Any] | None, + trainer_cls: Type[AgentLightningTrainer], + daemon_cls: Type[AgentModeDaemon], ) -> None: if not ray.is_initialized(): # this is for local ray cluster @@ -56,13 +71,15 @@ def run_ppo( runner = TaskRunner.remote() ray.get( - runner.run.remote( + runner.run.remote( # type: ignore config=config, train_dataset=train_dataset, val_dataset=val_dataset, store=store, llm_proxy=llm_proxy, adapter=adapter, + trainer_cls=trainer_cls, + daemon_cls=daemon_cls, ) ) @@ -72,11 +89,13 @@ class TaskRunner: def run( self, config: Any, - train_dataset: Dataset | None, - val_dataset: Dataset | None, + train_dataset: Dataset[Any] | None, + val_dataset: Dataset[Any] | None, store: LightningStore | None, llm_proxy: LLMProxy | None, - adapter: TraceAdapter | None, + adapter: TraceAdapter[Any] | None, + daemon_cls: Type[AgentModeDaemon], + trainer_cls: Type[AgentLightningTrainer], ): # print initial config from pprint import pprint @@ -91,7 +110,7 @@ def run( local_path = copy_to_local(config.actor_rollout_ref.model.path) # instantiate tokenizer - from verl.utils import hf_processor, hf_tokenizer + from verl.utils.tokenizer 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) @@ -112,7 +131,8 @@ def run( elif config.actor_rollout_ref.actor.strategy == "megatron": assert config.actor_rollout_ref.actor.strategy == config.critic.strategy - from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup + # FIXME: This import is outdated + from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup # type: ignore from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker actor_rollout_cls = ActorRolloutRefWorker @@ -121,9 +141,10 @@ def run( else: raise NotImplementedError - from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role + from verl.trainer.ppo.ray_trainer import ResourcePoolManager + from verl.trainer.ppo.utils import Role - role_worker_mapping = { + role_worker_mapping: dict[Role, ActorClass[Any]] = { Role.ActorRollout: ray.remote(actor_rollout_cls), Role.Critic: ray.remote(CriticWorker), } @@ -190,7 +211,7 @@ def run( val_dataset = LoadedDataset(val_dataset) train_sampler = create_rl_sampler(config.data, train_dataset) - trainer = AgentLightningTrainer( + trainer = trainer_cls( config=config, tokenizer=tokenizer, processor=processor, @@ -206,6 +227,7 @@ def run( store=store, llm_proxy=llm_proxy, adapter=adapter, + daemon_cls=daemon_cls, ) trainer.init_workers() trainer.fit() diff --git a/agentlightning/verl/trainer.py b/agentlightning/verl/trainer.py index b19d07124..950418598 100644 --- a/agentlightning/verl/trainer.py +++ b/agentlightning/verl/trainer.py @@ -174,12 +174,18 @@ class AgentLightningTrainer(RayPPOTrainer): """ def __init__( - self, store: LightningStore | None, llm_proxy: LLMProxy | None, adapter: TraceAdapter | None, **kwargs + self, + store: LightningStore | None, + llm_proxy: LLMProxy | None, + adapter: TraceAdapter | None, + daemon_cls: Type[AgentModeDaemon], + **kwargs, ): super().__init__(**kwargs) self.store = store self.llm_proxy = llm_proxy self.adapter = adapter + self.daemon_cls = daemon_cls def _validate(self): assert len(self.val_dataloader) == 1, "Please set val_batch_size to None for better throughput." @@ -444,7 +450,7 @@ def fit(self): else: # For other versions (e.g., 0.6.0), we use the full path to the model. model = self.config.actor_rollout_ref.model.path - self.agent_mode_daemon = AgentModeDaemon( + self.agent_mode_daemon = self.daemon_cls( self.config.agentlightning.port, self.config.actor_rollout_ref.rollout.n, train_information={ From 8aedda64fe22370e4caff3ba757486f420184f3f Mon Sep 17 00:00:00 2001 From: Yuge Zhang Date: Fri, 12 Dec 2025 09:54:04 +0800 Subject: [PATCH 2/4] missing annotation flag --- agentlightning/algorithm/verl/interface.py | 2 ++ agentlightning/verl/entrypoint.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/agentlightning/algorithm/verl/interface.py b/agentlightning/algorithm/verl/interface.py index 0cea87d61..f8a90dc35 100644 --- a/agentlightning/algorithm/verl/interface.py +++ b/agentlightning/algorithm/verl/interface.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + from typing import TYPE_CHECKING, Any, Optional, Type from hydra import compose, initialize diff --git a/agentlightning/verl/entrypoint.py b/agentlightning/verl/entrypoint.py index 8bcb3b4de..d23984234 100644 --- a/agentlightning/verl/entrypoint.py +++ b/agentlightning/verl/entrypoint.py @@ -4,6 +4,8 @@ # pyright: reportUnknownMemberType=false # pyright: reportUnknownArgumentType=false +from __future__ import annotations + from typing import TYPE_CHECKING, Any, Type import hydra From 9c98c58a05a28a49a83662282274fd42541d0e34 Mon Sep 17 00:00:00 2001 From: Yuge Zhang Date: Fri, 12 Dec 2025 10:22:33 +0800 Subject: [PATCH 3/4] resolve comments --- agentlightning/verl/entrypoint.py | 2 +- agentlightning/verl/trainer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/agentlightning/verl/entrypoint.py b/agentlightning/verl/entrypoint.py index d23984234..a0638ebc6 100644 --- a/agentlightning/verl/entrypoint.py +++ b/agentlightning/verl/entrypoint.py @@ -96,8 +96,8 @@ def run( store: LightningStore | None, llm_proxy: LLMProxy | None, adapter: TraceAdapter[Any] | None, - daemon_cls: Type[AgentModeDaemon], trainer_cls: Type[AgentLightningTrainer], + daemon_cls: Type[AgentModeDaemon], ): # print initial config from pprint import pprint diff --git a/agentlightning/verl/trainer.py b/agentlightning/verl/trainer.py index 950418598..3916597c3 100644 --- a/agentlightning/verl/trainer.py +++ b/agentlightning/verl/trainer.py @@ -8,7 +8,7 @@ from contextlib import contextmanager from copy import deepcopy from pprint import pprint -from typing import Dict, Tuple +from typing import Dict, Tuple, Type import numpy as np import torch From 17abbfc9f6827e385f3f26b6edda18fda1c8561d Mon Sep 17 00:00:00 2001 From: Yuge Zhang Date: Fri, 12 Dec 2025 12:18:28 +0800 Subject: [PATCH 4/4] update imports --- agentlightning/verl/entrypoint.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/agentlightning/verl/entrypoint.py b/agentlightning/verl/entrypoint.py index a0638ebc6..07150516b 100644 --- a/agentlightning/verl/entrypoint.py +++ b/agentlightning/verl/entrypoint.py @@ -144,7 +144,13 @@ def run( raise NotImplementedError from verl.trainer.ppo.ray_trainer import ResourcePoolManager - from verl.trainer.ppo.utils import Role + + try: + # verl >= 0.6.0 + from verl.trainer.ppo.utils import Role + except ImportError: + # Fallback for verl <= 0.5.0 + from verl.trainer.ppo.ray_trainer import Role # type: ignore role_worker_mapping: dict[Role, ActorClass[Any]] = { Role.ActorRollout: ray.remote(actor_rollout_cls),