Skip to content

Support customizing trainer and daemon in VERL #407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ultmaster merged 4 commits into from
Dec 12, 2025
Merged

Support customizing trainer and daemon in VERL #407

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4760e78
support
ultmaster Dec 12, 2025
8aedda6
missing annotation flag
ultmaster Dec 12, 2025
9c98c58
resolve comments
ultmaster Dec 12, 2025
17abbfc
update imports
ultmaster Dec 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 26 additions & 2 deletions agentlightning/algorithm/verl/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Copyright (c) Microsoft. All rights reserved.

from typing import Any, Optional
from __future__ import annotations

from typing import TYPE_CHECKING, Any, Optional, Type

from hydra import compose, initialize
from omegaconf import OmegaConf
Expand All @@ -10,6 +12,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.
Expand All @@ -23,6 +29,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
Expand Down Expand Up @@ -90,7 +98,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:
Expand All @@ -102,6 +115,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,
Expand All @@ -119,6 +134,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:
Expand All @@ -130,6 +150,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.")
Expand All @@ -142,6 +164,8 @@ def run(
store=store,
llm_proxy=llm_proxy,
adapter=adapter,
trainer_cls=trainer_cls,
daemon_cls=daemon_cls,
)

def get_client(self) -> AgentLightningClient:
Expand Down
62 changes: 46 additions & 16 deletions agentlightning/verl/entrypoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,16 @@
# 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 __future__ import annotations

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

Expand All @@ -17,7 +20,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",
Expand All @@ -27,8 +33,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,
Comment on lines +44 to +45

Copilot AI Dec 12, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The AgentLightningTrainer and AgentModeDaemon classes are imported inside TYPE_CHECKING but are used at runtime in the main function. This will cause a NameError at runtime because these names are not available outside of type checking. Move these imports outside the TYPE_CHECKING block or use string literals for the default values.

Copilot uses AI. Check for mistakes.
)


def run_ppo(
Expand All @@ -38,6 +53,8 @@ def run_ppo(
store: LightningStore | None,
llm_proxy: LLMProxy | None,
adapter: TraceAdapter[Any] | None,
trainer_cls: Type[AgentLightningTrainer],
daemon_cls: Type[AgentModeDaemon],
) -> None:

Copilot AI Dec 12, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The run_ppo function has new parameters trainer_cls and daemon_cls but lacks a docstring to document what these parameters are for. Add a docstring that describes all parameters, especially the new ones, to help users understand their purpose and usage.

Suggested change
) -> None:
) -> None:
"""
Run the PPO (Proximal Policy Optimization) training loop using the provided configuration and components.
Parameters:
config (Any): The configuration object for the training run, typically loaded via Hydra.
train_dataset (Dataset[Any] | None): The training dataset to use, or None if not provided.
val_dataset (Dataset[Any] | None): The validation dataset to use, or None if not provided.
store (LightningStore | None): The LightningStore instance for storing and retrieving data, or None.
llm_proxy (LLMProxy | None): The LLMProxy instance for model inference, or None.
adapter (TraceAdapter[Any] | None): The TraceAdapter for logging or tracing, or None.
trainer_cls (Type[AgentLightningTrainer]): The class to use for creating the PPO trainer. This allows customization of the training logic by providing a different trainer class.
daemon_cls (Type[AgentModeDaemon]): The class to use for creating the agent mode daemon. This allows customization of the agent's runtime behavior by providing a different daemon class.
Returns:
None
This function initializes Ray if necessary, then launches the PPO training process using the provided datasets,
store, LLM proxy, adapter, and customizable trainer and daemon classes.
"""

Copilot uses AI. Check for mistakes.
if not ray.is_initialized():
# this is for local ray cluster
Expand All @@ -56,13 +73,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,
)
)

Expand All @@ -72,11 +91,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,
trainer_cls: Type[AgentLightningTrainer],
daemon_cls: Type[AgentModeDaemon],
):

Copilot AI Dec 12, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The TaskRunner.run method has new parameters daemon_cls and trainer_cls but lacks a docstring to document what these parameters are for. Add a docstring that describes all parameters, especially the new ones.

Suggested change
):
):
"""
Run the main training or evaluation task using the provided configuration and components.
Args:
config (Any): The configuration object for the experiment, typically an OmegaConf config.
train_dataset (Dataset[Any] | None): The training dataset to use, or None if not provided.
val_dataset (Dataset[Any] | None): The validation dataset to use, or None if not provided.
store (LightningStore | None): The LightningStore instance for storing experiment data, or None.
llm_proxy (LLMProxy | None): The LLMProxy instance for model inference, or None.
adapter (TraceAdapter[Any] | None): The TraceAdapter for logging or tracing, or None.
daemon_cls (Type[AgentModeDaemon]): The class to use for creating the agent mode daemon. This should be a subclass of AgentModeDaemon and is responsible for managing agent modes during training or evaluation.
trainer_cls (Type[AgentLightningTrainer]): The class to use for creating the trainer. This should be a subclass of AgentLightningTrainer and is responsible for orchestrating the training or evaluation process.
Returns:
None
"""

Copilot uses AI. Check for mistakes.
# print initial config
from pprint import pprint
Expand All @@ -91,7 +112,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)
Expand All @@ -112,7 +133,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
Expand All @@ -121,9 +143,16 @@ def run(
else:
raise NotImplementedError

from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role
from verl.trainer.ppo.ray_trainer import ResourcePoolManager

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 = {
role_worker_mapping: dict[Role, ActorClass[Any]] = {
Role.ActorRollout: ray.remote(actor_rollout_cls),
Role.Critic: ray.remote(CriticWorker),
}
Expand Down Expand Up @@ -190,7 +219,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,
Expand All @@ -206,6 +235,7 @@ def run(
store=store,
llm_proxy=llm_proxy,
adapter=adapter,
daemon_cls=daemon_cls,
)
trainer.init_workers()
trainer.fit()
Expand Down
12 changes: 9 additions & 3 deletions agentlightning/verl/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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],

Copilot AI Dec 12, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Type import is missing from the typing imports. The daemon_cls parameter in line 181 uses Type[AgentModeDaemon] as a type hint, but Type is not imported. Add Type to the imports from typing.

Copilot uses AI. Check for mistakes.
**kwargs,
):

Copilot AI Dec 12, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docstring for the __init__ method is missing. The new daemon_cls parameter should be documented to help users understand its purpose and usage. Add a docstring that describes all parameters including the newly added daemon_cls.

Suggested change
):
):
"""
Initialize the AgentLightningTrainer.
Args:
store (LightningStore | None): The storage backend for logging and data persistence.
llm_proxy (LLMProxy | None): Proxy for interacting with the language model.
adapter (TraceAdapter | None): Adapter for converting traces to the required format.
daemon_cls (Type[AgentModeDaemon]): The class to use for creating the agent mode daemon responsible for server communication and agent orchestration.
**kwargs: Additional keyword arguments passed to the base RayPPOTrainer.
"""

Copilot uses AI. Check for mistakes.
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."
Expand Down Expand Up @@ -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={
Expand Down