diff --git a/code/RL_model/verl/verl_train/verl/experimental/agent_loop/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d43683df3e482e63c897bcbd8135064037f4de5d
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/__init__.py
@@ -0,0 +1,21 @@
+# Copyright 2024 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 .agent_loop import AgentLoopBase, AgentLoopManager, AgentLoopWorker, AsyncLLMServerManager
+from .single_turn_agent_loop import SingleTurnAgentLoop
+from .tool_agent_loop import ToolAgentLoop
+
+_ = [SingleTurnAgentLoop, ToolAgentLoop]
+
+__all__ = ["AgentLoopBase", "AgentLoopManager", "AsyncLLMServerManager", "AgentLoopWorker"]
diff --git a/code/RL_model/verl/verl_train/verl/experimental/agent_loop/agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..fe984d47e19e6c6f0994da00d87e7cc6d540cf45
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/agent_loop.py
@@ -0,0 +1,1022 @@
+# Copyright 2024 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.
+import asyncio
+import heapq
+import logging
+import os
+import random
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+from uuid import uuid4
+
+import hydra
+import numpy as np
+import ray
+import torch
+from cachetools import LRUCache
+from omegaconf import DictConfig, OmegaConf
+from PIL import Image
+from pydantic import BaseModel, ConfigDict
+from tensordict import TensorDict
+from transformers import AutoProcessor, AutoTokenizer
+
+from verl.experimental.agent_loop.prometheus_utils import update_prometheus_config
+from verl.experimental.agent_loop.utils import resolve_config_path
+from verl.experimental.reward_loop import RewardLoopWorker
+from verl.protocol import DataProto
+from verl.single_controller.ray.base import RayResourcePool, RayWorkerGroup
+from verl.utils import hf_processor, hf_tokenizer
+from verl.utils.chat_template import initialize_system_prompt
+from verl.utils.dataset.rl_dataset import RLHFDataset, get_dataset_class
+from verl.utils.fs import copy_to_local
+from verl.utils.model import compute_position_id_with_mask
+from verl.utils.ray_utils import get_event_loop
+from verl.utils.rollout_trace import (
+ RolloutTraceConfig,
+ rollout_trace_attr,
+ rollout_trace_op,
+)
+from verl.utils.transferqueue_utils import tqbridge
+from verl.workers.rollout.replica import TokenOutput, get_rollout_replica_class
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class AsyncLLMServerManager:
+ """
+ A class to manage multiple OpenAI compatible LLM servers. This class provides
+ - Load balance: least requests load balancing
+ - Sticky session: send multi-turn chat completions to same server for automatic prefix caching
+ """
+
+ def __init__(self, config: DictConfig, server_handles: list[ray.actor.ActorHandle], max_cache_size: int = 10000):
+ """Initialize the AsyncLLMServerManager.
+
+ Args:
+ config (DictConfig): YAML config.
+ server_handles (List[ray.actor.ActorHandle]): OpenAI compatible LLM server actor handles.
+ max_cache_size (int, optional): max cache size for request_id to server mapping. Defaults to 10000.
+ """
+ self.config = config
+ self.server_handles = server_handles
+ random.shuffle(self.server_handles)
+
+ # Least requests load balancing
+ self.weighted_serveres = [[0, idx, server] for idx, server in enumerate(self.server_handles)]
+ heapq.heapify(self.weighted_serveres)
+
+ # LRU cache to map request_id to server
+ self.request_id_to_server = LRUCache(maxsize=max_cache_size)
+
+ def _choose_server(self, request_id: str) -> ray.actor.ActorHandle:
+ # TODO: implement server pressure awareness load balancing
+ if request_id in self.request_id_to_server:
+ return self.request_id_to_server[request_id]
+
+ _, _, server = self.weighted_serveres[0]
+ self.weighted_serveres[0][0] += 1
+ heapq.heapreplace(self.weighted_serveres, self.weighted_serveres[0])
+ self.request_id_to_server[request_id] = server
+ return server
+
+ @rollout_trace_op
+ async def generate(
+ self,
+ request_id,
+ *,
+ prompt_ids: list[int],
+ sampling_params: dict[str, Any],
+ image_data: Optional[list[Any]] = None,
+ video_data: Optional[list[Any]] = None,
+ ) -> TokenOutput:
+ """Generate tokens from prompt ids.
+
+ Args:
+ request_id (str): request id for sticky session.
+ prompt_ids (List[int]): List of prompt token ids.
+ sampling_params (Dict[str, Any]): Sampling parameters for the chat completion.
+
+ Returns:
+ TokenOutput: token output
+ """
+ server = self._choose_server(request_id)
+ output = await server.generate.remote(
+ request_id=uuid4().hex, # use new request_id for each turn
+ prompt_ids=prompt_ids,
+ sampling_params=sampling_params,
+ image_data=image_data,
+ video_data=video_data,
+ )
+ return output
+
+
+class AgentLoopMetrics(BaseModel):
+ """Agent loop performance metrics."""
+
+ generate_sequences: float = 0.0
+ tool_calls: float = 0.0
+ num_preempted: int = -1 # -1 means not available
+
+
+class AgentLoopOutput(BaseModel):
+ """Agent loop output."""
+
+ prompt_ids: list[int]
+ """Prompt token ids."""
+ response_ids: list[int]
+ """Response token ids including LLM generated token, tool response token."""
+ response_mask: list[int]
+ """Response mask, 1 for LLM generated token, 0 for tool response token."""
+ response_logprobs: Optional[list[float]] = None
+ """Log probabilities for the response tokens."""
+ routed_experts: Optional[Any] = None
+ """Routed experts for the total tokens."""
+ multi_modal_data: Optional[dict[str, Any]] = None
+ """Multi-modal data for multi-modal tools."""
+ reward_score: Optional[float] = None
+ """Reward score for the trajectory."""
+ num_turns: int = 0
+ """Number of chat turns, including user, assistant, tool."""
+ metrics: AgentLoopMetrics
+ """Auxiliary performance metrics"""
+ extra_fields: dict[str, Any] = {}
+ """Extra fields for dynamic addition."""
+
+
+class _InternalAgentLoopOutput(AgentLoopOutput):
+ """Internal agent loop output with padded sequences."""
+
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+
+ prompt_ids: torch.Tensor
+ """Padded prompt token ids."""
+ response_ids: torch.Tensor
+ """Padded response token ids."""
+ input_ids: torch.Tensor
+ """Padded input ids(prompt_ids + response_ids)."""
+ position_ids: torch.Tensor
+ """Padded position ids."""
+ response_mask: torch.Tensor
+ """Padded response mask."""
+ attention_mask: torch.Tensor
+ """Padded attention mask."""
+ response_logprobs: Optional[torch.Tensor] = None
+ """Padded log probabilities for the response tokens."""
+ routed_experts: Optional[torch.Tensor] = None
+ """Padded routed experts for the total tokens."""
+ multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None
+ """Multi-modal inputs for processors (e.g., pixel_values, image_grid_thw)."""
+ extra_fields: dict[str, Any] = {}
+ """Extra fields for dynamic addition."""
+
+
+class DictConfigWrap:
+ """Wrapper for DictConfig to avoid hydra.utils.instantiate recursive resolve."""
+
+ def __init__(self, config: DictConfig):
+ self.config = config
+
+
+class AgentLoopBase(ABC):
+ """An agent loop takes an input message, chat with OpenAI compatible LLM server and interact with various
+ environments."""
+
+ def __init__(
+ self,
+ trainer_config: DictConfigWrap,
+ server_manager: AsyncLLMServerManager,
+ tokenizer: AutoTokenizer,
+ processor: AutoProcessor,
+ dataset_cls: type[RLHFDataset],
+ dataset_config: DictConfigWrap,
+ **kwargs,
+ ):
+ """Initialize agent loop, each sample will have its own loop instance.
+
+ Args:
+ trainer_config (DictConfigWrap): trainer config.
+ server_manager (AsyncLLMServerManager): OpenAI compatible LLM server manager.
+ tokenizer (AutoTokenizer): Tokenizer for tokenize messages.
+ processor (AutoProcessor): Processor for process messages.
+ dataset_cls (type[Dataset]): Dataset class for creating dataset, Defaults to RLHFDataset.
+ dataset_config (DictConfigWrap): Dataset config.
+ """
+ self.config = trainer_config.config
+ self.server_manager = server_manager
+ self.tokenizer = tokenizer
+ self.processor = processor
+ self.dataset_cls = dataset_cls
+ self.dataset_config = dataset_config.config
+ self.apply_chat_template_kwargs = self.dataset_config.get("apply_chat_template_kwargs", {})
+ self.system_prompt = initialize_system_prompt(self.tokenizer, **self.apply_chat_template_kwargs)
+ self.loop = get_event_loop()
+
+ async def process_vision_info(self, messages: list[dict]) -> dict:
+ """Extract images and videos from messages.
+
+ Args:
+ messages (list[dict]): Input messages.
+
+ Returns:
+ dict: Multi-modal data with keys "images" and "videos".
+ """
+ multi_modal_data = {}
+ if self.processor is not None:
+ images, videos = await self.dataset_cls.process_vision_info(
+ messages, image_patch_size=self.processor.image_processor.patch_size, config=self.dataset_config
+ )
+ if images is not None:
+ multi_modal_data["images"] = images
+ if videos is not None:
+ multi_modal_data["videos"] = videos
+
+ return multi_modal_data
+
+ async def apply_chat_template(
+ self,
+ messages: list[dict],
+ tools: list[dict] = None,
+ images: list[Image.Image] = None,
+ videos: list[tuple[torch.Tensor, dict]] = None,
+ remove_system_prompt: bool = False,
+ ):
+ """Apply chat template to messages with optional tools, images, and videos.
+
+ Args:
+ messages (list[dict]): Input messages.
+ tools (list[dict], optional): Tools schemas. Defaults to None.
+ images (list[Image.Image], optional): Input images. Defaults to None.
+ videos (list[tuple[torch.Tensor, dict]], optional): Input videos. Defaults to None.
+ remove_system_prompt (bool, optional): Whether to remove system prompt. Defaults to False.
+
+ Returns:
+ list[int]: Prompt token ids.
+ """
+ if self.processor is not None:
+ raw_prompt = await self.loop.run_in_executor(
+ None,
+ lambda: self.processor.apply_chat_template(
+ messages,
+ tools=tools,
+ add_generation_prompt=True,
+ tokenize=False,
+ **self.apply_chat_template_kwargs,
+ ),
+ )
+
+ # split the videos and according metadatas
+ if videos is not None:
+ videos, video_metadatas = zip(*videos, strict=False)
+ videos, video_metadatas = list(videos), list(video_metadatas)
+ else:
+ video_metadatas = None
+
+ model_inputs = self.processor(
+ text=[raw_prompt],
+ images=images,
+ videos=videos,
+ video_metadatas=video_metadatas,
+ return_tensors="pt",
+ do_sample_frames=False,
+ )
+ prompt_ids = model_inputs.pop("input_ids").squeeze(0).tolist()
+ else:
+ prompt_ids = await self.loop.run_in_executor(
+ None,
+ lambda: self.tokenizer.apply_chat_template(
+ messages,
+ tools=tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ **self.apply_chat_template_kwargs,
+ ),
+ )
+
+ if remove_system_prompt:
+ prompt_ids = prompt_ids[len(self.system_prompt) :]
+
+ return prompt_ids
+
+ @abstractmethod
+ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
+ """Run agent loop to interact with LLM server and environment.
+
+ Args:
+ sampling_params (Dict[str, Any]): LLM sampling params.
+ **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`.
+
+ Returns:
+ AgentLoopOutput: Agent loop output.
+ """
+ raise NotImplementedError
+
+
+"""Agent loop registry: key is agent_name, value is a dict of agent loop config
+used by hydra.utils.instantiate to initialize agent loop instance.
+
+https://hydra.cc/docs/advanced/instantiate_objects/overview/
+"""
+_agent_loop_registry: dict[str, dict] = {}
+
+
+def register(agent_name: str):
+ """Register agent loop class."""
+
+ def decorator(subclass: type[AgentLoopBase]) -> type[AgentLoopBase]:
+ fqdn = f"{subclass.__module__}.{subclass.__qualname__}"
+ _agent_loop_registry[agent_name] = {"_target_": fqdn}
+ return subclass
+
+ return decorator
+
+
+class AgentLoopWorker:
+ """Agent loop worker takes a batch of messages and run each message in an agent loop."""
+
+ def __init__(
+ self,
+ config: DictConfig,
+ server_handles: list[ray.actor.ActorHandle],
+ reward_router_address: str = None,
+ ):
+ """Initialize agent loop manager.
+ Args:
+ config (DictConfig): YAML config.
+ server_handles (List[ray.actor.ActorHandle]): OpenAI compatible LLM server actor handles.
+ reward_router_address (str): reward router address.
+ """
+ self.config = config
+
+ # for recipe to change
+ if not hasattr(self, "server_manager"):
+ self.server_manager = AsyncLLMServerManager(config, server_handles)
+
+ self.dataset_cls = get_dataset_class(config.data)
+ self.reward_router_address = reward_router_address
+
+ model_path = config.actor_rollout_ref.model.path
+ self.model_name = "/".join(model_path.split("/")[-2:])
+ local_path = copy_to_local(config.actor_rollout_ref.model.path)
+ self.tokenizer = hf_tokenizer(local_path, trust_remote_code=True)
+ self.processor = hf_processor(local_path, trust_remote_code=True)
+
+ agent_loop_config_path = config.actor_rollout_ref.rollout.agent.agent_loop_config_path
+ if agent_loop_config_path:
+ resolved_path = resolve_config_path(agent_loop_config_path)
+ agent_loop_configs = OmegaConf.load(resolved_path)
+ for agent_loop_config in agent_loop_configs:
+ _agent_loop_registry[agent_loop_config.name] = agent_loop_config
+ if self.config.actor_rollout_ref.model.get("custom_chat_template", None) is not None:
+ if self.processor is not None:
+ self.processor.chat_template = self.config.actor_rollout_ref.model.custom_chat_template
+ self.tokenizer.chat_template = self.config.actor_rollout_ref.model.custom_chat_template
+
+ use_reward_loop = True if self.config.reward_model.use_reward_loop else None
+ self.use_reward_loop = use_reward_loop
+ if use_reward_loop and not hasattr(self, "reward_loop_worker"):
+ self.reward_loop_worker = RewardLoopWorker.options(
+ scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+ node_id=ray.get_runtime_context().get_node_id(),
+ soft=False,
+ ),
+ ).remote(self.config, self.reward_router_address)
+
+ trace_config = self.config.actor_rollout_ref.rollout.get("trace", {})
+ RolloutTraceConfig.init(
+ self.config.trainer.project_name,
+ self.config.trainer.experiment_name,
+ trace_config.get("backend"),
+ trace_config.get("token2text", False),
+ trace_config.get("max_samples_per_step_per_worker", None),
+ )
+
+ @tqbridge()
+ async def generate_sequences(self, batch: DataProto) -> DataProto:
+ """Generate sequences from agent loop.
+
+ Args:
+ batch (DataProto): Input batch.
+
+ Returns:
+ DataProto: Output batch.
+ - prompts: [bsz, prompt_length], prompt token ids from dataset.
+ - responses: [bsz, response_length], output token ids include response tokens
+ from LLM generation and observation tokens from tool_calls.
+ - response_mask: [bsz, response_length], 1 for LLM generated tokens, 0 for observation/padding tokens.
+ - input_ids: [bsz, prompt_length + response_length], whole sequence token ids, including prompt tokens
+ and response tokens.
+ - attention_mask: [bsz, prompt_length + response_length], 0 for padding tokens, 1 for other tokens.
+ - position_ids: [bsz, prompt_length + response_length], incremental position ids.
+
+ For multi-turn conversations:
+ responses: |<- LLM generation ->|<- tool_calls ->|<- LLM generation ->|<- padding ->|
+ response_mask: | 1, 1, 1, ..., 1, 1 | 0, 0, .., 0, 0 | 1, 1, 1, ..., 1, 1 | 0, 0, ..., 0|
+ """
+ config = self.config.actor_rollout_ref.rollout
+ sampling_params = dict(
+ temperature=config.temperature,
+ top_p=config.top_p,
+ top_k=config.top_k,
+ repetition_penalty=1.0,
+ logprobs=config.calculate_log_probs,
+ )
+
+ # override sampling params for validation
+ if batch.meta_info.get("validate", False):
+ sampling_params["top_p"] = config.val_kwargs.top_p
+ sampling_params["top_k"] = config.val_kwargs.top_k
+ sampling_params["temperature"] = config.val_kwargs.temperature
+
+ # by default, we assume it's a single turn agent
+ if "agent_name" not in batch.non_tensor_batch:
+ default_agent_loop = config.agent.default_agent_loop
+ batch.non_tensor_batch["agent_name"] = np.array([default_agent_loop] * len(batch), dtype=object)
+
+ if "index" in batch.non_tensor_batch:
+ index = batch.non_tensor_batch["index"]
+ else:
+ index = np.arange(len(batch))
+
+ max_samples_per_worker = RolloutTraceConfig.get_instance().max_samples_per_step_per_worker
+
+ # For n rollouts per sample, we trace all n rollouts for selected samples
+ # Note: This sampling happens per-worker, so total traces = max_samples_per_worker * num_workers * n
+ if max_samples_per_worker is not None:
+ unique_sample_indices = np.unique(index)
+ if max_samples_per_worker < len(unique_sample_indices):
+ selected_samples = set(
+ np.random.choice(unique_sample_indices, max_samples_per_worker, replace=False).tolist()
+ )
+ traced_indices = set(i for i in range(len(batch)) if index[i] in selected_samples)
+ else:
+ traced_indices = set(range(len(batch)))
+ else:
+ traced_indices = set(range(len(batch)))
+
+ trajectory_info = await get_trajectory_info(
+ batch.meta_info.get("global_steps", -1), index.tolist(), batch.meta_info.get("validate", False)
+ )
+
+ tasks = []
+ for i in range(len(batch)):
+ trace_this_sample = i in traced_indices
+ kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items()}
+ tasks.append(
+ asyncio.create_task(
+ self._run_agent_loop(sampling_params, trajectory_info[i], trace=trace_this_sample, **kwargs)
+ )
+ )
+ outputs = await asyncio.gather(*tasks)
+
+ output = self._postprocess(outputs)
+
+ return output
+
+ async def _run_agent_loop(
+ self,
+ sampling_params: dict[str, Any],
+ trajectory: dict[str, Any],
+ *,
+ agent_name: str,
+ trace: bool = True,
+ **kwargs,
+ ) -> _InternalAgentLoopOutput:
+ with rollout_trace_attr(
+ step=trajectory["step"],
+ sample_index=trajectory["sample_index"],
+ rollout_n=trajectory["rollout_n"],
+ validate=trajectory["validate"],
+ name="agent_loop",
+ trace=trace,
+ ):
+ assert agent_name in _agent_loop_registry, (
+ f"Agent loop {agent_name} not registered, registered agent loops: {_agent_loop_registry.keys()}"
+ )
+
+ agent_loop_config = _agent_loop_registry[agent_name]
+ agent_loop = hydra.utils.instantiate(
+ config=agent_loop_config,
+ trainer_config=DictConfigWrap(config=self.config),
+ server_manager=self.server_manager,
+ tokenizer=self.tokenizer,
+ processor=self.processor,
+ dataset_cls=self.dataset_cls,
+ dataset_config=DictConfigWrap(self.config.data),
+ )
+ output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs)
+ return await self._agent_loop_postprocess(output, **kwargs)
+
+ async def _agent_loop_postprocess(self, output, **kwargs) -> _InternalAgentLoopOutput:
+ """Perform post-processing operations on the output of each individual agent loop."""
+ output.extra_fields["raw_prompt"] = kwargs["raw_prompt"]
+
+ # Some AgentLoop may have already computed the reward score, e.g SWE-agent.
+
+ # NOTE: consistent with the legacy batch version of generate_sequences that existed in the
+ # deprecated vLLM SPMD rollout implementation.
+ # prompt_ids: left padded with zeros (e.g., [0,0,0,0,1,2,3,4])
+ # response_ids: right padded with zeros (e.g., [5,6,7,8,0,0,0,0])
+ # input_ids: concatenation of prompt + response
+ # Mask:
+ # For example, if the prompt is [1,2,3,4] and the response is [5,6,7,(tool start)8,9(tool end),10,11,12]
+ # - prompt_attention_mask: 0s for padding, 1s for tokens
+ # e.g., [0,0,0,0,1,1,1,1]
+ # - response_attention_mask: 0s for padding, 1s for tokens
+ # e.g., [1,1,1,1,1,1,1,1,1,1,1,0,0,0,0]
+ # attention_mask: concatenation of prompt_attention_mask and response_attention_mask
+ # e.g., [0,0,0,0,1,1,1,1(prompt),1,1,1,1,1,1,1,1,1,1,1,0,0,0,0(response)]
+ # - response_mask: 1s for LLM generated tokens, 0 for tool response/padding tokens
+ # e.g., [1,1,1,1,1,1,1,(tool start),0,0(tool end),1,1,0,0,0,0]
+ # - position_ids: sequential positions for tokens, starting at 0
+ # e.g., [0,0,0,0,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,0,0,0,0]
+
+ # TODO(wuxibin): remove padding and use tensordict.
+ self.tokenizer.padding_side = "left"
+ prompt_output = self.tokenizer.pad(
+ {"input_ids": output.prompt_ids},
+ padding="max_length",
+ max_length=self.config.actor_rollout_ref.rollout.prompt_length,
+ return_tensors="pt",
+ return_attention_mask=True,
+ )
+ if prompt_output["input_ids"].dim() == 1:
+ prompt_output["input_ids"] = prompt_output["input_ids"].unsqueeze(0)
+ prompt_output["attention_mask"] = prompt_output["attention_mask"].unsqueeze(0)
+
+ self.tokenizer.padding_side = "right"
+ response_output = self.tokenizer.pad(
+ {"input_ids": output.response_ids},
+ padding="max_length",
+ max_length=self.config.actor_rollout_ref.rollout.response_length,
+ return_tensors="pt",
+ return_attention_mask=True,
+ )
+ if response_output["input_ids"].dim() == 1:
+ response_output["input_ids"] = response_output["input_ids"].unsqueeze(0)
+ response_output["attention_mask"] = response_output["attention_mask"].unsqueeze(0)
+
+ response_mask_output = self.tokenizer.pad(
+ {"input_ids": output.response_mask},
+ padding="max_length",
+ max_length=self.config.actor_rollout_ref.rollout.response_length,
+ return_tensors="pt",
+ return_attention_mask=False,
+ )
+ if response_mask_output["input_ids"].dim() == 1:
+ response_mask_output["input_ids"] = response_mask_output["input_ids"].unsqueeze(0)
+
+ response_logprobs = None
+ if output.response_logprobs is not None:
+ pad_size = self.config.actor_rollout_ref.rollout.response_length - len(output.response_logprobs)
+ response_logprobs = torch.tensor(output.response_logprobs + [0.0] * pad_size).unsqueeze(0)
+
+ response_mask = response_mask_output["input_ids"] * response_output["attention_mask"]
+ attention_mask = torch.cat([prompt_output["attention_mask"], response_output["attention_mask"]], dim=1)
+ input_ids = torch.cat([prompt_output["input_ids"], response_output["input_ids"]], dim=1)
+
+ routed_experts = None
+ if output.routed_experts is not None:
+ total_length = input_ids.shape[1]
+ length, layer_num, topk_num = output.routed_experts.shape
+ if isinstance(output.routed_experts, np.ndarray):
+ experts_tensor = torch.from_numpy(output.routed_experts)
+ elif isinstance(output.routed_experts, torch.Tensor):
+ experts_tensor = output.routed_experts
+ else:
+ raise TypeError(f"Unsupported type for routed_experts: {type(output.routed_experts)}")
+ routed_experts = torch.zeros(1, total_length, layer_num, topk_num, dtype=experts_tensor.dtype)
+
+ # Calculate start position: left padding means original prompt starts at the end
+ start_pos = prompt_output["input_ids"].shape[1] - len(output.prompt_ids)
+ end_pos = min(start_pos + length, total_length)
+
+ # Add boundary checks for robustness
+ if start_pos < 0 or end_pos > total_length:
+ raise ValueError(
+ f"Invalid position range: start_pos={start_pos}, end_pos={end_pos}, total_length={total_length}"
+ )
+
+ routed_experts[:, start_pos:end_pos] = experts_tensor.unsqueeze(0)
+
+ multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids)
+ position_ids = self._compute_position_ids(input_ids, attention_mask, multi_modal_inputs)
+ await self._compute_score(
+ output,
+ prompts=prompt_output["input_ids"],
+ responses=response_output["input_ids"],
+ attention_mask=attention_mask,
+ input_ids=input_ids,
+ position_ids=position_ids,
+ kwargs=kwargs,
+ )
+
+ return _InternalAgentLoopOutput(
+ prompt_ids=prompt_output["input_ids"],
+ response_ids=response_output["input_ids"],
+ input_ids=input_ids,
+ position_ids=position_ids,
+ response_mask=response_mask,
+ attention_mask=attention_mask,
+ response_logprobs=response_logprobs,
+ routed_experts=routed_experts,
+ multi_modal_inputs=multi_modal_inputs,
+ multi_modal_data=output.multi_modal_data,
+ reward_score=output.reward_score,
+ num_turns=output.num_turns,
+ metrics=output.metrics,
+ extra_fields=output.extra_fields,
+ )
+
+ def _compute_multi_modal_inputs(self, output, input_ids) -> dict[str, torch.Tensor]:
+ """Compute multi-modal inputs with image and video."""
+ multi_modal_inputs = {}
+ if self.processor is None:
+ return multi_modal_inputs
+
+ images = output.multi_modal_data.get("images")
+ videos = output.multi_modal_data.get("videos")
+ # split the videos and according metadatas
+ if videos is not None:
+ videos, video_metadatas = zip(*videos, strict=False)
+ videos, video_metadatas = list(videos), list(video_metadatas)
+ else:
+ video_metadatas = None
+ current_text = self.tokenizer.decode(input_ids.squeeze(0), skip_special_tokens=True)
+ multi_modal_inputs = self.processor(
+ text=[current_text],
+ images=images,
+ videos=videos,
+ video_metadatas=video_metadatas,
+ return_tensors="pt",
+ do_sample_frames=False,
+ )
+ multi_modal_inputs.pop("input_ids", None)
+ multi_modal_inputs.pop("attention_mask", None)
+
+ # We must use dict(multi_modal_inputs) to convert BatchFeature values to a new dict
+ # because np.array() only keeps the keys for BatchFeature.
+ multi_modal_inputs = dict(multi_modal_inputs.convert_to_tensors("pt"))
+ image_grid_thw = multi_modal_inputs.get("image_grid_thw")
+ if image_grid_thw is not None:
+ images_seqlens = torch.repeat_interleave(image_grid_thw[:, 1] * image_grid_thw[:, 2], image_grid_thw[:, 0])
+ multi_modal_inputs["images_seqlens"] = images_seqlens
+ return multi_modal_inputs
+
+ def _compute_position_ids(self, input_ids, attention_mask, multi_modal_inputs) -> torch.Tensor:
+ """Compute position ids for multi-modal inputs."""
+ if self.processor is None:
+ return compute_position_id_with_mask(attention_mask) # (1, seq_len)
+
+ image_grid_thw = multi_modal_inputs.get("image_grid_thw")
+ video_grid_thw = multi_modal_inputs.get("video_grid_thw")
+
+ # Model's get_rope_index has been dynamically bind to the processor.
+ vision_position_ids, _ = self.processor.get_rope_index(
+ input_ids=input_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ attention_mask=attention_mask,
+ )
+ vision_position_ids = vision_position_ids.transpose(0, 1) # (3, 1, seq_len) => (1, 3, seq_len)
+
+ valid_mask = attention_mask[0].bool()
+ text_position_ids = torch.ones((1, len(input_ids[0])), dtype=torch.long)
+ text_position_ids[0, valid_mask] = torch.arange(valid_mask.sum().item())
+ text_position_ids = text_position_ids.unsqueeze(0)
+ position_ids = torch.cat((text_position_ids, vision_position_ids), dim=1) # (1, 4, seq_length)
+ return position_ids
+
+ async def _compute_score(self, output, prompts, responses, attention_mask, input_ids, position_ids, kwargs):
+ """Compute reward score for single sample."""
+ enable_async_reward = (
+ self.reward_router_address is not None and self.config.reward_model.enable_resource_pool
+ ) or not self.config.reward_model.enable
+
+ if output.reward_score is None and enable_async_reward and self.use_reward_loop:
+ batch = TensorDict(
+ {
+ "prompts": prompts, # [1, prompt_length]
+ "responses": responses, # [1, response_length]
+ "attention_mask": attention_mask, # [1, prompt_length + response_length]
+ "input_ids": input_ids, # [1, prompt_length + response_length]
+ "position_ids": position_ids,
+ },
+ batch_size=1,
+ )
+ non_tensor_batch = {
+ **{k: np.array([v]) for k, v in kwargs.items()},
+ "__num_turns__": np.array([output.num_turns]),
+ "tool_extra_fields": np.array([output.extra_fields], dtype=object),
+ }
+
+ data = DataProto(
+ batch=batch,
+ non_tensor_batch=non_tensor_batch,
+ )
+ result = await self.reward_loop_worker.compute_score.remote(data)
+ output.reward_score = result["reward_score"]
+ output.extra_fields["reward_extra_info"] = result["reward_extra_info"]
+
+ def _postprocess(self, inputs: list[_InternalAgentLoopOutput]) -> DataProto:
+ """Process the padded outputs from _run_agent_loop and combine them into a batch."""
+ # Convert lists back to tensors and stack them to create a batch.
+ prompt_ids = torch.cat([input.prompt_ids for input in inputs], dim=0)
+ response_ids = torch.cat([input.response_ids for input in inputs], dim=0)
+ response_mask = torch.cat([input.response_mask for input in inputs], dim=0)
+ attention_mask = torch.cat([input.attention_mask for input in inputs], dim=0)
+ input_ids = torch.cat([input.input_ids for input in inputs], dim=0)
+ position_ids = torch.cat([input.position_ids for input in inputs], dim=0)
+ optional_outputs = {}
+ if inputs[0].response_logprobs is not None:
+ optional_outputs["rollout_log_probs"] = torch.cat([input.response_logprobs for input in inputs], dim=0)
+ if inputs[0].routed_experts is not None:
+ optional_outputs["routed_experts"] = torch.cat([input.routed_experts for input in inputs], dim=0)
+
+ batch = TensorDict(
+ {
+ "prompts": prompt_ids, # [bsz, prompt_length]
+ "responses": response_ids, # [bsz, response_length]
+ "response_mask": response_mask, # [bsz, response_length]
+ "input_ids": input_ids, # [bsz, prompt_length + response_length]
+ "attention_mask": attention_mask, # [bsz, prompt_length + response_length]
+ # position_ids: [bsz, 3, prompt_length + response_length] or [bsz, prompt_length + response_length]
+ "position_ids": position_ids,
+ **optional_outputs,
+ },
+ batch_size=len(inputs),
+ )
+
+ scores = [input.reward_score for input in inputs]
+ if all(score is not None for score in scores):
+ prompt_length = prompt_ids.size(1)
+ response_length = attention_mask[:, prompt_length:].sum(dim=1) - 1
+ rm_scores = torch.zeros_like(response_mask, dtype=torch.float32)
+ rm_scores[torch.arange(response_mask.size(0)), response_length] = torch.tensor(scores, dtype=torch.float32)
+ batch["rm_scores"] = rm_scores
+
+ non_tensor_batch = {
+ "__num_turns__": np.array([input.num_turns for input in inputs], dtype=np.int32),
+ }
+
+ # add reward_extra_info to non_tensor_batch
+ reward_extra_infos = [input.extra_fields.get("reward_extra_info", {}) for input in inputs]
+ reward_extra_keys = list(reward_extra_infos[0].keys())
+ for key in reward_extra_keys:
+ non_tensor_batch[key] = np.array([info[key] for info in reward_extra_infos])
+
+ # Add multi_modal_inputs to non_tensor_batch if any samples have them
+ multi_modal_inputs_list = [input.multi_modal_inputs for input in inputs]
+ if any(mmi is not None for mmi in multi_modal_inputs_list):
+ non_tensor_batch["multi_modal_inputs"] = np.array(multi_modal_inputs_list, dtype=object)
+
+ metrics = [input.metrics.model_dump() for input in inputs]
+ # Collect extra fields from all inputs and convert them to np.ndarray
+ extra_fields = {}
+ all_keys = set(key for input_item in inputs for key in input_item.extra_fields)
+ for key in all_keys:
+ temp_arr = np.empty(len(inputs), dtype=object)
+ temp_arr[:] = [input.extra_fields.get(key) for input in inputs]
+ extra_fields[key] = temp_arr
+
+ non_tensor_batch.update(extra_fields)
+
+ # Only include reward_extra_keys in meta_info if rm_scores is in batch
+ # This avoids conflicts when reward_tensor is merged later in ray_trainer.py
+ if "rm_scores" in batch.keys():
+ meta_info = {"metrics": metrics, "reward_extra_keys": reward_extra_keys}
+ else:
+ meta_info = {"metrics": metrics}
+
+ return DataProto(
+ batch=batch,
+ non_tensor_batch=non_tensor_batch,
+ meta_info=meta_info,
+ )
+
+ def create_transferqueue_client(
+ self,
+ ):
+ """Create a client for data system (TransferQueue)."""
+ from verl.single_controller.ray.base import get_random_string
+ from verl.utils.transferqueue_utils import create_transferqueue_client
+
+ client_name = get_random_string(length=6)
+
+ self.tq_client = create_transferqueue_client(
+ client_id=f"AgentLoopWorker_{client_name}",
+ config=self.config.transfer_queue,
+ )
+
+
+async def get_trajectory_info(step, index, validate):
+ """Get trajectory info.
+
+ Args:
+ step (int): global steps in the trainer.
+ index (list): form datastore extra_info.index column.
+ validate (bool): whether is a validate step.
+
+ Returns:
+ list: trajectory.
+ """
+ trajectory_info = []
+ rollout_n = 0
+ for i in range(len(index)):
+ if i > 0 and index[i - 1] == index[i]:
+ rollout_n += 1
+ else:
+ rollout_n = 0
+ trajectory_info.append({"step": step, "sample_index": index[i], "rollout_n": rollout_n, "validate": validate})
+ return trajectory_info
+
+
+class AgentLoopManager:
+ """Agent loop manager that manages a group of agent loop workers."""
+
+ def __init__(
+ self,
+ config: DictConfig,
+ worker_group: RayWorkerGroup = None,
+ rollout_resource_pool: RayResourcePool = None,
+ rm_resource_pool: RayResourcePool = None,
+ ):
+ """Initialize agent loop manager.
+
+ Args:
+ config (DictConfig): trainer config.
+ worker_group (RayWorkerGroup): ActorRolloutRef worker group for hybrid mode; None for standalone mode.
+ rollout_resource_pool (RayResourcePool): Resource pool for actor rollout (Colocate or Standalone mode).
+ rm_resource_pool (RayResourcePool): Resource pool for reward model (Standalone mode).
+ """
+ self.config = config
+ self.worker_group = worker_group
+ self.reward_model_manager = None
+ self.reward_router_address = None
+ if self.config.reward_model.enable and self.config.reward_model.enable_resource_pool:
+ from verl.experimental.reward_loop import RewardModelManager
+
+ self.reward_model_manager = RewardModelManager(config.reward_model, rm_resource_pool)
+ self.reward_router_address = self.reward_model_manager.get_router_address()
+
+ # for recipe to change
+ if not hasattr(self, "rollout_replica_class"):
+ self.rollout_replica_class = get_rollout_replica_class(self.config.actor_rollout_ref.rollout.name)
+ if not hasattr(self, "agent_loop_workers_class"):
+ self.agent_loop_workers_class = ray.remote(AgentLoopWorker)
+
+ self._initialize_llm_servers(rollout_resource_pool)
+ self._init_agent_loop_workers()
+
+ def _initialize_llm_servers(self, rollout_resource_pool: RayResourcePool):
+ rollout_world_size = (
+ self.config.actor_rollout_ref.rollout.tensor_model_parallel_size
+ * self.config.actor_rollout_ref.rollout.data_parallel_size
+ * self.config.actor_rollout_ref.rollout.pipeline_model_parallel_size
+ )
+ world_size = (
+ self.worker_group.world_size
+ if self.worker_group
+ else self.config.trainer.n_gpus_per_node * self.config.trainer.nnodes
+ )
+ num_replicas = world_size // rollout_world_size
+
+ rollout_config = self.config.actor_rollout_ref.rollout
+ model_config = self.config.actor_rollout_ref.model
+ self.rollout_replicas = [
+ self.rollout_replica_class(
+ replica_rank=replica_rank,
+ config=rollout_config,
+ model_config=model_config,
+ gpus_per_node=self.config.trainer.n_gpus_per_node,
+ )
+ for replica_rank in range(num_replicas)
+ ]
+
+ if self.worker_group and rollout_config.name != "trtllm":
+ self._run_all([server.init_hybrid(self.worker_group) for server in self.rollout_replicas])
+ elif self.worker_group and rollout_config.name == "trtllm":
+ self._run_all(
+ [
+ server.init_hybrid_colocated(self.worker_group, rollout_resource_pool)
+ for server in self.rollout_replicas
+ ]
+ )
+ else:
+ self._run_all([server.init_standalone() for server in self.rollout_replicas])
+
+ self.server_handles = [server._server_handle for server in self.rollout_replicas]
+ self.server_addresses = [server._server_address for server in self.rollout_replicas]
+
+ print(f"AgentLoopManager: {self.server_addresses}")
+
+ # Update Prometheus configuration with server addresses
+ if rollout_config.prometheus.enable:
+ if rollout_config.disable_log_stats:
+ raise ValueError("PROMETHEUS needs disable_log_stats==False, but it is currently True.")
+ update_prometheus_config(rollout_config.prometheus, self.server_addresses, rollout_config.name)
+
+ def _init_agent_loop_workers(self):
+ self.agent_loop_workers = []
+ num_workers = self.config.actor_rollout_ref.rollout.agent.num_workers
+
+ node_ids = [node["NodeID"] for node in ray.nodes() if node["Alive"] and node["Resources"].get("CPU", 0) > 0]
+ for i in range(num_workers):
+ # Round-robin scheduling over the all nodes
+ node_id = node_ids[i % len(node_ids)]
+ self.agent_loop_workers.append(
+ self.agent_loop_workers_class.options(
+ name=f"agent_loop_worker_{i}" + f"_{uuid4().hex[:8]}",
+ scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+ node_id=node_id, soft=True
+ ),
+ ).remote(self.config, self.server_handles, self.reward_router_address)
+ )
+
+ def generate_sequences(self, prompts: DataProto) -> DataProto:
+ """Split input batch and dispatch to agent loop workers.
+
+ Args:
+ prompts (DataProto): Input batch.
+
+ Returns:
+ DataProto: Output batch.
+ """
+
+ # TODO: move reward_model_manager out of agent_loop manager
+ if self.reward_model_manager:
+ self.reward_model_manager.wake_up()
+
+ chunkes = prompts.chunk(len(self.agent_loop_workers))
+ outputs = ray.get(
+ [
+ worker.generate_sequences.remote(chunk)
+ for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True)
+ ]
+ )
+ output = DataProto.concat(outputs)
+ if self.reward_model_manager:
+ self.reward_model_manager.sleep()
+
+ # calculate performance metrics
+ metrics = [output.meta_info.pop("metrics") for output in outputs] # List[List[Dict[str, str]]]
+ timing = self._performance_metrics(metrics, output)
+
+ output.meta_info = {"timing": timing, **outputs[0].meta_info}
+ return output
+
+ def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: DataProto) -> dict[str, float]:
+ timing = {}
+ t_generate_sequences = np.array([metric["generate_sequences"] for chunk in metrics for metric in chunk])
+ t_tool_calls = np.array([metric["tool_calls"] for chunk in metrics for metric in chunk])
+ num_preempted = np.array([metric["num_preempted"] for chunk in metrics for metric in chunk])
+ timing["agent_loop/num_preempted/min"] = num_preempted.min()
+ timing["agent_loop/num_preempted/max"] = num_preempted.max()
+ timing["agent_loop/num_preempted/mean"] = num_preempted.mean()
+ timing["agent_loop/generate_sequences/min"] = t_generate_sequences.min()
+ timing["agent_loop/generate_sequences/max"] = t_generate_sequences.max()
+ timing["agent_loop/generate_sequences/mean"] = t_generate_sequences.mean()
+ timing["agent_loop/tool_calls/min"] = t_tool_calls.min()
+ timing["agent_loop/tool_calls/max"] = t_tool_calls.max()
+ timing["agent_loop/tool_calls/mean"] = t_tool_calls.mean()
+
+ # batch sequence generation is bounded by the slowest sample
+ slowest = np.argmax(t_generate_sequences + t_tool_calls)
+ attention_mask = output.batch["attention_mask"][slowest]
+ prompt_length = output.batch["prompts"].shape[1]
+ timing["agent_loop/slowest/generate_sequences"] = t_generate_sequences[slowest]
+ timing["agent_loop/slowest/tool_calls"] = t_tool_calls[slowest]
+ timing["agent_loop/slowest/prompt_length"] = attention_mask[:prompt_length].sum().item()
+ timing["agent_loop/slowest/response_length"] = attention_mask[prompt_length:].sum().item()
+ timing["agent_loop/slowest/num_preempted"] = num_preempted[slowest]
+
+ return timing
+
+ def clear_kv_cache(self):
+ """Clear all rollout kv cache, but don`t sleep."""
+ self._run_all([replica.clear_kv_cache() for replica in self.rollout_replicas])
+
+ def start_profile(self, **kwargs):
+ """Start profiling on all rollout replicas."""
+ self._run_all([replica.start_profile(**kwargs) for replica in self.rollout_replicas])
+
+ def stop_profile(self):
+ """Stop profiling on all rollout replicas."""
+ self._run_all([replica.stop_profile() for replica in self.rollout_replicas])
+
+ def _run_all(self, tasks: list[asyncio.Task]):
+ async def run_all():
+ await asyncio.gather(*tasks)
+
+ asyncio.run(run_all())
diff --git a/code/RL_model/verl/verl_train/verl/experimental/agent_loop/prometheus_utils.py b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/prometheus_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..0ce582df61ed9cd27d63c6eb27a7885831ddc24a
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/prometheus_utils.py
@@ -0,0 +1,110 @@
+# Copyright 2025 Meituan 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.
+
+
+import logging
+import os
+
+import ray
+import yaml
+
+from verl.workers.config.rollout import PrometheusConfig
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+def update_prometheus_config(config: PrometheusConfig, server_addresses: list[str], rollout_name: str | None = None):
+ """
+ Update Prometheus configuration file with server addresses and reload on first node.
+
+ server_addresses: vllm or sglang server addresses
+
+ rollout_name: name of the rollout backend (e.g., "vllm", "sglang")
+ """
+
+ if not server_addresses:
+ logger.warning("No server addresses available to update Prometheus config")
+ return
+
+ try:
+ # Get Prometheus config file path from environment or use default
+ prometheus_config_json = {
+ "global": {"scrape_interval": "10s", "evaluation_interval": "10s"},
+ "scrape_configs": [
+ {
+ "job_name": "ray",
+ "file_sd_configs": [{"files": ["/tmp/ray/prom_metrics_service_discovery.json"]}],
+ },
+ {"job_name": "rollout", "static_configs": [{"targets": server_addresses}]},
+ ],
+ }
+
+ # Write configuration file to all nodes
+ @ray.remote(num_cpus=0)
+ def write_config_file(config_data, config_path):
+ os.makedirs(os.path.dirname(config_path), exist_ok=True)
+ with open(config_path, "w") as f:
+ yaml.dump(config_data, f, default_flow_style=False, indent=2)
+ return True
+
+ # Reload Prometheus on all nodes. Only master node should succeed, skip errors on other nodes.
+ @ray.remote(num_cpus=0)
+ def reload_prometheus(port):
+ import socket
+ import subprocess
+
+ hostname = socket.gethostname()
+ ip_address = socket.gethostbyname(hostname)
+
+ reload_url = f"http://{ip_address}:{port}/-/reload"
+
+ try:
+ subprocess.run(["curl", "-X", "POST", reload_url], capture_output=True, text=True, timeout=10)
+ print(f"Reloading Prometheus on node: {reload_url}")
+ except Exception:
+ # Skip errors on non-master nodes
+ pass
+
+ # Get all available nodes and schedule tasks on each node
+ nodes = ray.nodes()
+ alive_nodes = [node for node in nodes if node["Alive"]]
+
+ # Write config files on all nodes
+ write_tasks = []
+ for node in alive_nodes:
+ node_ip = node["NodeManagerAddress"]
+ task = write_config_file.options(
+ resources={"node:" + node_ip: 0.001} # Schedule to specific node
+ ).remote(prometheus_config_json, config.file)
+ write_tasks.append(task)
+
+ ray.get(write_tasks)
+
+ server_type = rollout_name.upper() if rollout_name else "rollout"
+ print(f"Updated Prometheus configuration at {config.file} with {len(server_addresses)} {server_type} servers")
+
+ # Reload Prometheus on all nodes
+ reload_tasks = []
+ for node in alive_nodes:
+ node_ip = node["NodeManagerAddress"]
+ task = reload_prometheus.options(
+ resources={"node:" + node_ip: 0.001} # Schedule to specific node
+ ).remote(config.port)
+ reload_tasks.append(task)
+
+ ray.get(reload_tasks)
+
+ except Exception as e:
+ logger.error(f"Failed to update Prometheus configuration: {e}")
diff --git a/code/RL_model/verl/verl_train/verl/experimental/agent_loop/single_turn_agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/single_turn_agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c479362aa4c9be50fb037afe18628d2122a51b4
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/single_turn_agent_loop.py
@@ -0,0 +1,84 @@
+# Copyright 2024 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.
+import logging
+import os
+from typing import Any
+from uuid import uuid4
+
+from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput, register
+from verl.tools.utils.tool_registry import initialize_tools_from_config
+from verl.utils.profiler import simple_timer
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+@register("single_turn_agent")
+class SingleTurnAgentLoop(AgentLoopBase):
+ """Naive agent loop that only do single turn chat completion."""
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.prompt_length = self.config.actor_rollout_ref.rollout.prompt_length
+ self.response_length = self.config.actor_rollout_ref.rollout.response_length
+
+ tool_config_path = self.config.data.tool_config_path
+ tool_list = initialize_tools_from_config(tool_config_path) if tool_config_path else []
+ self.tool_schemas = [tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list]
+
+ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
+ messages = list(kwargs["raw_prompt"])
+
+ # 1. extract images and videos from messages
+ multi_modal_data = await self.process_vision_info(messages)
+ images = multi_modal_data.get("images")
+ videos = multi_modal_data.get("videos")
+
+ # 2. apply chat template and tokenize
+ prompt_ids = await self.apply_chat_template(
+ messages,
+ tools=self.tool_schemas,
+ images=images,
+ videos=videos,
+ )
+
+ # 3. generate sequences
+ metrics = {}
+ with simple_timer("generate_sequences", metrics):
+ output = await self.server_manager.generate(
+ request_id=uuid4().hex,
+ prompt_ids=prompt_ids,
+ sampling_params=sampling_params,
+ image_data=images,
+ video_data=videos,
+ )
+ if metrics.get("num_preempted") is None:
+ metrics["num_preempted"] = output.num_preempted if output.num_preempted is not None else -1
+ response_mask = [1] * len(output.token_ids)
+
+ output = AgentLoopOutput(
+ prompt_ids=prompt_ids,
+ response_ids=output.token_ids[: self.response_length],
+ response_mask=response_mask[: self.response_length],
+ response_logprobs=output.log_probs[: self.response_length] if output.log_probs else None,
+ routed_experts=(
+ output.routed_experts[: len(prompt_ids) + self.response_length]
+ if output.routed_experts is not None
+ else None
+ ),
+ multi_modal_data=multi_modal_data,
+ num_turns=2,
+ metrics=metrics,
+ )
+ return output
diff --git a/code/RL_model/verl/verl_train/verl/experimental/agent_loop/tool_agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/tool_agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..f98485a6781f55592056d0df5a5923885693ab25
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/tool_agent_loop.py
@@ -0,0 +1,475 @@
+# Copyright 2025 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.
+import asyncio
+import json
+import logging
+import os
+from enum import Enum
+from typing import Any, Optional
+from uuid import uuid4
+
+import torch
+from PIL import Image
+from transformers import AutoProcessor, AutoTokenizer
+
+from verl.experimental.agent_loop.agent_loop import (
+ AgentLoopBase,
+ AgentLoopOutput,
+ AsyncLLMServerManager,
+ DictConfigWrap,
+ register,
+)
+from verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser
+from verl.experimental.agent_loop.utils import build_gpt_oss_tool_response_text
+from verl.interactions.base import BaseInteraction
+from verl.interactions.utils.interaction_registry import initialize_interactions_from_config
+from verl.tools.schemas import ToolResponse
+from verl.tools.utils.tool_registry import initialize_tools_from_config
+from verl.utils.profiler import simple_timer
+from verl.utils.rollout_trace import rollout_trace_op
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class AgentState(Enum):
+ PENDING = "pending"
+ GENERATING = "generating"
+ PROCESSING_TOOLS = "processing_tools"
+ TERMINATED = "terminated"
+ INTERACTING = "interacting"
+
+
+class AgentData:
+ """Encapsulates all state variables for the agent loop. AgentData is passed to tool calling in case that
+ tool may need to access full history state. User can store any tool session data in `extra_fields`."""
+
+ def __init__(
+ self,
+ messages: list[dict[str, Any]],
+ image_data: list[Image.Image],
+ video_data: list[tuple[torch.Tensor, dict[str, Any]]],
+ metrics: dict[str, Any],
+ request_id: str,
+ tools_kwargs: dict[str, Any],
+ interaction: Optional[BaseInteraction] = None,
+ interaction_kwargs: Optional[dict[str, Any]] = None,
+ ):
+ self.messages = messages
+ self.image_data = image_data
+ self.video_data = video_data
+ self.metrics = metrics
+ self.request_id = request_id
+ self.tools_kwargs = tools_kwargs
+ self.interaction = interaction
+ self.interaction_kwargs = interaction_kwargs or {}
+
+ # State variables
+ self.prompt_ids: list[int] = []
+ self.response_ids: list[int] = []
+ self.response_mask: list[int] = []
+ self.response_logprobs: list[float] = []
+ self.turn_scores: list[float] = []
+ self.tool_rewards: list[float] = []
+ self.user_turns = 0
+ self.assistant_turns = 0
+
+ # Temporary state for tool calls
+ self.tool_calls: list[FunctionCall] = []
+
+ # Extra fields for dynamic addition, e.g., tool session data
+ self.extra_fields: dict[str, Any] = {}
+
+
+@register("tool_agent")
+class ToolAgentLoop(AgentLoopBase):
+ def __init__(
+ self,
+ trainer_config: DictConfigWrap,
+ server_manager: AsyncLLMServerManager,
+ tokenizer: AutoTokenizer,
+ processor: AutoProcessor,
+ **kwargs,
+ ):
+ super().__init__(trainer_config, server_manager, tokenizer, processor, **kwargs)
+ config = trainer_config.config
+
+ # Initialize tools from config file
+ self.max_user_turns = config.actor_rollout_ref.rollout.multi_turn.max_user_turns
+ self.max_assistant_turns = config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns
+ self.max_parallel_calls = config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls
+ self.max_tool_response_length = config.actor_rollout_ref.rollout.multi_turn.max_tool_response_length
+ self.tool_response_truncate_side = config.actor_rollout_ref.rollout.multi_turn.tool_response_truncate_side
+ tool_config_path = config.actor_rollout_ref.rollout.multi_turn.tool_config_path
+ tool_list = initialize_tools_from_config(tool_config_path) if tool_config_path else []
+ self.tools = {tool.name: tool for tool in tool_list}
+ self.tool_schemas = [tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list]
+ self.tool_parser = ToolParser.get_tool_parser(
+ config.actor_rollout_ref.rollout.multi_turn.format, self.tokenizer
+ )
+ self.tool_parser_name = config.actor_rollout_ref.rollout.multi_turn.format
+
+ self.prompt_length = config.actor_rollout_ref.rollout.prompt_length
+ self.response_length = config.actor_rollout_ref.rollout.response_length
+
+ # Initialize interactions from config file
+ self.interaction_config_file = config.actor_rollout_ref.rollout.multi_turn.interaction_config_path
+ if self.interaction_config_file:
+ self.interaction_map: dict[str, BaseInteraction] = self._initialize_interactions(
+ self.interaction_config_file
+ )
+
+ @rollout_trace_op
+ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
+ messages = list(kwargs["raw_prompt"])
+
+ # extract images and videos from messages
+ multi_modal_data = await self.process_vision_info(messages)
+ images = multi_modal_data.get("images")
+ videos = multi_modal_data.get("videos")
+
+ metrics = {}
+ request_id = uuid4().hex
+ tools_kwargs = kwargs.get("tools_kwargs", {})
+
+ # Initialize interaction if needed
+ interaction = None
+ interaction_kwargs = {}
+ if self.interaction_config_file:
+ interaction_kwargs = kwargs["extra_info"]["interaction_kwargs"]
+ if "name" not in interaction_kwargs:
+ raise ValueError("'name' key is required in interaction_kwargs")
+ interaction_name = interaction_kwargs["name"]
+ if interaction_name not in self.interaction_map:
+ raise ValueError(
+ f"Interaction '{interaction_name}' not found in interaction_map. Available interactions: "
+ f"{list(self.interaction_map.keys())}"
+ )
+ interaction = self.interaction_map[interaction_name]
+ await interaction.start_interaction(request_id, **interaction_kwargs)
+ # Create AgentData instance to encapsulate all state
+ agent_data = AgentData(
+ messages=messages,
+ image_data=images,
+ video_data=videos,
+ metrics=metrics,
+ request_id=request_id,
+ tools_kwargs=tools_kwargs,
+ interaction=interaction,
+ interaction_kwargs=interaction_kwargs,
+ )
+
+ # State machine loop
+ state = AgentState.PENDING
+ while state != AgentState.TERMINATED:
+ if state == AgentState.PENDING:
+ state = await self._handle_pending_state(agent_data, sampling_params)
+ elif state == AgentState.GENERATING:
+ state = await self._handle_generating_state(agent_data, sampling_params)
+ elif state == AgentState.PROCESSING_TOOLS:
+ state = await self._handle_processing_tools_state(agent_data)
+ elif state == AgentState.INTERACTING:
+ state = await self._handle_interacting_state(agent_data)
+ else:
+ logger.error(f"Invalid state: {state}")
+ state = AgentState.TERMINATED
+
+ # Finalize output
+ response_ids = agent_data.prompt_ids[-len(agent_data.response_mask) :]
+ prompt_ids = agent_data.prompt_ids[: len(agent_data.prompt_ids) - len(agent_data.response_mask)]
+ multi_modal_data = {}
+ if agent_data.image_data is not None:
+ multi_modal_data["images"] = agent_data.image_data
+ if agent_data.video_data is not None:
+ multi_modal_data["videos"] = agent_data.video_data
+ output = AgentLoopOutput(
+ prompt_ids=prompt_ids,
+ response_ids=response_ids[: self.response_length],
+ response_mask=agent_data.response_mask[: self.response_length],
+ multi_modal_data=multi_modal_data,
+ response_logprobs=agent_data.response_logprobs[: self.response_length]
+ if agent_data.response_logprobs
+ else None,
+ num_turns=agent_data.user_turns + agent_data.assistant_turns + 1,
+ metrics=agent_data.metrics,
+ extra_fields={},
+ )
+ output.extra_fields.update({"turn_scores": agent_data.turn_scores, "tool_rewards": agent_data.tool_rewards})
+ return output
+
+ async def _handle_pending_state(self, agent_data: AgentData, sampling_params: dict[str, Any]) -> AgentState:
+ """Handle the pending state: prepare the prompt and start generation."""
+ prompt_ids = await self.apply_chat_template(
+ agent_data.messages,
+ tools=self.tool_schemas,
+ images=agent_data.image_data,
+ videos=agent_data.video_data,
+ )
+ agent_data.prompt_ids = prompt_ids
+ return AgentState.GENERATING
+
+ async def _handle_generating_state(
+ self, agent_data: AgentData, sampling_params: dict[str, Any], ignore_termination: bool = False
+ ) -> AgentState:
+ """Handle the generating state: generate model response and check for tool calls."""
+ add_messages: list[dict[str, Any]] = []
+
+ with simple_timer("generate_sequences", agent_data.metrics):
+ output = await self.server_manager.generate(
+ request_id=agent_data.request_id,
+ prompt_ids=agent_data.prompt_ids,
+ sampling_params=sampling_params,
+ image_data=agent_data.image_data,
+ video_data=agent_data.video_data,
+ )
+ # first time to set num_preempted
+ if agent_data.metrics.get("num_preempted") is None:
+ agent_data.metrics["num_preempted"] = output.num_preempted if output.num_preempted is not None else -1
+ # then add num_preempted to the metrics
+ else:
+ agent_data.metrics["num_preempted"] += output.num_preempted if output.num_preempted is not None else 0
+
+ agent_data.assistant_turns += 1
+ agent_data.response_ids = output.token_ids
+ agent_data.prompt_ids += agent_data.response_ids
+ agent_data.response_mask += [1] * len(agent_data.response_ids)
+ if output.log_probs:
+ agent_data.response_logprobs += output.log_probs
+
+ if output.routed_experts is not None:
+ agent_data.routed_experts = output.routed_experts
+
+ # Check termination conditions
+ if not ignore_termination and len(agent_data.response_mask) >= self.response_length:
+ return AgentState.TERMINATED
+ if self.max_assistant_turns and agent_data.assistant_turns >= self.max_assistant_turns:
+ return AgentState.TERMINATED
+ if self.max_user_turns and agent_data.user_turns >= self.max_user_turns:
+ return AgentState.TERMINATED
+
+ # Extract tool calls
+ _, agent_data.tool_calls = await self.tool_parser.extract_tool_calls(agent_data.response_ids)
+
+ # Handle interaction if needed
+ if self.interaction_config_file:
+ assistant_message = await self.loop.run_in_executor(
+ None, lambda: self.tokenizer.decode(agent_data.response_ids, skip_special_tokens=True)
+ )
+ add_messages.append({"role": "assistant", "content": assistant_message})
+ agent_data.messages.extend(add_messages)
+
+ # Determine next state
+ if agent_data.tool_calls:
+ return AgentState.PROCESSING_TOOLS
+ elif self.interaction_config_file:
+ return AgentState.INTERACTING
+ else:
+ return AgentState.TERMINATED
+
+ async def _handle_processing_tools_state(self, agent_data: AgentData) -> AgentState:
+ """Handle the processing tools state: execute tool calls and prepare tool responses."""
+ add_messages: list[dict[str, Any]] = []
+ new_images_this_turn: list[Any] = [] # Local variable instead of agent_data attribute
+
+ tasks = []
+ tool_call_names = []
+ for tool_call in agent_data.tool_calls[: self.max_parallel_calls]:
+ tasks.append(self._call_tool(tool_call, agent_data.tools_kwargs, agent_data))
+ tool_call_names.append(tool_call.name)
+
+ with simple_timer("tool_calls", agent_data.metrics):
+ responses = await asyncio.gather(*tasks)
+
+ # Process tool responses and update multi_modal_data
+ # Removed: agent_data.new_images_this_turn = []
+ for tool_response, tool_reward, _ in responses:
+ # Create message from tool response
+ if tool_response.image or tool_response.video:
+ # Multi-modal content with structured format
+ if not getattr(self.processor, "image_processor", None):
+ raise ValueError(
+ "Multimedia data can only be processed by `processor`, but the processor is None. "
+ "This error is often caused if you are using a LLM model but your tool returns multimodal "
+ "data. Plase use a vlm as the base model."
+ )
+ content = []
+ if tool_response.image:
+ content.append({"type": "image"})
+ if tool_response.video:
+ content.append({"type": "video"})
+ if tool_response.text:
+ content.append({"type": "text", "text": tool_response.text})
+ message = {"role": "tool", "content": content}
+ else:
+ # Text-only content
+ message = {"role": "tool", "content": tool_response.text or ""}
+
+ add_messages.append(message)
+
+ # Handle image data
+ if tool_response.image:
+ # Add new image data
+ if isinstance(tool_response.image, list):
+ # Ensure all elements in the list are valid image objects
+ for img in tool_response.image:
+ if img is not None: # Add a check to ensure the image is not None
+ new_images_this_turn.append(img) # Using local variable
+ else:
+ # Ensure the image is not None
+ if tool_response.image is not None:
+ new_images_this_turn.append(tool_response.image) # Using local variable
+
+ # Handle video data
+ if tool_response.video:
+ # Currently not supported, raise informative error
+ logger.warning("Multimedia type 'video' is not currently supported. Only 'image' is supported.")
+ raise NotImplementedError(
+ "Multimedia type 'video' is not currently supported. Only 'image' is supported."
+ )
+
+ if tool_reward is not None:
+ agent_data.tool_rewards.append(tool_reward)
+
+ agent_data.messages.extend(add_messages)
+
+ if self.tool_parser_name == "gpt-oss":
+ logger.info("manually format tool responses for gpt-oss")
+ tool_response_text = build_gpt_oss_tool_response_text(add_messages, tool_call_names)
+ response_ids = await self.loop.run_in_executor(
+ None, lambda: self.tokenizer.encode(tool_response_text, add_special_tokens=False)
+ )
+ else:
+ response_ids = await self.apply_chat_template(
+ add_messages,
+ images=new_images_this_turn, # Using local variable
+ videos=None,
+ remove_system_prompt=True,
+ )
+
+ if len(agent_data.response_mask) + len(response_ids) >= self.response_length:
+ return AgentState.TERMINATED
+ # Update prompt_ids and response_mask
+
+ if new_images_this_turn:
+ if agent_data.image_data is None:
+ agent_data.image_data = []
+ elif not isinstance(agent_data.image_data, list):
+ agent_data.image_data = [agent_data.image_data]
+ for img in new_images_this_turn:
+ agent_data.image_data.append(img)
+
+ agent_data.prompt_ids += response_ids
+ agent_data.response_mask += [0] * len(response_ids)
+ if agent_data.response_logprobs:
+ agent_data.response_logprobs += [0.0] * len(response_ids)
+ agent_data.user_turns += 1
+ return AgentState.GENERATING
+
+ async def _handle_interacting_state(self, agent_data: AgentData) -> AgentState:
+ """Handle the interacting state: get user input from interaction."""
+ (
+ should_terminate_sequence,
+ interaction_responses,
+ reward,
+ metrics,
+ ) = await agent_data.interaction.generate_response(
+ agent_data.request_id, agent_data.messages, **agent_data.interaction_kwargs
+ )
+ agent_data.user_turns += 1
+
+ add_messages: list[dict[str, Any]] = [{"role": "user", "content": interaction_responses}]
+ agent_data.messages.extend(add_messages)
+
+ if reward is not None:
+ agent_data.turn_scores.append(reward)
+
+ # Update prompt with user responses (similar to _handle_processing_tools_state)
+ response_ids = await self.apply_chat_template(
+ add_messages,
+ remove_system_prompt=True,
+ )
+
+ # Update prompt_ids and response_mask
+ agent_data.prompt_ids += response_ids
+ agent_data.response_mask += [0] * len(response_ids)
+ if agent_data.response_logprobs:
+ agent_data.response_logprobs += [0.0] * len(response_ids)
+
+ # double check prompt
+ # Check termination condition
+ if should_terminate_sequence:
+ return AgentState.TERMINATED
+ else:
+ return AgentState.GENERATING
+
+ async def _call_tool(
+ self, tool_call: FunctionCall, tools_kwargs: dict[str, Any], agent_data: AgentData
+ ) -> tuple[ToolResponse, float, dict]:
+ """Call tool and return tool response."""
+ tool, instance_id = None, None
+ try:
+ # TODO: append malformed tool_call to the prompt: invalid function name or arguments
+ tool_name = tool_call.name
+ tool_args = json.loads(tool_call.arguments)
+ tool = self.tools[tool_name]
+ kwargs = tools_kwargs.get(tool_name, {})
+ instance_id, _ = await tool.create(create_kwargs=kwargs.get("create_kwargs", {}))
+ tool_execution_response, tool_reward, res = await tool.execute(
+ instance_id, tool_args, agent_data=agent_data
+ )
+ except Exception as e:
+ logger.warning(f"Error when executing tool: {e}")
+ return (
+ ToolResponse(
+ text=f"Error when executing tool: {e}",
+ ),
+ 0.0,
+ {},
+ )
+ finally:
+ if tool and instance_id:
+ await tool.release(instance_id)
+
+ tool_response_text = tool_execution_response.text
+ if tool_response_text and len(tool_response_text) > self.max_tool_response_length:
+ if self.tool_response_truncate_side == "left":
+ tool_response_text = tool_response_text[: self.max_tool_response_length] + "...(truncated)"
+ elif self.tool_response_truncate_side == "right":
+ tool_response_text = "(truncated)..." + tool_response_text[-self.max_tool_response_length :]
+ else:
+ length = self.max_tool_response_length // 2
+ tool_response_text = tool_response_text[:length] + "...(truncated)..." + tool_response_text[-length:]
+
+ # Create ToolResponse from tool execution result
+ tool_response_kwargs = {"text": tool_response_text}
+
+ # Add multimedia data if present
+ for attr_name in ["image", "video"]:
+ if hasattr(tool_execution_response, attr_name):
+ attr_value = getattr(tool_execution_response, attr_name)
+ if attr_value is not None:
+ tool_response_kwargs[attr_name] = attr_value
+
+ return ToolResponse(**tool_response_kwargs), tool_reward, res
+
+ def _initialize_interactions(self, interaction_config_file):
+ """Initialize interactions from configuration.
+ Returns:
+ dict[str, BaseInteraction]: A dictionary mapping interaction names to interaction instances.
+ """
+ if interaction_config_file is None:
+ return {}
+
+ interaction_map = initialize_interactions_from_config(interaction_config_file)
+ return interaction_map
diff --git a/code/RL_model/verl/verl_train/verl/experimental/agent_loop/tool_parser.py b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/tool_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..67ad75e2bb8f1c1eaebe2e3a175e3ac55dcb10c3
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/tool_parser.py
@@ -0,0 +1,161 @@
+# Copyright 2024 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.
+import json
+import logging
+import os
+from abc import ABC, abstractmethod
+
+import regex
+from pydantic import BaseModel
+
+from verl.utils.ray_utils import get_event_loop
+from verl.utils.rollout_trace import rollout_trace_op
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class FunctionCall(BaseModel):
+ arguments: str
+ """
+ The arguments to call the function with, as generated by the model in JSON
+ format. Note that the model does not always generate valid JSON, and may
+ hallucinate parameters not defined by your function schema. Validate the
+ arguments in your code before calling your function.
+ """
+
+ name: str
+ """The name of the function to call."""
+
+
+class ToolParser(ABC):
+ _registry: dict[str, type["ToolParser"]] = {}
+
+ def __init__(self, tokenizer) -> None:
+ self.tokenizer = tokenizer
+
+ @abstractmethod
+ async def extract_tool_calls(self, responses_ids: list[int]) -> tuple[str, list[FunctionCall]]:
+ """Extract tool calls from the responses.
+
+ Args:
+ responses_ids (List[int]): The ids of the responses.
+
+ Returns:
+ Tuple[str, List[FunctionCall]]: Content and extracted tool calls.
+ """
+ raise NotImplementedError
+
+ @classmethod
+ def get_tool_parser(cls, name: str, tokenizer):
+ if name not in cls._registry:
+ raise ValueError(f"Unknown tool parser: {name}")
+ return cls._registry[name](tokenizer)
+
+ @classmethod
+ def register(cls, name: str):
+ def decorator(subclass: type[ToolParser]) -> type[ToolParser]:
+ cls._registry[name] = subclass
+ return subclass
+
+ return decorator
+
+
+@ToolParser.register("hermes")
+class HermesToolParser(ToolParser):
+ """Adapted from https://github.com/vllm-project/vllm/blob/v0.9.1/vllm/entrypoints/openai/tool_parsers/hermes_tool_parser.py"""
+
+ def __init__(self, tokenizer) -> None:
+ super().__init__(tokenizer)
+
+ self.tool_call_start_token: str = ""
+ self.tool_call_end_token: str = ""
+ self.tool_call_regex = regex.compile(r"(.*?)", regex.DOTALL)
+
+ @rollout_trace_op
+ async def extract_tool_calls(self, responses_ids: list[int]) -> tuple[str, list[FunctionCall]]:
+ loop = get_event_loop()
+ text = await loop.run_in_executor(None, self.tokenizer.decode, responses_ids)
+ if self.tool_call_start_token not in text or self.tool_call_end_token not in text:
+ return text, []
+
+ matches = self.tool_call_regex.findall(text)
+ function_calls = []
+ for match in matches:
+ try:
+ function_call = json.loads(match)
+ name, arguments = function_call["name"], function_call["arguments"]
+ function_calls.append(FunctionCall(name=name, arguments=json.dumps(arguments, ensure_ascii=False)))
+ except Exception as e:
+ logger.error(f"Failed to decode tool call: {e}")
+
+ # remaing text exclude tool call tokens
+ content = self.tool_call_regex.sub("", text)
+
+ return content, function_calls
+
+
+@ToolParser.register("gpt-oss")
+class GptOssToolParser(ToolParser):
+ """
+ Tool parser for gpt-oss model.
+ Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/function_call/gpt_oss_detector.py
+
+ Args:
+ tokenizer: The tokenizer to use.
+ """
+
+ def __init__(self, tokenizer) -> None:
+ super().__init__(tokenizer)
+ # check https://cookbook.openai.com/articles/openai-harmony for more details.
+ self.cot_pattern = regex.compile(
+ r"<\|start\|>assistant<\|channel\|>analysis<\|message\|>.*?<\|end\|>", regex.DOTALL
+ )
+ # <|start|>assistant may be pre-appended in prompts, so we need to remove it.
+ self.partial_cot_pattern = regex.compile(r"<\|channel\|>analysis<\|message\|>(.*?)<\|end\|>", regex.DOTALL)
+ self.tool_call_pattern = regex.compile(
+ r"<\|start\|>assistant<\|channel\|>[^<]* to=functions\.([^<]+) "
+ r"<\|constrain\|>json<\|message\|>(.*?)<\|call\|>",
+ regex.DOTALL,
+ )
+
+ @rollout_trace_op
+ async def extract_tool_calls(self, responses_ids: list[int]) -> tuple[str, list[FunctionCall]]:
+ loop = get_event_loop()
+ # We need to keep special tokens for gpt-oss model for better tool call extraction.
+ text = await loop.run_in_executor(None, lambda: self.tokenizer.decode(responses_ids, skip_special_tokens=False))
+ # Need to remove padding tokens for better tool call extraction.
+ text = text.replace(self.tokenizer.pad_token, "")
+ # Need to reomve COT since COT may contain tool call tokens.But they are not valid tool calls.
+ text = regex.sub(self.cot_pattern, "", text)
+ text = regex.sub(self.partial_cot_pattern, "", text)
+
+ # check if there are tool calls in the text by re.findall
+ matches = regex.findall(self.tool_call_pattern, text)
+ if not matches:
+ return text, []
+
+ function_calls = []
+ for match in matches:
+ try:
+ name, arguments = match[0], match[1]
+ # don't check if arguments is valid JSON and leave it to client
+ function_calls.append(FunctionCall(name=name, arguments=arguments))
+ except Exception as e:
+ logger.error(f"Failed to decode tool call: {e}")
+
+ # remaing text exclude tool call tokens
+ content = regex.sub(self.tool_call_pattern, "", text)
+
+ return content, function_calls
diff --git a/code/RL_model/verl/verl_train/verl/experimental/agent_loop/utils.py b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..68cb57d870f8fdb65aa892ba75598e4c66020239
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/agent_loop/utils.py
@@ -0,0 +1,108 @@
+# Copyright 2024 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.
+
+import os
+from typing import Any
+
+
+def resolve_config_path(config_path: str) -> str:
+ """Resolve agent loop configuration file path.
+
+ In multi-node Ray training, relative paths may not resolve correctly
+ because the working directory on remote nodes can differ from the driver node.
+ This function resolves relative paths by checking multiple locations in order:
+ 1. If already absolute, return as-is
+ 2. Try current working directory
+ 3. Try relative to verl package installation (project root)
+
+ Args:
+ config_path: Configuration file path (relative or absolute)
+
+ Returns:
+ Absolute path to the configuration file
+
+ Raises:
+ FileNotFoundError: If the configuration file cannot be found
+ """
+ # Return absolute paths unchanged
+ if os.path.isabs(config_path):
+ return config_path
+
+ # Try current working directory first
+ cwd = os.path.abspath(os.getcwd())
+ cwd_path = os.path.abspath(os.path.join(cwd, config_path))
+ if (cwd_path == cwd or cwd_path.startswith(cwd + os.sep)) and os.path.exists(cwd_path):
+ return cwd_path
+
+ # Try relative to verl project root (where verl package is installed)
+ try:
+ import verl
+
+ verl_package_dir = os.path.abspath(os.path.dirname(verl.__file__))
+
+ # Strategy 1: For development/editable installs.
+ project_root = os.path.dirname(verl_package_dir)
+ dev_path = os.path.abspath(os.path.join(project_root, config_path))
+ if (dev_path == project_root or dev_path.startswith(project_root + os.sep)) and os.path.exists(dev_path):
+ return dev_path
+
+ # Strategy 2: For standard package installations.
+ install_path = os.path.abspath(os.path.join(verl_package_dir, config_path))
+ if (install_path == verl_package_dir or install_path.startswith(verl_package_dir + os.sep)) and os.path.exists(
+ install_path
+ ):
+ return install_path
+ except (ImportError, AttributeError):
+ pass # verl not installed or __file__ not available
+
+ # File not found - raise clear error
+ raise FileNotFoundError(
+ f"Agent loop configuration file not found: {config_path}. Tried current directory and verl project root."
+ )
+
+
+# tokenizer.apply_chat_template is not working properly for gpt-oss model.
+# Because the chat template requires tool call messages to parse tool response messages
+# so we need to format the tool response manually.
+def format_gpt_oss_tool_response_manually(tool_response: str, tool_call_name: str) -> str:
+ """Format tool response for gpt-oss model.
+ Args:
+ tool_response: Tool response string
+ tool_call_name: Name of the tool that was called
+
+ Returns:
+ Formatted tool response string
+ """
+ return f"<|start|>functions.{tool_call_name} to=assistant<|channel|>commentary<|message|>{tool_response}<|end|>"
+
+
+def add_generation_prompt_for_gpt_oss(message_content: str) -> str:
+ """Add generation prompt for gpt-oss model.
+ Args:
+ message_content: Message content string
+
+ Returns:
+ Message content string with generation prompt
+ """
+ return message_content + "<|start|>assistant"
+
+
+def build_gpt_oss_tool_response_text(messages: list[dict[str, Any]], tool_call_names: list[str]) -> str:
+ """Build gpt-oss tool response text (manual formatting + generation prompt)."""
+ tool_response_texts: list[str] = []
+ for i, tool_msg in enumerate(messages):
+ actual_tool_name = tool_call_names[i]
+ formatted = format_gpt_oss_tool_response_manually(tool_msg["content"], actual_tool_name)
+ tool_response_texts.append(formatted)
+ return add_generation_prompt_for_gpt_oss("".join(tool_response_texts))
diff --git a/code/RL_model/verl/verl_train/verl/experimental/dataset/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/dataset/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ce90c5eb352d85c59105c0dc85b5f1dd576f095
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/dataset/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2024 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.
diff --git a/code/RL_model/verl/verl_train/verl/experimental/dataset/sampler.py b/code/RL_model/verl/verl_train/verl/experimental/dataset/sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..b7b15b422c823280c862397dd88c362aac213554
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/dataset/sampler.py
@@ -0,0 +1,40 @@
+# Copyright 2025 Amazon.com Inc 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 abc import abstractmethod
+from collections.abc import Sized
+
+from omegaconf import DictConfig
+from torch.utils.data import Sampler
+
+from verl import DataProto
+
+
+class AbstractSampler(Sampler[int]):
+ """Abstract interface for custom samplers."""
+
+ @abstractmethod
+ def __init__(
+ self,
+ data_source: Sized,
+ data_config: DictConfig,
+ ):
+ pass
+
+
+class AbstractCurriculumSampler(AbstractSampler):
+ """Experimental interface for curriculum learning samplers."""
+
+ @abstractmethod
+ def update(self, batch: DataProto) -> None:
+ pass
diff --git a/code/RL_model/verl/verl_train/verl/experimental/dynamic_dataset/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/dynamic_dataset/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ce90c5eb352d85c59105c0dc85b5f1dd576f095
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/dynamic_dataset/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2024 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.
diff --git a/code/RL_model/verl/verl_train/verl/experimental/dynamic_dataset/dynamicgen_dataset.py b/code/RL_model/verl/verl_train/verl/experimental/dynamic_dataset/dynamicgen_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e60b7836e02a7b2e3397eeb4b09b380ed3e03f5
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/dynamic_dataset/dynamicgen_dataset.py
@@ -0,0 +1,112 @@
+# Copyright 2025 Amazon.com Inc 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.
+"""
+Dataset class that enables dynamic data generation strategies between iterations of training.
+This class extends RLHFDataset and uses an AbstractDataGen instance to generate data.
+
+This is especially useful in settings where proposer model generates new tasks based
+on rollout data.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import datasets
+from omegaconf import DictConfig
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer, ProcessorMixin
+
+from verl import DataProto
+from verl.utils.dataset import RLHFDataset
+from verl.utils.import_utils import load_extern_object
+
+logger = logging.getLogger(__name__)
+
+
+class AbstractDataGenerator(ABC):
+ def __init__(self, config: DictConfig):
+ self.config = config
+
+ @abstractmethod
+ def generate(self, dataset: Dataset) -> datasets.Dataset:
+ """
+ Generate method must be implemented by subclasses.
+ Args:
+ dataset: The dataset to generate from.
+ Returns:
+ Processed data or result as implemented by the subclass.
+ """
+ pass
+
+
+class MockDataGenerator(AbstractDataGenerator):
+ """
+ A noop data gen class that only reappends the first datapoint.
+ This class is useful as a placeholder and testing.
+ """
+
+ def __init__(self, config: DictConfig = None):
+ super().__init__(config)
+
+ def generate(self, dataset: Dataset) -> datasets.Dataset:
+ print("MockDataGenerator: No operation performed on the dataset.")
+ return dataset.dataframe.select([0])
+
+
+class DynamicGenDataset(RLHFDataset):
+ """
+ A dataset class that uses a data generation strategy to process data.
+ This class extends RLHFDataset and uses an AbstractDataGen instance to generate data.
+ """
+
+ def __init__(
+ self,
+ data_files: str | list[str],
+ tokenizer: PreTrainedTokenizer,
+ config: DictConfig,
+ processor: Optional[ProcessorMixin] = None,
+ ):
+ super().__init__(data_files, tokenizer, config, processor)
+ self.datagen: AbstractDataGenerator = config.datagen
+ assert "datagen" in config and config.datagen.get("path", None) is not None, (
+ f"datagen path is not set in config: {config}"
+ )
+ # Dynamically load the custom datagen class
+ datagen_cls = load_extern_object(config.datagen.path, config.datagen.name)
+
+ # Verify that the custom datagen class inherits from AbstractDataGenerator
+ abs_cls = AbstractDataGenerator
+ if not issubclass(datagen_cls, abs_cls):
+ raise TypeError(
+ f"The custom datagen class '{config.datagen.name}' from '{config.datagen.path}'"
+ + " must inherit from {abs_cls}"
+ )
+
+ self.data_generator = datagen_cls(config.datagen)
+ self.on_batch_end()
+
+ def append_dataframe(self, new_dataframe: datasets.Dataset):
+ new_dataframe = self.maybe_filter_out_long_prompts(new_dataframe)
+ self.dataframe = datasets.concatenate_datasets([self.dataframe, new_dataframe])
+
+ logger.info(f"new dataset len: {len(self.dataframe)}")
+
+ def on_batch_end(self, batch: DataProto) -> None:
+ """
+ Generate data using the provided data generation strategy.
+ Note: This method is intended to change the dataset after each training batch.
+ """
+ new_data = self.data_generator.generate(self)
+ self.append_dataframe(new_data)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/README.md b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..b7406514f3b0b694ccbd5fb2a9d6789ed12fea15
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/README.md
@@ -0,0 +1,599 @@
+# Recipe: Fully Async Policy Trainer
+
+**Author:** `https://github.com/meituan-search`
+
+Last updated: 12/25/2025.
+
+This document introduces a fully asynchronous PPO training system that completely decouples the Trainer and Rollouter,
+supporting asynchronous sample generation and training.
+Under this system, we achieved a 2.35x-2.67x performance improvement when training the Qwen2.5-7B model with 128 GPUs,
+without significantly affecting the results.
+
+## Introduction
+
+### Background
+
+The separated rollout and train architecture, compared to the colocate architecture, can allocate resources more
+flexibly and design more flexible training logic, thereby addressing issues such as low GPU utilization and training
+efficiency caused by long-tail problems.
+The one_step_off_policy alleviates the problem of long rollout times and achieves some gains in training efficiency by
+designing a separated architecture and performing asynchronous training between rollout and train for one round.
+However, it forcibly uses data from one round of asynchronous training, which is not flexible enough and cannot
+completely eliminate the impact of long-tail on training efficiency.
+In other frameworks such as AReaL, Magistral, StreamRL, and AsyncFlow, asynchronous training and streaming training have
+been implemented based on the separated architecture and have achieved gains.
+We borrow from their methods and implemented them in VERL. The fully_async_policy supports asynchronous, streaming, and
+partial
+rollout training.
+By reasonably setting parameters such as resource allocation and parameter synchronization frequency, fully_async_policy
+can significantly improve training efficiency.
+
+> Magistral https://arxiv.org/abs/2506.10910
+>
+> AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language
+> Reasoning https://arxiv.org/abs/2505.24298
+>
+> StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream
+> Generation https://arxiv.org/abs/2504.15930
+>
+> AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training https://arxiv.org/abs/2507.01663
+>
+
+### Core Contributions
+
+* **Resource Isolation**: Unlike using hybrid_engine, Rollouter and Trainer use separate computing resources and need to
+ specify the resources they occupy separately.
+* **Parallel Generation and Training**: While the Trainer is training, the Rollouter is generating new samples.
+* **Multi-step Asynchronous**: Compared to one step off policy, it supports asynchronous settings from 0.x steps to
+ multiple steps, making the asynchronous solution more flexible.
+* **NCCL Parameter Synchronization**: Based on the nccl communication primitive, refer to [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine) to
+ achieve efficient parameter synchronization between Rollouter and Trainer.
+* **Stream Inference and Training**: Rollouter generates data sample by sample, and data transmission uses a single
+ sample as the minimum transmission unit.
+* **Asynchronous Training and Freshness Control**: By setting the parameter async_training.staleness_threshold, it
+ supports training with samples generated by old parameters.
+* **PartialRollout**: The Rollouter's inference process supports partial rollout logic. During parameter
+ synchronization, by adding `sleep() and resume()` logic, it
+ saves samples from ongoing rollouts and continues using them in the next rollout, reducing the time spent waiting for
+ ongoing tasks to finish during parameter synchronization.
+
+Currently, the supported usage mode is megatron/fsdp+vllm. vllm must use the server mode based on AgentLoop.
+
+## Design
+
+The overall architecture of fully_async_policy is shown in the figure below. fully_async_policy mainly consists of four
+parts: Rollouter, MessageQueue, Trainer, and ParameterSynchronizer.
+
+
+
+1. Rollouter generates sequences sample by sample and puts the generated samples into the MessageQueue, with the
+ production speed controlled by freshness.
+2. MessageQueue is used to temporarily store samples generated by Rollouter.
+3. Trainer fetches samples from MessageQueue sample by sample. After fetching `require_batches*ppo_mini_batch_size`
+ samples, it will perform training. After training for async_training.trigger_parameter_sync_step rounds, it triggers
+ a parameter synchronization with Rollouter.
+4. ParameterSynchronizer implements the NCCL synchronous parameter synchronization capability.
+
+The source of benefits compared to the base scheme lies in the fact that in the colocate case, using more resources for
+rollout cannot solve the idleness caused by long-tail samples.
+After we perform resource isolation, the time for rollout and train may be longer than before (because fewer resources
+are used),
+but the overlap in their time consumption reduces the end-to-end time consumption.
+
+
+
+## Usage
+
+### Parameter Description
+
+| super params | implication |
+|-----------------------------------------------|------------------------------------------------------------------------------------------------|
+| `trainer.nnodes` | Number of nodes for Trainer |
+| `trainer.n_gpus_per_node` | Number of GPUs per node for Trainer |
+| `rollout.nnodes` | Number of nodes for Rollouter |
+| `rollout.n_gpus_per_node` | Number of GPUs per node for Rollouter |
+| `data.train_batch_size` | In the fully async strategy, this value is not effective (default is 0) |
+| `data.gen_batch_size` | In the fully async strategy, uses streaming sample production logic (default is 1) |
+| `rollout.total_rollout_steps` | Total number of rollout samples |
+| `rollout.test_freq` | How many times Rollouter updates parameters before performing a validation |
+| `actor_rollout_ref.actor.ppo_mini_batch_size` | The ppo_mini_batch_size is a global num across all workers/gpus |
+| `async_training.require_batches` | Number of ppo_mini_batch_size that FullyAsyncTrainer fetches at once |
+| `async_training.trigger_parameter_sync_step` | Indicates how many local updates FullyAsyncTrainer performs before a parameter synchronization |
+| `async_training.staleness_threshold` | Freshness control |
+| `async_training.partial_rollout` | Whether to perform partial_rollout |
+| `async_training.use_rollout_log_probs` | Use log_probs generated by rollout |
+| `async_training.compute_prox_log_prob` | Whether to compute log_prob using the training model's parameters during the training phase. | |
+| `async_training.checkpoint_engine.enable`| Whether to use checkpoint_engine for accelerating, default `True`|
+| `async_training.checkpoint_engine.overlap_broadcast_and_consume` | When use checkpoint_engine, whether to overlap broadcast and load_weights, default `False`|
+| `async_training.checkpoint_engine.device_buffer_size_M` | When use checkpoint_engine, the user-specific bucket size (MB), default `4096`|
+| `async_training.use_trainer_do_validate` | Whether use trainer node to do validate process, default `False`|
+
+**Further Explanation:**
+
+* `rollout.total_rollout_steps`
+
+ Compared to colocate, the quantity can be aligned by multiplying train_batch_size and step:
+ `rollout.total_rollout_steps = data.train_batch_size * step`.
+
+* `async_training.trigger_parameter_sync_step`
+
+ In the fully async strategy, it indicates how many local updates the Trainer performs (i.e., how many times it fetches
+ `require_batches * ppo_mini_batch_size` samples) before a parameter synchronization with Rollouter.
+ Between every two parameter synchronizations between Rollouter and Trainer, the Trainer will process
+ `trigger_parameter_sync_step* require_batches*ppo_mini_batch_size` samples.
+ To fairly compare speed with colocate, `trigger_parameter_sync_step` should be set to
+ `data.train_batch_size / (require_batches * ppo_mini_batch_size)`.
+
+* `async_training.staleness_threshold`
+
+ In the fully async strategy, it indicates the maximum proportion of stale samples allowed to be used.
+
+ * `staleness_threshold`=0, indicates synchronous training.
+ Rollouter will generate a fixed number of samples between two parameter updates, the sample count is:
+
+ `rollout_num = (trigger_parameter_sync_step*require_batches*ppo_mini_batch_size)`
+ * `staleness_threshold`>0, indicates asynchronous training, can be set to a decimal for more flexible asynchronous
+ calls.
+ Rollouter will generate at most the following number of samples between two parameter updates:
+
+ `rollout_num = (1+staleness_threshold)*(trigger_parameter_sync_step*require_batches*ppo_mini_batch_size) - num_staleness_sample`
+
+ `num_staleness_sample` represents the number of stale samples generated in excess during the last rollout.
+
+ Since it's a streaming system, rollout continues to generate and trainer continues to consume. If rollouter is slower,
+ trainer will trigger parameter synchronization earlier, and rollouter will not actually produce rollout_num samples.
+ When rollout is fast enough, setting `staleness_threshold` to 1 is basically equivalent to one_step_off policy.
+ To avoid too many expired samples affecting training accuracy, it is recommended to set this value to less than 1.
+
+* `async_training.partial_rollout`
+
+ partial_rollout only actually takes effect when staleness_threshold>0.
+
+* `async_training.use_rollout_log_probs`
+
+ In reinforcement learning algorithms, log_probs have implicit correlations with parameter versions and tokens. Due to
+ the settings of algorithms like PPO/GRPO/DAPO, when calculating importance sampling,
+ old_log_prob must use the log_probs corresponding to the rollout parameters and tokens to ensure algorithm
+ correctness. In the fully
+ async strategy, we default to old_log_prob being calculated by rollout rather than by trainer.
+
+* `async_training.require_batches`
+
+ In streaming training, require_batches should be set to 1, indicating that training is performed after producing
+ enough ppo_mini_batch_size samples.
+ In actual testing, we found that if fewer samples are issued at once, due to the order of data distribution, it can
+ cause training instability and longer response lengths.
+ Here, we additionally provide require_batches for streaming distribution and control the number of samples
+ participating in training at once.
+
+* `async_training.compute_prox_log_prob` (experimental)
+
+ During the training process, we observed that metrics and response lengths may become unstable in the later
+ stages of training. To mitigate this issue, we can use
+ the [Rollout Importance Sampling](https://verl.readthedocs.io/en/latest/advance/rollout_is.html)
+ technique for importance sampling. To utilize Rollout Importance Sampling, we need to compute log_prob using
+ the training engine, which requires enabling this switch.
+ Additionally, when compute_prox_log_prob and Rollout Importance Sampling are enabled under mode d
+ (async stream pipeline with partial rollout), our implementation approximates `Areal's Decoupled PPO`.
+
+* `async_training.checkpoint_engine.enable`
+
+ Enabling the checkpoint engine generally reduces synchronization time overhead by more than 60% compared to
+ the original per-tensor parameter synchronization method. However, assembling buckets incurs additional
+ temporary GPU memory overhead.
+
+* `async_training.checkpoint_engine.overlap_broadcast_and_consume`
+
+ Enabling pipeline between the broadcast and load_weights parameters will allocate additional GPU memory.
+ Since the main time consumption for parameter synchronization is not in the broadcast and load_weights phases,
+ but in the parameter generation phase (by megatron or FSDP), this option is off by default.
+
+* `async_training.checkpoint_engine.device_buffer_size_M`
+
+ It controls the size of the memory buffer used for synchronization when the checkpoint-engine is enabled.
+ The actual `bucket_size` = `max(device_buffer_size_M, maximum parameter tensor size)`.
+ * When enable `overlap_broadcast_and_consume`, the additional device memory overhead of
+ trainer rank is `3 * bucket_size`and rollout rank is `2 * bucket_size`。
+ * When disable `overlap_broadcast_and_consume`, the additional device memory overhead of
+ trainer rank is `2 * bucket_size`and rollout rank is `1 * bucket_size`。
+
+* `async_training.use_trainer_do_validate`
+
+ It controls whether to use the trainer's `do_validate` method for validation.
+ If set to True, the trainer will perform validation after each parameter update. It can reduce the validation time
+ overhead and trainer node idle time.
+ If set to False, the trainer will not perform validation.
+
+### Supported Modes
+
+1. on policy pipeline:
+ 1. **trigger_parameter_sync_step=1, staleness_threshold=0**
+ 2. Rollouter produces `require_batches*ppo_mini_batch_size` samples at once, Trainer fetches these samples for
+ training, and after training completes, Trainer and Rollouter perform a parameter synchronization;
+ 3. During the rollout phase, if there are long-tail samples but few rollout samples, shorter samples cannot fill
+ idle resources, causing some resource waste.
+ 4. As shown in figure a;
+
+2. stream off policy pipeline:
+ 1. **trigger_parameter_sync_step>1, staleness_threshold=0**
+ 2. Synchronous streaming training will be performed. Rollouter produces
+ `require_batches*ppo_mini_batch_size*trigger_parameter_sync_step` samples at once, Trainer performs a local
+ training every time it fetches `require_batches*ppo_mini_batch_size` samples, and after training
+ trigger_parameter_sync_step times, Trainer and Rollouter perform a parameter synchronization;
+ 3. Compared to a, since more samples are generated at once, resource idleness will be lower.
+ 4. In one step training, there will be two periods of resource idleness: when fetching the first batch of samples,
+ train waits for `require_batches*ppo_mini_batch_size` samples to be produced, and during the last parameter
+ update, rollout waits for training to complete.
+ 5. As shown in figure b;
+
+3. async stream pipeline with stale samples:
+ 1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=False**
+ 2. After each parameter update, Rollouter will plan to produce at most rollout_num samples (in practice, the number
+ of samples generated may be less than this value depending on rollout speed).
+ 3. If the rollout process is relatively fast, Rollouter will generate some additional samples num_stale_samples
+ before parameter synchronization for immediate use by Trainer after synchronization.
+ When triggering parameter synchronization, if Rollouter has ongoing tasks, it will wait for the tasks to complete
+ and not add new tasks;
+ 4. Compared to b, except for the first step training, subsequent training will not have the time to wait for the
+ first batch rollout to finish, but will have the time to wait for active tasks to finish.
+ 5. As shown in figure c;
+
+4. async stream pipeline with partial rollout:
+ 1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=True**
+ 2. Compared to c, when triggering parameter synchronization, if Rollouter has samples being produced, it will
+ interrupt the rollout process and perform parameter synchronization. The interrupted samples will continue to be
+ generated after synchronization. This reduces the time to wait for active tasks to finish.
+ 3. As shown in figure d;
+
+
+
+### Key Metrics
+
+| metrics | implication |
+|------------------------------------------------|--------------------------------------------------------------------------------------------------------|
+| `trainer/idle_ratio` | Trainer idle rate |
+| `rollouter/idle_ratio` | Rollouter idle rate |
+| `fully_async/count/stale_samples_processed` | Total number of old samples used in training |
+| `fully_async/count/stale_trajectory_processed` | Total number of old trajectories used in training (one sample produces rollout.n trajectories) |
+| `fully_async/partial/total_partial_num` | Number of partial samples processed by Trainer between two trigger_parameter_sync_step |
+| `fully_async/partial/partial_ratio` | Ratio of partial samples processed by Trainer between two trigger_parameter_sync_step |
+| `fully_async/partial/max_partial_span` | Maximum parameter span of partial samples processed by Trainer between two trigger_parameter_sync_step |
+
+### Parameter Tuning Recommendations
+
+* Resource Allocation and Adjustment:
+ * Reasonable resource allocation is the prerequisite for achieving good training efficiency. The ideal resource
+ allocation should make the rollout time and train time close, thereby minimizing pipeline bubbles in the entire
+ training process,
+ avoiding resource idleness, and ensuring Trainer does not use old samples. In real training scenarios, resource
+ allocation can be adjusted based on the idle time of rollout and train during actual training,
+ which can be obtained from rollouter/idle_ratio and trainer/idle_ratio. If rollouter/idle_ratio is high and
+ trainer/idle_ratio is low,
+ Trainer resources should be increased and Rollouter resources should be reduced, and vice versa.
+
+* Key Parameters:
+ * staleness_threshold: Setting it too high will cause more old samples to be used, affecting model performance. It
+ is recommended to set it to less than 1.
+ * require_batches: The closer to 1, the closer to a pure streaming process, the smaller the training bubbles, and
+ the faster the acceleration effect that can be achieved in terms of speed, but it will affect the order of sample
+ processing;
+ * trigger_parameter_sync_step: The smaller the setting, the closer to on policy, but it will cause frequent
+ parameter synchronization. Long-tail samples waste resources that cannot be filled by short samples, resulting in
+ low resource utilization.
+ The larger the setting, the higher the computational efficiency, but the accuracy will be affected by off policy.
+ * rollout.test_freq: It will occupy Rollouter resources and is not recommended to be set too small.
+
+* Mode Selection: By adjusting different parameters, the Fully Async architecture supports optimization acceleration at
+ different levels, suitable for tasks in different scenarios.
+ * For small-scale tasks that need to ensure training stability and on-policy nature, and have low speed
+ requirements, the on policy pipeline mode (Mode 1) can be tried.
+ * For scenarios that need to improve training throughput but are sensitive to staleness, the stream off policy
+ pipeline mode can be tried. That is, by
+ setting trigger_parameter_sync_step>1 to improve training efficiency, but still maintaining the synchronization
+ mechanism (staleness_threshold=0) (Mode 2).
+ * For large-scale tasks with high training speed requirements and can tolerate a certain degree of off-policy and
+ staleness, setting staleness_threshold>
+ 0 and partial_rollout=True can improve training efficiency, using the async stream pipeline mode (Mode 3 or 4).
+
+### Quick Start
+
+```shell
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*400)))
+test_freq=10
+staleness_threshold=0
+trigger_parameter_sync_step=16
+partial_rollout=False
+
+
+python -m recipe.fully_async_policy.fully_async_main \
+ train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.partial_rollout="${partial_rollout}"
+```
+
+## Experiments
+
+### Asynchronous Training on 7B Model
+
+We used Qwen2.5-Math-7B to verify the benefits of the fully async strategy under long candidates and multiple resources.
+Using the `async stream pipeline with stale samples` strategy, we achieved about 2x performance improvement on 32 cards,
+64 cards, and 128 cards without significantly affecting experimental results.
+
+* Machine: H20
+* Model: Qwen2.5-Math-7B
+* Rollout length: max_response_length FSDP2: 28K tokens;
+* Algorithm: DAPO
+* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet
+* Engine: vllm+FSDP2
+* rollout.n: 16
+* ppo_mini_batch_size: 32
+* test_freq: 20
+
+* colocate sync:
+ * step: 400
+ * train_batch_size: 512
+
+* fully_async_policy
+ * total_rollout_steps: 512*400
+ * require_batches: 4
+ * trigger_parameter_sync_step: 4
+ * staleness_threshold: 0.5
+ * partial_rollout: True
+
+| training mode | resource allocation | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 |
+|:--------------------:|:---------------------:|:--------:|:--------:|:--------------:|:---------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-------------------------------:|
+| colocate sync | 32 | 790.10 | 357.41 | 107.71 | 269.80 | 13h 44m | 1d 3h 43m | 2d 9h 22m | 3d 17h 5m | max: 0.3313
last: 0.2448 |
+| fully_async_policy | 16:16 | 294.77 | 21.26 | \ | 313.81 | 7h 58m
(1.72x) | 16h 21m
(1.70x) | 1d 0h 53m
(2.31x) | 1d 9h 26m
(2.66x) | max: 0.3302
last: 0.2333 |
+| colocate sync | 64 | 365.28 | 150.72 | 70.26 | 133.41 | 10h 22m | 20h 45m | 1d 7h 6m | 1d 17h 32m | max: 0.3365
last: 0.2333 |
+| fully_async_policy | 32:32 | 189.26 | 28.46 | \ | 156.98 | 4h 57m
(2.09x) | 10h 14m
(2.03x) | 16h 58m
(1.83x) | 21h 40m
(1.92x) | max: 0.3677
last: 0.3406 |
+| colocate sync | 128 | 356.30 | 177.85 | 53.92 | 113.81 | 8h 36m | 17h 56m | 1d 5h 6m | 1d 16h 48m | max: 0.3573
last: 0.2958 |
+| fully_async_policy | 64:64 | 150.63 | 33.14 | \ | 113.16 | 3h 13m
(2.67x) | 6h 46m
(2.65x) | 10h 53m
(2.67x) | 17h 22m
(2.35x) | max: 0.3521
last: 0.3094 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-colocate_async?nw=nwuserhouzg
+
+### 128-card 7B Asynchronous Mode Experiment
+
+We used Qwen2.5-Math-7B to verify the effects of various modes supported by fully async.
+We can see that the benefit brought by streaming is approximately 1.6x, and after combining staleness and
+partial_rollout, the benefit reaches 2.35x.
+
+| mode | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 |
+|:-------------------------------------------------------------------------------------------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------------:|
+| colocate sync | 356.30 | 177.85 | 53.92 | 113.81 | 8h 36m | 17h 56m | 1d 5h 6m | 1d 16h 48m | max: 0.3573
last: 0.2958 |
+| `stream off policy pipeline`
(+fully async: trigger_parameter_sync_step= 4,
require_batches= 4) | 231.34 | 128.47 | \ | 98.77 | 4h 25m | 9h 41m | 15h 2m | 1d 1h 53m | max: 0.2844
last: 0.2604 |
+| `async stream pipeline with stale samples`
(+staleness_threshold=0.5) | | | | | | | | | |
+| `async stream pipeline with partial rollout`
(+partial_rollout=True) | 150.63 | 33.14 | \ | 113.16 | 3h 13m | 6h 46m | 10h 53m | 17h 22m | max: 0.3521
last: 0.3094 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg
+
+### 128-card Stale Ablation Experiment
+
+Under the `async stream pipeline with partial rollout` mode, we verified the impact of staleness settings on training
+efficiency.
+We found that the larger the staleness, the more obvious the final gains.
+We also noticed that the times for staleness values of 0.3 and 0.5 are quite close, because as the training steps
+increase, the response length changes significantly, causing training instability.
+Further analysis and optimization are needed for this issue.
+
+| staleness_threshold | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 |
+|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|
+| 0 | 231.34 | 128.47 | \ | 98.77 | 4h 25m | 9h 41m | 15h 2m | 1d 1h 53m | max: 0.2844
last: 0.2604 |
+| 0.1 | 171.30 | 58.17 | \ | 109.12 | 3h 53m | 8h 37m | 14h 25m | 19h 59m | max: 0.3542
last: 0.2979 |
+| 0.3 | 146.11 | 38.88 | \ | 103.22 | 3h 18m | 6h 49m | 11h 40m | 17h 20m | max: 0.3469
last: 0.2865 |
+| 0.5 | 150.63 | 33.14 | \ | 113.16 | 3h 13m | 6h 46m | 10h 53m | 17h 22m | max: 0.3521
last: 0.3094 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg
+
+### 128-card 7B require_batches Ablation Experiment
+
+In multiple tests, we found that the number of samples issued each time in streaming affects the response length during
+training, which in turn affects training time. We verified the impact on results by modifying
+`async_training.require_batches`.
+
+| require_batches | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | acc/mean@1 |
+|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|
+| 1 | 203.47 | 30.88 | \ | 181.08 | 3h 31m | 8h 29m | 17h 36m | max: 0.349
last: 0.326 |
+| 2 | 158.72 | 26.32 | \ | 128.08 | 3h 35m | 7h 38m | 13h 57m | max: 0.351
last: 0.3406 |
+| 4 | 124.64 | 25.62 | \ | 95.06 | 3h 13m | 6h 46m | 10h 53m | max: 0.3521
last: 0.3521 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_require_batches?nw=nwuserhouzg
+
+### 30B Model Mode Experiment
+
+We achieved a 1.7x performance improvement with `async stream pipeline with staleness samples` strategy on the
+Qwen3-30B-A3B-Base model compared to the colocate setup. It is worth noting that this is far from the upper limit of
+performance gains achievable through asynchrony. Firstly, the comparative experiments used a maximum response length of
+only 8k, which is much shorter than the 20k sequence length in previous experiments, resulting in a less pronounced
+rollout tail effect. Secondly, we adopted a highly skewed resource allocation, with rollout using 96 GPUs and trainer
+using 32 GPUs, which is not an optimal configuration. During the experiments, we observed that the current verl
+implementation imposes certain constraints, such as requiring data to be evenly divisible by the number of GPUs, making
+resource adjustment less flexible. Additionally, as asynchronous training and deployment accelerate, the performance gap
+is gradually narrowing. Therefore, enabling more flexible resource allocation and dynamic resource adjustment in the
+future will be our next focus.
+
+* Machine: H20
+* Model: Qwen3-30B-A3B-Base
+* Rollout length: max_response_length : 8K tokens;
+* Algorithm: GRPO
+* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet
+* Engine: vllm+Megatron
+* rollout.n: 16
+* ppo_mini_batch_size: 128
+* test_freq: 20
+
+* colocate sync:
+ * step:400
+ * train_batch_size: 512
+
+* fully_async_policy
+ * total_rollout_steps: 512*400
+ * trigger_parameter_sync_step: 512/128 = 4
+ * staleness_threshold: 0.5
+ * partial_rollout: True
+
+| Training Mode | Resource Allocation | Step | Gen | Old Log Prob | Ref | Update Actor | Total Time 100 Step | Total Time 200 Step | Total Time 300 Step | Total Time 400 Step | Acc/Mean@1 |
+|--------------------|---------------------|--------|--------|--------------|-------|--------------|---------------------|---------------------|---------------------|---------------------|-----------------------------|
+| Colocate Sync | 128 | 497.89 | 348.05 | 28.73 | 20.86 | 86.27 | 13h 36m | 1d 3h 48m | 1d 19h 4m | 2d 11h 39m | max: 0.3500
last: 0.3208 |
+| Fully Async Policy | 96:32 | 282.75 | 22.06 | \ | 50.05 | 206.63 | 6h 45m (2.01x) | 14h 48m (1.88x) | 1d 0h 9m (1.78x) | 1d 10h 41m (1.72x) | max: 0.3813
last: 0.3448 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-30B?nw=nwuserhouzg | | |
+
+### checkpoint-engine Ablation Experiment
+We tested the single-step parameter synchronization time of the checkpoint-engine on three models: Qwen2.5-Math-7B, Qwen3-30B-A3B, and Qwen3-235B-A22B, using default checkpoint-engine configurations. All experiments were performed on H20 machines, and the Megatron engine was used for training.
+| model | trainer rank | rollout rank | checkpoint-engine | total sync time |
+|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|
+| Qwen2.5-Math-7B | 4 | 4 | False | 0.12s |
+| Qwen2.5-Math-7B | 4 | 4 | True | 0.02s |
+| Qwen3-30B-A3B | 16 | 16 | False | 15.76s |
+| Qwen3-30B-A3B | 16 | 16 | True | 4.38s |
+| Qwen3-235B-A22B | 64 | 64 | False | 58.57s |
+| Qwen3-235B-A22B | 64 | 64 | True | 23.70s |
+
+
+### use_trainer_do_validate Experiment
+We tested the effect of setting `use_trainer_do_validate=True` on the training process. The results show that setting
+this parameter to True can reduce the validation time overhead and trainer node idle time.
+We used Qwen2.5-Math-7B to verify the benefits of `use_trainer_do_validate=True` on the training process, we achieved about 2x performance improvement on validation time, and the trainer node idle time is reduced by about 40%.
+
+* Machine: H20
+* Model: Qwen2.5-Math-7B
+* Rollout length: max_response_length FSDP2: 10K tokens;
+* Algorithm: DAPO
+* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet
+* Engine: vllm+FSDP2
+* rollout.n: 16
+* ppo_mini_batch_size: 32
+* test_freq: 10
+
+* fully_async_policy
+ * total_rollout_steps: 512*400
+ * require_batches: 4
+ * trigger_parameter_sync_step: 4
+ * staleness_threshold: 0.5
+ * partial_rollout: True
+
+| training mode | resource allocation | step | gen | old_log_prob | update_actor | validate time | total time
50 step | acc/mean@2 |
+|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|
+| colocate sync | 16 | 484.623 | 52.939 | 0 | 430.263 | 205.080 | 7h9m | 22.6 |
+| fully_async_policy | 8:8 | 489.953 | 52.622 | 0 | 435.874 | 95.699 | 7h2m | 21.0 |
+
+
+## Multi-Turn Tool Calling
+
+Referencing **recipe/retool** and **ToolAgentLoop**, we implemented **AsyncPartialToolAgentLoop**, a multi-turn
+tool-calling loop that supports partial_rollout for **fully_async_policy**.
+
+### Core Design
+
+`AsyncPartialToolAgentLoop` inherits from `ToolAgentLoop` and is adapted for the asynchronous training mode of
+`fully_async_policy`. When `partial_rollout=True`, the Rollouter interrupts ongoing generation tasks before
+synchronizing parameters with the Trainer. `AsyncPartialToolAgentLoop` is capable of:
+
+1. **Interrupting Tasks**: Responding to an interrupt signal to save the current state. Currently, interruptions occur
+ during the `GENERATING` process or after other states have completed.
+2. **Resuming Tasks**: Resuming execution from the saved state after parameter synchronization is complete, rather than
+ starting over.
+
+### How to Use
+
+RL training with multi-turn tool calling in `fully_async_policy` is similar to `recipe/retool`. It is enabled by
+specifying `multi_turn` configurations in the config file.
+
+1. **SFT Stage**: First, the model should undergo SFT to learn how to follow tool-calling format instructions.
+2. **Multi-turn Configuration**: In the `fully_async_policy` training configuration, set the following parameters:
+ ```yaml
+ actor_rollout_ref:
+ rollout:
+ multi_turn:
+ enable: True # AsyncPartialToolAgentLoop will be used by default in fully_async_policy mode
+ # Other multi_turn related configurations
+ ```
+3. **Async Parameters**: To improve efficiency, enable `partial_rollout` and `staleness_threshold` when using multi-turn
+ tool calling:
+ ```yaml
+ async_training:
+ partial_rollout: True
+ staleness_threshold: 0.5
+ # Other async parameters
+ ```
+4. **Example**: See `recipe/fully_async_policy/shell/dapo_7b_async_retool.sh`.
+
+### Experimental Results
+
+To validate the performance of `fully_async_policy` on multi-turn tool-calling tasks, we compared it with the standard
+`colocate` synchronous mode. Key parameter settings are as follows.
+
+* **SFT Model**: Based on `Qwen2.5-7B-Instruct`, trained for 6 epochs on the `ReTool-SFT` dataset
+* **RL Algorithm**: DAPO
+* **Dataset**:
+ * Train: `DAPO-Math-17k`
+ * Test: `aime_2025`
+* **Resource and Mode Comparison**:
+ * `colocate sync`: 32 H20 gpus
+ * `fully_async_policy`: 16 gpus for Trainer + 16 gpus for Rollouter
+* **Key Configurations**:
+ 1. **Tool Calling Configuration**:
+ * `multi_turn.enable: True`
+ * `multi_turn.max_user_turns: 16`
+ * `multi_turn.max_assistant_turns: 16`
+ * `multi_turn.tool_config_path: recipe/retool/sandbox_fusion_tool_config.yaml`
+ 2. **`colocate sync` Configuration**:
+ * `ppo_mini_batch_size: 16`
+ * `train_batch_size: 64`
+ 3. **`fully_async_policy` Configuration**:
+ * `ppo_mini_batch_size: 16`
+ * `trigger_parameter_sync_step: 4`
+ * `require_batches: 1`
+ * `staleness_threshold: 1`
+ * `partial_rollout: True`
+
+| training mode | Resource allocation | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | aime_2025
acc/mean@30 |
+|:--------------------:|:---------------------:|:---------:|:---------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:-------------------------------:|
+| colocate | 32 | 375.47 | 228.03 | 35.19 | 111.84 | 9h 46m | 22h 28m | start:0.1078
last:0.2056 |
+| fully_async_policy | 16: 16 | 221.36 | 40.59 | \ | 179.58 | 6h 19m
(1.55x) | 14h 4m
(1.60x) | start:0.11
last:0.2044 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-multiturn-tool?nw=nwuserhouzg
+
+## Future Plans
+
+* GRPO experiments
+* Megatron adaptation
+* SGLang integration
+* Transfer queue integration
+* Asynchronous parameter synchronization
+* AReaL asynchronous algorithm implementation
+* TPPO algorithm implementation
+* Multi-turn and Tool support
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/README_zh.md b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/README_zh.md
new file mode 100644
index 0000000000000000000000000000000000000000..b6b5eb5344a4b4d7236db1263f8445bf68c35f84
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/README_zh.md
@@ -0,0 +1,517 @@
+# Recipe: Fully Async Policy Trainer
+
+**Author:** `https://github.com/meituan-search`
+
+Last updated: 12/15/2025.
+
+本文档介绍了完全异步PPO训练系统，该系统实现了 Trainer 和 Rollouter 的完全解耦，支持异步样本生成和训练。
+在该系统下，我们使用128卡训练qwen2.5-7B模型取得了2.35x-2.67x的性能提升,同时效果没有显著受到影响。
+
+## Introduction
+
+### Background
+
+rollout和train分离架构相较于colocate的架构能够更加灵活地分配资源，设计更加灵活的训练逻辑，从而处理长尾等问题带来的GPU利用率低，训练效率低的问题。
+one_step_off_policy通过分离架构的设计并进行rollout和train一轮异步的训练方法，缓解了rollout时间过长的问题，并在训练效率上取得了一些收益，
+但其强制使用一轮异步的数据，存在不够灵活等问题，而且并不能完全去除长尾对训练效率带来的的影响；在其他框架如areal、Magistral、streamrl、asyncflow上，
+已经基于分离架构实现了异步训练、流式训练，并取得了收益；我们借鉴其方法，在verl上进行了实现。fully_async_policy支持异步、流式、partial
+rollout的训练， 通过合理设置资源分配情况、参数同步频率等参数，fully_async_policy能够显著提高训练效率。
+
+> Magistral https://arxiv.org/abs/2506.10910
+>
+> AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language
+> Reasoning https://arxiv.org/abs/2505.24298
+>
+> StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream
+> Generation https://arxiv.org/abs/2504.15930
+>
+> AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training https://arxiv.org/abs/2507.01663
+>
+
+### 核心贡献
+
+* **资源隔离**：与使用hybrid_engine不同，Rollouter和Trainer使用分离的计算资源，需要分别指定所占用的资源。
+* **生成与训练并行**：Trainer在训练的同时，Rollouter在生成新的样本。
+* **多步异步**: 相比 one step off policy 支持0.x步到多步的异步设定，异步方案更加灵活。
+* **nccl参数同步**：基于nccl通信原语，参考[checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine)实现Rollouter与Trainer间的高效参数同步。
+* **Stream推理与训练**：Rollouter逐样本生成数据，同时数据传输以单个sample为最小传输单位。
+* **异步训练与新鲜度控制**：通过设置参数async_training.staleness_threshold，支持使用旧参数生成的样本进行训练。
+* **PartialRollout**: Rollouter推理过程支持partial rollout逻辑，通过参数同步时，添加`sleep()`和`resume()`
+ 逻辑，保存进行中的rollout的样本，并在下一次rollout中继续使用，减少参数同步等待进行中的任务结束时间。
+
+目前支持使用模式为 megatron/fsdp+vllm。vllm必须使用基于AgentLoop的server模式。
+
+## 设计
+
+fully_async_policy的整体架构如下图所示，fully_async_policy主要由Rollouter、MessageQueue、Trainer、ParameterSynchronizer四部分组成。
+
+
+
+1. Rollouter逐样本生成序列，并将生成的sample放入MessageQueue中，生产的速度受新鲜度控制。
+2. MessageQueue用于暂存Rollouter生成的sample。
+3. Trainer逐样本从MessageQueue中获取，获取到`require_batches*ppo_mini_batch_size`
+ 数量的样本后，就会进行训练，训练async_training.trigger_parameter_sync_step轮后，触发与Rollouter的一次参数同步。
+4. ParameterSynchronizer 实现了Nccl的同步参数同步能力。
+
+当前方案对比base的收益来源，在于colocate情况下，rollout使用更多的资源无法解决长尾样本带来的空闲，
+当我们进行资源隔离后，rollout的时间和train的时间都可能相较于之前更长（因为使用的资源变少了），
+但是相互之间的耗时overlap，端到端的耗时反而有所缩减。
+
+
+
+## 使用方式
+
+### 参数说明
+
+| super params | implication |
+|------------------------------------------------------|-----------------------------------------------------------------|
+| `trainer.nnodes` | Trainer的node数量 |
+| `trainer.n_gpus_per_node` | Trainer每个node上gpu的数量 |
+| `rollout.nnodes` | Rollouter的node数量 |
+| `rollout.n_gpus_per_node` | Rollouter每个node上gpu的数量 |
+| `data.train_batch_size` | 在fully async策略中，该值不生效（默认设置为0） |
+| `data.gen_batch_size` | 在fully async策略中，使用流式的样本生产逻辑（默认设置为1) |
+| `rollout.total_rollout_steps` | 总的rollout的sample数量 |
+| `rollout.test_freq` | Rollouter每更新多少次参数，进行一次validation |
+| `actor_rollout_ref.actor.ppo_mini_batch_size` | The ppo_mini_batch_size is a global num across all workers/gpus |
+| `async_training.require_batches` | FullyAsyncTrainer一次性获取的ppo_mini_batch_size的数量 |
+| `async_training.trigger_parameter_sync_step` | 表示FullyAsyncTrainer进行多少次本地更新后,进行一次参数同步 |
+| `async_training.staleness_threshold` | 新鲜度控制 |
+| `async_training.partial_rollout` | 是否进行partial_rollout |
+| `async_training.use_rollout_log_probs` | 使用rollout产生的log_probs |
+| `async_training.compute_prox_log_prob`（experimental） | 是否在train阶段，使用train模型的参数计算token的 log_prob |
+| `async_training.checkpoint_engine.enable`| 是否开启checkpoint_engine模式的加速，默认值True |
+| `async_training.checkpoint_engine.overlap_broadcast_and_consume` | 启动checkpoint_engine时，是否在参数同步时在broadcast和加载之间使用流水，默认值False|
+| `async_training.checkpoint_engine.device_buffer_size_M` | 启动checkpoint_engine时，组装的bucket的大小(MB)，默认为4096 |
+| `async_training.use_trainer_do_validate` | 是否使用Trainer的do_validate方法进行validation，默认值False |
+
+**进一步的解释：**
+
+* `rollout.total_rollout_steps`
+
+ 与 colocate 相比，数量可以通过 train_batch_size 与 step 相乘对齐:
+ `rollout.total_rollout_steps = data.train_batch_size * step`。
+
+* `async_training.trigger_parameter_sync_step`
+
+ 在fully async策略中，表示Trainer进行多少次本地更新后（也就是获取多少次`require_batches * ppo_mini_batch_size`数量样本），
+ 与Rollouter之间进行一次参数同步。
+ 每两次Rollouter和Trainer参数同步之间，Trainer将会处理`trigger_parameter_sync_step* require_batches\
+ ppo_mini_batch_size`份sample。
+ 如果为了与colocate在公平的情况下对比速度，trigger_parameter_sync_step应该设置为 `data.train_batch_size / (
+ require_batches * ppo_mini_batch_size)`。
+
+* `async_training.staleness_threshold`
+
+ 在fully async策略中，表示最大允许使用的staleness样本的比例。
+
+ * staleness_threshold=0，表示同步训练。
+ Rollouter两次参数更新之间将会生成固定数量的样本，样本数为：
+ $$rollout\_num = (trigger\_parameter\_sync\_step*require\_batches*ppo\_mini\_batch\_size)$$
+ * staleness_threshold>0，表示异步训练， 可以设置为小数，支持更灵活的异步调用。
+ Rollouter两次参数更新之间将会最多生成的样本数为：
+ $$rollout\_num = (1+staleness\_threshold)*(trigger\_parameter\_sync\_step*require\_batches*ppo\_mini\_batch\_size) - num\_staleness\_sample $$
+
+ num_staleness_sample 表示上一次rollout多生成的陈旧样本数。
+
+ 由于是流式系统，rollout持续生成，trainer持续消费。如果rollouter较慢，trainer会更早触发参数同步，rollouter并不会实际生产rollout_num个样本。
+ 当rollout 足够快时，staleness_threshold设置为1，基本上等价于one_step_off policy。
+ 为了避免过期样本太多影响训练精度，建议该值设置小于1。
+
+* `async_training.partial_rollout`
+
+ partial_rollout只会在staleness_threshold>0时才实际上起作用。
+
+* `async_training.use_rollout_log_probs`
+
+ 在强化学习算法中，log_probs与参数版本，token都存在隐性的相关性。由于PPO/GRPO/DAPO等算法的设定，我们在计算重要性采样时，
+ 即 old_log_prob必须使用rollout参数及token所对应log_probs，才能保证算法的正确性。在fully
+ async策略中，我们默认old_log_prob是有rollout所计算的，而不是由trainer所计算。
+
+* `async_training.require_batches`
+
+ 在流式训练中，require_batches 应该设置为1，表示生产够ppo_mini_batch_size样本后，就进行训练。
+ 在实际测试中，我们发现，如果单次下发的样本较少，由于数据分发的顺序，会导致训练不稳定，response 长度变长。
+ 在这里，我们额外提供 require_batches 进行流式分发，单次参与训练的样本数量控制。
+
+* `async_training.compute_prox_log_prob` （experimental）
+
+ 我们在训练过程中，观测到随着训练的进行，训练后期指标和response长度可能会出现不稳定的情况，
+ 这里我们可以使用 [Rollout Importance Sampling](https://verl.readthedocs.io/en/latest/advance/rollout_is.html) 的技术进行
+ 重要性采样，缓解这一问题。为了使用 `Rollout Importance Sampling` 我们需要使用训练引擎使用当前的参数版本计算old_log_prob，此开关需要打开。
+ 此外，在 mode d (async stream pipeline with partial rollout) 的情况下开启 `compute_prox_log_prob` 以及
+ `Rollout Importance Sampling` 后，我们的实现已近似Areal的 `Decoupled PPO`。
+
+* `async_training.checkpoint_engine.enable`
+
+ 开启checkpoint engine后，相较于原始的逐tensor的参数同步方式，同步时间开销普遍可以降低60%以上。但是组装bucket会带来额外的临时显存开销。
+
+* `async_training.checkpoint_engine.overlap_broadcast_and_consume`
+
+ 开启参数broadcast和load_weights之间的流水后，会进一步额外申请更多显存。由于目前分析参数同步的主要耗时并非来自broadcast和load_weights阶段，而是在参数生成阶段（由megatron或FSDP），因此该开关默认关闭。
+
+* `async_training.checkpoint_engine.device_buffer_size_M`
+
+ 控制开启checkpoint engine后，用于同步的显存buffer大小。实际的`bucket_size` = `max(device_buffer_size_M, 最大参数tensor size)`
+ * 在开启`overlap_broadcast_and_consume`时，trainer节点的临时额外显存开销为 `3 * bucket_size`, rollout节点的临时额外显存开销为`2 * bucket_size`。
+ * 在关闭`overlap_broadcast_and_consume`时，trainer节点的临时额外显存开销为 `2 * bucket_size`, rollout节点的临时额外显存开销为`1 * bucket_size`。
+
+* `async_training.use_trainer_do_validate`
+
+ 控制是否使用trainer的`do_validate`方法进行validation。
+ 如果设置为True，trainer会在每次参数更新后，调用`do_validate`方法进行validation。
+ 如果设置为False，trainer不会调用`do_validate`方法。
+
+### 模式支持
+
+1. on policy pipeline:
+ 1. **trigger_parameter_sync_step=1，staleness_threshold=0**
+ 2. Rollouter一次生产`require_batches*ppo_mini_batch_size`
+ 的samples，Trainer获取这些samples后进行训练，训练完后Trainer和Rollouter之间进行一次参数同步;
+ 3. 在rollout阶段，如果存在长尾的样本，但是rollout样本数较少时，较短的样本无法填充到空闲的资源中，会造成一定的资源浪费。
+ 4. 如图a所示；
+
+2. stream off policy pipeline:
+ 1. **trigger_parameter_sync_step>1，staleness_threshold=0**
+ 2. 将会进行同步的流式训练，Rollouter一次生产`require_batches*ppo_mini_batch_size*trigger_parameter_sync_step`
+ 的samples，Trainer每获取`require_batches*ppo_mini_batch_size`
+ 就进行一次本地训练，训练trigger_parameter_sync_step次后，Trainer和Rollouter之间进行一次参数同步;
+ 3. 相较于a，由于一次生成的样本更多，资源的空闲会更低。
+ 4. 在一次step训练中，会存在两次资源闲置的时间，分别是在第一次获取样本时，train等待`require_batches*ppo_mini_batch_size`
+ 个样本生产，以及最后一次参数更新时，rollout等待训练完成。
+ 5. 如图b所示；
+
+3. async stream pipeline with staleness samples:
+ 1. **trigger_parameter_sync_step>=1，staleness_threshold>0，partial_rollout=Flase**
+ 2. Rollouter在每次参数更新后将计划最多生产rollout_num个样本（实际根据rollout速度，生成的样本可能会少与这个值）。
+ 3. 如果rollout过程比较快，Rollouter将会在参数同步前额外生成一部分样本num_stale_samples，用于参数同步后立即给Trainer使用。
+ 触发参数同步时，如果Rollouter有正在生产的任务，将会等待任务完成，同时不会添加新的任务；
+ 4. 相较于b，除第一次step训练外，后续的训练都不会有wait first batch rollout finish的时间，但是会有wait active task
+ finish的时间。
+ 5. 如图c所示；
+
+4. async stream pipeline with partial rollout:
+ 1. **trigger_parameter_sync_step>=1，staleness_threshold>0，partial_rollout=True**
+ 2. 相较于c，触发参数同步时，Rollouter如果有正在生产的sample，会打断rollout过程并进行参数同步，被中断的sample会在参数同步后继续生成。减少了wait
+ active task finish的时间。
+ 3. 如图d所示；
+
+
+
+### 关键指标
+
+| metrics | implication |
+|------------------------------------------------|-----------------------------------------------------------|
+| `trainer/idle_ratio` | Trainer闲置率 |
+| `rollouter/idle_ratio` | Rollouter闲置率 |
+| `fully_async/count/stale_samples_processed` | 训练使用的旧sample总数 |
+| `fully_async/count/stale_trajectory_processed` | 训练使用的旧trajectory总数(一个sample会生产rollout.n条trajectory) |
+| `fully_async/partial/total_partial_num` | 两次trigger_parameter_sync_step之间Trainer处理的partial样本数 |
+| `fully_async/partial/partial_ratio` | 两次trigger_parameter_sync_step之间Trainer处理的partial样本的比例 |
+| `fully_async/partial/max_partial_span` | 两次trigger_parameter_sync_step之间Trainer处理的partial样本的最大参数跨度 |
+
+### 调参建议
+
+* 资源分配与调整:
+ * 合理的资源分配是获得好的训练效率的前提。理想的资源分配情况应该是使得Rollout的时间和Train的时间接近，从而使得整个训练过程流水气泡最小，
+ 避免资源闲置，同时Trainer不会使用旧样本。在真实训练场景下，可以根据实际训练过程中rollout和train的空闲时间调整资源分配，
+ 可从rollouter/idle_ratio和trainer/idle_ratio获得，如果rollouter/idle_ratio较高trainer/idle_ratio较低，
+ 应该增多Trainer的资源减少Rollouter的资源，反之亦然。
+
+* 关键参数：
+ * staleness_threshold: 设置太大会导致较多的旧样本使用，影响模型效果，建议设置小于1。
+ * require_batches：越接近1，越接近纯流式过程，训练过程中bubble越小，能够在速度上获得更快的加速效果，但会对样本的处理顺序产生影响；
+ * trigger_parameter_sync_step: 设置的越小越接近on policy，但会导致频繁的参数同步，长尾样本浪费的资源无法被短样本填充，资源利用率低。
+ 设置的越大有更高的计算效率，但是精度上会受到off policy的影响。
+ * rollout.test_freq: 会占用Rollouter资源，不建议设置太小。
+
+* 模式选择：通过调整不同的参数，Fully Async架构支持不同程度上的优化加速，适用于不同场景的任务。
+ * 对于小规模任务，需要保证训练的稳定性和 on-policy 性，对速度要求不高的场景，可以尝试使用on policy pipeline的模式（模式1）。
+ * 对于需要提高训练吞吐量，但对 staleness 敏感的场景，可以尝试使用 stream off policy pipeline 的模式。即通过
+ 设置trigger_parameter_sync_step>1 ，提高 训练效率，但仍保持同步机制 (staleness_threshold=0 )（模式2）。
+ * 对于大规模任务，对训练速度有较高要求，且可以容忍一定 off-policy 程度、staleness的场景，可以设置staleness_threshold>
+ 0、partial_rollout=True提高训练效率，使用 async stream pipeline 模式（模式 3 或 4）。
+
+### 快速开始
+
+```shell
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*400)))
+test_freq=10
+staleness_threshold=0
+trigger_parameter_sync_step=16
+partial_rollout=False
+
+
+python -m recipe.fully_async_policy.fully_async_main \
+ train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.partial_rollout="${partial_rollout}"
+```
+
+## 实验
+
+### 在7B模型上进行异步训练
+
+我们使用 Qwen2.5-Math-7B 验证 fully async 策略在长候选下，多种资源下的收益情况。
+使用`async stream pipeline with staleness samples` 策略，我们在32卡，64卡，128卡都取得2x左右的性能提升，同时没有显著影响实验效果。
+
+* 机器：H20
+* 模型：Qwen2.5-Math-7B
+* rollout长度：max_response_length FSDP2: 28K tokens;
+* 算法：DAPO
+* 数据集： TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet
+* engine: vllm+FSDP2
+* rollout.n: 16
+* ppo_mini_batch_size: 32
+* test_freq: 20
+
+* colocate sync:
+ * step: 400
+ * train_batch_size: 512
+
+* fully_async_policy
+ * total_rollout_steps: 512*400
+ * require_batches: 4
+ * trigger_parameter_sync_step: 4
+ * staleness_threshold: 0.5
+ * partial_rollout: True
+
+| training mode | resource allocation | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 |
+|:--------------------:|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-------------------------------:|
+| colocate sync | 32 | 790.10 | 357.41 | 107.71 | 269.80 | 13h 44m | 1d 3h 43m | 2d 9h 22m | 3d 17h 5m | max: 0.3313
last: 0.2448 |
+| fully_async_policy | 16:16 | 294.77 | 21.26 | \ | 313.81 | 7h 58m
(1.72x) | 16h 21m
(1.70x) | 1d 0h 53m
(2.31x) | 1d 9h 26m
(2.66x) | max: 0.3302
last: 0.2333 |
+| colocate sync | 64 | 365.28 | 150.72 | 70.26 | 133.41 | 10h 22m | 20h 45m | 1d 7h 6m | 1d 17h 32m | max: 0.3365
last: 0.2333 |
+| fully_async_policy | 32:32 | 189.26 | 28.46 | \ | 156.98 | 4h 57m
(2.09x) | 10h 14m
(2.03x) | 16h 58m
(1.83x) | 21h 40m
(1.92x) | max: 0.3677
last: 0.3406 |
+| colocate sync | 128 | 356.30 | 177.85 | 53.92 | 113.81 | 8h 36m | 17h 56m | 1d 5h 6m | 1d 16h 48m | max: 0.3573
last: 0.2958 |
+| fully_async_policy | 64:64 | 150.63 | 33.14 | \ | 113.16 | 3h 13m
(2.67x) | 6h 46m
(2.65x) | 10h 53m
(2.67x) | 17h 22m
(2.35x) | max: 0.3521
last: 0.3094 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-colocate_async?nw=nwuserhouzg
+
+### 128卡 7B 异步模式实验
+
+我们使用 Qwen2.5-Math-7B 验证 fully async 所支持的各个模式的效果。
+我们可以看到 stream 带来的收益大约1.6x，叠加 staleness 和 partial_rollout 后，收益为2.35x。
+
+| mode | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 |
+|:-------------------------------------------------------------------------------------------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------------:|
+| colocate sync | 356.30 | 177.85 | 53.92 | 113.81 | 8h 36m | 17h 56m | 1d 5h 6m | 1d 16h 48m | max: 0.3573
last: 0.2958 |
+| `stream off policy pipeline`
(+fully async: trigger_parameter_sync_step= 4,
require_batches= 4) | 231.34 | 128.47 | \ | 98.77 | 4h 25m | 9h 41m | 15h 2m | 1d 1h 53m | max: 0.2844
last: 0.2604 |
+| `async stream pipeline with staleness samples`
(+staleness_threshold=0.5) | | | | | | | | | |
+| `async stream pipeline with partial rollout`
(+partial_rollout=True) | 150.63 | 33.14 | \ | 113.16 | 3h 13m | 6h 46m | 10h 53m | 17h 22m | max: 0.3521
last: 0.3094 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg
+
+### 128卡 stale 消融实验
+
+在 `async stream pipeline with partial rollout` 模式下，我们验证 staleness 的设置对于训练效率的影响。
+我们可以发现，staleness 越大，最终取得的收益越明显。
+同时我们也注意到 staleness 取 0.3 和 0.5 的时间比较接近，原因是随着训练步数的增量，response 长度变化较大，训练出现了不稳定的问题。
+后续还需要针对该问题进行进一步的分析和优化。
+
+| staleness_threshold | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | total time
400 step | acc/mean@1 |
+|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|
+| 0 | 231.34 | 128.47 | \ | 98.77 | 4h 25m | 9h 41m | 15h 2m | 1d 1h 53m | max: 0.2844
last: 0.2604 |
+| 0.1 | 171.30 | 58.17 | \ | 109.12 | 3h 53m | 8h 37m | 14h 25m | 19h 59m | max: 0.3542
last: 0.2979 |
+| 0.3 | 146.11 | 38.88 | \ | 103.22 | 3h 18m | 6h 49m | 11h 40m | 17h 20m | max: 0.3469
last: 0.2865 |
+| 0.5 | 150.63 | 33.14 | \ | 113.16 | 3h 13m | 6h 46m | 10h 53m | 17h 22m | max: 0.3521
last: 0.3094 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_stale?nw=nwuserhouzg
+
+### 128卡 7B require_batches 消融实验
+
+在多次测试下，我们发现流式每次下发样本的数量会影响训练的response长度，进而影响训练时长，我们通过修改
+`async_training.require_batches` 验证对与结果的影响。
+
+| require_batches | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | total time
300 step | acc/mean@1 |
+|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|
+| 1 | 203.47 | 30.88 | \ | 181.08 | 3h 31m | 8h 29m | 17h 36m | max: 0.349
last: 0.326 |
+| 2 | 158.72 | 26.32 | \ | 128.08 | 3h 35m | 7h 38m | 13h 57m | max: 0.351
last: 0.3406 |
+| 4 | 124.64 | 25.62 | \ | 95.06 | 3h 13m | 6h 46m | 10h 53m | max: 0.3521
last: 0.3521 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_require_batches?nw=nwuserhouzg
+
+### 30B模型模式实验
+
+我们在 Qwen3-30B-A3B-Base 模型上通过`async stream pipeline with staleness samples` 策略，相比于 colocate 方案取得了 1.7
+倍的性能提升。值得说明的是，这距离异步方式所能带来的性能提升上限还有很大空间。首先，对比实验中使用的最大响应长度仅为
+8k，这远低于此前实验的 20k 序列长度，因此 rollout 的长尾效应并不明显。其次，我们采用了极为倾斜的资源分配方案，rollout 使用了
+96 张 GPU，而 trainer 仅使用了 32 张 GPU，这并不是最优的配置。在实验过程中，我们观察到当前的 verl 实现存在一些限制，比如要求数据必须能被
+GPU 数量整除，这使得资源调整的灵活性受到影响。此外，随着异步训练和部署的加速，性能差距也在逐渐缩小。因此，未来我们将重点关注如何实现更灵活的资源分配和动态调整资源。
+
+* 机器：H20
+* 模型：Qwen3-30B-A3B-Base
+* rollout长度：max_response_length : 8K tokens;
+* 算法： GRPO
+* 数据集： TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet
+* Engine: vllm+Megatron
+* rollout.n: 16
+* ppo_mini_batch_size: 128
+* test_freq: 20
+
+* colocate sync:
+ * step:400
+ * train_batch_size: 512
+
+* fully_async_policy
+ * total_rollout_steps: 512*400
+ * trigger_parameter_sync_step: 512/128 = 4
+ * staleness_threshold: 0.5
+ * partial_rollout: True
+
+| Training Mode | Resource Allocation | Step | Gen | Old Log Prob | Ref | Update Actor | Total Time 100 Step | Total Time 200 Step | Total Time 300 Step | Total Time 400 Step | Acc/Mean@1 |
+|----------------------|--------------------|---------|--------|--------------|--------|--------------|---------------------|---------------------|---------------------|---------------------|-----------------------------|
+| Colocate Sync | 128 | 497.89 | 348.05 | 28.73 | 20.86 | 86.27 | 13h 36m | 1d 3h 48m | 1d 19h 4m | 2d 11h 39m | max: 0.3500
last: 0.3208 |
+| Fully Async Policy | 96:32 | 282.75 | 22.06 | \ | 50.05 | 206.63 | 6h 45m (2.01x) | 14h 48m (1.88x) | 1d 0h 9m (1.78x) | 1d 10h 41m (1.72x) | max: 0.3813
last: 0.3448 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-30B?nw=nwuserhouzg
+
+### checkpoint-engine参数同步消融实验
+我们在Qwen2.5-Math-7B，Qwen3-30B-A3B和Qwen3-235B-A22B三个模型上测试了checkpoint-engine参数同步的单步参数同步耗时，使用的参数均为默认参数配置。实验均在H20机器上完成，并使用megatron训练引擎。
+| model | trainer rank | rollout rank | checkpoint-engine | total sync time |
+|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|
+| Qwen2.5-Math-7B | 4 | 4 | False | 0.12s |
+| Qwen2.5-Math-7B | 4 | 4 | True | 0.02s |
+| Qwen3-30B-A3B | 16 | 16 | False | 15.76s |
+| Qwen3-30B-A3B | 16 | 16 | True | 4.38s |
+| Qwen3-235B-A22B | 64 | 64 | False | 58.57s |
+| Qwen3-235B-A22B | 64 | 64 | True | 23.70s |
+
+### use_trainer_do_validate 实验测试
+我们在Qwen2.5-Math-7B模型上测试了`use_trainer_do_validate`参数的影响。这个结果展示使用`use_trainer_do_validate=True`可以减少验证时间开销，并且训练器节点的空闲时间也减少了。
+
+* Machine: H20
+* Model: Qwen2.5-Math-7B
+* Rollout length: max_response_length FSDP2: 10K tokens;
+* Algorithm: DAPO
+* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet
+* Engine: vllm+FSDP2
+* rollout.n: 16
+* ppo_mini_batch_size: 32
+* test_freq: 10
+
+* fully_async_policy
+ * total_rollout_steps: 512*400
+ * require_batches: 4
+ * trigger_parameter_sync_step: 4
+ * staleness_threshold: 0.5
+ * partial_rollout: True
+
+| training mode | resource allocation | step | gen | old_log_prob | update_actor | validate time | total time
50 step | acc/mean@2 |
+|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|
+| colocate sync | 16 | 484.623 | 52.939 | 0 | 430.263 | 205.080 | 7h9m | 22.6 |
+| fully_async_policy | 8:8 | 489.953 | 52.622 | 0 | 435.874 | 95.699 | 7h2m | 21.0 |
+
+
+## 多轮工具调用
+
+参考 **recipe/retool** 和 **ToolAgentLoop**，我们为 **fully_async_policy** 实现了支持partial rollout的多轮工具调用循环 *
+*AsyncPartialToolAgentLoop**。
+
+### 核心设计
+
+`AsyncPartialToolAgentLoop` 继承自 `ToolAgentLoop`，其核心是适配了 `fully_async_policy` 的异步训练模式。当
+`partial_rollout=True` 时，Rollouter 在与 Trainer 同步参数前会中断正在进行的生成任务。`AsyncPartialToolAgentLoop` 能够：
+
+1. **中断任务**: 响应中断信号，保存当前的生成状态。目前，中断会发生在GENERATING过程中，或其他状态结束后；
+2. **恢复任务**: 在参数同步完成后，从保存的状态恢复，继续执行，而不是从头开始。
+
+### 使用方法
+
+`fully_async_policy`多轮与工具调用的RL训练与 `recipe/retool` 类似，通过在配置文件中指定 `multi_turn` 相关配置来启用。
+
+1. **SFT 阶段**: 首先，需要对模型进行 SFT训练，使其具备遵循工具调用格式指令的能力。
+2. **配置启用**: 在 `fully_async_policy` 的训练配置中，设置以下参数:
+ ```yaml
+ actor_rollout_ref:
+ rollout:
+ multi_turn:
+ enable: True # 在fully_async_policy模式下将默认使用AsyncPartialToolAgentLoop
+ # 其他 multi_turn 相关配置
+ ```
+3. **配置async参数**: 为提高效率，在启用多轮工具调用时，同时开启 `partial_rollout`和`staleness_threshold`：
+ ```yaml
+ async_training:
+ partial_rollout: True
+ staleness_threshold: 0.5
+ # 其他async参数
+ ```
+4. **example**: 参考`recipe/fully_async_policy/shell/dapo_7b_async_retool.sh`
+
+### 实验结果
+
+为验证 `fully_async_policy` 在多轮工具调用任务中的性能，我们将其与标准 `colocate` 同步模式进行了对比。实验具体设置如下。
+
+* **SFT模型**: 实验基于 `Qwen2.5-7B-Instruct` 模型，使用`ReTool-SFT`数据集训练6个epoch；
+* **RL算法**: DAPO
+* **数据集**:
+ * 训练集: `DAPO-Math-17k`
+ * 测试集: `aime_2025`
+* **资源与模式对比**:
+ * `colocate sync`: 32卡 H20
+ * `fully_async_policy`: 16卡 Trainer + 16卡 Rollouter
+* **关键配置**:
+ 1. **工具调用配置**:
+ * `multi_turn.enable: True`
+ * `multi_turn.max_user_turns: 16`
+ * `multi_turn.max_assistant_turns: 16`
+ * `multi_turn.tool_config_path: recipe/retool/sandbox_fusion_tool_config.yaml`
+ 2. **`colocate sync`配置**:
+ * `ppo_mini_batch_size: 16`
+ * `train_batch_size: 64`
+ 3. **`fully_async_policy`配置**:
+ * `ppo_mini_batch_size: 16`
+ * `trigger_parameter_sync_step: 4`
+ * `require_batches: 1`
+ * `staleness_threshold: 1`
+ * `partial_rollout: True`
+
+| training mode | Resource allocation | step | gen | old_log_prob | update_actor | total time
100 step | total time
200 step | aime_2025
acc/mean@30 |
+|:------------------: |:-------------------: |:-------: |:-------: |:------------: |:------------: |:----------------------: |:----------------------: |:---------------------------: |
+| colocate | 32 | 375.47 | 228.03 | 35.19 | 111.84 | 9h 46m | 22h 28m | start:0.1078
last:0.2056 |
+| fully_async_policy | 16: 16 | 221.36 | 40.59 | \ | 179.58 | 6h 19m
(1.55x) | 14h 4m
(1.60x) | start:0.11
last:0.2044 |
+
+> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-multiturn-tool?nw=nwuserhouzg
+
+## 后续计划
+
+* GRPO实验
+* megatron 适配
+* sglang 集成
+* transfer queue 集成
+* 异步参数同步
+* Areal异步算法实现
+* TPPO算法实现
+* 多轮及Tool的支持
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef46df0e529be7ea447c1fbb2554122428dc7147
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/__init__.py
@@ -0,0 +1,20 @@
+# Copyright 2025 Meituan 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 .agent_loop import FullyAsyncAgentLoopManager
+from .partial_single_turn_agent_loop import PartialSingleTurnAgentLoop
+from .partial_tool_agent_loop import AsyncPartialToolAgentLoop
+
+_ = [PartialSingleTurnAgentLoop, AsyncPartialToolAgentLoop]
+__all__ = [FullyAsyncAgentLoopManager]
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..a21796de79b2f151244cb3d79d68289cb8c120b9
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/agent_loop.py
@@ -0,0 +1,370 @@
+# Copyright 2025 Meituan 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.
+import asyncio
+import logging
+import os
+from typing import Any, Optional, Sequence
+
+import hydra
+import numpy as np
+import ray
+from omegaconf import DictConfig
+
+from verl.experimental.agent_loop.agent_loop import (
+ AgentLoopManager,
+ AgentLoopOutput,
+ AgentLoopWorker,
+ AsyncLLMServerManager,
+ DictConfigWrap,
+ _agent_loop_registry,
+ get_trajectory_info,
+)
+from verl.experimental.agent_loop.prometheus_utils import update_prometheus_config
+from verl.protocol import DataProto
+from verl.single_controller.ray import RayResourcePool, RayWorkerGroup
+from verl.utils.rollout_trace import (
+ rollout_trace_attr,
+ rollout_trace_op,
+)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class FullyAsyncLLMServerManager(AsyncLLMServerManager):
+ @rollout_trace_op
+ async def generate_for_partial(
+ self,
+ request_id,
+ *,
+ prompt_ids: list[int],
+ sampling_params: dict[str, Any],
+ image_data: Optional[list[Any]] = None,
+ ) -> tuple[list[Any], list[Any], Any] | tuple[Sequence[int], list[float], bool]:
+ """Generate tokens from prompt ids, used for async partial.
+
+ Args:
+ request_id (str): request id for sticky session.
+ prompt_ids (List[int]): List of prompt token ids.
+ sampling_params (Dict[str, Any]): Sampling parameters for the chat completion.
+
+ Returns:
+ output: A tuple representing the generation output.
+ - Element 0 (Sequence[int]): Generated response token IDs.
+ - Element 1 (list[float]): Log probabilities for the response token IDs.
+ - Element 2 (bool): A flag or status indicating cancellation.
+ """
+ server = self._choose_server(request_id)
+ output = await server.generate_for_partial.remote(
+ request_id=request_id,
+ prompt_ids=prompt_ids,
+ sampling_params=sampling_params,
+ image_data=image_data,
+ )
+ return output
+
+
+@ray.remote
+class FullyAsyncAgentLoopWorker(AgentLoopWorker):
+ def __init__(
+ self, config: DictConfig, server_handles: list[ray.actor.ActorHandle], reward_router_address: str = None
+ ):
+ self.server_manager = FullyAsyncLLMServerManager(config, server_handles)
+ super().__init__(config, server_handles, reward_router_address)
+ # A shared cancellation event for all agent loops running on this worker.
+ self.cancellation_event = asyncio.Event()
+
+ async def generate_sequences_no_post(
+ self, batch: DataProto, partial_output_list: Optional[list[AgentLoopOutput]]
+ ) -> tuple[list[AgentLoopOutput], bool] | tuple[DataProto, bool]:
+ """Generate sequences from agent loop.
+
+ Args:
+ batch (DataProto): Input batch.
+ partial_output_list: Optional[List[AgentLoopOutput]]: already rollout result.
+
+ Returns:
+ list[AgentLoopOutput]: List of agent loop outputs, one per sample in the batch.
+ """
+ config = self.config.actor_rollout_ref.rollout
+ sampling_params = dict(
+ temperature=config.temperature,
+ top_p=config.top_p,
+ repetition_penalty=1.0,
+ logprobs=config.calculate_log_probs,
+ )
+
+ # override sampling params for validation
+ if batch.meta_info.get("validate", False):
+ sampling_params["top_p"] = config.val_kwargs.top_p
+ sampling_params["temperature"] = config.val_kwargs.temperature
+
+ if "agent_name" not in batch.non_tensor_batch:
+ default_agent_loop = config.agent.default_agent_loop
+ batch.non_tensor_batch["agent_name"] = np.array([default_agent_loop] * len(batch), dtype=object)
+
+ if "index" in batch.non_tensor_batch:
+ index = batch.non_tensor_batch["index"]
+ else:
+ index = np.arange(len(batch))
+
+ trajectory_info = await get_trajectory_info(
+ batch.meta_info.get("global_steps", -1), index, batch.meta_info.get("validate", False)
+ )
+
+ if not partial_output_list:
+ partial_output_list = [None] * len(batch)
+ try:
+ tasks = []
+ for i in range(len(batch)):
+ kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items()}
+ kwargs["output"] = partial_output_list[i]
+ tasks.append(
+ asyncio.create_task(self._partial_run_agent_loop(sampling_params, trajectory_info[i], **kwargs))
+ )
+ outputs = await asyncio.gather(*tasks)
+ except Exception:
+ logger.exception("_partial_run_agent_loop failed")
+ raise
+
+ is_cancel = any(output.extra_fields.get("is_cancel", False) for output in outputs)
+ if not is_cancel:
+ output = self._postprocess(outputs)
+ output = self._addition_process(output)
+ return output, is_cancel
+ return outputs, is_cancel
+
+ def _addition_process(self, output: DataProto):
+ """collect metirics"""
+ metrics = output.meta_info.pop("metrics") # List[Dict[str, str]]
+ processing_times_list = [item["generate_sequences"] for item in metrics]
+ tool_calls_times_list = [item["tool_calls"] for item in metrics]
+ output.non_tensor_batch["processing_times"] = processing_times_list
+ output.non_tensor_batch["tool_calls_times"] = tool_calls_times_list
+ return output
+
+ async def _partial_run_agent_loop(
+ self,
+ sampling_params: dict[str, Any],
+ trajectory: dict[str, Any],
+ *,
+ agent_name: str,
+ **kwargs,
+ ) -> AgentLoopOutput:
+ # Completed, return directly
+ if kwargs["output"] is not None and not kwargs["output"].extra_fields.get("is_cancel", False):
+ logger.info("In _partial_run_agent_loop, already completed, return derictly!")
+ return kwargs["output"]
+ try:
+ with rollout_trace_attr(
+ step=trajectory["step"],
+ sample_index=trajectory["sample_index"],
+ rollout_n=trajectory["rollout_n"],
+ validate=trajectory["validate"],
+ name="agent_loop",
+ ):
+ assert agent_name in _agent_loop_registry, (
+ f"Agent loop {agent_name} not registered, registered agent loops: {_agent_loop_registry.keys()}"
+ )
+
+ agent_loop_config = _agent_loop_registry[agent_name]
+ agent_loop = hydra.utils.instantiate(
+ config=agent_loop_config,
+ trainer_config=DictConfigWrap(config=self.config),
+ server_manager=self.server_manager,
+ tokenizer=self.tokenizer,
+ processor=self.processor,
+ dataset_cls=self.dataset_cls,
+ dataset_config=self.config.data,
+ )
+ output: AgentLoopOutput = await agent_loop.run(
+ sampling_params, cancellation_event=self.cancellation_event, **kwargs
+ )
+ if not output.extra_fields.get("is_cancel", False):
+ kwargs.pop("output", None)
+ output = await self._agent_loop_postprocess(output, **kwargs)
+
+ return output
+ except Exception:
+ logger.exception("Agent_loop run failed")
+ raise
+
+ async def cancel_agent_loops(self):
+ """Set the shared cancellation event to stop all agent loops."""
+ self.cancellation_event.set()
+
+ async def resume_agent_loops(self):
+ """Clear the shared cancellation event."""
+ self.cancellation_event.clear()
+
+
+class FullyAsyncAgentLoopManager(AgentLoopManager):
+ def __init__(
+ self, config: DictConfig, worker_group: RayWorkerGroup = None, rm_resource_pool: RayResourcePool = None
+ ):
+ self.config = config
+ self.worker_group = worker_group
+ self.reward_model_manager = None
+ self.reward_router_address = None
+ self.agent_loop_workers_class = FullyAsyncAgentLoopWorker
+
+ # Select rollout replica class based on rollout name
+ rollout_name = config.actor_rollout_ref.rollout.name
+ if rollout_name == "sglang":
+ from verl.experimental.fully_async_policy.sglang_rollout.sglang_async_server import FullyAsyncSGLangReplica
+
+ self.rollout_replica_class = FullyAsyncSGLangReplica
+ print("[FullyAsyncAgentLoopManager] SGLang replica class selected")
+ elif rollout_name == "vllm":
+ from verl.experimental.fully_async_policy.vllm_rollout.vllm_async_server import FullyAsyncvLLMReplica
+
+ self.rollout_replica_class = FullyAsyncvLLMReplica
+ print("[FullyAsyncAgentLoopManager] vLLM replica class selected")
+ else:
+ raise ValueError(f"Unsupported rollout name: {rollout_name}. Supported values are 'sglang' and 'vllm'.")
+
+ self.rm_resource_pool = rm_resource_pool
+ self.rollout_replicas = None
+ self.server_handles = None
+ self.server_addresses = None
+ self.agent_loop_workers = None
+
+ @classmethod
+ async def create(
+ cls, config: DictConfig, worker_group: RayWorkerGroup = None, rm_resource_pool: RayResourcePool = None
+ ):
+ instance = cls(config, worker_group, rm_resource_pool)
+ await instance._async_init()
+ return instance
+
+ async def _async_init(self):
+ if self.config.reward_model.enable and self.config.reward_model.enable_resource_pool:
+ from verl.experimental.reward_loop import RewardModelManager
+
+ self.reward_model_manager = RewardModelManager(self.config.reward_model, self.rm_resource_pool)
+ self.reward_router_address = self.reward_model_manager.get_router_address()
+
+ await self._initialize_llm_servers_async()
+ self._init_agent_loop_workers()
+
+ async def _initialize_llm_servers_async(self):
+ rollout_world_size = (
+ self.config.actor_rollout_ref.rollout.tensor_model_parallel_size
+ * self.config.actor_rollout_ref.rollout.data_parallel_size
+ * self.config.actor_rollout_ref.rollout.pipeline_model_parallel_size
+ )
+ world_size = (
+ self.worker_group.world_size
+ if self.worker_group
+ else self.config.trainer.n_gpus_per_node * self.config.trainer.nnodes
+ )
+ num_replicas = world_size // rollout_world_size
+
+ rollout_config = self.config.actor_rollout_ref.rollout
+ model_config = self.config.actor_rollout_ref.model
+ self.rollout_replicas = [
+ self.rollout_replica_class(
+ replica_rank=replica_rank,
+ config=rollout_config,
+ model_config=model_config,
+ gpus_per_node=self.config.trainer.n_gpus_per_node,
+ )
+ for replica_rank in range(num_replicas)
+ ]
+
+ if self.worker_group:
+ await asyncio.gather(*[server.init_hybrid(self.worker_group) for server in self.rollout_replicas])
+ else:
+ await asyncio.gather(*[server.init_standalone() for server in self.rollout_replicas])
+
+ self.server_handles = [server._server_handle for server in self.rollout_replicas]
+ self.server_addresses = [server._server_address for server in self.rollout_replicas]
+
+ print(f"AgentLoopManager: {self.server_addresses}")
+ # Update Prometheus configuration with server addresses
+ if rollout_config.prometheus.enable:
+ if rollout_config.disable_log_stats:
+ raise ValueError("PROMETHEUS needs disable_log_stats==False, but it is currently True.")
+ await asyncio.to_thread(
+ update_prometheus_config, rollout_config.prometheus, self.server_addresses, rollout_config.name
+ )
+
+ async def generate_single_sample_async(
+ self,
+ sample: DataProto,
+ partial_output_list: Optional[list[AgentLoopOutput]],
+ ) -> tuple[list[AgentLoopOutput], bool] | tuple[DataProto, bool]:
+ """
+ Asynchronously process a single sample
+
+ Args:
+ sample: Single sample data
+ partial_output_list: Optional[List[AgentLoopOutput]]: already rollout result.
+
+ Returns:
+ list[AgentLoopOutput]: Processing results
+ """
+ worker = self._select_best_worker()
+ output_future = worker.generate_sequences_no_post.remote(sample, partial_output_list)
+ return await asyncio.wrap_future(output_future.future())
+
+ def _select_best_worker(self):
+ """Select the best worker, simple round-robin load balancing"""
+ if not hasattr(self, "_worker_index"):
+ self._worker_index = 0
+
+ worker = self.agent_loop_workers[self._worker_index]
+ self._worker_index = (self._worker_index + 1) % len(self.agent_loop_workers)
+ return worker
+
+ async def cancel(self):
+ worker_cancel_tasks = [worker.cancel_agent_loops.remote() for worker in self.agent_loop_workers]
+ rollout_cancel_tasks = [replica.cancel() for replica in self.rollout_replicas]
+ await asyncio.gather(*rollout_cancel_tasks, *worker_cancel_tasks)
+
+ async def resume(self):
+ rollout_resume_tasks = [replica.resume() for replica in self.rollout_replicas]
+ worker_resume_tasks = [worker.resume_agent_loops.remote() for worker in self.agent_loop_workers]
+ await asyncio.gather(*rollout_resume_tasks, *worker_resume_tasks)
+
+ async def wake_up(self):
+ await asyncio.gather(*[replica.wake_up() for replica in self.rollout_replicas])
+
+ async def sleep(self):
+ await asyncio.gather(*[replica.sleep() for replica in self.rollout_replicas])
+
+ async def reset_prefix_cache(self):
+ print("[FullyAsyncAgentLoopManager] Reset prefix cache ...")
+ # await asyncio.gather(*[replica.reset_prefix_cache() for replica in self.rollout_replicas])
+ # Note: debug
+ timeout = 5.0
+
+ async def reset_one(idx, replica):
+ print(f"[reset_prefix_cache] start replica={idx}")
+ try:
+ await asyncio.wait_for(replica.reset_prefix_cache(), timeout=timeout)
+ except asyncio.TimeoutError:
+ print(f"[reset_prefix_cache] TIMEOUT replica={idx} after {timeout}s")
+ return
+ except Exception as e:
+ print(f"[reset_prefix_cache] ERROR replica={idx}: {e!r}")
+ return
+ print(f"[reset_prefix_cache] done replica={idx}")
+
+ tasks = [reset_one(i, replica) for i, replica in enumerate(self.rollout_replicas)]
+ await asyncio.gather(*tasks, return_exceptions=True)
+ print("[FullyAsyncAgentLoopManager] Reset prefix cache finished")
+
+ async def clear_kv_cache(self):
+ await asyncio.gather(*[replica.clear_kv_cache() for replica in self.rollout_replicas])
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/partial_single_turn_agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/partial_single_turn_agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..30f3fb9220ce36c06c6c6e9a5380db1e52a5adee
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/partial_single_turn_agent_loop.py
@@ -0,0 +1,115 @@
+# Copyright 2025 Meituan 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.
+import logging
+import os
+from typing import Any, Optional
+from uuid import uuid4
+
+from verl.experimental.agent_loop import AgentLoopBase
+from verl.experimental.agent_loop.agent_loop import AgentLoopOutput, register
+from verl.utils.profiler import simple_timer
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+@register("partial_single_turn_agent")
+class PartialSingleTurnAgentLoop(AgentLoopBase):
+ """Naive agent loop that only do single turn chat completion."""
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.prompt_length = self.config.actor_rollout_ref.rollout.prompt_length
+ self.response_length = self.config.actor_rollout_ref.rollout.response_length
+ self.apply_chat_template_kwargs = self.config.data.get("apply_chat_template_kwargs", {})
+
+ async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
+ output: Optional[AgentLoopOutput] = kwargs.get("output", None)
+ messages = list(kwargs["raw_prompt"])
+ param_version = kwargs.get("param_version", 0)
+
+ metrics = {}
+ request_id = uuid4().hex
+ image_data = (kwargs.get("multi_modal_data") or {}).get("image", None)
+
+ param_version_start = param_version
+ param_version_end = param_version
+
+ if not output:
+ # TODO(baiyan): it is supposed to use the correct processor,
+ # but I found the async training would hang if use_correct_processor=True.
+ # so we use the tokenizer to tokenize the prompt for now.
+ use_correct_processor = False
+ if self.processor is not None and use_correct_processor:
+
+ def get_prompt_ids():
+ raw_prompt = self.processor.apply_chat_template(
+ messages,
+ add_generation_prompt=True,
+ tokenize=False,
+ **self.apply_chat_template_kwargs,
+ )
+ model_inputs = self.processor(text=[raw_prompt], images=image_data, return_tensors="pt")
+ return model_inputs.pop("input_ids").squeeze(0).tolist()
+
+ prompt_ids = await self.loop.run_in_executor(None, get_prompt_ids)
+ else:
+ prompt_ids = await self.loop.run_in_executor(
+ None,
+ lambda: self.tokenizer.apply_chat_template(
+ messages, add_generation_prompt=True, tokenize=True, **self.apply_chat_template_kwargs
+ ),
+ )
+ else:
+ if output.extra_fields.get("is_cancel", False):
+ # Resume the paused sample,
+ # add the result directly after prompt_ids,
+ # and reset generate_sequences metric
+ prompt_ids = output.prompt_ids + output.response_ids
+ metrics["generate_sequences"] = output.metrics.generate_sequences
+ param_version_start = output.extra_fields.get("param_version_start", param_version)
+ else:
+ # In the same batch of samples,
+ # some are canceled and some are not.
+ # The samples without partial rollout are returned directly.
+ return output
+ with simple_timer("generate_sequences", metrics):
+ response_ids, response_logprobs, is_cancel = await self.server_manager.generate_for_partial(
+ request_id=request_id, prompt_ids=prompt_ids, sampling_params=sampling_params, image_data=image_data
+ )
+ if not output:
+ response_mask = [1] * len(response_ids)
+ else:
+ # Pause the sample to be resumed, add the output result to response_ids, and reset response_mask
+ prompt_ids = output.prompt_ids
+ response_logprobs = output.response_logprobs + response_logprobs
+ response_ids = output.response_ids + response_ids
+ response_mask = [1] * len(response_ids)
+ if len(response_ids) >= self.response_length:
+ is_cancel = False
+
+ return AgentLoopOutput(
+ prompt_ids=prompt_ids,
+ response_ids=response_ids[: self.response_length],
+ response_mask=response_mask[: self.response_length],
+ response_logprobs=response_logprobs[: self.response_length],
+ num_turns=2,
+ metrics=metrics,
+ extra_fields={
+ "is_cancel": is_cancel,
+ "param_version_start": param_version_start,
+ "param_version_end": param_version_end,
+ },
+ # multi_modal_data={"image": image_data} if image_data is not None else {},
+ )
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/partial_tool_agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/partial_tool_agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed404586f290096a848e5b4694afc766dc8ff248
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/agent_loop/partial_tool_agent_loop.py
@@ -0,0 +1,281 @@
+# Copyright 2025 Meituan 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.
+
+import asyncio
+import copy
+import logging
+import os
+from typing import Any, Optional
+from uuid import uuid4
+
+from verl.experimental.agent_loop.agent_loop import AgentLoopOutput, register
+from verl.experimental.agent_loop.tool_agent_loop import AgentData, AgentState, ToolAgentLoop
+from verl.utils.profiler import simple_timer
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+@register("async_partial_tool_agent")
+class AsyncPartialToolAgentLoop(ToolAgentLoop):
+ """
+ Support for partial rollout with multiple tool invocations in Agent Loop
+
+ """
+
+ def __init__(self, trainer_config, **kwargs):
+ super().__init__(trainer_config, **kwargs)
+ self.enable_partial_rollout = trainer_config.config.async_training.get("partial_rollout", False)
+
+ # async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:
+ async def run(
+ self, sampling_params: dict[str, Any], *, cancellation_event: asyncio.Event = None, **kwargs
+ ) -> AgentLoopOutput:
+ """
+ Main entrance, supports interruption/recovery
+
+ Args:
+ sampling_params: Sampling parameters
+ cancellation_event: cancellationn sginal
+ **kwargs: Contains output (for recovery), raw_prompt, param_version, etc.
+
+ Returns:
+ AgentLoopOutput: Include the is_cancel flag
+ """
+ param_version = kwargs.get("param_version", 0)
+ agent_data = None
+ state = None
+
+ # 1. check whether is the partial task
+ output: Optional[AgentLoopOutput] = kwargs.get("output", None)
+ if output and output.extra_fields.get("is_cancel", False):
+ agent_data, state = self._restore_from_output(output)
+
+ logger.info(f"[PartialToolAgent] Resuming from {state.value}")
+ else:
+ if output and not output.extra_fields.get("is_cancel", False):
+ # Completed, return directly
+ return output
+
+ agent_data = await self._init_agent_data(kwargs, param_version)
+ state = AgentState.PENDING
+ logger.info("[PartialToolAgent] Start from scratch")
+ # 2. run state machine
+ state = await self._run_state_machine(agent_data, state, sampling_params, cancellation_event)
+
+ # 3. bulid output
+ if state == AgentState.TERMINATED:
+ return self._build_completed_output(agent_data, param_version)
+ else:
+ # build cancelled output
+ return self._build_cancelled_output(agent_data, state)
+
+ async def _init_agent_data(self, kwargs: dict, param_version: int) -> AgentData:
+ messages = list(kwargs["raw_prompt"])
+ image_data = copy.deepcopy(kwargs.get("multi_modal_data", {}).get("image", None))
+ video_data = copy.deepcopy(kwargs.get("multi_modal_data", {}).get("video", None))
+ metrics = {}
+ request_id = uuid4().hex
+ tools_kwargs = kwargs.get("tools_kwargs", {})
+
+ # Initialize interaction if needed
+ interaction = None
+ interaction_kwargs = {}
+ if self.interaction_config_file:
+ interaction_kwargs = kwargs["extra_info"]["interaction_kwargs"]
+ if "name" not in interaction_kwargs:
+ raise ValueError("'name' key is required in interaction_kwargs")
+ interaction_name = interaction_kwargs["name"]
+ if interaction_name not in self.interaction_map:
+ raise ValueError(
+ f"Interaction '{interaction_name}' not found in interaction_map. Available interactions: "
+ f"{list(self.interaction_map.keys())}"
+ )
+ interaction = self.interaction_map[interaction_name]
+ await interaction.start_interaction(request_id, **interaction_kwargs)
+ # Create AgentData instance to encapsulate all state
+ agent_data = AgentData(
+ messages=messages,
+ image_data=image_data,
+ video_data=video_data,
+ metrics=metrics,
+ request_id=request_id,
+ tools_kwargs=tools_kwargs,
+ interaction=interaction,
+ interaction_kwargs=interaction_kwargs,
+ )
+
+ # additional param version record
+ agent_data.extra_fields["param_version_start"] = param_version
+ agent_data.extra_fields["param_version_end"] = param_version
+
+ return agent_data
+
+ def _restore_from_output(self, output: AgentLoopOutput) -> tuple[AgentData, AgentState]:
+ """restore AgentState and AgentData from output"""
+ agent_data = output.extra_fields.get("agent_data", None)
+ agent_state = output.extra_fields.get("agent_state", None)
+ if agent_data is None or agent_state is None:
+ raise ValueError(f"Unexpected situation: agent_data is {agent_data}, agent_state is {agent_state}")
+ return agent_data, agent_state
+
+ async def _run_state_machine(
+ self,
+ agent_data: AgentData,
+ state: AgentState,
+ sampling_params: dict[str, Any],
+ cancellation_event: asyncio.Event = None,
+ ) -> AgentState:
+ """
+ State machine.
+ Currently, interruptions are only supported to occur in the GENERATING state or other states have ended.
+ """
+ # State machine loop
+ while state != AgentState.TERMINATED:
+ if cancellation_event and cancellation_event.is_set():
+ logger.info(f"[PartialToolAgent] Cancellation detected. Interrupted before/at state: {state.value}")
+ return state
+ if state == AgentState.PENDING:
+ state = await self._handle_pending_state(agent_data, sampling_params)
+ elif state == AgentState.GENERATING:
+ state = await self._handle_generating_state_partial(agent_data, sampling_params)
+ elif state == AgentState.PROCESSING_TOOLS:
+ state = await self._handle_processing_tools_state(agent_data)
+ elif state == AgentState.INTERACTING:
+ state = await self._handle_interacting_state(agent_data)
+ else:
+ logger.error(f"[PartialToolAgent] Invalid state: {state}")
+ return AgentState.TERMINATED
+
+ return AgentState.TERMINATED
+
+ async def _handle_generating_state_partial(
+ self, agent_data: AgentData, sampling_params: dict[str, Any], ignore_termination: bool = False
+ ) -> AgentState:
+ """
+ Handle GENERATING state, support partial rollout
+ """
+ add_messages: list[dict[str, Any]] = []
+
+ with simple_timer("generate_sequences", agent_data.metrics):
+ # partial interface
+ if self.enable_partial_rollout:
+ response_ids, log_probs, is_cancel = await self.server_manager.generate_for_partial(
+ request_id=agent_data.request_id,
+ prompt_ids=agent_data.prompt_ids,
+ sampling_params=sampling_params,
+ image_data=agent_data.image_data,
+ )
+
+ if is_cancel:
+ # Save the generated parts
+ agent_data.response_ids = response_ids
+ agent_data.prompt_ids += agent_data.response_ids
+ agent_data.response_mask += [1] * len(response_ids)
+ if log_probs:
+ agent_data.response_logprobs += log_probs
+ if not ignore_termination and len(agent_data.response_mask) >= self.response_length:
+ # If response_length has reached the limit,
+ # it is considered to have ended normally.
+ agent_data.assistant_turns += 1
+ return AgentState.TERMINATED
+ return AgentState.GENERATING
+ else:
+ # original generate interface
+ output = await self.server_manager.generate(
+ request_id=agent_data.request_id,
+ prompt_ids=agent_data.prompt_ids,
+ sampling_params=sampling_params,
+ image_data=agent_data.image_data,
+ )
+ response_ids = output.token_ids
+ log_probs = output.log_probs
+
+ agent_data.assistant_turns += 1
+ agent_data.response_ids = response_ids
+ agent_data.prompt_ids += agent_data.response_ids
+ agent_data.response_mask += [1] * len(agent_data.response_ids)
+ if log_probs:
+ agent_data.response_logprobs += log_probs
+
+ if not ignore_termination and len(agent_data.response_mask) >= self.response_length:
+ return AgentState.TERMINATED
+ if self.max_assistant_turns and agent_data.assistant_turns >= self.max_assistant_turns:
+ return AgentState.TERMINATED
+ if self.max_user_turns and agent_data.user_turns >= self.max_user_turns:
+ return AgentState.TERMINATED
+
+ # Extract tool calls
+ _, agent_data.tool_calls = await self.tool_parser.extract_tool_calls(agent_data.response_ids)
+
+ # Handle interaction if needed
+ if self.interaction_config_file:
+ assistant_message = await self.loop.run_in_executor(
+ None, lambda: self.tokenizer.decode(agent_data.response_ids, skip_special_tokens=True)
+ )
+ add_messages.append({"role": "assistant", "content": assistant_message})
+ agent_data.messages.extend(add_messages)
+
+ # Determine next state
+ if agent_data.tool_calls:
+ return AgentState.PROCESSING_TOOLS
+ elif self.interaction_config_file:
+ return AgentState.INTERACTING
+ else:
+ return AgentState.TERMINATED
+
+ def _build_completed_output(self, agent_data: AgentData, param_version: int) -> AgentLoopOutput:
+ """build completed output"""
+ response_ids = agent_data.prompt_ids[-len(agent_data.response_mask) :]
+ prompt_ids = agent_data.prompt_ids[: len(agent_data.prompt_ids) - len(agent_data.response_mask)]
+ multi_modal_data = {"image": agent_data.image_data} if agent_data.image_data is not None else {}
+ output = AgentLoopOutput(
+ prompt_ids=prompt_ids,
+ response_ids=response_ids[: self.response_length],
+ response_mask=agent_data.response_mask[: self.response_length],
+ multi_modal_data=multi_modal_data,
+ response_logprobs=agent_data.response_logprobs[: self.response_length]
+ if agent_data.response_logprobs
+ else None,
+ num_turns=agent_data.user_turns + agent_data.assistant_turns + 1,
+ metrics=agent_data.metrics,
+ extra_fields={},
+ )
+ output.extra_fields.update(
+ {
+ "turn_scores": agent_data.turn_scores,
+ "tool_rewards": agent_data.tool_rewards,
+ "is_cancel": False,
+ "param_version_start": agent_data.extra_fields["param_version_start"],
+ "param_version_end": param_version,
+ }
+ )
+ return output
+
+ def _build_cancelled_output(self, agent_data: AgentData, state: AgentState) -> AgentLoopOutput:
+ """build cancelled output"""
+ return AgentLoopOutput(
+ prompt_ids=[],
+ response_ids=[],
+ response_mask=[],
+ multi_modal_data={},
+ response_logprobs=None,
+ num_turns=0,
+ metrics=agent_data.metrics,
+ extra_fields={
+ "is_cancel": True,
+ "agent_data": agent_data,
+ "agent_state": state,
+ },
+ )
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/base_detach_sync.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/base_detach_sync.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0924417d78413c74fd20b10cab281c1618664ec
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/base_detach_sync.py
@@ -0,0 +1,238 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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.
+
+import asyncio
+import logging
+import os
+import threading
+
+import torch
+from omegaconf import DictConfig
+from ray.util.collective import collective
+
+from verl.single_controller.base.decorator import Dispatch, register
+from verl.utils.device import get_torch_device, is_npu_available
+from verl.utils.distributed import stateless_init_process_group
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class BaseDetachNcclSync:
+ _bucket_size_mb = 1024.0
+ _sync_history = []
+ _max_history_size = 20
+ _last_avg_bucket_size = 1024.0
+
+ def __init__(self, config: DictConfig, role: str):
+ self._bg_loop = asyncio.new_event_loop()
+ self._bg_thread = threading.Thread(
+ target=self._start_background_loop, args=(self._bg_loop,), name="rollout_actor_async_worker", daemon=True
+ )
+ self._bg_thread.start()
+ logger.info(f"[DetachNcclSync] Background thread for SGLang sync started. PID: {os.getpid()}")
+
+ @classmethod
+ def get_bucket_size_mb(cls):
+ return cls._bucket_size_mb
+
+ @classmethod
+ def get_last_avg_bucket_size(cls):
+ return cls._last_avg_bucket_size
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=True)
+ def get_last_avg_bucket_size_remote(self):
+ return BaseDetachNcclSync._last_avg_bucket_size
+
+ @classmethod
+ def record_sync_metrics(cls, bucket_size_mb, sync_time):
+ """Dynamically adjust the bucket size based on past synchronization times."""
+ bucket_size_mb_value = bucket_size_mb[0] if isinstance(bucket_size_mb, list) else bucket_size_mb
+ print(f"[DetachNcclSync] sync_metrics: bucket_size_mb={bucket_size_mb_value:.2f}MB, sync_time={sync_time:.2f}s")
+ cls._sync_history.append((bucket_size_mb_value, sync_time))
+ if len(cls._sync_history) > cls._max_history_size:
+ cls._sync_history.pop(0)
+
+ MIN_BUCKET_SIZE_MB = 512
+ MAX_BUCKET_SIZE_MB = 8192 # 8GB
+
+ if len(cls._sync_history) < 4:
+ cls._bucket_size_mb = min(MAX_BUCKET_SIZE_MB, cls._bucket_size_mb * 1.5)
+ else:
+ times = [t for _, t in cls._sync_history]
+ buckets = [b for b, _ in cls._sync_history]
+ recent_avg_time = sum(times[-2:]) / 2
+ previous_avg_time = sum(times[-4:-2]) / 2
+ recent_avg_bucket = sum(buckets[-2:]) / 2
+ previous_avg_bucket = sum(buckets[-4:-2]) / 2
+
+ performance_improved = recent_avg_time < previous_avg_time
+ bucket_increased = recent_avg_bucket > previous_avg_bucket
+ time_change_ratio = (
+ abs(recent_avg_time - previous_avg_time) / previous_avg_time if previous_avg_time > 0 else 0.0
+ )
+
+ if time_change_ratio > 0.2:
+ increase_step, decrease_step = 1.2, 0.8
+ elif time_change_ratio > 0.1:
+ increase_step, decrease_step = 1.1, 0.9
+ elif time_change_ratio > 0.05:
+ increase_step, decrease_step = 1.05, 0.95
+ else:
+ increase_step, decrease_step = 1.02, 0.98
+
+ should_increase = (performance_improved and bucket_increased) or (
+ not performance_improved and not bucket_increased
+ )
+ step = increase_step if should_increase else decrease_step
+ new_size = cls._bucket_size_mb * step
+ cls._bucket_size_mb = min(MAX_BUCKET_SIZE_MB, max(MIN_BUCKET_SIZE_MB, new_size))
+
+ def _start_background_loop(self, loop):
+ asyncio.set_event_loop(loop)
+ try:
+ loop.run_forever()
+ except Exception as e:
+ logger.error(f"[DetachNcclSync] Background loop crashed: {e}")
+
+ def _run_async_safely(self, coro):
+ if not self._bg_thread.is_alive():
+ raise RuntimeError("Background thread for SGLang sync is not running!")
+
+ future = asyncio.run_coroutine_threadsafe(coro, self._bg_loop)
+ return future.result()
+
+ def __del__(self):
+ if hasattr(self, "_bg_loop") and self._bg_loop.is_running():
+ self._bg_loop.call_soon_threadsafe(self._bg_loop.stop)
+ if hasattr(self, "_bg_thread") and self._bg_thread.is_alive():
+ self._bg_thread.join(timeout=1.0)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def init_checkpoint_engine(self, rank_offset: int, actor_num: int, rollout_num: int):
+ from .checkpoint_engine import CheckpointEngine
+
+ current_rank = torch.distributed.get_rank() + rank_offset
+ actor_ranks = list(range(actor_num))
+ rollout_ranks = [rank + actor_num for rank in range(rollout_num)]
+ assert rank_offset == 0 or rank_offset == actor_num
+
+ self.checkpoint_engine = CheckpointEngine(
+ current_rank, actor_ranks, rollout_ranks, self.config.checkpoint_engine.device_buffer_size_M
+ )
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def create_weight_sync_group(self, master_address, master_port, rank_offset, world_size):
+ rank = torch.distributed.get_rank() + rank_offset
+ self._weight_sync_group = stateless_init_process_group(
+ master_address,
+ master_port,
+ rank,
+ world_size,
+ get_torch_device().current_device(),
+ )
+
+ @staticmethod
+ def get_inference_model(rollout):
+ """
+ Get models according to different types of inference_engine
+ Args:
+ rollout: rollout object
+ Returns:
+ model: model object (for vllm) or rollout object itself (for sglang)
+ """
+ inference_engine = rollout.inference_engine
+ if hasattr(inference_engine, "llm_engine"):
+ inference_model = inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model
+ elif hasattr(inference_engine, "worker"):
+ inference_model = inference_engine.worker.model_runner.model
+ else:
+ raise AttributeError(
+ f"Unsupported inference_engine type: {type(inference_engine)}. "
+ f"Expected LLM (with llm_engine attribute) or WorkerWrapperBase (with worker attribute)."
+ )
+ return inference_model
+
+ def _sync_sglang_weights(self, inference_model, params, sync_group_name):
+ bucket_size_bytes = int(self.get_bucket_size_mb() * 1024 * 1024)
+ actual_bucket_sizes = []
+ current_batch = []
+ current_batch_size = 0
+
+ def flush_batch():
+ if current_batch:
+ actual_bucket_sizes.append(current_batch_size / (1024 * 1024))
+ self._run_async_safely(self.update_weights(inference_model, iter(current_batch)))
+ get_torch_device().synchronize()
+ current_batch.clear()
+
+ for key, shape, dtype in self._weights_info:
+ tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device())
+ if self._is_actor:
+ assert key in params
+ origin_data = params[key]
+ if hasattr(origin_data, "full_tensor"):
+ origin_data = origin_data.full_tensor()
+ if torch.distributed.get_rank() == 0:
+ tensor.copy_(origin_data)
+ collective.broadcast(tensor, src_rank=0, group_name=sync_group_name)
+
+ tensor_size = tensor.numel() * tensor.element_size()
+ current_batch.append((key, tensor))
+ current_batch_size += tensor_size
+
+ if current_batch_size >= bucket_size_bytes:
+ flush_batch()
+ current_batch_size = 0
+
+ flush_batch()
+ cls = type(self)
+ cls._last_avg_bucket_size = (
+ sum(actual_bucket_sizes) / len(actual_bucket_sizes) if actual_bucket_sizes else self.get_bucket_size_mb()
+ )
+
+ # Resume kv_cache after weights sync to restore GPU memory released during pause
+ if self._is_rollout and self.rollout_device_mesh["infer_tp"].get_local_rank() == 0:
+ self._run_async_safely(inference_model.resume_memory_occupation(tags=["kv_cache"]))
+
+ def _sync_vllm_weights(self, inference_model, params, sync_group_name):
+ for key, shape, dtype in self._weights_info:
+ tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device())
+ if self._is_actor:
+ assert key in params
+ origin_data = params[key]
+ if hasattr(origin_data, "full_tensor"):
+ origin_data = origin_data.full_tensor()
+ if torch.distributed.get_rank() == 0:
+ tensor.copy_(origin_data)
+ if is_npu_available:
+ self._weight_sync_group.broadcast(tensor, src=0, stream=get_torch_device().current_stream())
+ else:
+ collective.broadcast(tensor, src_rank=0, group_name=sync_group_name)
+ if self._is_rollout:
+ inference_model.load_weights([(key, tensor)])
+
+ async def update_weights(self, inference_engine, params):
+ from sglang.srt.weight_sync.utils import update_weights as sgl_update_weights
+
+ await sgl_update_weights(
+ engine=inference_engine,
+ params_batch=params,
+ device_mesh_key="infer_tp",
+ device_mesh=self.rollout_device_mesh,
+ )
+
+ if self.rollout_device_mesh["infer_tp"].get_local_rank() == 0:
+ await inference_engine.flush_cache()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/checkpoint_engine.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/checkpoint_engine.py
new file mode 100644
index 0000000000000000000000000000000000000000..28f932d61b3a46f9a01d2d454a0b5d66d932509b
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/checkpoint_engine.py
@@ -0,0 +1,522 @@
+# Copyright 2025 Meituan 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.
+"""
+This logic is largely copied from:
+- https://github.com/MoonshotAI/checkpoint-engine
+"""
+
+import concurrent.futures
+import os
+import re
+import socket
+import subprocess
+import threading
+from collections.abc import Callable
+from functools import lru_cache
+from typing import TYPE_CHECKING, Annotated, Any, TypedDict
+
+import torch
+import zmq
+from pydantic import BaseModel, PlainSerializer, PlainValidator, WithJsonSchema
+from ray.util.collective import collective
+
+from verl.utils.device import (
+ get_device_name,
+ get_torch_device,
+)
+
+if TYPE_CHECKING:
+ from typing import TypeVar
+
+ from typing_extensions import TypedDict
+
+ class FileMeta(TypedDict):
+ key: str # parameter name
+ dtype: torch.dtype
+ shape: torch.Size
+ type: type
+ tp_concat_dim: int
+
+ T = TypeVar("T")
+
+
+def _dt_validate(value: Any) -> torch.dtype:
+ """Validate the input value to ensure it is a valid torch.dtype"""
+ if isinstance(value, str):
+ if not value.startswith("torch."):
+ raise ValueError(f"dtype {value} should start with torch.")
+ try:
+ value = getattr(torch, value.split(".")[1])
+ except AttributeError as e:
+ raise ValueError(f"unknown dtype: {value}") from e
+ if not isinstance(value, torch.dtype):
+ raise TypeError(f"dtype {value} should be torch.dtype, got {type(value)}")
+ return value
+
+
+# Annotated type for torch.dtype with validation and serialization
+_TorchDtype = Annotated[
+ torch.dtype,
+ PlainValidator(_dt_validate),
+ PlainSerializer(lambda x: str(x), return_type=str),
+ WithJsonSchema({"type": "string"}, mode="serialization"),
+]
+
+
+def _size_validate(value: Any) -> torch.Size:
+ """Validate the input value to ensure it is a valid torch.Size"""
+ if isinstance(value, list | tuple):
+ return torch.Size(value)
+ if not isinstance(value, torch.Size):
+ raise TypeError(f"size {value} should be torch.Size, got {type(value)}")
+ return value
+
+
+# Annotated type for torch.Size with validation and serialization
+_TorchSize = Annotated[
+ torch.Size,
+ PlainValidator(_size_validate),
+ PlainSerializer(lambda x: tuple(x), return_type=tuple),
+ WithJsonSchema({"type": "array", "items": {"type": "integer"}}, mode="serialization"),
+]
+
+
+def _tensor_validate(value: Any) -> torch.Tensor:
+ """Validate the input value to ensure it is a valid torch.Tensor"""
+ if isinstance(value, torch.Tensor):
+ return value
+ raise TypeError(f"tensor {value} should be torch.Tensor, got {type(value)}")
+
+
+# Annotated type for torch.Tensor with validation
+_TorchTensor = Annotated[
+ torch.Tensor,
+ PlainValidator(_tensor_validate),
+]
+
+
+class ParameterMeta(BaseModel):
+ """Metadata for a parameter including name, dtype, and shape"""
+
+ name: str
+ dtype: _TorchDtype
+ shape: _TorchSize
+
+
+class MemoryBuffer(BaseModel):
+ """
+ MemoryBuffer assembles a group of parameter tensors into a single buffer,
+ and records the meta information of each original parameter.
+ """
+
+ buffer: _TorchTensor
+ size: int # size of buffer in bytes
+ metas: list[ParameterMeta]
+
+
+class MemoryBufferMeta(BaseModel):
+ """The meta info of MemoryBuffer, but not store the buffer data"""
+
+ size: int
+ metas: list[ParameterMeta]
+
+
+# 256 bytes alignment when flatten torch tensors to uint8 buffer
+_ALIGN_SIZE = 256
+
+
+def _align_size(dtype: torch.dtype, shape: torch.Size) -> int:
+ """
+ Calculate the aligned size of a torch tensor
+
+ If the tensor's size (in bytes) cannot be evenly divided by _ALIGN_SIZE,
+ it will be rounded up to the nearest multiple of _ALIGN_SIZE.
+
+ Args:
+ dtype (torch.dtype): The data type of the tensor (e.g., torch.float32, torch.int64).
+ shape (torch.Size): The shape of the tensor, representing its dimensions.
+
+ Returns:
+ int: The aligned size of the tensor in bytes.
+ """
+ return (dtype.itemsize * shape.numel() + _ALIGN_SIZE - 1) // _ALIGN_SIZE * _ALIGN_SIZE
+
+
+@lru_cache(maxsize=1)
+def get_ip() -> str:
+ try:
+ # try to get ip from network interface
+ with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+ s.connect(("8.8.8.8", 80))
+ return s.getsockname()[0]
+ except Exception as e: # noqa: BLE001
+ # fallback to get ip from hostname
+ print(f"fail to get ip from network interface, fallback to get ip from hostname: {e}")
+ return socket.gethostbyname(socket.gethostname())
+
+
+def npu_generate_uuid() -> str:
+ """Generate uuid for each npu device"""
+ str_pid = str(os.getpid())
+ npu_num = 8
+ try:
+ for npu_id in range(npu_num):
+ cmd = ["npu-smi", "info", "-t", "proc-mem", "-i", str(npu_id)]
+ result = subprocess.run(cmd, check=True, capture_output=True, text=True) # noqa: S603
+ str_result = str(result.stdout)
+ if str_pid in str_result:
+ # In A3 server, one NPU has two chips.
+ match_chip_count = re.search(r"Chip Count[^\d]*(\d+)", str_result)
+ chip_count = int(match_chip_count.group(1))
+ search_after_pid = str_result[str_result.find(str_pid) + len(str_pid) :]
+ match_chip_id = re.search(r"Chip ID[^\d]*(\d+)", search_after_pid)
+ chip_id = int(match_chip_id.group(1))
+ return f"{get_ip()}-{npu_id * chip_count + chip_id}"
+ raise ValueError("The current process is not running on the npu device")
+ except subprocess.CalledProcessError as e:
+ raise ValueError("The current process is not running on the npu device") from e
+
+
+def _get_physical_device_id(device_index: int | None = None) -> str:
+ """
+ Get the physical device (GPU or NPU) uuid of the current device
+ """
+ try:
+ if get_device_name() == "npu":
+ return f"NPU-{npu_generate_uuid()}"
+ else:
+ return f"GPU-{get_torch_device().get_device_properties(device_index).uuid!s}"
+ except AssertionError as e:
+ raise ValueError(f"fail to get physical gpu id {device_index}") from e
+
+
+class FlattenedTensorMetadata(TypedDict):
+ name: str
+ shape: torch.Size
+ dtype: torch.dtype
+ # specify the start offset of this tensor in shared ipc_buffer tensor
+ offset: int
+
+
+def _to_flattened_tensor_meta(metas: list[ParameterMeta], offset: int = 0) -> list[FlattenedTensorMetadata]:
+ """
+ compute the offset of each parameter in the buffer
+
+ Args:
+ metas (list[ParameterMeta]): The list of parameter metas info
+ offset (int): The start offset of the buffer. Defaults to 0.
+
+ Returns:
+ list[FlattenedTensorMetadata]: The list of FlattenedTensorMetadata:
+ """
+ ret = []
+ for meta in metas:
+ size = _align_size(meta.dtype, meta.shape)
+ ret.append(
+ {
+ "name": meta.name,
+ "dtype": meta.dtype,
+ "shape": meta.shape,
+ "offset": offset,
+ }
+ )
+ offset += size
+ return ret
+
+
+def _extract_weights(
+ flatten_metas: list[FlattenedTensorMetadata], buffer: torch.Tensor
+) -> list[tuple[str, torch.Tensor]]:
+ """
+ According to the flatten_metas and buffer, extract the weights
+ """
+
+ assert buffer is not None
+ weights: list[tuple[str, torch.Tensor]] = []
+ for item in flatten_metas:
+ shape = item["shape"]
+ if isinstance(shape, list | tuple):
+ shape = torch.Size(shape)
+ assert isinstance(shape, torch.Size)
+ dtype, offset = item["dtype"], item["offset"]
+ size = dtype.itemsize * shape.numel()
+ tensor = buffer[offset : offset + size].view(dtype=dtype).view(shape)
+ weights.append((item["name"], tensor))
+ return weights
+
+
+class CheckpointEngine:
+ """
+ CheckpointEngine class for control parameters synchronization.
+ Each trainer/rollout rank has a CheckpointEngine instance.
+ """
+
+ def __init__(
+ self, current_rank: int, actor_ranks: list[int], rollout_ranks: list[int], device_buffer_size_M: int
+ ) -> None:
+ self.current_rank = current_rank
+ self.actor_ranks = actor_ranks
+ self.rollout_ranks = rollout_ranks
+ # global_buckets saves the global MemoryBufferMeta infos.
+ # Thus each CheckpointEngine instance can control their operations in SPMD
+ self.global_buckets: dict[int, list[MemoryBufferMeta]] = None
+ # min device_buffer_size for h2d and broadcast
+ self.device_buffer_size_M = device_buffer_size_M
+
+ # ipc config for broadcast in pipeline mode
+ self._zmq_ctx = zmq.Context()
+ self._zmq_addr_counter: int = 0
+ device_index = self.current_rank % get_torch_device().device_count()
+ self._device_uuid = _get_physical_device_id(device_index)
+
+ def register_checkpoint(
+ self, weights_info: list[tuple[str, torch.Size, torch.dtype]], cpu_named_params: dict[str, torch.Tensor]
+ ):
+ """
+ Register checkpoint information and prepare memory buffers for parameter synchronization.
+
+ This function organizes the parameters into memory buckets for efficient synchronization
+ and prepares pinned memory buffers for faster data transfer between CPU and device.
+
+ Args:
+ weights_info (list[tuple[str, torch.Size, torch.dtype]]):
+ A list of tuples containing parameter name, shape, and data type.
+ cpu_named_params (dict[str, torch.Tensor]):
+ A dictionary mapping parameter names to their corresponding CPU tensors.
+
+ Steps:
+ 1. Calculate the bucket size based on the largest parameter tensor size and the device buffer size.
+ 2. Organize parameters into global buckets for each actor rank, ensuring that the total size of each bucket
+ does not exceed the bucket size.
+ 3. For actor ranks, allocate pinned memory buffers for each bucket and copy the parameter tensors
+ into these buffers.
+
+ Notes:
+ Each CheckpointEngine instance maintains the global buckets metas,
+ but stores part of parmas data in host memory
+ """
+ bucket_size = max(
+ self.device_buffer_size_M << 20, max(_align_size(dtype, shape) for _, shape, dtype in weights_info)
+ )
+ print(
+ f"set checkpoint_engine device buffer size: {self.device_buffer_size_M}M, "
+ f"and finally set it to {bucket_size >> 20}M considering the largest parameter tensor size"
+ )
+ self.bucket_size = bucket_size
+
+ # global_buckets saves the global MemoryBufferMeta infos.
+ if self.global_buckets is None:
+ self.global_buckets = {rank: [MemoryBufferMeta(size=0, metas=[])] for rank in self.actor_ranks}
+
+ actor_ranks_size = len(self.actor_ranks)
+ assert actor_ranks_size > 0, f"actor_ranks:{self.actor_ranks} should not be empty"
+ for param_idx, (param_name, param_shape, param_dtype) in enumerate(weights_info):
+ # Each parameter is assigned to an actor rank, and only this rank will store it
+ assgin_rank = self.actor_ranks[param_idx % actor_ranks_size]
+ param_size = _align_size(param_dtype, param_shape)
+
+ if self.global_buckets[assgin_rank][-1].size + param_size > bucket_size:
+ assert self.global_buckets[assgin_rank][-1].size, (
+ f"global_buckets[{assgin_rank}][-1].size:{self.global_buckets[assgin_rank][-1].size}"
+ " should not be 0"
+ )
+ self.global_buckets[assgin_rank].append(MemoryBufferMeta(size=0, metas=[]))
+ self.global_buckets[assgin_rank][-1].metas.append(
+ ParameterMeta(name=param_name, dtype=param_dtype, shape=param_shape)
+ )
+ self.global_buckets[assgin_rank][-1].size += param_size
+
+ def register_pin_memory(idx: int, size: int) -> tuple[int, torch.Tensor]:
+ """Allocate pinned memory for a bucket."""
+ buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
+ return idx, buffer
+
+ def register_tensor(buffer: torch.Tensor, offset: int, tensor: torch.Tensor):
+ """Copy a tensor into a pinned memory buffer."""
+ buffer[offset : offset + tensor.nbytes] = tensor.view(-1).view(dtype=torch.uint8)
+
+ memory_buffers = [] # for rollout rank, return empty buffer
+ if self.current_rank in self.actor_ranks: # is_actor
+ local_buckets = self.global_buckets[self.current_rank]
+ memory_buffers = [
+ MemoryBuffer(buffer=torch.empty(0), size=bucket.size, metas=bucket.metas) for bucket in local_buckets
+ ]
+
+ # Use thread pool to accelerate organize parameters into buckets
+ with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
+ futures = [
+ executor.submit(register_pin_memory, idx, bucket.size) for idx, bucket in enumerate(local_buckets)
+ ]
+ new_futures = []
+ for future in concurrent.futures.as_completed(futures):
+ idx, buffer = future.result()
+ assert buffer.numel() == local_buckets[idx].size, (
+ f"buffer numel {buffer.numel()} should be equal to bucket size {local_buckets[idx].size}"
+ )
+ memory_buffers[idx].buffer = buffer
+ print(
+ f"[rank{self.current_rank}] register pin_memory for "
+ f" bucket {idx + 1}/{len(local_buckets)} finished, "
+ f"size {buffer.numel() / 1024 / 1024:.2f}MiB, start to copy tensors to buffer"
+ )
+ offset = 0
+ for meta in local_buckets[idx].metas:
+ name = meta.name
+ tensor = cpu_named_params[name]
+ size = _align_size(tensor.dtype, tensor.shape)
+ assert size == _align_size(meta.dtype, meta.shape), (
+ f"tensor {name} size {size} should be equal to "
+ f"meta size {_align_size(meta.dtype, meta.shape)}"
+ )
+ new_futures.append(executor.submit(register_tensor, buffer, offset, tensor))
+ offset += size
+ for future in concurrent.futures.as_completed(new_futures):
+ future.result()
+
+ self.memory_buffers = memory_buffers
+
+ def get_max_buckets_num_per_rank(self):
+ """
+ Get the maximum number of buckets for all rank.
+ """
+ assert self.global_buckets is not None
+ return max(len(buckets) for buckets in self.global_buckets.values())
+
+ def _bind_zmq_socket(self) -> tuple[zmq.Socket, list[tuple[str, str]]]:
+ """
+ Bind zmq socket for broadcast.
+ """
+
+ def zmq_handle(device_uuid: str) -> str:
+ return f"ipc://@checkpoint-engine-{device_uuid}-{self._zmq_addr_counter}.sock"
+
+ socket_path = zmq_handle(self._device_uuid)
+ socket = self._zmq_ctx.socket(zmq.REQ)
+ socket.bind(socket_path)
+ self._zmq_addr_counter += 1
+ return socket, socket_path
+
+ def update_checkpoint(self, inference_model, group_name: str, overlap_broadcast_and_consume: bool = False):
+ """
+ Update the checkpoint by broadcasting and loading weights.
+
+ This function handles the synchronization of parameters across ranks by:
+ 1. Copying data from memory buffers to device buffers (h2d_buffer).
+ 2. Broadcasting the data to all ranks using collective communication.
+ 3. Loading the weights into the inference model if provided.
+ 4. Optionally, use a pipeline approach for broadcasting and loading weights.
+
+ Args:
+ inference_model: The model to load weights into. If None (trainer rank), weights are only broadcasted.
+ group_name (str): The name of the collective communication group.
+ overlap_broadcast_and_consume (bool): Whether to use the pipeline approach
+ for broadcasting and loading weights.
+ """
+ try:
+ h2d_buffer: torch.Tensor | None = (
+ None
+ if self.current_rank in self.rollout_ranks
+ else torch.empty(self.bucket_size, dtype=torch.uint8, device=get_torch_device().current_device())
+ )
+ # for pipeline mode, we need to allocate 2x buffer size
+ broadcast_load_buffer = torch.empty(
+ self.bucket_size * (2 if overlap_broadcast_and_consume else 1),
+ dtype=torch.uint8,
+ device=get_torch_device().current_device(),
+ )
+ except Exception:
+ print(
+ "allocate buffer for update_checkpoint failed, "
+ "you may need to reduce "
+ "config.async_training.checkpoint_engine.device_buffer_size_M"
+ )
+ raise
+
+ max_h2d_iter = self.get_max_buckets_num_per_rank()
+
+ if overlap_broadcast_and_consume:
+ socket, socket_path = self._bind_zmq_socket()
+
+ # Define a function to update weights from IPC
+ def update_weights_from_ipc_(socket_path):
+ zmq_ctx = zmq.Context()
+ socket = zmq_ctx.socket(zmq.REP)
+ socket.connect(socket_path)
+ socket.recv_pyobj()
+ socket.send(b"")
+
+ while True:
+ payload: tuple[Callable, tuple] | list[FlattenedTensorMetadata] | None = socket.recv_pyobj()
+ if payload is None:
+ # means the update is done
+ get_torch_device().synchronize()
+ socket.send(b"")
+ break
+ assert isinstance(payload, list)
+ if inference_model is not None:
+ inference_model.load_weights(_extract_weights(payload, broadcast_load_buffer))
+ get_torch_device().synchronize()
+ socket.send(b"")
+
+ req_thread = threading.Thread(
+ target=update_weights_from_ipc_,
+ args=(socket_path,),
+ )
+ req_thread.start()
+ socket.send_pyobj(b"")
+ get_torch_device().synchronize()
+
+ gidx = 0
+ local_buckets = self.global_buckets.get(self.current_rank, [])
+
+ for i in range(max_h2d_iter):
+ # Step 1: Each actor rank copy the parameter tensor into device memory
+ if i < len(self.memory_buffers):
+ h2d_buffer[: local_buckets[i].size].data.copy_(self.memory_buffers[i].buffer)
+
+ # Step 2: Broadcast the device data in turn
+ for broadcast_rank, _buckets in self.global_buckets.items():
+ if i >= len(_buckets):
+ continue
+ bucket = _buckets[i]
+
+ # Prepare the broadcast buffer
+ start = gidx % 2 * self.bucket_size if overlap_broadcast_and_consume else 0
+ buffer_b: torch.Tensor = broadcast_load_buffer[start : start + bucket.size]
+ if broadcast_rank == self.current_rank:
+ buffer_b.data.copy_(h2d_buffer[: bucket.size])
+
+ # Broadcast the buffer to all ranks
+ collective.broadcast(buffer_b, src_rank=broadcast_rank, group_name=group_name)
+
+ if overlap_broadcast_and_consume:
+ socket.recv()
+ collective.barrier(group_name=group_name)
+ socket.send_pyobj(_to_flattened_tensor_meta(bucket.metas, start))
+ elif inference_model is not None:
+ named_tensor = _to_flattened_tensor_meta(bucket.metas, 0)
+ inference_model.load_weights(_extract_weights(named_tensor, buffer_b))
+
+ gidx += 1
+
+ if overlap_broadcast_and_consume:
+ socket.recv()
+ socket.send_pyobj(None)
+ socket.recv()
+ req_thread.join()
+ socket.close()
+
+ collective.barrier(group_name=group_name)
+ # clear host memory cache
+ self.memory_buffers = []
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/config/fully_async_ppo_megatron_trainer.yaml b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/config/fully_async_ppo_megatron_trainer.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..85b8307ee0c23f35ba1ff50436212b49d37ef3f6
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/config/fully_async_ppo_megatron_trainer.yaml
@@ -0,0 +1,76 @@
+hydra:
+ searchpath:
+ - file://verl/trainer/config
+
+defaults:
+ - ppo_megatron_trainer
+ - _self_
+
+async_training:
+
+ # Maximum samples staleness threshold
+ staleness_threshold: 0.1
+
+ # Frequency of parameter synchronization between rollouter and trainer,
+ # One step means trainer obtains a batch of required samples
+ trigger_parameter_sync_step: 4
+
+ # The number of ppo_mini_batches that the FullyAsyncTrainer obtains once
+ require_batches: 1
+
+ # When synchronizing parameters, whether to interrupt rollouter and perform partial rollout
+ partial_rollout: True
+
+ # Whether to use rollout log probs for training
+ use_rollout_log_probs: True
+
+ # compute_prox_log_prob
+ compute_prox_log_prob: False
+
+ # whether to use trainer do_validate
+ use_trainer_do_validate: False
+
+
+ # checkpoint_engine config for accelerating parameter synchronization between rollouter and trainer
+ checkpoint_engine:
+ # Whether to use checkpoint_engine
+ enable: True
+
+ # Device buffer size for checkpoint_engine, default is 4096 MB
+ device_buffer_size_M: 4096
+
+ # Enable the pipeline for broadcasting and updating parameters, but it requires more device memory
+ overlap_broadcast_and_consume: False
+
+# Rollout config
+rollout:
+
+ # Number of nodes used in the rollout
+ nnodes: 1
+
+ # Number of GPUs per node
+ n_gpus_per_node: 8
+
+ # number of responses (i.e. num sample times). > 1 for grpo
+ n: 4
+
+ # total rollout samples # TODO rename to total_rollout_samples
+ total_rollout_steps: 100
+
+ # Number of epochs in training
+ total_epochs: 10
+
+ # Test frequency, how many times a parameter update triggers a validation
+ test_freq: 1
+
+data:
+ # Number of samples generated, currently only support 1
+ gen_batch_size: 1
+
+actor_rollout_ref:
+ # checkpoint_engine config for accelerating parameter synchronization between rollouter and trainer
+ checkpoint_engine: ${oc.select:async_training.checkpoint_engine, null}
+
+ actor:
+ # Whether to use rollout log probs for training
+ use_rollout_log_probs: ${oc.select:async_training.use_rollout_log_probs, True}
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/config/fully_async_ppo_trainer.yaml b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/config/fully_async_ppo_trainer.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..c5692b4a931ec8452b015dda0db9e7d78678165f
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/config/fully_async_ppo_trainer.yaml
@@ -0,0 +1,76 @@
+hydra:
+ searchpath:
+ - file://verl/trainer/config
+
+defaults:
+ - ppo_trainer
+ - _self_
+
+async_training:
+
+ # Maximum samples staleness threshold
+ staleness_threshold: 0.1
+
+ # Frequency of parameter synchronization between rollouter and trainer,
+ # One step means trainer obtains a batch of required samples
+ trigger_parameter_sync_step: 4
+
+ # The number of ppo_mini_batches that the FullyAsyncTrainer obtains once
+ require_batches: 1
+
+ # When synchronizing parameters, whether to interrupt rollouter and perform partial rollout
+ partial_rollout: True
+
+ # Whether to use rollout log probs for training
+ use_rollout_log_probs: True
+
+ # compute_prox_log_prob
+ compute_prox_log_prob: False
+
+ # whether to use trainer do_validate
+ use_trainer_do_validate: False
+
+
+ # checkpoint_engine config for accelerating parameter synchronization between rollouter and trainer
+ checkpoint_engine:
+ # Whether to use checkpoint_engine
+ enable: True
+
+ # Device buffer size for checkpoint_engine, default is 4096 MB
+ device_buffer_size_M: 4096
+
+ # Enable the pipeline for broadcasting and updating parameters, but it requires more device memory
+ overlap_broadcast_and_consume: False
+
+# Rollout config
+rollout:
+
+ # Number of nodes used in the rollout
+ nnodes: 1
+
+ # Number of GPUs per node
+ n_gpus_per_node: 8
+
+ # number of responses (i.e. num sample times). > 1 for grpo
+ n: 4
+
+ # total rollout samples # TODO rename to total_rollout_samples
+ total_rollout_steps: 100
+
+ # Number of epochs in training
+ total_epochs: 10
+
+ # Test frequency, how many times a parameter update triggers a validation
+ test_freq: 1
+
+data:
+ # Number of samples generated, currently only support 1
+ gen_batch_size: 1
+
+actor_rollout_ref:
+ # checkpoint_engine config for accelerating parameter synchronization between rollouter and trainer
+ checkpoint_engine: ${oc.select:async_training.checkpoint_engine, null}
+
+ actor:
+ # Whether to use rollout log probs for training
+ use_rollout_log_probs: ${oc.select:async_training.use_rollout_log_probs, True}
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/detach_utils.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/detach_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..c8d2c02ebcae9ad723f88a5110984a615e4336cd
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/detach_utils.py
@@ -0,0 +1,363 @@
+# Copyright 2025 Meituan 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.
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import numpy as np
+import torch
+
+from verl import DataProto
+from verl.experimental.agent_loop.agent_loop import AgentLoopOutput
+from verl.trainer.ppo.ray_trainer import compute_response_mask
+
+
+@dataclass
+class RolloutSample:
+ """Enhanced rollout sample containing both original batch info and AgentLoopOutput"""
+
+ # Original batch information
+ full_batch: Any
+
+ # AgentLoopOutput from generation
+ agent_loop_output_list: list[AgentLoopOutput]
+
+ # Metadata
+ sample_id: str
+ epoch: int
+
+ # Processing metadata
+ processing_times: list[float]
+ tool_calls: list[float]
+ param_version: int
+ param_version_start: list[int]
+ param_version_end: list[int]
+ rollout_status: dict[str, Any]
+
+
+@dataclass
+class ValidateMetrics:
+ """Metrics for validation"""
+
+ timing_raw: dict[str, Any]
+ metrics: Optional[dict[str, Any]] = None
+ global_steps: Optional[int] = None
+ param_version: Optional[int] = None
+
+
+def prepare_single_generation_data(batch_dict, config) -> DataProto:
+ """
+ Similar to the logic of ray_trainer._prepare_generate_batch, but for a single sample.
+ Separate the data used for generation from the original data.
+
+ Returns:
+ tuple: (original_batch_dict, gen_data_for_single_sample)
+ """
+
+ full_batch = DataProto.from_single_dict(batch_dict)
+
+ batch_keys_to_pop = []
+ non_tensor_batch_keys_to_pop = []
+
+ existing_batch_keys = [k for k in batch_keys_to_pop if k in full_batch.batch.keys()]
+ existing_non_tensor_keys = [k for k in non_tensor_batch_keys_to_pop if k in full_batch.non_tensor_batch.keys()]
+
+ if existing_batch_keys or existing_non_tensor_keys:
+ full_batch.pop(
+ batch_keys=existing_batch_keys,
+ non_tensor_batch_keys=existing_non_tensor_keys,
+ )
+
+ # Setting selected agent, that supports partial
+ if config.actor_rollout_ref.rollout.multi_turn.enable:
+ full_batch.non_tensor_batch["agent_name"] = np.array(
+ ["async_partial_tool_agent"] * len(full_batch), dtype=object
+ )
+ else:
+ full_batch.non_tensor_batch["agent_name"] = np.array(
+ ["partial_single_turn_agent"] * len(full_batch), dtype=object
+ )
+
+ # Add global step count to generated data
+ full_batch = full_batch.repeat(repeat_times=config.actor_rollout_ref.rollout.n, interleave=True)
+ return full_batch
+
+
+def assemble_batch_from_rollout_samples(
+ rollout_samples: list[RolloutSample], tokenizer, config, balance_batch=None
+) -> DataProto:
+ """
+ Assemble gen_batch_output from RolloutSample objects
+ Assembles batches from RolloutSample objects, similar to the _post_generate_batch logic in ray_trainer.
+
+ Args:
+ rollout_samples: List of RolloutSample objects
+ tokenizer: Tokenizer instance
+ config: Configuration object containing trainer settings
+ balance_batch: Whether to balance the batch (simplified version)
+
+ Returns:
+ DataProto: Assembled gen_batch_output
+
+ Raises:
+ ValueError: If rollout_samples is empty
+ """
+ start_time = time.time()
+
+ if not rollout_samples:
+ raise ValueError("Empty rollout_samples provided for batch assembly")
+
+ print(f"[BatchUtils] Assembling batch from {len(rollout_samples)} RolloutSample objects")
+
+ rollout_samples_batch = []
+ processing_times = []
+ tool_calls = []
+ rollout_status = rollout_samples[0].rollout_status
+ # Add a prefix to all rollout_status keys
+ rollout_status = {f"fully_async/{key}": value for key, value in rollout_status.items()}
+
+ for rs in rollout_samples:
+ rollout_samples_batch.append(rs.full_batch)
+ final_batch = DataProto.concat(rollout_samples_batch)
+
+ # Calculate response_mask (if not present)
+ if "response_mask" not in final_batch.batch.keys():
+ final_batch.batch["response_mask"] = compute_response_mask(final_batch)
+
+ if balance_batch:
+ balance_batch(final_batch, metrics={})
+
+ # Calculate the global valid token number
+ if "attention_mask" in final_batch.batch:
+ final_batch.meta_info["global_token_num"] = torch.sum(final_batch.batch["attention_mask"], dim=-1).tolist()
+
+ processing_times = final_batch.non_tensor_batch["processing_times"]
+ tool_calls = final_batch.non_tensor_batch["tool_calls_times"]
+ # Collect statistics
+
+ processing_time_stats = {
+ "processing_time/avg": np.mean(processing_times),
+ "processing_time/max": np.max(processing_times),
+ "processing_time/min": np.min(processing_times),
+ "processing_time/tp50": np.percentile(processing_times, 50),
+ "processing_time/tp99": np.percentile(processing_times, 99),
+ "processing_time/tp95": np.percentile(processing_times, 95),
+ }
+ tool_calls_stats = {}
+ if len(tool_calls) > 0:
+ tool_calls_stats = {
+ "timing_s/agent_loop/tool_calls/max": np.max(tool_calls),
+ "timing_s/agent_loop/tool_calls/min": np.min(tool_calls),
+ "timing_s/agent_loop/tool_calls/mean": np.mean(tool_calls),
+ }
+ processing_time_stats = {f"fully_async/{key}": value for key, value in processing_time_stats.items()}
+
+ param_version_start = final_batch.non_tensor_batch["param_version_start"]
+ param_version_end = final_batch.non_tensor_batch["param_version_end"]
+ param_version_diff = [abs(a - b) for a, b in zip(param_version_end, param_version_start, strict=False)]
+ num_diff0 = param_version_diff.count(0)
+ partial_stats = {
+ "fully_async/partial/total_partial_num": len(param_version_diff) - num_diff0,
+ "fully_async/partial/partial_ratio": (len(param_version_diff) - num_diff0) / len(param_version_diff),
+ "fully_async/partial/max_partial_span": max(param_version_diff),
+ }
+ # add meta_info
+ param_versions = [rs.param_version for rs in rollout_samples]
+ trajectorys_param_versions = final_batch.non_tensor_batch["param_version_end"]
+
+ final_batch.meta_info.update(
+ {
+ "rollout_param_versions": param_versions,
+ "param_version_diversity": len(set(param_versions)) if param_versions else 0,
+ "trajectory_param_versions": trajectorys_param_versions,
+ **processing_time_stats,
+ **rollout_status,
+ **partial_stats,
+ **tool_calls_stats,
+ }
+ )
+
+ print(f"[BatchUtils] Batch assembly completed in {time.time() - start_time:.2f}s")
+
+ return final_batch
+
+
+class MetricsAggregator:
+ """Metrics aggregator, used to combine metrics from multiple training steps"""
+
+ def __init__(self, total_gpus: int):
+ # Store all values for each metric
+ self.metric_values: dict[str, list[float]] = defaultdict(list)
+ # Store the number of samples at each step for weighted averaging
+ self.sample_counts: list[int] = []
+ # Store the timestamp of each step for time-related calculations
+ self.timestamps: list[float] = []
+ # Step Count
+ self.step_count = 0
+ # total num gpus used
+ self.total_gpus = total_gpus
+
+ # Metric aggregation rule configuration
+ self.aggregation_rules = self._init_aggregation_rules()
+
+ def _init_aggregation_rules(self) -> dict[str, dict[str, list[str]]]:
+ """Initialize metrics aggregation rules"""
+ return {
+ # Time-Based metrics, can add metrics here
+ "time_sum": ["perf/time_per_step"],
+ "min": ["timing_s/agent_loop/tool_calls/min"],
+ "avg": ["timing_s/agent_loop/tool_calls/mean"],
+ "max": ["timing_s/agent_loop/tool_calls/max"],
+ "last": [
+ "fully_async/count/total_generated_samples",
+ "fully_async/count/stale_samples_processed",
+ "fully_async/count/stale_trajectory_processed",
+ "fully_async/count/current_param_version",
+ "fully_async/count/dropped_stale_samples",
+ "training/global_step", # TODO change name to: total_step
+ ],
+ }
+
+ def add_step_metrics(self, metrics: dict[str, Any], sample_count: int, timestamp: float = None):
+ """Adding a single-step metrics"""
+ if timestamp is None:
+ timestamp = time.time()
+
+ self.sample_counts.append(sample_count)
+ self.timestamps.append(timestamp)
+ self.step_count += 1
+
+ # Store all metrics values
+ for key, value in metrics.items():
+ if isinstance(value, int | float | np.number):
+ self.metric_values[key].append(float(value))
+ elif isinstance(value, torch.Tensor):
+ self.metric_values[key].append(float(value.item()))
+
+ def _get_aggregation_type(self, metric_name: str) -> str:
+ """Determine the aggregation type based on the metric name"""
+ for agg_type, metric_list in self.aggregation_rules.items():
+ if metric_name in metric_list:
+ return agg_type
+
+ metric_lower = metric_name.lower()
+ if any(keyword in metric_lower for keyword in ["timing_s/"]):
+ return "time_sum"
+ if any(keyword in metric_lower for keyword in ["mean", "avg", "average"]):
+ return "avg"
+ if any(keyword in metric_lower for keyword in ["max", "maximum"]):
+ return "max"
+ if any(keyword in metric_lower for keyword in ["min", "minimum"]):
+ return "min"
+ if any(keyword in metric_lower for keyword in ["sum", "total"]):
+ return "sum"
+ if any(keyword in metric_lower for keyword in ["weighted_avg"]):
+ return "weighted_avg"
+
+ return "avg"
+
+ def _aggregate_single_metric(self, metric_name: str, values: list[float]) -> float:
+ """Aggregating a single metric"""
+ if not values:
+ return 0.0
+
+ agg_type = self._get_aggregation_type(metric_name)
+
+ if agg_type == "last":
+ return values[-1]
+
+ elif agg_type == "weighted_avg":
+ # Weighted average
+ if len(values) != len(self.sample_counts):
+ # If the lengths do not match, use a simple average
+ return sum(values) / len(values)
+
+ total_samples = sum(self.sample_counts)
+ if total_samples == 0:
+ return sum(values) / len(values)
+
+ weighted_sum = sum(v * c for v, c in zip(values, self.sample_counts, strict=False))
+ return weighted_sum / total_samples
+
+ elif agg_type == "sum" or agg_type == "time_sum":
+ return sum(values)
+
+ elif agg_type == "avg":
+ return sum(values) / len(values)
+
+ elif agg_type == "max":
+ return max(values)
+
+ elif agg_type == "min":
+ return min(values)
+
+ else:
+ # Default average
+ return sum(values) / len(values)
+
+ def get_aggregated_metrics(self) -> dict[str, Any]:
+ """aggregated metrics"""
+ t = time.time()
+ if self.step_count == 0:
+ return {}
+
+ aggregated = {}
+
+ # Aggregate all metrics
+ for metric_name, values in self.metric_values.items():
+ aggregated[metric_name] = self._aggregate_single_metric(metric_name, values)
+
+ # Aggregate special metrics
+ aggregated = self._special_metrics_aggergate(aggregated)
+
+ print(f"aggregated metrics done. cost {time.time() - t}")
+
+ return aggregated
+
+ def _special_metrics_aggergate(self, aggregated: dict[str, Any]) -> dict[str, Any]:
+ """calculate special metrics"""
+
+ # global_seqlen/minmax_diff
+ if "global_seqlen/minmax_diff" in aggregated.keys():
+ aggregated["global_seqlen/minmax_diff"] = aggregated["global_seqlen/max"] - aggregated["global_seqlen/min"]
+
+ # perf/throughput
+ REQUIRED_PERF_KEYS = {"perf/throughput", "perf/total_num_tokens", "perf/time_per_step"}
+ if REQUIRED_PERF_KEYS.issubset(aggregated):
+ aggregated["perf/throughput"] = aggregated["perf/total_num_tokens"] / (
+ aggregated["perf/time_per_step"] * self.total_gpus
+ )
+
+ # trainer/idle_ratio
+ if "timing_s/gen" in aggregated.keys() and "timing_s/step" in aggregated.keys():
+ aggregated["trainer/idle_ratio"] = aggregated["timing_s/gen"] / aggregated["timing_s/step"]
+
+ return aggregated
+
+ def reset(self):
+ """Reset Aggregator"""
+ self.metric_values.clear()
+ self.sample_counts.clear()
+ self.timestamps.clear()
+ self.step_count = 0
+
+ def get_current_stats(self) -> dict[str, Any]:
+ """Get statistics about the current aggregation state (for debugging)"""
+ return {
+ "step_count": self.step_count,
+ "metric_count": len(self.metric_values),
+ "total_samples": sum(self.sample_counts),
+ "metric_names": list(self.metric_values.keys()),
+ }
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fsdp2_utils.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fsdp2_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f1856596fb4ac27e572a4290fc9b6c7117ffffc
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fsdp2_utils.py
@@ -0,0 +1,125 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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 typing import Optional
+
+import torch
+import torch.distributed as dist
+from packaging import version
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor._dtensor_spec import DTensorSpec
+
+if version.parse(torch.__version__) < version.parse("2.6"):
+ raise RuntimeError("PyTorch 2.6 or higher is required to use fstp_utils.")
+
+
+def fsdp2_sharded_save_to_cpu(
+ model: torch.nn.Module,
+) -> tuple[dict[str, tuple[torch.Tensor, DTensorSpec]], DTensorSpec]:
+ """
+ Sharded Save: Each process only saves the local DTensor shard from its own GPU to CPU memory.
+
+ Args:
+ model: FSDP2-wrapped model whose parameters are of DTensor type.
+
+ Returns:
+ cpu_sharded_state: Dictionary of CPU shards for the current process.
+ Key = parameter name, Value = (CPU shard tensor, original DTensorSpec)
+ global_spec: DTensorSpec of the first parameter (used to verify global rules during loading)
+ """
+ cpu_sharded_state = {}
+ global_spec = None # Record global sharding rules (all parameters follow the same spec)
+
+ for param_name, param in model.named_parameters():
+ # Only process sharded parameters of DTensor type (core parameters of FSDP2)
+ if not isinstance(param, DTensor):
+ # Save non-sharded parameters (e.g., running_mean of BatchNorm) as local data
+ cpu_tensor = param.detach().cpu()
+ cpu_sharded_state[param_name] = (cpu_tensor, None)
+ continue
+
+ # Record global sharding rules (take spec of the first DTensor to ensure consistency)
+ if global_spec is None:
+ global_spec = param._spec
+ assert hasattr(global_spec, "device_mesh"), "DTensorSpec must contain 'device_mesh' attribute"
+ assert hasattr(global_spec, "placements"), "DTensorSpec must contain 'placements' attribute"
+
+ # 1. Extract local shard data from the current GPU (_local_tensor)
+ local_gpu_tensor = param._local_tensor # Local shard attribute defined in your DTensor class
+ # 2. Move to CPU memory and detach from computation graph
+ local_cpu_tensor = local_gpu_tensor.detach().cpu()
+ # 3. Save CPU shard + original DTensorSpec (ensure sharding rules remain unchanged)
+ cpu_sharded_state[param_name] = (local_cpu_tensor, param._spec)
+
+ assert global_spec is not None, "No DTensor-type parameters found in the model. FSDP2 sharding may not be enabled."
+ return cpu_sharded_state, global_spec
+
+
+def fsdp2_sharded_load_from_cpu(
+ model: torch.nn.Module,
+ cpu_sharded_state: dict[str, tuple[torch.Tensor, Optional[DTensorSpec]]],
+ target_spec: DTensorSpec,
+) -> None:
+ """
+ Sharded Load: Each process only loads the CPU shard it is responsible for to the GPU,
+ keeping sharding rules unchanged.
+
+ Args:
+ model: FSDP2 model to be restored (must have the same structure as when saved)
+ cpu_sharded_state: Shard data read from CPU memory by the current process
+ (from fsdp2_sharded_save_to_cpu)
+ target_spec: Global DTensorSpec from saving (used to verify sharding rule consistency)
+ """
+ # Verify device_mesh consistency (core: ensure loaded shards map to original GPUs)
+ current_device_mesh = None
+ for param in model.parameters():
+ if isinstance(param, DTensor):
+ current_device_mesh = param._spec.device_mesh
+ break
+ assert current_device_mesh is not None, "DTensor parameters not initialized in the model to be loaded"
+ assert current_device_mesh == target_spec.device_mesh, (
+ f"device_mesh mismatch during loading! Original: {target_spec.device_mesh}, Current: {current_device_mesh}"
+ )
+
+ for param_name, param in model.named_parameters():
+ # Skip parameters not in the saved state (e.g., newly added parameters)
+ if param_name not in cpu_sharded_state:
+ continue
+
+ # Extract CPU shard data and original Spec
+ local_cpu_tensor, saved_spec = cpu_sharded_state[param_name]
+
+ # Handle different parameter types: DTensor sharded parameters vs. regular parameters
+ if isinstance(param, DTensor):
+ # 1. Verify sharding rule consistency (placements must match original Spec)
+ assert saved_spec is not None, f"DTensorSpec missing in saved state for parameter {param_name}"
+ assert saved_spec.placements == target_spec.placements, (
+ f"Sharding strategy mismatch for parameter {param_name} (conflicts with global rules)!"
+ )
+
+ # 2. Move CPU shard data to the current GPU (device of param._local_tensor)
+ target_device = param._local_tensor.device
+ local_gpu_tensor = local_cpu_tensor.to(target_device)
+
+ # 3. Restore to DTensor's local shard (directly copy to _local_tensor, keep spec unchanged)
+ param._local_tensor.copy_(local_gpu_tensor)
+
+ else:
+ # Regular parameters: load directly to original device
+ target_device = param.device
+ param.data.copy_(local_cpu_tensor.to(target_device))
+
+ # Process synchronization: ensure all processes complete loading before proceeding
+ dist.barrier()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fsdp_workers.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fsdp_workers.py
new file mode 100644
index 0000000000000000000000000000000000000000..86d15d63a49136b3ccf3cf37d36ae0df659312bd
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fsdp_workers.py
@@ -0,0 +1,247 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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.
+
+import logging
+import os
+import time
+
+import torch
+import torch.distributed
+from omegaconf import DictConfig
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+from verl.experimental.fully_async_policy.base_detach_sync import BaseDetachNcclSync
+from verl.experimental.fully_async_policy.fsdp2_utils import fsdp2_sharded_load_from_cpu, fsdp2_sharded_save_to_cpu
+from verl.single_controller.base.decorator import Dispatch, register
+from verl.utils.device import (
+ get_device_name,
+ get_torch_device,
+)
+from verl.utils.fsdp_utils import (
+ fsdp_version,
+ load_fsdp_model_to_gpu,
+ offload_fsdp_model_to_cpu,
+)
+from verl.workers.fsdp_workers import AsyncActorRolloutRefWorker, CriticWorker
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+device_name = get_device_name()
+
+__all__ = ["DetachActorWorker", "DetachAsyncRolloutWorker", "CriticWorker"]
+
+
+class DetachNcclSync(BaseDetachNcclSync, AsyncActorRolloutRefWorker):
+ def __init__(self, config: DictConfig, role: str):
+ BaseDetachNcclSync.__init__(self, config, role)
+ AsyncActorRolloutRefWorker.__init__(self, config, role)
+
+ def _get_actor_params(self):
+ pass
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def sync_rollout_weights(self, sync_group_name="actor_rollout"):
+ assert (self._is_actor or self._is_rollout) and not self.config.hybrid_engine
+ assert hasattr(self, "_weights_info") and self._weights_info is not None
+
+ if self._is_actor and self._is_offload_param:
+ load_fsdp_model_to_gpu(self.actor_module_fsdp)
+ params = self._get_actor_params() if self._is_actor else None
+ rollout_name = self.config.rollout.name
+
+ inference_model = None
+ if self._is_rollout:
+ if rollout_name == "vllm":
+ inference_model = BaseDetachNcclSync.get_inference_model(self.rollout)
+
+ from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+
+ patch_vllm_moe_model_weight_loader(inference_model)
+ elif rollout_name == "sglang":
+ inference_model = self.rollout._engine
+ # For ServerAdapter, _engine might be None and needs async initialization
+ if inference_model is None:
+ # Initialize the server adapter engine
+ print("[sync_rollout_weights] Initialize server adapter engine")
+
+ async def init_engine():
+ if hasattr(self.rollout, "_init_server_adapter"):
+ await self.rollout._init_server_adapter()
+ else:
+ print("[sync_rollout_weights] No _init_server_adapter method found")
+ return self.rollout._engine
+
+ inference_model = self._run_async_safely(init_engine())
+ if inference_model is None:
+ raise RuntimeError(
+ f"Failed to initialize rollout engine. "
+ f"rollout type: {type(self.rollout)}, "
+ f"has _init_server_adapter: {hasattr(self.rollout, '_init_server_adapter')}"
+ )
+ else:
+ raise NotImplementedError(f"Unknown rollout name: {rollout_name}")
+
+ if rollout_name == "sglang" and self._is_rollout:
+ self._sync_sglang_weights(inference_model, params, sync_group_name)
+ else:
+ self._sync_vllm_weights(inference_model, params, sync_group_name)
+
+ if self._is_actor and self._is_offload_param:
+ offload_fsdp_model_to_cpu(self.actor_module_fsdp)
+ get_torch_device().empty_cache()
+
+ def cache_actor_weights_to_cpu(self):
+ self.cpu_named_params = {}
+ if self._is_actor:
+ params = self._get_actor_params()
+ local_rank = torch.distributed.get_rank()
+ world_size = torch.distributed.get_world_size()
+
+ for tensor_idx, (key, _, _) in enumerate(self._weights_info):
+ origin_data = params[key]
+ if hasattr(origin_data, "full_tensor"):
+ origin_data = origin_data.full_tensor()
+
+ if tensor_idx % world_size == local_rank:
+ self.cpu_named_params[key] = origin_data.to("cpu", non_blocking=True)
+ get_torch_device().synchronize()
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def sync_rollout_weights_by_checkpoint(self, sync_group_name="actor_rollout"):
+ assert (self._is_actor or self._is_rollout) and not self.config.hybrid_engine
+ assert hasattr(self, "_weights_info") and self._weights_info is not None
+
+ # Load model to GPU
+ load_start_time = time.time()
+ if self._is_actor and self._is_offload_param:
+ load_fsdp_model_to_gpu(self.actor_module_fsdp)
+ load_duration = time.time() - load_start_time
+
+ from ray.util.collective import collective
+
+ # Cache actor weights to CPU and measure the time taken
+ cache_start_time = time.time()
+ self.cache_actor_weights_to_cpu()
+ cache_end_time = time.time()
+ cache_duration = cache_end_time - cache_start_time
+
+ # Register the cached weights into the checkpoint engine
+ self.checkpoint_engine.register_checkpoint(self._weights_info, self.cpu_named_params)
+ register_end_time = time.time()
+ register_duration = register_end_time - cache_end_time
+ self.cpu_named_params = {}
+
+ collective.barrier(group_name=sync_group_name)
+ update_start_time = time.time()
+
+ inference_model = None
+ if self._is_rollout:
+ inference_model = BaseDetachNcclSync.get_inference_model(self.rollout)
+ from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+
+ patch_vllm_moe_model_weight_loader(inference_model)
+
+ # Update the checkpoint with the inference model and broadcast weights
+ self.checkpoint_engine.update_checkpoint(
+ inference_model=inference_model,
+ group_name=sync_group_name,
+ overlap_broadcast_and_consume=self.config.checkpoint_engine.overlap_broadcast_and_consume,
+ )
+
+ update_end_time = time.time()
+ update_duration = update_end_time - update_start_time
+
+ offload_start_time = time.time()
+ if self._is_actor and self._is_offload_param:
+ offload_fsdp_model_to_cpu(self.actor_module_fsdp)
+ offload_duration = time.time() - offload_start_time
+
+ print(
+ f"sync_rollout_weights_by_checkpoint finish!, rank:{torch.distributed.get_rank()},"
+ f" is_actor:{self._is_actor}, is_rollout:{self._is_rollout},"
+ f" total cost:{update_end_time - cache_start_time} seconds, while cache cost {cache_duration} seconds, "
+ f" register cost {register_duration} seconds, update cost {update_duration} seconds"
+ )
+
+ if self._is_actor and self._is_offload_param:
+ print(
+ f"sync_rollout_weights_by_checkpoint load model to gpu cost {load_duration} seconds,"
+ f" offload model to cpu cost {offload_duration} seconds"
+ )
+
+
+class DetachActorWorker(DetachNcclSync):
+ def __init__(self, config: DictConfig, role: str):
+ print("[DetachAsyncRolloutWorker] Initializing via DetachNcclSync...")
+ DetachNcclSync.__init__(self, config, role)
+
+ def _get_actor_params(self):
+ assert self._is_actor
+ params = self.actor_module_fsdp.state_dict()
+ from verl.utils.model import convert_weight_keys
+
+ params = convert_weight_keys(
+ params, getattr(self.actor_module_fsdp, "_fsdp_wrapped_module", self.actor_module_fsdp)
+ )
+ return params
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def get_actor_weights_info(self):
+ assert self._is_actor
+ if hasattr(self, "_weights_info"):
+ return self._weights_info
+ if fsdp_version(self.actor_module_fsdp) == 1:
+ from torch.distributed.fsdp.api import ShardedStateDictConfig, StateDictType
+
+ FSDP.set_state_dict_type(
+ self.actor_module_fsdp,
+ state_dict_type=StateDictType.SHARDED_STATE_DICT,
+ state_dict_config=ShardedStateDictConfig(),
+ )
+ params = self._get_actor_params()
+ ret = []
+ for key, tensor in params.items():
+ ret.append((key, tensor.size(), tensor.dtype))
+ self._weights_info = ret
+ return ret
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def save_model_to_cpu(self, n):
+ if not hasattr(self, "cpu_saved_models"):
+ self.cpu_saved_models = {}
+ self.cpu_saved_models[n] = fsdp2_sharded_save_to_cpu(self.actor_module_fsdp)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def restore_model_from_cpu(self, n):
+ if n in self.cpu_saved_models:
+ cpu_sharded_state, global_spec = self.cpu_saved_models[n]
+ fsdp2_sharded_load_from_cpu(self.actor_module_fsdp, cpu_sharded_state, global_spec)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def clear_cpu_model(self, n):
+ if n in self.cpu_saved_models:
+ del self.cpu_saved_models[n]
+
+
+class DetachAsyncRolloutWorker(DetachNcclSync):
+ def __init__(self, config: DictConfig, role: str):
+ print(f"[DetachAsyncRolloutWorker] {DetachAsyncRolloutWorker.__mro__}")
+ DetachNcclSync.__init__(self, config, role)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def set_actor_weights_info(self, weights_info):
+ assert self._is_rollout
+ self._weights_info = weights_info
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_main.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_main.py
new file mode 100644
index 0000000000000000000000000000000000000000..685af1a2eaae6bf0af9aa318c89513c216fab686
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_main.py
@@ -0,0 +1,312 @@
+# Copyright 2025 Meituan 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.
+
+import asyncio
+import os
+import socket
+import threading
+from pprint import pprint
+
+import hydra
+import ray
+from omegaconf import OmegaConf
+
+from verl.experimental.fully_async_policy.fully_async_rollouter import FullyAsyncRollouter
+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, need_reference_policy
+from verl.utils.fs import copy_to_local
+
+
+def create_resource_pool_manager(config, roles: list) -> ResourcePoolManager:
+ """
+ Create resource pool manager
+
+ Args:
+ config: Configuration object
+ roles: List of roles that need to create resource pools
+
+ Returns:
+ ResourcePoolManager: Resource pool manager
+ """
+ resource_pool_spec = {}
+ mapping = {}
+
+ # Actor/Critic resource pool
+ if any(role in roles for role in [Role.Actor, Role.ActorRollout, Role.Critic, Role.RefPolicy, Role.RewardModel]):
+ assert config.trainer.n_gpus_per_node > 0, "config.trainer.n_gpus_per_node must be greater than 0"
+ assert config.trainer.nnodes > 0, "config.trainer.nnodes must be greater than 0"
+
+ trainer_pool = [config.trainer.n_gpus_per_node] * config.trainer.nnodes
+ resource_pool_spec["trainer_pool"] = trainer_pool
+
+ # Map training-related roles to the same resource pool
+ for role in [Role.Actor, Role.ActorRollout, Role.Critic, Role.RefPolicy, Role.RewardModel]:
+ if role in roles:
+ mapping[role] = "trainer_pool"
+
+ # Rollout resource pool
+ if Role.Rollout in roles:
+ assert config.rollout.n_gpus_per_node > 0, "config.rollout.n_gpus_per_node must be greater than 0"
+ assert config.rollout.nnodes > 0, "config.rollout.nnodes must be greater than 0"
+
+ rollout_pool = [config.rollout.n_gpus_per_node] * config.rollout.nnodes
+ resource_pool_spec["rollout_pool"] = rollout_pool
+ mapping[Role.Rollout] = "rollout_pool"
+
+ return ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
+
+
+def create_role_worker_mapping(config):
+ """
+ Create mapping from roles to worker classes
+
+ Args:
+ config: Configuration object
+
+ Returns:
+ dict: Mapping from roles to worker classes
+ """
+ # Select worker class based on strategy
+ if config.actor_rollout_ref.actor.strategy in ["fsdp", "fsdp2"]:
+ assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
+ from verl.experimental.fully_async_policy.fsdp_workers import (
+ CriticWorker,
+ DetachActorWorker,
+ DetachAsyncRolloutWorker,
+ )
+ from verl.single_controller.ray import RayWorkerGroup
+
+ ray_worker_group_cls = RayWorkerGroup
+
+ elif config.actor_rollout_ref.actor.strategy == "megatron":
+ assert config.critic.strategy == "megatron"
+ from verl.experimental.fully_async_policy.megatron_worker import (
+ CriticWorker,
+ DetachActorWorker,
+ DetachAsyncRolloutWorker,
+ )
+ from verl.single_controller.ray import RayWorkerGroup
+
+ ray_worker_group_cls = RayWorkerGroup
+ else:
+ raise NotImplementedError(f"Unsupported strategy: {config.actor_rollout_ref.actor.strategy}")
+
+ train_role = Role.ActorRollout if config.async_training.use_trainer_do_validate else Role.Actor
+ role_worker_mapping = {
+ train_role: ray.remote(DetachActorWorker),
+ Role.Rollout: ray.remote(DetachAsyncRolloutWorker),
+ Role.Critic: ray.remote(CriticWorker),
+ }
+
+ if config.reward_model.enable:
+ if config.reward_model.strategy in ["fsdp", "fsdp2"]:
+ from verl.workers.fsdp_workers import RewardModelWorker
+ elif config.reward_model.strategy == "megatron":
+ from verl.workers.megatron_workers import RewardModelWorker
+ else:
+ raise NotImplementedError
+
+ role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker)
+
+ # Add reference policy (if KL loss or reward is required)
+ if need_reference_policy(config):
+ role_worker_mapping[Role.RefPolicy] = ray.remote(DetachActorWorker)
+
+ return role_worker_mapping, ray_worker_group_cls
+
+
+@ray.remote(num_cpus=1)
+class FullyAsyncTaskRunner:
+ """
+ Ray remote class for executing distributed PPO training tasks.
+ """
+
+ def __init__(self):
+ self.running = False
+ self.components = {}
+ self.shutdown_event = threading.Event()
+
+ def run(self, config):
+ print("[ASYNC MAIN] Starting fully async PPO training...")
+ self._initialize_components(config)
+ self._run_training_loop()
+
+ def _initialize_components(self, config) -> None:
+ print(f"[ASYNC MAIN] TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
+ pprint(OmegaConf.to_container(config, resolve=True))
+ OmegaConf.resolve(config)
+
+ print("[ASYNC MAIN] Initializing model and tokenizer...")
+ local_path = copy_to_local(
+ config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get("use_shm", False)
+ )
+ from verl.utils 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)
+
+ # Used for multimodal LLM, could be None
+ processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)
+
+ self.components["tokenizer"] = tokenizer
+ self.components["processor"] = processor
+ self.components["config"] = config
+
+ print("[ASYNC MAIN] Creating worker mapping and resource pools...")
+ role_worker_mapping, ray_worker_group_cls = create_role_worker_mapping(config)
+ self.components["role_worker_mapping"] = role_worker_mapping
+ self.components["ray_worker_group_cls"] = ray_worker_group_cls
+
+ print("[ASYNC MAIN] Creating FullyAsyncRollouter...")
+ self._create_rollouter(config)
+
+ print("[ASYNC MAIN] Creating FullyAsyncTrainer...")
+ self._create_trainer(config)
+
+ # sync total_train_steps between rollouter and trainer
+ total_train_steps = ray.get(self.components["rollouter"].get_total_train_steps.remote())
+ print(f"total_train_steps {total_train_steps}")
+ ray.get(self.components["trainer"].set_total_train_steps.remote(total_train_steps))
+
+ # max_queue_size
+ max_queue_size = ray.get(self.components["rollouter"].get_max_queue_size.remote())
+ print(f"[ASYNC MAIN] Creating MessageQueue... max_queue_size {max_queue_size}")
+ message_queue = MessageQueue.remote(config, max_queue_size)
+ message_queue_client = MessageQueueClient(message_queue)
+ self.components["message_queue"] = message_queue
+ self.components["message_queue_client"] = message_queue_client
+
+ ray.get(self.components["rollouter"].set_message_queue_client.remote(self.components["message_queue_client"]))
+ ray.get(self.components["trainer"].set_message_queue_client.remote(self.components["message_queue_client"]))
+
+ print("[ASYNC MAIN] Setting up parameter synchronization...")
+ from verl.experimental.fully_async_policy.param_sync import ParameterSynchronizer
+
+ param_synchronizer = ParameterSynchronizer.remote(
+ config=config,
+ trainer=self.components["trainer"],
+ rollouter=self.components["rollouter"],
+ mq=self.components["message_queue_client"],
+ )
+ ray.get(self.components["trainer"].set_parameter_synchronizer.remote(param_synchronizer))
+
+ # load checkpoint and sync parameter before doing anything
+ val_before_train = config.trainer.get("val_before_train", True)
+ # param_version resume from ckpt or default 0
+ param_version = ray.get(self.components["trainer"].load_checkpoint.remote())
+ ray.get(self.components["rollouter"].load_checkpoint.remote())
+ ray.get(
+ param_synchronizer.sync_weights.remote(
+ version=param_version,
+ validate=val_before_train,
+ use_trainer_do_validate=config.async_training.use_trainer_do_validate,
+ )
+ )
+ ray.get(param_synchronizer.wait_last_valid.remote())
+
+ self.components["param_synchronizer"] = param_synchronizer
+ print("[ASYNC MAIN] All components initialized successfully")
+
+ def _create_rollouter(self, config) -> None:
+ rollouter = FullyAsyncRollouter.remote(
+ config=config,
+ tokenizer=self.components["tokenizer"],
+ role_worker_mapping={Role.Rollout: self.components["role_worker_mapping"][Role.Rollout]},
+ resource_pool_manager=create_resource_pool_manager(config, roles=[Role.Rollout]),
+ ray_worker_group_cls=self.components["ray_worker_group_cls"],
+ processor=self.components["processor"],
+ device_name=config.trainer.device,
+ )
+
+ ray.get(rollouter.init_workers.remote())
+ ray.get(rollouter.set_max_required_samples.remote())
+
+ self.components["rollouter"] = rollouter
+ print("[ASYNC MAIN] Rollouter created and initialized successfully")
+
+ def _create_trainer(self, config) -> None:
+ trainer_role_mapping = {
+ role: worker_cls
+ for role, worker_cls in self.components["role_worker_mapping"].items()
+ if role != Role.Rollout
+ }
+
+ trainer = FullyAsyncTrainer.remote(
+ config=config,
+ tokenizer=self.components["tokenizer"],
+ role_worker_mapping=trainer_role_mapping,
+ resource_pool_manager=create_resource_pool_manager(config, roles=list(trainer_role_mapping.keys())),
+ ray_worker_group_cls=self.components["ray_worker_group_cls"],
+ processor=self.components["processor"],
+ device_name=config.trainer.device,
+ )
+
+ ray.get(trainer.init_workers.remote())
+ self.components["trainer"] = trainer
+ print("[ASYNC MAIN] FullyAsyncTrainer created and initialized successfully")
+
+ def _run_training_loop(self):
+ self.running = True
+
+ print("[ASYNC MAIN] Starting Rollouter and Trainer...")
+ rollouter_future = self.components["rollouter"].fit.remote()
+ trainer_future = self.components["trainer"].fit.remote()
+
+ futures = [rollouter_future, trainer_future]
+
+ try:
+ while futures:
+ # Use ray.wait to monitor all futures and return when any one is completed.
+ done_futures, remaining_futures = ray.wait(futures, num_returns=1, timeout=None)
+
+ for future in done_futures:
+ try:
+ ray.get(future)
+ print("[ASYNC MAIN] One component completed successfully")
+ except Exception as e:
+ print(f"[ASYNC MAIN] Component failed with error: {e}")
+ for remaining_future in remaining_futures:
+ ray.cancel(remaining_future)
+ raise e
+
+ futures = remaining_futures
+
+ except Exception as e:
+ print(f"[ASYNC MAIN] Training failed: {e}")
+ for future in futures:
+ ray.cancel(future)
+ raise
+ finally:
+ asyncio.run(self.components["message_queue_client"].clear_queue())
+ print("[ASYNC MAIN] Training completed or interrupted")
+
+
+@hydra.main(config_path="config", config_name="fully_async_ppo_trainer", version_base=None)
+def main(config):
+ from verl.trainer.main_ppo import run_ppo
+
+ # Ensure async training config exists
+ if not hasattr(config, "async_training"):
+ raise RuntimeError("must set async_training config")
+ from time import time
+
+ start_time = time()
+ run_ppo(config, task_runner_class=FullyAsyncTaskRunner)
+ print(f"total time: {time() - start_time:.2f} seconds")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_rollouter.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_rollouter.py
new file mode 100644
index 0000000000000000000000000000000000000000..757432f4cf06f22fd91544503ec78442354f0cf9
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_rollouter.py
@@ -0,0 +1,793 @@
+# Copyright 2025 Meituan 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.
+
+import asyncio
+import functools
+import multiprocessing
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pprint import pformat
+
+import numpy as np
+import ray
+import torch
+from ray import ObjectRef
+
+from verl.experimental.fully_async_policy.detach_utils import (
+ RolloutSample,
+ ValidateMetrics,
+ prepare_single_generation_data,
+)
+from verl.experimental.fully_async_policy.message_queue import MessageQueueClient
+from verl.experimental.fully_async_policy.ray_trainer import FullyAsyncRayPPOTrainer
+from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup
+from verl.trainer.ppo.ray_trainer import ResourcePoolManager
+from verl.trainer.ppo.reward import load_reward_manager
+from verl.trainer.ppo.utils import Role, WorkerType
+from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path
+from verl.utils.profiler import marked_timer
+from verl.utils.tracking import ValidationGenerationsLogger
+
+
+@ray.remote(num_cpus=10, max_concurrency=100)
+class FullyAsyncRollouter(FullyAsyncRayPPOTrainer):
+ """
+ Asynchronous sample generator, responsible for continuously generating training samples
+ and putting them into MessageQueue
+ Based on the mature implementation improvements of OneStepOffRayTrainer
+ """
+
+ def __init__(
+ self,
+ config,
+ tokenizer,
+ role_worker_mapping: dict[Role, WorkerType],
+ resource_pool_manager: ResourcePoolManager,
+ ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,
+ processor=None,
+ reward_fn=None,
+ val_reward_fn=None,
+ device_name=None,
+ ):
+ # Store the tokenizer for text processing
+ self.tokenizer = tokenizer
+ self.processor = processor
+ self.config = config
+ self.reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})
+ )
+ self.val_reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {})
+ )
+ self.hybrid_engine = config.actor_rollout_ref.hybrid_engine
+
+ assert not self.hybrid_engine
+ assert self.config.data.train_batch_size == 0, "train_batch_size must be zero"
+ assert self.config.data.gen_batch_size == 1, "gen_batch_size must be one"
+ assert self.config.async_training.staleness_threshold >= 0, "staleness_threshold must larger than 0"
+ assert self.config.async_training.trigger_parameter_sync_step >= 1, (
+ "trigger_parameter_sync_step must larger than 1"
+ )
+
+ self.role_worker_mapping = role_worker_mapping
+ self.resource_pool_manager = resource_pool_manager
+ self.ray_worker_group_cls = ray_worker_group_cls
+ self.device_name = device_name if device_name else self.config.trainer.device
+ self.validation_generations_logger = ValidationGenerationsLogger(
+ project_name=self.config.trainer.project_name,
+ experiment_name=self.config.trainer.experiment_name,
+ )
+
+ self.ref_in_actor = False
+ self.kl_ctrl_in_reward = False
+ self.use_critic = False
+ self.use_reference_policy = False
+ self.use_rm = False
+
+ print("[FullyAsyncRollouter] Creating datasets...")
+ from verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler
+ from verl.utils.dataset.rl_dataset import collate_fn
+
+ train_dataset = create_rl_dataset(config.data.train_files, config.data, tokenizer, processor)
+ val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor)
+ train_sampler = create_rl_sampler(config.data, train_dataset)
+
+ self._validate_config()
+ if self.config.async_training.use_trainer_do_validate:
+ rollout_gpus = config.rollout.nnodes * config.rollout.n_gpus_per_node
+ train_gpus = config.trainer.nnodes * config.trainer.n_gpus_per_node
+ total_gpus = rollout_gpus + train_gpus
+ print(f"[FullyAsyncRollouter] split before val_dataset total len: {len(val_dataset)}")
+ split_dataset = val_dataset.split(total_gpus)
+ rollout_val_dataset0 = split_dataset[:rollout_gpus]
+ from torch.utils.data import ConcatDataset
+
+ val_dataset = ConcatDataset(rollout_val_dataset0)
+ print(f"[FullyAsyncRollouter] split after val_dataset total len: {len(val_dataset)}")
+ print(f"[FullyAsyncRollouter] Rollouter _create_dataloader...\n{train_dataset}\n{val_dataset}")
+
+ self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)
+
+ # ==================== fully async config ====================
+
+ self.total_rollout_steps = len(self.train_dataloader) * self.config.trainer.total_epochs
+ if self.config.rollout.total_rollout_steps is not None:
+ self.total_rollout_steps = min(self.config.rollout.total_rollout_steps, self.total_rollout_steps)
+ print(f"[FullyAsyncRollouter] Total rollout steps: {self.total_rollout_steps}")
+ self.total_train_steps = None
+
+ # Rollouter parameter configuration
+ self.message_queue_client = None
+
+ # Worker groups: rollout_wg is same to actor_rollout_wg
+ self.rollout_wg = None
+ self.actor_rollout_wg = None
+ self.async_rollout_manager = None
+
+ # Config
+ self.staleness_threshold: float = config.async_training.get("staleness_threshold", 1)
+ # required_samples use ppo_mini_batch_size*require_batches as the minimum number of samples.
+ self.require_batches = config.async_training.require_batches
+ self.required_samples = config.actor_rollout_ref.actor.ppo_mini_batch_size * self.require_batches
+ self.max_required_samples = None
+ self.max_concurrent_samples = None
+ # queue size
+ self.max_queue_size = None
+
+ # Statistics
+ self.current_param_version = 0
+ self.total_generated_samples = 0
+ self.staleness_samples = 0
+ self.dropped_stale_samples = 0
+ self.processed_sample_count = 0
+ # we start from step 1
+ self.global_steps = 1
+ self.idle_start_time = None
+ self.version_start_time = None
+
+ # Concurrency control
+ # Modified by self.pause() or self._should_pause_generation()
+ self.paused = False
+ self.running = True
+ self.monitor_loop_trigger = True
+
+ # Add dataloader lock
+ self.dataloader_lock = asyncio.Lock()
+
+ # Initialize async queues
+ self.pending_queue = asyncio.Queue(maxsize=128)
+ self.active_tasks = set()
+ self.cancel_queue = asyncio.Queue()
+
+ cpu_cores = multiprocessing.cpu_count()
+ # cpu case use cpu_cores; io case use cpu_cores*2
+ self.validate_executor = ThreadPoolExecutor(max_workers=cpu_cores)
+ self.parallel_validate_and_rollout = config.async_training.get("parallel_validate_and_rollout", False)
+ self.validate_task = None
+
+ def _init_async_objects(self):
+ # Initialize asyncio synchronization primitives.
+ # We let asyncio.Condition create the Lock internally to ensure they share the same Event Loop.
+ # This avoids 'ValueError: loop argument must agree with lock' which can occur in Ray environments
+ # where the lock's captured loop (get_running_loop) differs from Condition's default loop check.
+ # Explicitly passing the loop is deprecated/removed in Python 3.10+, so this reverse-initialization
+ # is the most robust workaround.
+ self.condition = asyncio.Condition()
+ self.lock = self.condition._lock
+
+ async def set_message_queue_client(self, message_queue_client: MessageQueueClient):
+ """Set message queue client"""
+ async with self.lock:
+ self.message_queue_client = message_queue_client
+
+ async def set_max_required_samples(self):
+ async with self.lock:
+ self.max_required_samples = int(
+ self.required_samples
+ * (self.staleness_threshold + 1)
+ * self.config.async_training.trigger_parameter_sync_step
+ )
+ self.total_train_steps = int(
+ self.total_rollout_steps
+ / (self.required_samples * self.config.async_training.trigger_parameter_sync_step)
+ )
+
+ self.max_concurrent_samples = len(self.async_rollout_manager.server_handles) * 16
+ self.max_concurrent_samples = min(self.max_concurrent_samples, self.max_required_samples)
+ self.max_queue_size = self.max_required_samples
+
+ print(
+ f"[FullyAsyncRollouter] required_samples : {self.required_samples} "
+ f"max_required_samples: {self.max_required_samples} "
+ f"max_queue_size: {self.max_queue_size} "
+ f"total_train_steps: {self.total_train_steps} "
+ f"total_rollout_steps: {self.total_rollout_steps} "
+ f"max_concurrent_samples: {self.max_concurrent_samples} "
+ )
+
+ def get_rollout_wg(self):
+ """Get rollout worker group"""
+ return self.rollout_wg
+
+ def get_max_queue_size(self):
+ return self.max_queue_size
+
+ def get_total_train_steps(self):
+ return self.total_train_steps
+
+ async def update_param_version(
+ self, version: int, validate: bool = False, global_steps: int = 0, use_trainer_do_validate: bool = False
+ ):
+ """Update current parameter version"""
+ async with self.lock:
+ old_version = self.current_param_version
+ self.current_param_version = version
+ # every time param change, reset staleness_samples
+ self.staleness_samples = (
+ len(self.active_tasks) + self.cancel_queue.qsize() + await self.message_queue_client.get_queue_size()
+ )
+ timing_raw = {}
+ idle_ratio = None
+ if self.idle_start_time is not None and self.version_start_time is not None:
+ rollout_active_time = self.idle_start_time - self.version_start_time
+ rollout_version_time = time.time() - self.version_start_time
+ idle_ratio = 1 - rollout_active_time / rollout_version_time
+ timing_raw["rollouter/active_time"] = rollout_active_time
+ timing_raw["rollouter/version_time"] = rollout_version_time
+ timing_raw["rollouter/idle_ratio"] = idle_ratio
+ self.idle_start_time = None
+ print(
+ f"[FullyAsyncRollouter][Public][update_param_version] "
+ f"Parameter version updated from {old_version} to {version} "
+ f",reset staleness_samples to: {self.staleness_samples}"
+ f",idle_ratio: {idle_ratio}"
+ )
+ need_validate = (
+ (
+ self.val_reward_fn is not None
+ and self.config.rollout.test_freq > 0
+ and self.current_param_version % self.config.rollout.test_freq == 0
+ and self.current_param_version > 0
+ ) # don't test here in the initial parameter sync
+ or (validate and self.val_reward_fn is not None)
+ )
+ print(
+ f"[FullyAsyncRollouter] need_validate: {need_validate},"
+ f"parallel_validate_and_rollout: {self.parallel_validate_and_rollout}"
+ )
+ if not need_validate:
+ data = ValidateMetrics(
+ timing_raw=timing_raw, metrics=None, global_steps=global_steps, param_version=version
+ )
+ elif need_validate and not self.parallel_validate_and_rollout:
+ data = self._validate_wrapper(timing_raw, version, global_steps, use_trainer_do_validate)
+
+ if not need_validate or not self.parallel_validate_and_rollout:
+ await self.message_queue_client.put_validate(ray.cloudpickle.dumps(data))
+
+ self.version_start_time = time.time()
+
+ if need_validate and self.parallel_validate_and_rollout:
+ if self.validate_task and not self.validate_task.done():
+ print("[FullyAsyncRollouter] validate_task is running, wait last validate_task to finish")
+ self.validate_task.get()
+ self.validate_task = asyncio.create_task(
+ self.do_validate_async(timing_raw, version, global_steps, use_trainer_do_validate)
+ )
+
+ def _validate_wrapper(
+ self, timing_raw: dict, version: int, global_steps: int = 0, use_trainer_do_validate: bool = False
+ ):
+ val_metrics = None
+ with marked_timer("rollouter/validate_time", timing_raw, color="green"):
+ val_metrics: dict = self._validate(use_trainer_do_validate)
+ data = ValidateMetrics(
+ timing_raw=timing_raw, metrics=val_metrics, global_steps=global_steps, param_version=version
+ )
+ return data
+
+ async def do_validate_async(
+ self,
+ timing_raw: dict,
+ version: int,
+ global_steps: int = 0,
+ use_trainer_do_validate: bool = False,
+ ):
+ loop = asyncio.get_running_loop()
+
+ data = await loop.run_in_executor(
+ self.validate_executor,
+ functools.partial(
+ self._validate_wrapper,
+ timing_raw=timing_raw,
+ version=version,
+ global_steps=global_steps,
+ use_trainer_do_validate=use_trainer_do_validate,
+ ),
+ )
+ await self.message_queue_client.put_validate(ray.cloudpickle.dumps(data))
+
+ async def save_checkpoint(self, local_global_step_folder: str):
+ # WARNING!: Due to the asynchronous nature, there are some in-flight samples
+ # (pending/cancel/result queue and message queue).
+ # Therefore, directly saving the state of the dataloader will result in losing these
+ # samples when resuming training.
+ # TODO: Implement dataloader recovery without losing in-flight samples.
+ from verl.utils.fs import local_mkdir_safe
+
+ # save dataloader
+ local_mkdir_safe(local_global_step_folder)
+ dataloader_local_path = os.path.join(local_global_step_folder, "data.pt")
+ async with self.dataloader_lock:
+ dataloader_state_dict = self.train_dataloader.state_dict()
+ torch.save(dataloader_state_dict, dataloader_local_path)
+ print(f"[FullyAsyncRollouter] Saved dataloader checkpoint to {dataloader_local_path}")
+
+ def load_checkpoint(self):
+ """Load checkpoint including dataloader state based on resume mode"""
+
+ if self.config.trainer.resume_mode == "disable":
+ print("[FullyAsyncRollouter] Resume mode is disabled, starting from scratch")
+ return 0
+
+ # Determine checkpoint folder path
+ if self.config.trainer.default_hdfs_dir is not None:
+ raise NotImplementedError("[FullyAsyncRollouter] Load from hdfs is not implemented yet")
+ else:
+ checkpoint_folder = self.config.trainer.default_local_dir
+ if not os.path.isabs(checkpoint_folder):
+ working_dir = os.getcwd()
+ checkpoint_folder = os.path.join(working_dir, checkpoint_folder)
+
+ global_step_folder = find_latest_ckpt_path(checkpoint_folder)
+
+ # Find and validate global_step_folder based on resume mode
+ if self.config.trainer.resume_mode == "auto":
+ if global_step_folder is None:
+ print("[FullyAsyncRollouter] Training from scratch (no checkpoint found)")
+ return 0
+ elif self.config.trainer.resume_mode == "resume_path":
+ assert isinstance(self.config.trainer.resume_from_path, str), (
+ "[FullyAsyncRollouter] resume_from_path must be str type"
+ )
+ assert "global_step_" in self.config.trainer.resume_from_path, (
+ "[FullyAsyncRollouter] resume_from_path must specify the global_steps"
+ )
+ global_step_folder = self.config.trainer.resume_from_path
+ if not os.path.isabs(global_step_folder):
+ working_dir = os.getcwd()
+ global_step_folder = os.path.join(working_dir, global_step_folder)
+ else:
+ raise ValueError(f"[FullyAsyncRollouter] Unknown resume_mode: {self.config.trainer.resume_mode}")
+
+ print(f"[FullyAsyncRollouter] Loading checkpoint from: {global_step_folder}")
+
+ # Extract and set global step
+ trainer_global_steps = int(global_step_folder.split("global_step_")[-1])
+ self.global_steps = (
+ trainer_global_steps * self.required_samples * self.config.async_training.trigger_parameter_sync_step + 1
+ )
+ print(f"[FullyAsyncRollouter] Setting global_steps to {self.global_steps}")
+
+ # Load dataloader state
+ dataloader_local_path = os.path.join(global_step_folder, "data.pt")
+ if os.path.exists(dataloader_local_path):
+ dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False)
+ self.train_dataloader.load_state_dict(dataloader_state_dict)
+ print(f"[FullyAsyncRollouter] Loaded dataloader state from {dataloader_local_path}")
+ else:
+ print(
+ f"[FullyAsyncRollouter] Warning: No dataloader state found at {dataloader_local_path}, "
+ f"will start from scratch"
+ )
+
+ def _validate_config(self):
+ # Validate asynchronous training configuration
+ if not hasattr(self.config, "async_training"):
+ raise ValueError("[FullyAsyncRollouter] Missing async_training configuration")
+ assert self.config.actor_rollout_ref.rollout.calculate_log_probs, "must rollout calculate log_probs"
+
+ async def init_workers(self):
+ """Initialize distributed training workers using Ray backend.
+
+ Creates:
+ 1. Ray resource pools from configuration
+ 2. Worker groups for each role (actor, critic, etc.)
+ """
+ self._init_async_objects()
+ self._init_resource_pools()
+ self._create_worker_classes()
+ self._init_worker_groups()
+ self._init_models()
+ await self._init_async_rollout_manager()
+
+ def _create_actor_rollout_classes(self):
+ # only create rollout
+ for role in [Role.Rollout]:
+ resource_pool = self.resource_pool_manager.get_resource_pool(role)
+ role_cls = RayClassWithInitArgs(
+ cls=self.role_worker_mapping[role],
+ config=self.config.actor_rollout_ref,
+ role=str(role),
+ )
+ self.resource_pool_to_cls[resource_pool][str(role)] = role_cls
+
+ def _init_models(self):
+ self.rollout_wg = self.all_wg[str(Role.Rollout)]
+ self.rollout_wg.init_model()
+ self.actor_rollout_wg = self.rollout_wg
+
+ def _create_continuous_iterator(self):
+ """
+ Create a continuous data iterator across epoch
+ """
+ for epoch in range(self.config.rollout.total_epochs):
+ iterator = iter(self.train_dataloader)
+ for batch_dict in iterator:
+ yield epoch, batch_dict
+
+ async def _init_async_rollout_manager(self):
+ # create async rollout manager and request scheduler
+ assert self.config.actor_rollout_ref.rollout.mode == "async"
+ from verl.experimental.fully_async_policy.agent_loop import FullyAsyncAgentLoopManager
+
+ self.async_rollout_mode = True
+ self.async_rollout_manager = await FullyAsyncAgentLoopManager.create(
+ config=self.config,
+ worker_group=self.rollout_wg,
+ )
+
+ # Add samples to the pending_queue
+ async def _feed_samples(self):
+ continuous_iterator = self._create_continuous_iterator()
+
+ for epoch, batch_dict in continuous_iterator:
+ # Similar to _prepare_generate_batch: Separate data
+ full_batch = prepare_single_generation_data(batch_dict, self.config)
+
+ sample_id = f"sample_{epoch}_{self.global_steps}"
+
+ rollout_sample = RolloutSample(
+ full_batch=full_batch,
+ agent_loop_output_list=[None] * self.config.actor_rollout_ref.rollout.n,
+ sample_id=sample_id,
+ epoch=epoch,
+ param_version=0,
+ param_version_start=[],
+ param_version_end=[],
+ processing_times=[],
+ tool_calls=[],
+ rollout_status={},
+ )
+
+ await self.pending_queue.put(rollout_sample)
+
+ # Check if have reached the last step
+ if self.global_steps >= self.total_rollout_steps:
+ print(
+ f"[FullyAsyncRollouter][Feed] "
+ f"Maximum count has been reached, stop adding new samples"
+ f"{self.global_steps} >= {self.total_rollout_steps}"
+ )
+ break
+
+ self.global_steps += 1
+
+ # End signal
+ await self.pending_queue.put("DONE")
+ print(f"[FullyAsyncRollouter][Feed] Sample addition is complete, {self.global_steps} samples have been added")
+
+ async def _processor_worker(self):
+ """
+ Streaming worker coroutines, a sample is submitted for processing without waiting for batches
+ """
+ while True:
+ if self.paused or await self._should_pause_generation():
+ print(
+ "[FullyAsyncRollouter][Processor] Received pause signal, waiting for remaining tasks to return..."
+ )
+ async with self.lock:
+ self.paused = True
+ while self.active_tasks:
+ async with self.lock:
+ # After acquiring the lock, the number of active_tasks may change, need to be verified again
+ if self.active_tasks:
+ done_tasks, self.active_tasks = await asyncio.wait(
+ self.active_tasks, return_when=asyncio.FIRST_COMPLETED
+ )
+ for task in done_tasks:
+ await task
+
+ async with self.lock:
+ while self.paused:
+ self.idle_start_time = time.time()
+ await self.condition.wait()
+ continue
+
+ simple_from_cancel_queue = False
+ if not self.cancel_queue.empty():
+ rollout_sample = await self.cancel_queue.get()
+ simple_from_cancel_queue = True
+ else:
+ rollout_sample = await self.pending_queue.get()
+ self.staleness_samples += 1
+
+ if rollout_sample == "DONE":
+ print(
+ "[FullyAsyncRollouter][Processor] Received end signal, waiting for remaining tasks to complete..."
+ )
+ while self.active_tasks:
+ async with self.lock:
+ if self.active_tasks:
+ done_tasks, self.active_tasks = await asyncio.wait(
+ self.active_tasks, return_when=asyncio.FIRST_COMPLETED
+ )
+ for task in done_tasks:
+ await task
+ break
+
+ # Check whether the number of concurrent tasks exceeds the limit
+ while len(self.active_tasks) >= self.max_concurrent_samples:
+ async with self.lock:
+ if self.active_tasks:
+ done_tasks, self.active_tasks = await asyncio.wait(
+ self.active_tasks, return_when=asyncio.FIRST_COMPLETED
+ )
+ for task in done_tasks:
+ await task
+
+ # Submit single sample processing
+ async with self.lock:
+ # After the pause is over, the lock is acquired and it is necessary
+ # to determine whether it is the pause phase, otherwise continue to wait
+ while self.paused:
+ await self.condition.wait()
+ task = asyncio.create_task(
+ self._process_single_sample_streaming(rollout_sample),
+ name=rollout_sample.sample_id,
+ )
+ self.active_tasks.add(task)
+
+ if simple_from_cancel_queue:
+ self.cancel_queue.task_done()
+ else:
+ self.pending_queue.task_done()
+
+ async def _process_single_sample_streaming(self, rollout_sample: RolloutSample):
+ """Process a single sample streamingly"""
+ # Calling asynchronous generation methods
+ rollout_sample.full_batch.non_tensor_batch["param_version"] = [self.current_param_version] * len(
+ rollout_sample.full_batch
+ )
+ ret, is_cancel = await self.async_rollout_manager.generate_single_sample_async(
+ rollout_sample.full_batch, rollout_sample.agent_loop_output_list
+ )
+ if not is_cancel:
+ rollout_sample.full_batch = ret
+ rollout_sample.full_batch.non_tensor_batch["uid"] = np.array(
+ [f"uid_{rollout_sample.sample_id}"] * len(rollout_sample.full_batch), dtype=object
+ )
+ rollout_sample.param_version = self.current_param_version
+ rollout_sample.rollout_status = await self.get_statistics()
+ rollout_sample.agent_loop_output_list = []
+
+ success = await self.message_queue_client.put_sample(
+ sample=ray.cloudpickle.dumps(rollout_sample),
+ param_version=rollout_sample.param_version,
+ )
+ if success:
+ self.total_generated_samples += 1
+ else:
+ self.dropped_stale_samples += 1
+ else:
+ rollout_sample.agent_loop_output_list = ret
+ await self.cancel_queue.put(rollout_sample)
+
+ self.processed_sample_count += 1
+
+ async def _streaming_generation_main(self):
+ """The main entry method for stream processing"""
+
+ if self.async_rollout_manager is None:
+ await self._init_async_rollout_manager()
+
+ # Start the streaming loop
+ print(f"[FullyAsyncRollouter] Start streaming mode, maximum concurrent samples: {self.max_concurrent_samples}")
+
+ # Start sample feed coroutine, streaming process coroutine
+ self.feed_task = asyncio.create_task(self._feed_samples())
+ self.processor_task = asyncio.create_task(self._processor_worker())
+
+ try:
+ # Wait for sample feed to complete
+ # Use asyncio.wait to monitor all tasks. If processor exits early,
+ # detect it instead of blocking on feed_task (it might be stuck on a full queue).
+ done, pending = await asyncio.wait(
+ [self.feed_task, self.processor_task], return_when=asyncio.FIRST_COMPLETED
+ )
+
+ for task in done:
+ if task.exception():
+ raise task.exception()
+
+ if self.feed_task not in done:
+ raise RuntimeError("Processor task exited prematurely")
+
+ print("[FullyAsyncRollouter] Sample feed completed")
+
+ # Wait for streaming to complete
+ await self.processor_task
+ print("[FullyAsyncRollouter] Streaming process completed")
+
+ except Exception as e:
+ print(f"[FullyAsyncRollouter] Streaming process exception:{e}")
+
+ finally:
+ if self.processor_task:
+ self.processor_task.cancel()
+
+ await asyncio.gather(self.processor_task, return_exceptions=True)
+
+ # Send a finish signal
+ await self.message_queue_client.put_sample(
+ sample=None,
+ param_version=self.current_param_version,
+ )
+
+ async with self.lock:
+ self.running = False
+
+ async def fit(self):
+ """
+ Start the async rollouter - entry point that sets up and runs async tasks
+ Main async fit method that coordinates all coroutines
+ """
+
+ print("[FullyAsyncRollouter] Starting FullyAsyncRollouter...")
+
+ if self.message_queue_client is None:
+ raise ValueError("MessageQueue client not set. Call set_message_queue_client() first.")
+
+ # Set the running status flag
+ async with self.lock:
+ self.paused = False
+ self.running = True
+
+ # Create the main asynchronous task
+ generation_task = asyncio.create_task(self._streaming_generation_main())
+ monitor_task = asyncio.create_task(self._async_monitor_loop())
+
+ try:
+ # Run build and monitoring tasks concurrently
+ await asyncio.gather(generation_task, monitor_task, return_exceptions=True)
+ except Exception as e:
+ print(f"[FullyAsyncRollouter] Asynchronous task execution error: {e}")
+ finally:
+ if not generation_task.done():
+ generation_task.cancel()
+ if not monitor_task.done():
+ monitor_task.cancel()
+
+ # Wait for the task to complete
+ await asyncio.gather(generation_task, monitor_task, return_exceptions=True)
+
+ print("[FullyAsyncRollouter] Rollouter fit completed")
+
+ async def _async_monitor_loop(self):
+ """
+ Async coroutine for monitoring:
+ Function 1: Log information output
+ Function 2: Trigger rollout recovery
+ """
+ last_stats_time = time.time()
+ stats_interval = 60.0
+ check_interval = 10.0
+
+ while True:
+ async with self.lock:
+ if not self.running:
+ break
+ await asyncio.sleep(check_interval)
+ # Print statistics periodically
+ current_time = time.time()
+ if current_time - last_stats_time >= stats_interval:
+ stats = await self.get_statistics()
+ print(f"[FullyAsyncRollouter][MonitorLoop][Statistics] {pformat(stats)}")
+ last_stats_time = current_time
+
+ # Trigger rollout recovery
+ if self.monitor_loop_trigger:
+ if not await self._should_pause_generation():
+ async with self.lock:
+ self.paused = False
+ self.condition.notify_all()
+
+ async def _should_pause_generation(self) -> bool:
+ """Determine whether the build should be paused"""
+ queue_stats = self.message_queue_client.get_statistics_sync()
+ queue_size = queue_stats["queue_size"]
+
+ if queue_size >= self.max_queue_size:
+ if not self.paused:
+ print(
+ f"[FullyAsyncRollouter][ShouldPause] "
+ f"due to full queue: size={queue_size}, max={self.max_queue_size}"
+ )
+ return True
+
+ if self.staleness_samples >= self.max_required_samples:
+ if not self.paused:
+ print(
+ "[FullyAsyncRollouter][ShouldPause] "
+ f"due to "
+ f"staleness_samples {self.staleness_samples} >= max_required_samples {self.max_required_samples} "
+ )
+ return True
+
+ return False
+
+ async def pause(self):
+ """pause rollout"""
+ print("[FullyAsyncRollouter][Public][Pause] partial rollout:", self.config.async_training.partial_rollout)
+ async with self.lock:
+ self.paused = True
+ # Cancel all rollout tasks
+ if self.config.async_training.partial_rollout:
+ await self.async_rollout_manager.cancel()
+ print("[FullyAsyncRollouter][Public][Pause] Unfinished rollout tasks canceled")
+ if self.active_tasks:
+ await asyncio.gather(*self.active_tasks, return_exceptions=True)
+ self.active_tasks.clear()
+ print("[FullyAsyncRollouter][Public][Pause] All active tasks completed")
+ print("[FullyAsyncRollouter][Public][Pause] Prefix cache reset")
+ # Always clear KV cache to release GPU memory during weight synchronization,
+ # regardless of partial_rollout setting.
+ await self.async_rollout_manager.clear_kv_cache()
+ self.monitor_loop_trigger = False
+
+ async def resume(self, dependency_ref: ObjectRef = None):
+ if dependency_ref is not None:
+ ray.get(dependency_ref)
+ print("[FullyAsyncRollouter][Public][Resume]")
+ async with self.lock:
+ if self.config.async_training.partial_rollout:
+ await self.async_rollout_manager.resume()
+ self.paused = False
+ self.monitor_loop_trigger = True
+ self.condition.notify_all()
+
+ async def get_statistics(self) -> dict:
+ queue_stats = self.message_queue_client.get_statistics_sync()
+
+ stats = {
+ # monitor stats
+ "monitor/active_tasks_size": len(self.active_tasks),
+ "monitor/queue/pending_queue_size": self.pending_queue.qsize(),
+ "monitor/queue/cancel_queue_size": self.cancel_queue.qsize(),
+ "monitor/queue/mq_queue_size": queue_stats["queue_size"],
+ # counting stats
+ "count/current_param_version": self.current_param_version,
+ "count/total_generated_samples": self.total_generated_samples,
+ "count/staleness_samples": self.staleness_samples,
+ "count/dropped_stale_samples": self.dropped_stale_samples,
+ # static stats
+ "static/max_required_samples": self.max_required_samples,
+ "static/required_samples": self.required_samples,
+ "static/staleness_threshold": self.staleness_threshold,
+ "static/max_queue_size": self.max_queue_size,
+ "static/max_concurrent_samples": self.max_concurrent_samples,
+ }
+
+ return stats
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_trainer.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb2721854230c187f70658d3946eb5f19b57b2c3
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/fully_async_trainer.py
@@ -0,0 +1,612 @@
+# Copyright 2025 Meituan 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.
+
+import os
+import time
+from datetime import datetime
+from pprint import pprint
+from typing import Any
+
+import ray
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from verl.experimental.fully_async_policy.detach_utils import (
+ MetricsAggregator,
+ ValidateMetrics,
+ assemble_batch_from_rollout_samples,
+)
+from verl.experimental.fully_async_policy.message_queue import MessageQueueClient
+from verl.experimental.fully_async_policy.ray_trainer import FullyAsyncRayPPOTrainer
+from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup
+from verl.trainer.ppo import core_algos
+from verl.trainer.ppo.ray_trainer import ResourcePoolManager
+from verl.trainer.ppo.reward import load_reward_manager
+from verl.trainer.ppo.utils import Role, WorkerType, need_critic, need_reference_policy, need_reward_model
+from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi
+from verl.utils.debug import marked_timer
+
+
+@ray.remote(num_cpus=10)
+class FullyAsyncTrainer(FullyAsyncRayPPOTrainer):
+ """
+ A fully asynchronous PPO trainer that obtains samples from a MessageQueue for training.
+ Based on an improved implementation of OneStepOffRayTrainer
+ """
+
+ def __init__(
+ self,
+ config,
+ tokenizer,
+ role_worker_mapping: dict[Role, WorkerType],
+ resource_pool_manager: ResourcePoolManager,
+ ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,
+ processor=None,
+ reward_fn=None,
+ val_reward_fn=None,
+ device_name=None,
+ ):
+ # Store the tokenizer for text processing
+ self.tokenizer = tokenizer
+ self.processor = processor
+ self.config = config
+ self.reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})
+ )
+ self.val_reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {})
+ )
+
+ self.hybrid_engine = config.actor_rollout_ref.hybrid_engine
+ assert not self.hybrid_engine
+
+ self.role_worker_mapping = role_worker_mapping
+ self.resource_pool_manager = resource_pool_manager
+ 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
+ self.device_name = device_name if device_name else self.config.trainer.device
+
+ 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)
+ # if ref_in_actor is True, the reference policy will be actor without lora applied
+ self.ref_in_actor = lora_rank > 0
+
+ # define in-reward KL control
+ # kl loss control currently not suppoorted
+ if self.config.algorithm.use_kl_in_reward:
+ self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl)
+
+ # ==================== fully async config ====================
+
+ self.message_queue_client = None
+ self.param_synchronizer = None
+
+ # Statistics
+ # we start from step 1
+ self.global_steps = 1
+ self.local_trigger_step = 1
+ self.processed_samples = 0
+ self.stale_samples_processed = 0
+ self.stale_trajectory_processed = 0
+ self.current_param_version = 0
+ self.total_train_steps = None
+ self.progress_bar = None
+ self.trigger_parameter_sync_step = config.async_training.trigger_parameter_sync_step
+ self.last_ckpt_version = 0
+ self.train_val_metrics = None
+ self.train_role = Role.ActorRollout if config.async_training.use_trainer_do_validate else Role.Actor
+
+ # required_samples use ppo_mini_batch_size*require_batches as the minimum number of samples.
+ self.require_batches = config.async_training.require_batches
+ self.required_samples = config.actor_rollout_ref.actor.ppo_mini_batch_size * self.require_batches
+ self.compute_prox_log_prob = self.config.async_training.compute_prox_log_prob
+ total_gpus = (
+ config.trainer.nnodes * config.trainer.n_gpus_per_node
+ + config.rollout.nnodes * config.rollout.n_gpus_per_node
+ )
+ self.metrics_aggregator = MetricsAggregator(total_gpus=total_gpus)
+
+ # use trainer to do validation
+ if self.config.async_training.use_trainer_do_validate:
+ from verl.trainer.main_ppo import create_rl_dataset
+ from verl.utils.dataset.rl_dataset import collate_fn
+
+ val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor)
+ rollout_gpus = config.rollout.nnodes * config.rollout.n_gpus_per_node
+ print(f"[FullyAsyncTrainer] split before val_dataset total len: {len(val_dataset)}")
+ split_dataset = val_dataset.split(total_gpus)
+ rollout_val_dataset0 = split_dataset[rollout_gpus:]
+ from torch.utils.data import ConcatDataset
+
+ val_dataset = ConcatDataset(rollout_val_dataset0)
+ print(f"[FullyAsyncTrainer] split after val_dataset total len: {len(val_dataset)}")
+ self.val_dataset = val_dataset
+ # update val_dataloader
+ val_batch_size = self.config.data.val_batch_size # Prefer config value if set
+ if val_batch_size is None:
+ val_batch_size = len(val_dataset)
+ from torchdata.stateful_dataloader import StatefulDataLoader
+
+ print(f"[FullyAsyncTrainer] create val_dataloader with batch_size: {val_batch_size}")
+ self.val_dataloader = StatefulDataLoader(
+ dataset=val_dataset,
+ batch_size=val_batch_size,
+ num_workers=self.config.data["dataloader_num_workers"],
+ shuffle=self.config.data.get("validation_shuffle", True),
+ drop_last=False,
+ collate_fn=collate_fn,
+ )
+
+ def set_message_queue_client(self, message_queue_client: MessageQueueClient):
+ """Set message queue client"""
+ self.message_queue_client = message_queue_client
+
+ def set_parameter_synchronizer(self, param_synchronizer):
+ """Set parameter synchronizer"""
+ self.param_synchronizer = param_synchronizer
+
+ def set_total_train_steps(self, total_train_steps):
+ self.total_train_steps = total_train_steps
+ self.progress_bar = tqdm(total=self.total_train_steps, initial=0, desc="Training Progress")
+
+ def get_actor_wg(self):
+ """Get actor worker group"""
+ return self.actor_wg
+
+ def _get_samples_from_queue(self) -> tuple[None, None] | tuple[int, Any]:
+ """
+ Get samples from message queue and compose gen_batch_output
+ Uses a loop to continuously collect samples until enough are gathered
+
+ Returns:
+ tuple: (epoch, batch_dict, gen_batch_output)
+ """
+ print(
+ f"[FullyAsyncTrainer] Requesting {self.required_samples} samples from queue",
+ flush=True,
+ )
+
+ # Collect samples using a simple loop calling get_sample
+ consumer_start = time.time()
+ queue_samples = []
+ queue_len = 0
+ while len(queue_samples) < self.required_samples:
+ # Get a single sample and wait until there is a sample or None is received
+ sample, queue_len = self.message_queue_client.get_sample_sync()
+
+ if sample is None:
+ print(
+ f"[FullyAsyncTrainer] Detected termination signal (None), stopping sample collection. "
+ f"Collected {len(queue_samples)}/{self.required_samples} samples"
+ )
+ break
+
+ queue_samples.append(sample)
+
+ if len(queue_samples) % 64 == 0:
+ print(
+ f"[FullyAsyncTrainer] Collected {len(queue_samples)}/{self.required_samples} samples. "
+ f"mq_len: {queue_len}"
+ )
+
+ consumer_end = time.time()
+
+ if not queue_samples or len(queue_samples) < self.required_samples:
+ print("[FullyAsyncTrainer] not enough samples collected after loop")
+ return None, None
+ total_wait_time = consumer_end - consumer_start
+
+ print(
+ f"[FullyAsyncTrainer] Loop collection completed: {len(queue_samples)}/{self.required_samples} samples, "
+ f"total wait time: {total_wait_time:.2f} seconds."
+ f"mq_len: {queue_len}"
+ )
+
+ queue_samples = [ray.cloudpickle.loads(x) for x in queue_samples]
+ # Assemble batch - now working directly with RolloutSample objects
+ if self.config.trainer.balance_batch:
+ batch = assemble_batch_from_rollout_samples(queue_samples, self.tokenizer, self.config, self._balance_batch)
+ else:
+ batch = assemble_batch_from_rollout_samples(queue_samples, self.tokenizer, self.config, None)
+
+ batch.meta_info["fully_async/total_wait_time"] = total_wait_time
+ return 0, batch
+
+ def _create_actor_rollout_classes(self):
+ # create actor
+ for role in [self.train_role]:
+ resource_pool = self.resource_pool_manager.get_resource_pool(role)
+ role_cls = RayClassWithInitArgs(
+ cls=self.role_worker_mapping[role],
+ config=self.config.actor_rollout_ref,
+ role=str(role),
+ )
+ self.resource_pool_to_cls[resource_pool][str(role)] = role_cls
+
+ def _init_models(self):
+ if self.use_critic:
+ self.critic_wg = self.all_wg[str(Role.Critic)]
+ self.critic_wg.init_model()
+
+ if self.use_reference_policy and not self.ref_in_actor:
+ self.ref_policy_wg = self.all_wg[str(Role.RefPolicy)]
+ self.ref_policy_wg.init_model()
+
+ if self.use_rm:
+ self.rm_wg = self.all_wg[str(Role.RewardModel)]
+ self.rm_wg.init_model()
+
+ self.actor_wg = self.all_wg[str(self.train_role)]
+ self.actor_wg.init_model()
+ self.actor_rollout_wg = self.actor_wg # to be compatible with the functions that not be modified
+
+ async def init_workers(self):
+ """Initialize distributed training workers using Ray backend.
+ Creates:
+ 1. Ray resource pools from configuration
+ 2. Worker groups for each role (actor, critic, etc.)
+ """
+ # self._init_async_objects()
+ self._init_resource_pools()
+ self._create_worker_classes()
+ self._init_worker_groups()
+ self._init_models()
+ await self._init_async_rollout_manager()
+
+ async def _init_async_rollout_manager(self):
+ # use async rollout do validate
+ print(f"[FullyAsyncTrainer] use_trainer_do_validate: {self.config.async_training.use_trainer_do_validate}")
+ if self.config.async_training.use_trainer_do_validate:
+ assert self.config.actor_rollout_ref.rollout.mode == "async"
+ self.async_rollout_mode = True
+ print("[FullyAsyncTrainer] Init async rollout manager")
+ from verl.experimental.fully_async_policy.agent_loop import FullyAsyncAgentLoopManager
+
+ self.async_rollout_manager = await FullyAsyncAgentLoopManager.create(
+ config=self.config, worker_group=self.actor_rollout_wg
+ )
+ print("[FullyAsyncTrainer] async_rollout_manager sleep")
+ await self.async_rollout_manager.sleep()
+ else:
+ print("[FullyAsyncTrainer] Skip async rollout manager (use_trainer_do_validate=False)")
+
+ async def fit(self):
+ """
+ The training loop of PPO.
+ The driver process only need to call the compute functions of the worker group through RPC
+ to construct the PPO dataflow.
+ The light-weight advantage computation is done on the driver process.
+ """
+ print("[FullyAsyncTrainer] Starting FullyAsyncTrainer...")
+ if self.message_queue_client is None:
+ raise ValueError("MessageQueue client not set. Call set_message_queue_client() first.")
+ if self.param_synchronizer is None:
+ raise ValueError("param_synchronizer client not set. Call set_parameter_synchronizer() first.")
+
+ from verl.utils.tracking import Tracking
+
+ self.logger = Tracking(
+ project_name=self.config.trainer.project_name,
+ experiment_name=self.config.trainer.experiment_name,
+ default_backend=self.config.trainer.logger,
+ config=OmegaConf.to_container(self.config, resolve=True),
+ )
+
+ self.max_steps_duration = 0
+
+ # get validate data before training
+ self._log_validation_data()
+
+ # Use queue mode, no need for traditional dataloader iterator
+ # Initialize to get the first batch of data
+ while True:
+ metrics = {}
+ timing_raw = {}
+
+ with marked_timer("step", timing_raw):
+ with marked_timer("gen", timing_raw, color="red"):
+ epoch, batch = self._get_samples_from_queue()
+ if batch is None:
+ break
+ self._collect_metrics_from_samples(batch, metrics)
+ batch, reward_extra_infos_dict = self._process_batch_common(
+ batch, metrics, timing_raw, self.local_trigger_step if self.compute_prox_log_prob else None
+ )
+ self._log_rollout(batch, reward_extra_infos_dict, timing_raw)
+
+ self._collect_metrics(batch, 0, metrics, timing_raw)
+ self.metrics_aggregator.add_step_metrics(
+ metrics=metrics, sample_count=self.required_samples, timestamp=time.time()
+ )
+ # Trigger parameter synchronization after training step
+ time_str = datetime.now().strftime("%H:%M:%S.%f")[:-3]
+ print(
+ f"[FullyAsyncTrainer] global_steps: {self.global_steps} "
+ f"local_trigger_step: {self.local_trigger_step} "
+ f"trigger_parameter_sync_step: {self.trigger_parameter_sync_step} "
+ f"{time_str}"
+ )
+ await self._trigger_parameter_sync_after_step(global_steps=self.global_steps)
+ self._log_validation_data()
+ self._check_save_checkpoint(timing_raw)
+ self.global_steps += 1
+
+ # final parameter sync and validate
+ # 1. waiting remaining validate task
+ ray.get(self.param_synchronizer.wait_last_valid.remote())
+ self._log_validation_data()
+ # 2. perform addtional parameter_sync and validate if trainer already updated
+ if self.current_param_version % self.config.rollout.test_freq != 0 or self.local_trigger_step > 1:
+ await self._trigger_parameter_sync_after_step(validate=True, global_steps=self.global_steps)
+ ray.get(self.param_synchronizer.wait_last_valid.remote())
+ self._log_validation_data()
+ self.progress_bar.close()
+
+ self._check_save_checkpoint(timing_raw)
+
+ def _check_save_checkpoint(self, timing_raw):
+ if self.current_param_version == self.last_ckpt_version:
+ return
+ # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.
+ esi_close_to_expiration = should_save_ckpt_esi(
+ max_steps_duration=self.max_steps_duration,
+ redundant_time=self.config.trainer.esi_redundant_time,
+ )
+ # Check if the conditions for saving a checkpoint are met.
+ # The conditions include a mandatory condition (1) and
+ # one of the following optional conditions (2/3/4):
+ # 1. The save frequency is set to a positive value.
+ # 2. The current step number is a multiple of the save frequency.
+ # 3. The ESI(Elastic Server Instance)/training plan is close to expiration.
+ if self.config.trainer.save_freq > 0 and (
+ self.current_param_version % self.config.trainer.save_freq == 0 or esi_close_to_expiration
+ ):
+ if esi_close_to_expiration:
+ print("Force saving checkpoint: ESI instance expiration approaching.")
+ with marked_timer("save_checkpoint", timing_raw, color="green"):
+ self._save_checkpoint()
+ self.last_ckpt_version = self.current_param_version
+
+ def _save_checkpoint(self):
+ # Warning: Currently, to align the training process and metrics of colocate,
+ # we use current_param_version instead of global step.
+ # This can be logically aligned with the original self.global_steps of colocate
+ # and is used for metrics and ckpt. which means that the parameter synchronization
+ # from trainer to rollouter will increase by 1 each time.
+
+ # path: given_path + `/global_step_{global_steps}` + `/actor`
+ local_global_step_folder = os.path.join(
+ self.config.trainer.default_local_dir, f"global_step_{self.current_param_version}"
+ )
+
+ print(f"[FullyAsyncTrainer] local_global_step_folder: {local_global_step_folder}")
+ actor_local_path = os.path.join(local_global_step_folder, "actor")
+
+ actor_remote_path = (
+ None
+ if self.config.trainer.default_hdfs_dir is None
+ else os.path.join(
+ self.config.trainer.default_hdfs_dir, f"global_step_{self.current_param_version}", "actor"
+ )
+ )
+
+ remove_previous_ckpt_in_save = self.config.trainer.get("remove_previous_ckpt_in_save", False)
+ if remove_previous_ckpt_in_save:
+ print(
+ "[FullyAsyncTrainer] Warning: remove_previous_ckpt_in_save is deprecated,"
+ + " set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead"
+ )
+ max_actor_ckpt_to_keep = (
+ self.config.trainer.get("max_actor_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1
+ )
+ max_critic_ckpt_to_keep = (
+ self.config.trainer.get("max_critic_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1
+ )
+
+ self.actor_rollout_wg.save_checkpoint(
+ actor_local_path, actor_remote_path, self.current_param_version, max_ckpt_to_keep=max_actor_ckpt_to_keep
+ )
+
+ if self.use_critic:
+ critic_local_path = os.path.join(local_global_step_folder, str(Role.Critic))
+ critic_remote_path = (
+ None
+ if self.config.trainer.default_hdfs_dir is None
+ else os.path.join(
+ self.config.trainer.default_hdfs_dir, f"global_step_{self.current_param_version}", str(Role.Critic)
+ )
+ )
+ self.critic_wg.save_checkpoint(
+ critic_local_path,
+ critic_remote_path,
+ self.current_param_version,
+ max_ckpt_to_keep=max_critic_ckpt_to_keep,
+ )
+ ray.get(self.param_synchronizer.rollouter_save_checkpoint.remote(local_global_step_folder))
+ # latest checkpointed iteration tracker (for atomic usage)
+ local_latest_checkpointed_iteration = os.path.join(
+ self.config.trainer.default_local_dir, "latest_checkpointed_iteration.txt"
+ )
+ with open(local_latest_checkpointed_iteration, "w") as f:
+ f.write(str(self.current_param_version))
+
+ def load_checkpoint(self):
+ if self.config.trainer.resume_mode == "disable":
+ # NOTE: while there is no checkpoint to load, we still need to offload the model and optimizer to CPU
+ self.actor_rollout_wg.load_checkpoint(None)
+ return 0
+
+ # load from hdfs
+ if self.config.trainer.default_hdfs_dir is not None:
+ raise NotImplementedError("load from hdfs is not implemented yet")
+ else:
+ checkpoint_folder = self.config.trainer.default_local_dir # TODO: check path
+ if not os.path.isabs(checkpoint_folder):
+ working_dir = os.getcwd()
+ checkpoint_folder = os.path.join(working_dir, checkpoint_folder)
+ global_step_folder = find_latest_ckpt_path(checkpoint_folder) # None if no latest
+
+ # find global_step_folder
+ if self.config.trainer.resume_mode == "auto":
+ if global_step_folder is None:
+ print("[FullyAsyncTrainer] Training from scratch")
+ self.actor_rollout_wg.load_checkpoint(None)
+ return 0
+ else:
+ if self.config.trainer.resume_mode == "resume_path":
+ assert isinstance(self.config.trainer.resume_from_path, str), "resume ckpt must be str type"
+ assert "global_step_" in self.config.trainer.resume_from_path, (
+ "resume ckpt must specify the global_steps"
+ )
+ global_step_folder = self.config.trainer.resume_from_path
+ if not os.path.isabs(global_step_folder):
+ working_dir = os.getcwd()
+ global_step_folder = os.path.join(working_dir, global_step_folder)
+ print(f"[FullyAsyncTrainer] Load from checkpoint folder: {global_step_folder}")
+ # set global step
+ self.current_param_version = int(global_step_folder.split("global_step_")[-1])
+ self.global_steps = self.current_param_version * self.trigger_parameter_sync_step + 1
+ self.last_ckpt_version = self.current_param_version
+ print(
+ f"[FullyAsyncTrainer] Setting global step to {self.global_steps}, "
+ f"current_param_version to {self.current_param_version}"
+ )
+ print(f"[FullyAsyncTrainer] Resuming from {global_step_folder}")
+
+ actor_path = os.path.join(global_step_folder, "actor")
+ critic_path = os.path.join(global_step_folder, str(Role.Critic))
+ # load actor
+ self.actor_rollout_wg.load_checkpoint(
+ actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load
+ )
+ # load critic
+ if self.use_critic:
+ self.critic_wg.load_checkpoint(
+ critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load
+ )
+ return self.current_param_version
+
+ def _collect_metrics_from_samples(self, batch, metrics):
+ """
+ Collect metrics from samples
+ """
+ if hasattr(batch, "meta_info") and batch.meta_info:
+ samples_param_versions = batch.meta_info["rollout_param_versions"]
+ stale_count = sum(1 for v in samples_param_versions if self.current_param_version - v >= 1)
+ self.stale_samples_processed += stale_count
+ trajectory_param_versions = batch.meta_info["trajectory_param_versions"]
+ stale_traj_count = sum(1 for v in trajectory_param_versions if self.current_param_version - v >= 1)
+ self.stale_trajectory_processed += stale_traj_count
+ metrics.update(
+ {
+ "fully_async/count/stale_samples_processed": self.stale_samples_processed,
+ "fully_async/count/stale_trajectory_processed": self.stale_trajectory_processed,
+ "fully_async/count/current_param_version": self.current_param_version,
+ }
+ )
+ for key, value in batch.meta_info.items():
+ if key.startswith("fully_async") or key.startswith("timing_s"):
+ metrics[key] = value
+
+ async def _trigger_parameter_sync_after_step(self, validate: bool = False, global_steps: int = None):
+ """
+ Trigger parameter synchronization after training step
+ This ensures rollouter always uses the latest trained parameters
+ """
+ if self.local_trigger_step < self.trigger_parameter_sync_step and not validate:
+ self.local_trigger_step += 1
+ return
+
+ self.current_param_version += 1
+ self.local_trigger_step = 1
+ self.logger.log(
+ data=self.metrics_aggregator.get_aggregated_metrics(),
+ step=self.current_param_version,
+ )
+ self.progress_bar.update(1)
+ self.metrics_aggregator.reset()
+ timing_param_sync = {}
+ with marked_timer("timing_s/wait_last_valid", timing_param_sync):
+ ray.get(self.param_synchronizer.wait_last_valid.remote())
+ with marked_timer("timing_s/param_sync", timing_param_sync):
+ ray.get(
+ self.param_synchronizer.sync_weights.remote(
+ self.current_param_version,
+ validate=validate,
+ global_steps=global_steps,
+ use_trainer_do_validate=self.config.async_training.use_trainer_do_validate,
+ )
+ )
+
+ # do trainer validate
+ do_validate_param = (
+ self.config.rollout.test_freq > 0
+ and self.current_param_version % self.config.rollout.test_freq == 0
+ and self.current_param_version > 0
+ )
+ print(f"do_validate_param: {do_validate_param}")
+ if do_validate_param and self.reward_fn is not None and self.config.async_training.use_trainer_do_validate:
+ print(f"[FullyAsyncTrainer] validate param version: {self.current_param_version}")
+ await self._validate_process()
+ else:
+ self.train_val_metrics = None
+ self.logger.log(data=timing_param_sync, step=self.current_param_version)
+
+ def _log_validation_data(self):
+ """
+ Log validation data
+ """
+ val_data = self.message_queue_client.get_validate_sync()
+ if not val_data:
+ return
+
+ val_metrics: ValidateMetrics = ray.cloudpickle.loads(val_data)
+ if self.train_val_metrics and self.config.async_training.use_trainer_do_validate:
+ # merge info
+ timing_param_sync = {}
+ with marked_timer("timing_s/merge_val", timing_param_sync):
+ new_metrics = self._merge_validation_results(self.train_val_metrics, val_metrics.metrics)
+ if new_metrics:
+ self.logger.log(data=new_metrics, step=val_metrics.param_version)
+ pprint(
+ f"[FullyAsyncTrainer] parameter version: {val_metrics.param_version} "
+ f"Validation metrics: {new_metrics}, timing_param_sync: {timing_param_sync['timing_s/merge_val']}"
+ )
+ self.logger.log(data=val_metrics.timing_raw, step=val_metrics.param_version)
+ else:
+ if val_metrics.metrics:
+ self.logger.log(data=val_metrics.metrics, step=val_metrics.param_version)
+ pprint(
+ f"[FullyAsyncTrainer] parameter version: {val_metrics.param_version} "
+ f"Validation metrics: {val_metrics.metrics}"
+ )
+ self.logger.log(data=val_metrics.timing_raw, step=val_metrics.param_version)
+
+ async def _validate_process(self):
+ if self.config.async_training.use_trainer_do_validate:
+ print("[FullyAsyncTrainer] _validate_process")
+ from verl.utils.profiler import marked_timer
+
+ timing_raw = {}
+ await self.async_rollout_manager.wake_up()
+ with marked_timer("trainer/validate_time", timing_raw):
+ self.train_val_metrics = self._validate(True)
+ await self.async_rollout_manager.sleep()
+ print(f"[FullyAsyncTrainer] validate timing_raw validate: {timing_raw['trainer/validate_time']}")
+ else:
+ self.train_val_metrics = None
+ print("[FullyAsyncTrainer] _validate_process without async_rollout_manager")
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/megatron_utils.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/megatron_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f5380f25c54ad7e7b28da04cc54f96313405448
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/megatron_utils.py
@@ -0,0 +1,99 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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.
+
+import torch
+from megatron.core.distributed import DistributedDataParallel as DDP
+
+
+@torch.no_grad()
+def copy_megatron_model_to_cpu(models):
+ """
+ Copy Megatron model parameters to CPU memory (non-destructive copy).
+ Unlike offload_megatron_model_to_cpu which moves data, this function creates
+ independent copies on CPU while keeping GPU data intact.
+
+ Args:
+ models: List of model chunks (DDP-wrapped or unwrapped)
+
+ Returns:
+ dict: CPU state containing copied parameters and buffers
+ """
+ cpu_state = {}
+
+ for model_idx, model_chunk in enumerate(models):
+ if isinstance(model_chunk, DDP):
+ # Handle DDP-wrapped models
+ model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
+ buffer_states = []
+
+ for buffers in model_chunk_all_buffers:
+ buffer_list = []
+ for buffer in buffers:
+ buffer_state = {}
+
+ # Copy parameter data to CPU
+ if buffer.param_data.storage().size() > 0:
+ buffer_state["param_data"] = buffer.param_data.data.cpu().clone().pin_memory()
+
+ buffer_list.append(buffer_state)
+ buffer_states.append(buffer_list)
+
+ cpu_state[f"model_chunk_{model_idx}"] = {"buffer_states": buffer_states, "is_ddp": True}
+ else:
+ # Handle non-DDP models (ref module)
+ model_state = {}
+ for name, param in model_chunk.named_parameters():
+ param_state = {"data": param.data.cpu().clone().pin_memory()}
+ model_state[name] = param_state
+
+ cpu_state[f"model_chunk_{model_idx}"] = {"model_state": model_state, "is_ddp": False}
+
+ return cpu_state
+
+
+@torch.no_grad()
+def restore_megatron_model_from_cpu(models, cpu_state):
+ """
+ Restore Megatron model parameters from CPU memory back to GPU.
+
+ Args:
+ models: List of model chunks to restore to
+ cpu_state: CPU state dict returned from copy_megatron_model_to_cpu
+ """
+ for model_idx, model_chunk in enumerate(models):
+ chunk_key = f"model_chunk_{model_idx}"
+ if chunk_key not in cpu_state:
+ continue
+
+ chunk_state = cpu_state[chunk_key]
+
+ if chunk_state["is_ddp"] and isinstance(model_chunk, DDP):
+ # Restore DDP buffers
+ model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
+ buffer_states = chunk_state["buffer_states"]
+
+ for buffers, buffer_list in zip(model_chunk_all_buffers, buffer_states, strict=False):
+ for buffer, buffer_state in zip(buffers, buffer_list, strict=False):
+ # Restore parameter data
+ if "param_data" in buffer_state:
+ buffer.param_data.data.copy_(buffer_state["param_data"].to(buffer.param_data.device))
+
+ elif not chunk_state["is_ddp"] and not isinstance(model_chunk, DDP):
+ # Restore non-DDP models
+ model_state = chunk_state["model_state"]
+ for name, param in model_chunk.named_parameters():
+ if name in model_state:
+ param_state = model_state[name]
+ param.data.copy_(param_state["data"].to(param.device))
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/megatron_worker.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/megatron_worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..44e63a94ce6ace763cfa4adbea9c7fd508344252
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/megatron_worker.py
@@ -0,0 +1,267 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan Ltd. and/or its affiliates
+# Copyright 2025 NVIDIA 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.
+
+import logging
+import os
+import time
+
+import torch
+import torch.distributed
+from omegaconf import DictConfig
+
+from verl.experimental.fully_async_policy.base_detach_sync import BaseDetachNcclSync
+from verl.experimental.fully_async_policy.megatron_utils import (
+ copy_megatron_model_to_cpu,
+ restore_megatron_model_from_cpu,
+)
+from verl.single_controller.base.decorator import Dispatch, register
+from verl.utils.device import (
+ get_device_name,
+ get_torch_device,
+)
+from verl.utils.megatron_utils import load_megatron_model_to_gpu, offload_megatron_model_to_cpu, per_tensor_generator
+from verl.workers.megatron_workers import AsyncActorRolloutRefWorker, CriticWorker
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+device_name = get_device_name()
+
+__all__ = ["DetachActorWorker", "DetachAsyncRolloutWorker", "CriticWorker"]
+
+
+class DetachNcclSync(BaseDetachNcclSync, AsyncActorRolloutRefWorker):
+ def __init__(self, config: DictConfig, role: str):
+ BaseDetachNcclSync.__init__(self, config, role)
+
+ AsyncActorRolloutRefWorker.__init__(self, config, role)
+
+ def _get_actor_params(self):
+ pass
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def sync_rollout_weights(self, sync_group_name="actor_rollout"):
+ assert (self._is_actor or self._is_rollout) and not self.config.hybrid_engine
+ assert hasattr(self, "_weights_info") and self._weights_info is not None
+ if self._is_actor and self._is_offload_param:
+ load_megatron_model_to_gpu(self.actor_module, False)
+ params_generator = self._get_actor_params_generator() if self._is_actor else None
+ params = {key: tensor for key, tensor in params_generator} if params_generator is not None else None
+
+ rollout_name = self.config.rollout.name
+ inference_model = None
+ if self._is_rollout and (not self._is_actor):
+ if rollout_name == "vllm":
+ inference_model = BaseDetachNcclSync.get_inference_model(self.rollout)
+ from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+
+ patch_vllm_moe_model_weight_loader(inference_model)
+ elif rollout_name == "sglang":
+ inference_model = self.rollout._engine
+ if inference_model is None:
+ print("[sync_rollout_weights] Initialize server adapter engine")
+
+ async def init_engine():
+ if hasattr(self.rollout, "_init_server_adapter"):
+ await self.rollout._init_server_adapter()
+ else:
+ print("[sync_rollout_weights] No _init_server_adapter method found")
+ return self.rollout._engine
+
+ inference_model = self._run_async_safely(init_engine())
+ if inference_model is None:
+ raise RuntimeError(
+ f"Failed to initialize rollout engine. "
+ f"rollout type: {type(self.rollout)}, "
+ f"has _init_server_adapter: {hasattr(self.rollout, '_init_server_adapter')}"
+ )
+ else:
+ raise NotImplementedError(f"Unknown rollout name: {rollout_name}")
+
+ if rollout_name == "sglang" and self._is_rollout:
+ self._sync_sglang_weights(inference_model, params, sync_group_name)
+ else:
+ self._sync_vllm_weights(inference_model, params, sync_group_name)
+
+ if self._is_actor and self._is_offload_param:
+ offload_megatron_model_to_cpu(self.actor_module)
+ get_torch_device().empty_cache()
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def save_model_to_cpu(self, n):
+ if not hasattr(self, "cpu_saved_models"):
+ self.cpu_saved_models = {}
+ self.cpu_saved_models[n] = copy_megatron_model_to_cpu(self.actor.actor_module)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def restore_model_from_cpu(self, n):
+ if n in self.cpu_saved_models:
+ restore_megatron_model_from_cpu(self.actor.actor_module, self.cpu_saved_models[n])
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def clear_cpu_model(self, n):
+ if n in self.cpu_saved_models:
+ del self.cpu_saved_models[n]
+
+ def cache_actor_weights_to_cpu(self):
+ self.cpu_named_params = {}
+ if self._is_actor:
+ params_generator = self._get_actor_params_generator()
+ local_rank = torch.distributed.get_rank()
+ world_size = torch.distributed.get_world_size()
+ print(f"cache_actor_weights_to_cpu, local_rank:{local_rank}, world_size:{world_size}")
+ for tensor_idx, (key, tensor) in enumerate(params_generator):
+ if tensor_idx % world_size == local_rank:
+ self.cpu_named_params[key] = tensor.to("cpu", non_blocking=True)
+ get_torch_device().synchronize()
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def sync_rollout_weights_by_checkpoint(self, sync_group_name="actor_rollout"):
+ assert (self._is_actor or self._is_rollout) and not self.config.hybrid_engine
+ assert hasattr(self, "_weights_info") and self._weights_info is not None
+
+ # Load model to GPU
+ load_start_time = time.time()
+ if self._is_actor and self._is_offload_param:
+ load_megatron_model_to_gpu(self.actor_module, False)
+ load_duration = time.time() - load_start_time
+
+ from ray.util.collective import collective
+
+ # Cache actor weights to CPU and measure the time taken
+ cache_start_time = time.time()
+ self.cache_actor_weights_to_cpu()
+ cache_end_time = time.time()
+ cache_duration = cache_end_time - cache_start_time
+
+ # Register the cached weights into the checkpoint engine
+ self.checkpoint_engine.register_checkpoint(self._weights_info, self.cpu_named_params)
+ register_end_time = time.time()
+ register_duration = register_end_time - cache_end_time
+ self.cpu_named_params = {}
+
+ collective.barrier(group_name=sync_group_name)
+ update_start_time = time.time()
+
+ rollout_name = self.config.rollout.name
+ inference_model = None
+ if self._is_rollout and (not self._is_actor):
+ if rollout_name == "vllm":
+ inference_model = BaseDetachNcclSync.get_inference_model(self.rollout)
+ from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+
+ patch_vllm_moe_model_weight_loader(inference_model)
+ elif rollout_name == "sglang":
+ inference_model = self.rollout._engine
+ # For ServerAdapter, _engine might be None and needs async initialization
+ if inference_model is None:
+ # Initialize the server adapter engine
+ print("[sync_rollout_weights] Initialize server adapter engine")
+
+ async def init_engine():
+ if hasattr(self.rollout, "_init_server_adapter"):
+ await self.rollout._init_server_adapter()
+ else:
+ print("[sync_rollout_weights] No _init_server_adapter method found")
+ return self.rollout._engine
+
+ inference_model = self._run_async_safely(init_engine())
+ if inference_model is None:
+ raise RuntimeError(
+ f"Failed to initialize rollout engine. "
+ f"rollout type: {type(self.rollout)}, "
+ f"has _init_server_adapter: {hasattr(self.rollout, '_init_server_adapter')}"
+ )
+ else:
+ raise NotImplementedError(f"Unknown rollout name: {rollout_name}")
+ # Update the checkpoint with the inference model and broadcast weights
+ self.checkpoint_engine.update_checkpoint(
+ inference_model=inference_model,
+ group_name=sync_group_name,
+ overlap_broadcast_and_consume=self.config.checkpoint_engine.overlap_broadcast_and_consume,
+ )
+
+ update_end_time = time.time()
+ update_duration = update_end_time - update_start_time
+
+ collective.barrier(group_name=sync_group_name)
+ offload_start_time = time.time()
+ if self._is_actor and self._is_offload_param:
+ offload_megatron_model_to_cpu(self.actor_module)
+ offload_duration = time.time() - offload_start_time
+
+ print(
+ f"sync_rollout_weights_by_checkpoint finish!, rank:{torch.distributed.get_rank()},"
+ f" is_actor:{self._is_actor}, is_rollout:{self._is_rollout},"
+ f" total cost:{update_end_time - cache_start_time} seconds, while cache cost {cache_duration} seconds, "
+ f" register cost {register_duration} seconds, update cost {update_duration} seconds"
+ )
+
+ if self._is_actor and self._is_offload_param:
+ print(
+ f"sync_rollout_weights_by_checkpoint load model to gpu cost {load_duration} seconds,"
+ f" offload model to cpu cost {offload_duration} seconds"
+ )
+
+
+class DetachActorWorker(DetachNcclSync):
+ def __init__(self, config: DictConfig, role: str):
+ print("[DetachAsyncRolloutWorker] Initializing via DetachNcclSync...")
+ DetachNcclSync.__init__(self, config, role)
+
+ def _get_actor_params_generator(self):
+ assert self._is_actor
+ if self.bridge is not None:
+ generator = self.bridge.export_weights(self.actor.actor_module)
+ else:
+ generator = per_tensor_generator(
+ self.actor.actor_module,
+ self.actor_model_config,
+ self.weight_converter,
+ self.tf_config,
+ self.layer_name_mapping,
+ )
+
+ return generator
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def get_actor_weights_info(self):
+ assert self._is_actor
+ if hasattr(self, "_weights_info"):
+ return self._weights_info
+ if self._is_offload_param:
+ load_megatron_model_to_gpu(self.actor_module, False)
+ params_generator = self._get_actor_params_generator()
+ ret = []
+ for key, tensor in params_generator:
+ ret.append((key, tensor.size(), tensor.dtype))
+
+ self._weights_info = ret
+ # Here, we only call this function at the beginning,
+ # and immediately afterwards we call sync_rollout_weights.
+ # So we no longer call offload in this.
+ return ret
+
+
+class DetachAsyncRolloutWorker(DetachNcclSync):
+ def __init__(self, config: DictConfig, role: str):
+ print(f"[DetachAsyncRolloutWorker] {DetachAsyncRolloutWorker.__mro__}")
+ DetachNcclSync.__init__(self, config, role)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def set_actor_weights_info(self, weights_info):
+ assert self._is_rollout
+ self._weights_info = weights_info
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/message_queue.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/message_queue.py
new file mode 100644
index 0000000000000000000000000000000000000000..85860c6f2a0d4ee711e80d6e696c2c2430a48a6b
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/message_queue.py
@@ -0,0 +1,265 @@
+# Copyright 2025 Meituan 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.
+
+import asyncio
+import logging
+from collections import deque
+from typing import Any
+
+import ray
+from omegaconf import DictConfig
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote(num_cpus=2, max_concurrency=20)
+class MessageQueue:
+ """
+ Simplified Ray-based asynchronous message queue for communication between Rollouter and Trainer
+ """
+
+ def __init__(self, config: DictConfig, max_queue_size: int = 1000):
+ self.config = config
+ if max_queue_size is None:
+ raise ValueError(f"max_queue_size cannot be None, got: {max_queue_size}")
+ self.max_queue_size = int(max_queue_size)
+ self.queue = deque(maxlen=self.max_queue_size)
+ self.current_param_version = 0
+
+ self.val_queue = deque()
+
+ try:
+ if hasattr(config, "async_training") and config.async_training is not None:
+ self.staleness_threshold = getattr(config.async_training, "staleness_threshold", 3)
+ else:
+ self.staleness_threshold = 3
+ except (AttributeError, RecursionError):
+ self.staleness_threshold = 3
+
+ # Asyncio for message handling
+ self.running = True
+
+ # async safe
+ self._lock = asyncio.Lock()
+ self._consumer_condition = asyncio.Condition(self._lock)
+
+ # statistic message
+ self.total_produced = 0
+ self.total_consumed = 0
+ self.dropped_samples = 0
+
+ print(
+ f"[MessageQueue] initialized with max_queue_size={max_queue_size},"
+ f"staleness_threshold={self.staleness_threshold}"
+ )
+
+ async def put_sample(self, sample: Any, param_version: int) -> bool:
+ """
+ Put a batch sample into the queue
+
+ Args:
+ sample: Sample data
+ param_version: Parameter version number
+
+ Returns:
+ bool: Whether the sample was successfully put into the queue
+ """
+ async with self._lock:
+ # If queue is full, remove the oldest sample (rarely happens)
+ is_drop = False
+ if len(self.queue) >= self.max_queue_size:
+ self.queue.popleft()
+ self.dropped_samples += 1
+ is_drop = True
+ logger.warning("Queue full, dropped sample")
+ self.queue.append(sample)
+ self.total_produced += 1
+
+ # Notify waiting consumers
+ self._consumer_condition.notify_all()
+
+ if self.total_produced % 100 == 0:
+ print(f"MessageQueue stats: produced={self.total_produced}, queue_size={len(self.queue)}")
+ if is_drop:
+ return False
+ return True
+
+ async def get_sample(self) -> Any | None:
+ """
+ Get a single sample from the queue, wait until one is available
+
+ Returns:
+ Any: Single sample data or None if queue is closed
+ """
+ async with self._lock:
+ while len(self.queue) == 0 and self.running:
+ await self._consumer_condition.wait()
+
+ # If queue is closed and empty, return None
+ if not self.running and len(self.queue) == 0:
+ return None
+
+ # Get one sample
+ data = self.queue.popleft()
+ self.total_consumed += 1
+ return data, len(self.queue)
+
+ async def update_param_version(self, version: int):
+ """Update current parameter version"""
+ async with self._lock:
+ old_version = self.current_param_version
+ self.current_param_version = version
+ print(f"Parameter version updated from {old_version} to {version}")
+
+ async def get_queue_size(self) -> int:
+ """Get current queue length"""
+ async with self._lock:
+ return len(self.queue)
+
+ async def get_statistics(self) -> dict[str, Any]:
+ """Get queue statistics"""
+ async with self._lock:
+ return {
+ "queue_size": len(self.queue),
+ "total_produced": self.total_produced,
+ "total_consumed": self.total_consumed,
+ "dropped_samples": self.dropped_samples,
+ "current_param_version": self.current_param_version,
+ "staleness_threshold": self.staleness_threshold,
+ "max_queue_size": self.max_queue_size,
+ }
+
+ async def clear_queue(self):
+ """Clear the queue"""
+ async with self._lock:
+ cleared_count = len(self.queue)
+ self.queue.clear()
+ logger.info(f"Cleared {cleared_count} samples from queue")
+
+ async def shutdown(self):
+ """Shutdown the message queue"""
+ async with self._lock:
+ self.running = False
+ # Notify all waiting coroutines so they can exit
+ self._consumer_condition.notify_all()
+ logger.info("MessageQueue shutdown")
+
+ async def get_memory_usage(self) -> dict:
+ """Get memory usage statistics"""
+ async with self._lock:
+ # Estimate memory usage of samples in queue
+ import sys
+
+ total_size = 0
+ sample_count = len(self.queue)
+
+ if sample_count > 0:
+ # Estimate size of a single sample (simplified estimation)
+ sample = list(self.queue)[0]
+ try:
+ sample_size = sys.getsizeof(sample)
+ # Since we now store RolloutSample directly, estimate based on its components
+ if hasattr(sample, "original_batch_dict") and sample.original_batch_dict:
+ # Estimate batch data size
+ batch_data = sample.original_batch_dict.get("batch", {})
+ sample_size += len(batch_data) * 1000 # Roughly estimate 1KB per batch entry
+ if hasattr(sample, "agent_loop_output"):
+ # Estimate AgentLoopOutput size
+ sample_size += 5000 # Roughly estimate 5KB for AgentLoopOutput
+ total_size = sample_size * sample_count
+ except Exception:
+ total_size = sample_count * 15000 # Roughly estimate 15KB per RolloutSample
+
+ return {
+ "queue_samples": sample_count,
+ "estimated_memory_bytes": total_size,
+ "estimated_memory_mb": total_size / (1024 * 1024),
+ }
+
+ async def put_validate(self, data):
+ async with self._lock:
+ self.val_queue.append(data)
+
+ async def get_validate(self):
+ async with self._lock:
+ if self.val_queue:
+ return self.val_queue.popleft()
+ else:
+ return None
+
+
+class MessageQueueClient:
+ """Asyncio-compatible MessageQueue client for communicating with MessageQueue Actor"""
+
+ def __init__(self, queue_actor: Any):
+ self.queue_actor = queue_actor
+
+ async def put_sample(self, sample: Any, param_version: int) -> bool:
+ """Put batch into queue (async)"""
+ future = self.queue_actor.put_sample.remote(sample, param_version)
+ return await asyncio.wrap_future(future.future())
+
+ async def put_validate(self, data: Any) -> bool:
+ future = self.queue_actor.put_validate.remote(data)
+ return await asyncio.wrap_future(future.future())
+
+ def get_validate_sync(self) -> Any | None:
+ return ray.get(self.queue_actor.get_validate.remote())
+
+ async def get_sample(self) -> Any | None:
+ """Get single sample from queue, wait until one is available (async)"""
+ future = self.queue_actor.get_sample.remote()
+ return await asyncio.wrap_future(future.future())
+
+ async def get_queue_size(self) -> int:
+ """Get queue size (async)"""
+ future = self.queue_actor.get_queue_size.remote()
+ return await asyncio.wrap_future(future.future())
+
+ async def get_statistics(self) -> dict[str, Any]:
+ """Get statistics (async)"""
+ future = self.queue_actor.get_statistics.remote()
+ return await asyncio.wrap_future(future.future())
+
+ async def clear_queue(self):
+ """Clear queue (async)"""
+ future = self.queue_actor.clear_queue.remote()
+ await asyncio.wrap_future(future.future())
+
+ async def shutdown(self):
+ """Shutdown queue (async)"""
+ future = self.queue_actor.shutdown.remote()
+ await asyncio.wrap_future(future.future())
+
+ async def get_memory_usage(self) -> dict:
+ """Get memory usage statistics (async)"""
+ future = self.queue_actor.get_memory_usage.remote()
+ return await asyncio.wrap_future(future.future())
+
+ # Synchronous version of the method (deprecated)
+ def put_sample_sync(self, sample: Any, param_version: int) -> bool:
+ """Put batch into queue (sync - deprecated, use put_sample instead)"""
+ return ray.get(self.queue_actor.put_sample.remote(sample, param_version))
+
+ def get_sample_sync(self) -> Any | None:
+ """Get single sample from queue (sync - deprecated, use get_sample instead)"""
+ return ray.get(self.queue_actor.get_sample.remote())
+
+ def get_statistics_sync(self) -> dict[str, Any]:
+ """Get statistics (sync - deprecated, use get_statistics instead)"""
+ return ray.get(self.queue_actor.get_statistics.remote())
+
+ def update_param_version_sync(self, version: int):
+ """Update parameter version (async)"""
+ return ray.get(self.queue_actor.update_param_version.remote(version))
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/param_sync.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/param_sync.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a9ac167aa33cf21bbd2c06afcea0757c3a90d61
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/param_sync.py
@@ -0,0 +1,173 @@
+# Copyright 2025 Meituan 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.
+
+import logging
+import time
+
+import ray
+from ray.util.collective import collective
+
+from verl.utils.device import get_nccl_backend
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class ParameterSynchronizer:
+ """
+ Unified parameter synchronizer, responsible for synchronizing model parameters between actor and rollout
+ Based on the mature synchronization mode implementation of one_step_off_policy
+ Merges the functions of the original multiple synchronizer classes
+ """
+
+ def __init__(self, config, trainer, rollouter, mq):
+ self.config = config
+ self.trainer = trainer
+ self.rollouter = rollouter
+ self.mq_client = mq
+ self.actor_wg = ray.get(trainer.get_actor_wg.remote())
+ self.rollout_wg = ray.get(rollouter.get_rollout_wg.remote())
+
+ # Basic attributes
+ self.weights_info = None
+ self.sync_group_initialized = False
+ self.sync_group_name = "actor_rollout"
+ self.wait_last_update = None
+ self.wait_last_resume = None
+ self.validate_task = None
+
+ # Statistics
+ self.current_version = 0
+
+ self._init_weights_info()
+ self._init_sync_group()
+
+ if self.config.async_training.checkpoint_engine.enable:
+ self._init_actor_rollout_checkpoint_engine()
+
+ def get_current_param_version(self) -> int:
+ """Get current parameter version number"""
+ return self.current_version
+
+ def get_weights_info(self):
+ """Get weights info"""
+ return self.weights_info
+
+ def _init_weights_info(self):
+ self.weights_info = self.actor_wg.get_actor_weights_info()[0]
+ self.rollout_wg.set_actor_weights_info(self.weights_info)
+
+ def _init_sync_group(self):
+ print("[ParameterSynchronizer] Initializing parameter synchronization group...")
+ actor_rollout_workers = self.actor_wg.workers + self.rollout_wg.workers
+ n_workers = len(self.actor_wg.workers + self.rollout_wg.workers)
+ if self.config.trainer.device == "npu":
+ master_address = ray.get(self.actor_wg.workers[0]._get_node_ip.remote()).strip("[]")
+ master_port = ray.get(self.actor_wg.workers[0]._get_free_port.remote())
+ self.actor_wg.create_weight_sync_group(
+ master_address,
+ master_port,
+ 0,
+ n_workers,
+ )
+ ray.get(
+ self.rollout_wg.create_weight_sync_group(
+ master_address,
+ master_port,
+ len(self.actor_wg.workers),
+ n_workers,
+ )
+ )
+ else:
+ collective.create_collective_group(
+ actor_rollout_workers,
+ n_workers,
+ list(range(0, n_workers)),
+ backend=get_nccl_backend(),
+ group_name=self.sync_group_name,
+ )
+
+ def _init_actor_rollout_checkpoint_engine(self):
+ ray.get(
+ self.actor_wg.init_checkpoint_engine(
+ rank_offset=0,
+ actor_num=len(self.actor_wg.workers),
+ rollout_num=len(self.rollout_wg.workers),
+ )
+ )
+ ray.get(
+ self.rollout_wg.init_checkpoint_engine(
+ rank_offset=len(self.actor_wg.workers),
+ actor_num=len(self.actor_wg.workers),
+ rollout_num=len(self.rollout_wg.workers),
+ )
+ )
+
+ def sync_weights(self, version, validate=False, global_steps=0, use_trainer_do_validate=False):
+ """Sync weights between trainer and rollouter, and update parameter version"""
+ start_time = time.time()
+
+ self.current_version = version
+ ray.get(self.rollouter.pause.remote())
+
+ print(f"[ParameterSynchronizer] rollout paused. cost {time.time() - start_time:.2f} seconds")
+ # Update MQ version
+ self.mq_client.update_param_version_sync(version)
+
+ pause_time = time.time()
+
+ # sync weights
+ # For sglang, always use sync_rollout_weights instead of sync_rollout_weights_by_checkpoint
+ rollout_name = getattr(self.config.actor_rollout_ref.rollout, "name", None)
+ use_checkpoint_engine = self.config.async_training.checkpoint_engine.enable and rollout_name != "sglang"
+
+ if use_checkpoint_engine:
+ self.actor_wg.sync_rollout_weights_by_checkpoint(self.sync_group_name)
+ ray.get(self.rollout_wg.sync_rollout_weights_by_checkpoint(self.sync_group_name))
+ else:
+ self.actor_wg.sync_rollout_weights(self.sync_group_name)
+ ray.get(self.rollout_wg.sync_rollout_weights(self.sync_group_name))
+ end_time = time.time()
+ print(
+ f"[ParameterSynchronizer] sync_weights success. cost {end_time - start_time:.2f} seconds, "
+ f"pause:{pause_time - start_time:.2f}s, sync:{end_time - pause_time:.2f}s"
+ )
+ # async train do validate
+ print(f"[ParameterSynchronizer] validate: {validate}, use_trainer_do_validate: {use_trainer_do_validate}")
+ if validate and use_trainer_do_validate:
+ print("[ParameterSynchronizer] use trainer to do validate")
+ self.validate_task = self.trainer._validate_process.remote()
+ else:
+ self.validate_task = None
+ # Async Update rollout version & validation
+ self.wait_last_update = self.rollouter.update_param_version.remote(
+ version, validate, global_steps, use_trainer_do_validate
+ )
+ self.wait_last_resume = self.rollouter.resume.remote(self.wait_last_update)
+
+ def wait_last_valid(self):
+ print("[ParameterSynchronizer] Waiting last sync and validate...")
+ start_time = time.time()
+ if self.wait_last_update:
+ ray.get(self.wait_last_update)
+ if self.wait_last_resume:
+ ray.get(self.wait_last_resume)
+ if self.validate_task:
+ ray.get(self.validate_task)
+ print(f"[ParameterSynchronizer] Wait last validate cost: {time.time() - start_time:.2f} seconds")
+
+ def rollouter_save_checkpoint(self, local_global_step_folder: str):
+ """Trigger rollout to save checkpoint(dataloader)"""
+ print(f"[ParameterSynchronizer] Triggering checkpoint save at {local_global_step_folder} ...")
+ return ray.get(self.rollouter.save_checkpoint.remote(local_global_step_folder))
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/ray_trainer.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/ray_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..f31e55d1388a08df7171c136c737bbaf8abf46a0
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/ray_trainer.py
@@ -0,0 +1,538 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023-2024 SGLang Team
+# Copyright 2025 ModelBest Inc. 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.
+"""
+PPO Trainer with Ray-based single controller.
+This trainer supports model-agonistic model initialization with huggingface
+"""
+
+import uuid
+from copy import deepcopy
+from pprint import pprint
+
+import numpy as np
+import ray
+import torch
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from verl import DataProto
+from verl.experimental.dataset.sampler import AbstractCurriculumSampler
+from verl.single_controller.ray import RayClassWithInitArgs
+from verl.single_controller.ray.base import create_colocated_worker_cls
+from verl.trainer.ppo.core_algos import AdvantageEstimator, agg_loss
+from verl.trainer.ppo.metric_utils import (
+ compute_data_metrics,
+ compute_throughout_metrics,
+ compute_timing_metrics,
+)
+from verl.trainer.ppo.ray_trainer import RayPPOTrainer, apply_kl_penalty, compute_advantage, compute_response_mask
+from verl.trainer.ppo.reward import compute_reward, compute_reward_async
+from verl.trainer.ppo.utils import Role
+from verl.utils.config import omega_conf_to_dataclass
+from verl.utils.debug import marked_timer
+from verl.utils.metric import (
+ reduce_metrics,
+)
+from verl.utils.rollout_skip import RolloutSkip
+
+
+class FullyAsyncRayPPOTrainer(RayPPOTrainer):
+ def init_workers(self):
+ """Initialize distributed training workers using Ray backend.
+
+ Creates:
+ 1. Ray resource pools from configuration
+ 2. Worker groups for each role (actor, critic, etc.)
+ """
+ self._init_resource_pools()
+ self._create_worker_classes()
+ self._init_worker_groups()
+ self._init_models()
+ self._init_async_rollout_manager()
+
+ def _init_resource_pools(self):
+ self.resource_pool_manager.create_resource_pool()
+
+ self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
+
+ def _create_worker_classes(self):
+ self._create_actor_rollout_classes()
+ self._create_critic_class()
+ self._create_reference_policy_class()
+ self._create_reward_model_class()
+
+ def _create_actor_rollout_classes(self):
+ raise NotImplementedError
+
+ def _create_critic_class(self):
+ # create critic
+ if self.use_critic:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)
+ critic_cfg = omega_conf_to_dataclass(self.config.critic)
+ critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg)
+ self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls
+
+ def _create_reference_policy_class(self):
+ # create reference policy if needed
+ if self.use_reference_policy:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)
+ ref_policy_cls = RayClassWithInitArgs(
+ self.role_worker_mapping[Role.RefPolicy],
+ config=self.config.actor_rollout_ref,
+ role=str(Role.RefPolicy),
+ # profile_option=self.config.trainer.npu_profile.options,
+ )
+ self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls
+
+ def _create_reward_model_class(self):
+ # create a reward model if reward_fn is None
+ 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
+
+ def _init_worker_groups(self):
+ # 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`.
+ # Instead, directly pass different resource pool to different worker groups.
+ # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.
+ all_wg = {}
+ wg_kwargs = {} # Setting up kwargs for RayWorkerGroup
+ if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None:
+ wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout
+ if OmegaConf.select(self.config.global_profiler, "steps") is not None:
+ wg_kwargs["profile_steps"] = OmegaConf.select(self.config.global_profiler, "steps")
+ # Only require nsight worker options when tool is nsys
+ if OmegaConf.select(self.config.global_profiler, "tool") == "nsys":
+ assert (
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ is not None
+ ), "worker_nsight_options must be set when using nsys with profile_steps"
+ wg_kwargs["worker_nsight_options"] = OmegaConf.to_container(
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ )
+ wg_kwargs["device_name"] = self.device_name
+
+ for resource_pool, class_dict in self.resource_pool_to_cls.items():
+ worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
+ wg_dict = self.ray_worker_group_cls(
+ resource_pool=resource_pool,
+ ray_cls_with_init=worker_dict_cls,
+ **wg_kwargs,
+ )
+ spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
+ all_wg.update(spawn_wg)
+ self.all_wg = all_wg
+
+ def _init_models(self):
+ if self.use_critic:
+ self.critic_wg = self.all_wg[str(Role.Critic)]
+ self.critic_wg.init_model()
+
+ if self.use_reference_policy and not self.ref_in_actor:
+ self.ref_policy_wg = self.all_wg[str(Role.RefPolicy)]
+ self.ref_policy_wg.init_model()
+
+ if self.use_rm:
+ self.rm_wg = self.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 = self.all_wg[str(Role.ActorRollout)]
+ self.actor_rollout_wg.init_model()
+
+ def _init_async_rollout_manager(self):
+ pass
+
+ def fit(self):
+ """
+ The training loop of PPO.
+ The driver process only need to call the compute functions of the worker group through RPC
+ to construct the PPO dataflow.
+ The light-weight advantage computation is done on the driver process.
+ """
+ from omegaconf import OmegaConf
+
+ from verl.utils.tracking import Tracking
+
+ logger = Tracking(
+ project_name=self.config.trainer.project_name,
+ experiment_name=self.config.trainer.experiment_name,
+ default_backend=self.config.trainer.logger,
+ config=OmegaConf.to_container(self.config, resolve=True),
+ )
+
+ self.global_steps = 0
+
+ # load checkpoint before doing anything
+ self._load_checkpoint()
+
+ # perform validation before training
+ # currently, we only support validation using the reward_function.
+ if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
+ val_metrics = self._validate()
+ assert val_metrics, f"{val_metrics=}"
+ pprint(f"Initial validation metrics: {val_metrics}")
+ logger.log(data=val_metrics, step=self.global_steps)
+ if self.config.trainer.get("val_only", False):
+ return
+
+ if self.config.actor_rollout_ref.rollout.get("skip_rollout", False):
+ rollout_skip = RolloutSkip(self.config, self.actor_rollout_wg)
+ rollout_skip.wrap_generate_sequences()
+
+ # add tqdm
+ progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress")
+
+ # we start from step 1
+ self.global_steps += 1
+ last_val_metrics = None
+ self.max_steps_duration = 0
+
+ prev_step_profile = False
+ curr_step_profile = (
+ self.global_steps in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ next_step_profile = False
+
+ for epoch in range(self.config.trainer.total_epochs):
+ for batch_dict in self.train_dataloader:
+ metrics = {}
+ timing_raw = {}
+
+ with marked_timer("start_profile", timing_raw):
+ self._start_profiling(
+ not prev_step_profile and curr_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+
+ batch, gen_batch = self._prepare_generate_batch(batch_dict)
+
+ is_last_step = self.global_steps >= self.total_training_steps
+
+ with marked_timer("step", timing_raw):
+ # generate a batch
+ with marked_timer("gen", timing_raw, color="red"):
+ if not self.async_rollout_mode:
+ gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)
+ else:
+ gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch)
+ timing_raw.update(gen_batch_output.meta_info["timing"])
+ gen_batch_output.meta_info.pop("timing", None)
+
+ if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
+ if self.reward_fn is None:
+ raise ValueError("A reward_fn is required for REMAX advantage estimation.")
+
+ with marked_timer("gen_max", timing_raw, color="purple"):
+ gen_baseline_batch = deepcopy(gen_batch)
+ gen_baseline_batch.meta_info["do_sample"] = False
+ if not self.async_rollout_mode:
+ gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)
+ else:
+ gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)
+ batch = batch.union(gen_baseline_output)
+ reward_baseline_tensor = self.reward_fn(batch)
+ reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
+
+ batch.pop(batch_keys=list(gen_baseline_output.batch.keys()))
+
+ batch.batch["reward_baselines"] = reward_baseline_tensor
+
+ del gen_baseline_batch, gen_baseline_output
+
+ batch = self._post_generate_batch(batch, gen_batch_output, metrics)
+ batch, reward_extra_infos_dict = self._process_batch_common(batch, metrics, timing_raw)
+ self._log_rollout(batch, reward_extra_infos_dict, timing_raw)
+
+ last_val_metrics = self._validate_metrics(is_last_step, last_val_metrics, metrics, timing_raw)
+ self._check_save_checkpoint(is_last_step, timing_raw)
+
+ with marked_timer("stop_profile", timing_raw):
+ next_step_profile = (
+ self.global_steps + 1 in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ self._stop_profiling(
+ curr_step_profile and not next_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+ prev_step_profile = curr_step_profile
+ curr_step_profile = next_step_profile
+
+ self._collect_metrics(batch, epoch, metrics, timing_raw)
+ self._post_batch_processing(batch)
+
+ # TODO: make a canonical logger that supports various backend
+ logger.log(data=metrics, step=self.global_steps)
+
+ progress_bar.update(1)
+ self.global_steps += 1
+
+ if (
+ hasattr(self.config.actor_rollout_ref.actor, "profiler")
+ and self.config.actor_rollout_ref.actor.profiler.tool == "torch_memory"
+ ):
+ self.actor_rollout_wg.dump_memory_snapshot(
+ tag=f"post_update_step{self.global_steps}", sub_dir=f"step{self.global_steps}"
+ )
+
+ if is_last_step:
+ pprint(f"Final validation metrics: {last_val_metrics}")
+ progress_bar.close()
+ return
+
+ def _prepare_generate_batch(self, batch_dict):
+ batch: DataProto = DataProto.from_single_dict(batch_dict)
+
+ # add uid to batch
+ batch.non_tensor_batch["uid"] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object)
+
+ gen_batch = self._get_gen_batch(batch)
+
+ # pass global_steps to trace
+ gen_batch.meta_info["global_steps"] = self.global_steps
+ gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
+ return batch, gen_batch
+
+ def _post_generate_batch(self, batch, gen_batch_output, metrics):
+ # repeat to align with repeated responses in rollout
+ batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
+ batch = batch.union(gen_batch_output)
+
+ if "response_mask" not in batch.batch.keys():
+ batch.batch["response_mask"] = compute_response_mask(batch)
+ # Balance the number of valid tokens across DP ranks.
+ # NOTE: This usually changes the order of data in the `batch`,
+ # which won't affect the advantage calculation (since it's based on uid),
+ # but might affect the loss calculation (due to the change of mini-batching).
+ # TODO: Decouple the DP balancing and mini-batching.
+ if self.config.trainer.balance_batch:
+ self._balance_batch(batch, metrics=metrics)
+
+ # compute global_valid tokens
+ batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
+
+ return batch
+
+ def _process_batch_common(self, batch, metrics, timing_raw, local_trigger_step=None):
+ with marked_timer("reward", timing_raw, color="yellow"):
+ # compute reward model score
+ if self.use_rm:
+ reward_tensor = self.rm_wg.compute_rm_score(batch)
+ batch = batch.union(reward_tensor)
+
+ if self.config.reward_model.launch_reward_fn_async:
+ future_reward = compute_reward_async.remote(data=batch, reward_fn=self.reward_fn)
+ else:
+ reward_tensor, reward_extra_infos_dict = compute_reward(batch, self.reward_fn)
+
+ with marked_timer("old_log_prob", timing_raw, color="blue"):
+
+ def compute_old_log_prob(batch):
+ old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
+ entropys = old_log_prob.batch["entropys"]
+ response_masks = batch.batch["response_mask"]
+ actor_config = self.config.actor_rollout_ref.actor
+ entropy_agg = agg_loss(
+ loss_mat=entropys,
+ loss_mask=response_masks,
+ loss_agg_mode=actor_config.loss_agg_mode,
+ loss_scale_factor=actor_config.loss_scale_factor,
+ )
+ old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()}
+ metrics.update(old_log_prob_metrics)
+ old_log_prob.batch.pop("entropys")
+ batch = batch.union(old_log_prob)
+ if "rollout_log_probs" in batch.batch.keys():
+ # TODO: we may want to add diff of probs too.
+ from verl.utils.debug.metrics import calculate_debug_metrics
+
+ metrics.update(calculate_debug_metrics(batch))
+ return batch
+
+ async_training = self.config.get("async_training", None)
+ if async_training and async_training.use_rollout_log_probs:
+ # If local_triger_step == 1, load the training engine's parameters to the CPU
+ # and save a copy for subsequent MIS use.
+ # If local_trigger_step == 2, 3, ..., restore the parameters of version 1 to calculate the old_log_prob,
+ # then restore the parameters of the current version.
+ if local_trigger_step == 1:
+ self.actor_rollout_wg.save_model_to_cpu(1)
+ batch = compute_old_log_prob(batch)
+ elif local_trigger_step is not None:
+ self.actor_rollout_wg.save_model_to_cpu(local_trigger_step)
+ self.actor_rollout_wg.restore_model_from_cpu(1)
+ batch = compute_old_log_prob(batch)
+ self.actor_rollout_wg.restore_model_from_cpu(local_trigger_step)
+ self.actor_rollout_wg.clear_cpu_model(local_trigger_step)
+ else:
+ batch.batch["old_log_probs"] = batch.batch["rollout_log_probs"]
+ batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature
+
+ else:
+ batch = compute_old_log_prob(batch)
+
+ if self.use_reference_policy:
+ # compute reference log_prob
+ with marked_timer("ref", timing_raw, color="olive"):
+ if not self.ref_in_actor:
+ ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
+ else:
+ ref_log_prob = self.actor_rollout_wg.compute_ref_log_prob(batch)
+ batch = batch.union(ref_log_prob)
+
+ # compute values
+ if self.use_critic:
+ with marked_timer("values", timing_raw, color="cyan"):
+ values = self.critic_wg.compute_values(batch)
+ batch = batch.union(values)
+
+ with marked_timer("adv", timing_raw, color="brown"):
+ # we combine with rule-based rm
+ reward_extra_infos_dict: dict[str, list]
+ if self.config.reward_model.launch_reward_fn_async:
+ reward_tensor, reward_extra_infos_dict = ray.get(future_reward)
+ batch.batch["token_level_scores"] = reward_tensor
+
+ if reward_extra_infos_dict:
+ batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})
+
+ # compute rewards. apply_kl_penalty if available
+ if self.config.algorithm.use_kl_in_reward:
+ batch, kl_metrics = apply_kl_penalty(
+ batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
+ )
+ metrics.update(kl_metrics)
+ else:
+ batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
+
+ # Compute rollout correction weights centrally (once per batch)
+ # This corrects for off-policy issues (policy mismatch, model staleness, etc.)
+ # Also computes off-policy diagnostic metrics (KL, PPL, etc.)
+ from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch
+
+ rollout_corr_config = self.config.algorithm.get("rollout_correction", None)
+ if rollout_corr_config is not None and "rollout_log_probs" in batch.batch:
+ batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config)
+ # IS and off-policy metrics already have rollout_corr/ prefix
+ metrics.update(is_metrics)
+
+ # compute advantages, executed on the driver process
+ norm_adv_by_std_in_grpo = self.config.algorithm.get(
+ "norm_adv_by_std_in_grpo", True
+ ) # GRPO adv normalization factor
+
+ batch = compute_advantage(
+ batch,
+ adv_estimator=self.config.algorithm.adv_estimator,
+ gamma=self.config.algorithm.gamma,
+ lam=self.config.algorithm.lam,
+ num_repeat=self.config.actor_rollout_ref.rollout.n,
+ norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+ config=self.config.algorithm,
+ )
+
+ # update critic
+ if self.use_critic:
+ with marked_timer("update_critic", timing_raw, color="pink"):
+ critic_output = self.critic_wg.update_critic(batch)
+ critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"])
+ metrics.update(critic_output_metrics)
+
+ # implement critic warmup
+ if self.config.trainer.critic_warmup <= self.global_steps:
+ # update actor
+ with marked_timer("update_actor", timing_raw, color="red"):
+ batch.meta_info["multi_turn"] = self.config.actor_rollout_ref.rollout.multi_turn.enable
+ actor_output = self.actor_rollout_wg.update_actor(batch)
+ actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])
+ metrics.update(actor_output_metrics)
+ return batch, reward_extra_infos_dict
+
+ def _log_rollout(self, batch, reward_extra_infos_dict, timing_raw):
+ # Log rollout generations if enabled
+ rollout_data_dir = self.config.trainer.get("rollout_data_dir", None)
+ if rollout_data_dir:
+ with marked_timer("dump_rollout_generations", timing_raw, color="green"):
+ inputs = self.tokenizer.batch_decode(batch.batch["prompts"], skip_special_tokens=True)
+ outputs = self.tokenizer.batch_decode(batch.batch["responses"], skip_special_tokens=True)
+ scores = batch.batch["token_level_scores"].sum(-1).cpu().tolist()
+ sample_gts = [item.non_tensor_batch.get("reward_model", {}).get("ground_truth", None) for item in batch]
+
+ if "request_id" in batch.non_tensor_batch:
+ reward_extra_infos_dict.setdefault(
+ "request_id",
+ batch.non_tensor_batch["request_id"].tolist(),
+ )
+
+ self._dump_generations(
+ inputs=inputs,
+ outputs=outputs,
+ gts=sample_gts,
+ scores=scores,
+ reward_extra_infos_dict=reward_extra_infos_dict,
+ dump_path=rollout_data_dir,
+ )
+
+ def _validate_metrics(self, is_last_step, last_val_metrics, metrics, timing_raw):
+ if (
+ self.val_reward_fn is not None
+ and self.config.trainer.test_freq > 0
+ and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
+ ):
+ with marked_timer("testing", timing_raw, color="green"):
+ val_metrics: dict = self._validate()
+ if is_last_step:
+ last_val_metrics = val_metrics
+ metrics.update(val_metrics)
+ return last_val_metrics
+
+ def _collect_metrics(self, batch, epoch, metrics, timing_raw):
+ steps_duration = timing_raw["step"]
+ self.max_steps_duration = max(self.max_steps_duration, steps_duration)
+
+ # training metrics
+ metrics.update(
+ {
+ "training/global_step": self.global_steps,
+ "training/epoch": epoch,
+ }
+ )
+ # collect metrics
+ metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
+ metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
+ # TODO: implement actual tflpo and theoretical tflpo
+ n_gpus = self.resource_pool_manager.get_n_gpus()
+ metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))
+
+ def _post_batch_processing(self, batch: DataProto):
+ # this is experimental and may be changed/removed in the future in favor of a general-purpose one
+ if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):
+ self.train_dataloader.sampler.update(batch=batch)
+
+ # this is experimental and may be changed/removed in the future
+ # in favor of a general-purpose data buffer pool
+ if hasattr(self.train_dataset, "on_batch_end"):
+ # The dataset may be changed after each training batch
+ self.train_dataset.on_batch_end(batch=batch)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/sglang_rollout/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/sglang_rollout/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9cd3ed5b8e9f967b0e91ce33ffa01d4902e69a38
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/sglang_rollout/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2025 Meituan 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.
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/sglang_rollout/sglang_async_server.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/sglang_rollout/sglang_async_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..d52434f8b4ced8972b42e6f982a510ed433120b2
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/sglang_rollout/sglang_async_server.py
@@ -0,0 +1,189 @@
+# Copyright 2025 Meituan 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.
+import asyncio
+import logging
+from typing import Any, Optional
+
+import ray
+import torch
+from ray.actor import ActorHandle
+
+from verl.workers.config import HFModelConfig, RewardModelConfig, RolloutConfig
+from verl.workers.rollout.replica import RolloutMode
+from verl.workers.rollout.sglang_rollout.async_sglang_server import (
+ SGLangHttpServer,
+ SGLangReplica,
+)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(logging.INFO)
+
+
+class SGLangHttpServerForPartial(SGLangHttpServer):
+ def __init__(
+ self,
+ config: RolloutConfig | RewardModelConfig,
+ model_config: HFModelConfig,
+ rollout_mode: RolloutMode,
+ workers: list[ActorHandle],
+ replica_rank: int,
+ node_rank: int,
+ nnodes: int,
+ cuda_visible_devices: str,
+ base_gpu_id: int,
+ ):
+ super().__init__(
+ config=config,
+ model_config=model_config,
+ rollout_mode=rollout_mode,
+ workers=workers,
+ replica_rank=replica_rank,
+ node_rank=node_rank,
+ nnodes=nnodes,
+ cuda_visible_devices=cuda_visible_devices,
+ base_gpu_id=base_gpu_id,
+ )
+
+ # for cancel LLMServer
+ self.paused = False
+ self.lock = asyncio.Lock()
+ self.cancel_event: dict[str, asyncio.Event] = {}
+ self.req_output: dict[str, Optional[dict[str, Any]]] = {}
+
+ async def _generate_step(
+ self,
+ prompt_ids: torch.Tensor,
+ sampling_params: dict[str, Any],
+ request_id: str,
+ image_data: Optional[list[Any]] = None,
+ ) -> None:
+ sampling_params = dict(sampling_params)
+
+ max_new_tokens = min(
+ self.config.response_length,
+ self.config.max_model_len - len(prompt_ids) - 1,
+ )
+ sampling_params["max_new_tokens"] = max_new_tokens
+
+ sampling_params.setdefault(
+ "repetition_penalty",
+ self.config.get("repetition_penalty", 1.0),
+ )
+
+ sampling_params.pop("logprobs", None)
+ return_logprob = True
+ from sglang.srt.managers.io_struct import GenerateReqInput
+
+ request = GenerateReqInput(
+ rid=request_id,
+ input_ids=prompt_ids,
+ sampling_params=sampling_params,
+ return_logprob=return_logprob,
+ image_data=image_data,
+ )
+ generator = self.tokenizer_manager.generate_request(request, None)
+ async for output in generator:
+ self.req_output[request_id] = output
+
+ assert self.req_output[request_id] is not None
+
+ async def generate_for_partial(
+ self,
+ prompt_ids: torch.Tensor,
+ sampling_params: dict[str, Any],
+ request_id: str,
+ image_data: Optional[list[Any]] = None,
+ ) -> tuple[list[int], list[float], bool]:
+ async with self.lock:
+ if self.paused:
+ return [], [], True
+ self.req_output[request_id] = None
+ self.cancel_event[request_id] = asyncio.Event()
+ cancel_handle = asyncio.create_task(self.cancel_event[request_id].wait())
+ generation_handle = asyncio.create_task(
+ self._generate_step(prompt_ids, sampling_params, request_id, image_data)
+ )
+ done, pending = await asyncio.wait(
+ [generation_handle, cancel_handle],
+ return_when=asyncio.FIRST_COMPLETED,
+ )
+ for task in done:
+ await task
+
+ for task in pending:
+ task.cancel()
+ async with self.lock:
+ output = self.req_output.get(request_id)
+ if output is None:
+ self.cancel_event.pop(request_id, None)
+ self.req_output.pop(request_id, None)
+ return [], [], True
+ meta_info = output.get("meta_info", {})
+ output_token_logprobs = meta_info.get("output_token_logprobs")
+
+ token_ids: list[int] = []
+ log_probs: list[float] = []
+
+ if output_token_logprobs is not None:
+ for log_prob, token_id, _ in output_token_logprobs:
+ token_ids.append(int(token_id))
+ log_probs.append(float(log_prob))
+ else:
+ token_ids = list(output["output_ids"])
+ log_probs = []
+ is_cancel = generation_handle not in done
+ self.cancel_event.pop(request_id, None)
+ self.req_output.pop(request_id, None)
+
+ return token_ids, log_probs, is_cancel
+
+ async def cancel(self):
+ async with self.lock:
+ self.paused = True
+ for request_id in self.cancel_event:
+ self.cancel_event[request_id].set()
+
+ async def resume(self):
+ async with self.lock:
+ self.paused = False
+
+ async def reset_prefix_cache(self):
+ async with self.lock:
+ print("Reset prefix cache ...")
+ await self.tokenizer_manager.flush_cache()
+
+
+class FullyAsyncSGLangReplica(SGLangReplica):
+ def __init__(
+ self,
+ replica_rank: int,
+ config: RolloutConfig | RewardModelConfig,
+ model_config: HFModelConfig,
+ gpus_per_node: int = 8,
+ is_reward_model: bool = False,
+ ):
+ super().__init__(replica_rank, config, model_config, gpus_per_node, is_reward_model)
+ self.server_class = ray.remote(SGLangHttpServerForPartial)
+
+ async def cancel(self):
+ """Cancel each rollout server."""
+ await asyncio.gather(*[server.cancel.remote() for server in self.servers])
+
+ async def resume(self):
+ """Resume each rollout server."""
+ await asyncio.gather(*[server.resume.remote() for server in self.servers])
+
+ async def reset_prefix_cache(self):
+ """reset kv cache in each rollout server."""
+ await asyncio.gather(*[server.reset_prefix_cache.remote() for server in self.servers])
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_30b_a3b_base_math_fsdp.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_30b_a3b_base_math_fsdp.sh
new file mode 100644
index 0000000000000000000000000000000000000000..09b22145e2665c13f63c977126fc53bccb9cf78a
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_30b_a3b_base_math_fsdp.sh
@@ -0,0 +1,191 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO-Qwen3-30B-A3B-Base-Async'
+exp_name='Fsdp2-tp4sp4'
+
+# Ray
+RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+DATA_PATH=${RAY_DATA_HOME:-"${HOME}/verl"}
+DATA_PATH=${DATA_PATH:-"/mnt/bn/${BYTENAS}"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${DATA_PATH}/shared/models/Qwen3-30B-A3B-Base"}
+CKPTS_DIR=${CKPTS_DIR:-"${DATA_PATH}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${DATA_PATH}/shared/data/dapo-math/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${DATA_PATH}/shared/data/dapo-math/aime-2024.parquet"}
+
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 20))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+enable_filter_groups=True
+filter_groups_metric=acc
+max_num_gen_batches=10
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+
+NNODES=${NNODES:-4}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+# Fully async specific parameters
+n_gpus_rollout=8
+n_gpus_training=8
+n_nodes_rollout=2
+n_nodes_train=2 # $((NNODES - n_nodes_rollout))
+
+train_bsz=512
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((train_bsz * 400)))
+test_freq=25
+staleness_threshold=0.6 # 0 0.3 1
+require_batches=1
+total_train_gpus=$((n_gpus_training * n_nodes_train))
+total_rollout_gpus=$((n_gpus_rollout * n_nodes_rollout))
+trigger_parameter_sync_step=$((train_bsz / ( train_prompt_mini_bsz * require_batches))) # 8 16 32
+partial_rollout=True
+enforce_eager=False
+nccl_timeout=72000
+enable_sleep_mode=False
+
+# Performance Related Parameter
+sp_size=4
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$((max_prompt_length + max_response_length))
+infer_ppo_max_token_len=$((max_prompt_length + max_response_length))
+ref_offload=True
+actor_offload=False
+gen_tp=4
+fsdp_size=-1
+
+
+ray job submit --no-wait --runtime-env="${RUNTIME_ENV}" \
+ --working-dir "${WORKING_DIR}" \
+ --address "${RAY_ADDRESS}" \
+ -- python3 -m verl.experimental.fully_async_policy.fully_async_main \
+ --config-path=config \
+ --config-name='fully_async_dapo_trainer.yaml' \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ actor_rollout_ref.actor.strategy=fsdp \
+ critic.strategy=fsdp \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ actor_rollout_ref.nccl_timeout=${nccl_timeout} \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.50 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ +actor_rollout_ref.rollout.enable_sleep_mode=${enable_sleep_mode} \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.enforce_eager=${enforce_eager} \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ reward_model.reward_manager=dapo \
+ reward_model.overlong_buffer.enable=${enable_overlong_buffer} \
+ reward_model.overlong_buffer.len=${overlong_buffer_len} \
+ reward_model.overlong_buffer.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','wandb'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}-i${total_rollout_gpus}_t${total_train_gpus}_s${staleness_threshold}" \
+ trainer.val_before_train=True \
+ trainer.test_freq="${test_freq}" \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${n_nodes_train}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${n_nodes_rollout}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.test_freq=${test_freq} \
+ rollout.total_epochs=10 \
+ async_training.require_batches=${require_batches} \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_async_retool.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_async_retool.sh
new file mode 100644
index 0000000000000000000000000000000000000000..b11705d8eca3fe378e11bb82ed76d5dd211cb6bc
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_async_retool.sh
@@ -0,0 +1,141 @@
+set -x
+
+export VLLM_USE_V1=1
+
+# ================= data/model/tool =================
+HDFS_ROOT=${HDFS_ROOT:-$PWD}
+DATA_ROOT=${DATA_ROOT:-$PWD}
+
+dapo_math_17k=$DATA_ROOT/dataset/BytedTsinghua-SIA/DAPO-Math-17k
+aime_2024=$DATA_ROOT/dataset/Maxwell-Jia/AIME_2024
+aime_2025=$DATA_ROOT/dataset/yentinglin/aime_2025
+model_path=$HDFS_ROOT/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372
+
+train_files="['$dapo_math_17k']"
+test_files="['$aime_2025', '$aime_2024']"
+
+# tool
+tool_config_path=recipe/retool/sandbox_fusion_tool_config.yaml
+retool_path=recipe/retool/retool.py
+
+# wandb / tensorboard
+project_name=retool
+experiment_name=qwen2.5-7b_dapo_async_tool
+default_local_dir=$DATA_ROOT/checkpoint/$experiment_name
+
+# ================= algorithm =================
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+max_turns=16
+max_prompt_length=2048
+max_response_length=16384
+actor_lr=1e-6
+
+# ================= perfomance =================
+infer_tp=4 # vllm
+train_sp=4 # train
+fsdp_size=4 # train
+offload=False
+
+actor_max_token_len_per_gpu=$(( (max_prompt_length + max_response_length) * 1 ))
+log_prob_max_token_len_per_gpu=$(( actor_max_token_len_per_gpu * 4 ))
+
+# ================= async policy =================
+rollout_name="vllm"
+rollout_mode="async"
+
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+n_gpus_rollout=4
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+train_batch_size=0
+ppo_mini_batch_size=16
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+n_resp_per_prompt_val=30
+total_rollout_steps=$(((64*250)))
+test_freq=10
+staleness_threshold=0.5
+trigger_parameter_sync_step=4
+require_batches=1
+partial_rollout=True
+
+python3 -m verl.experimental.fully_async_policy.fully_async_main \
+ algorithm.adv_estimator=$adv_estimator \
+ algorithm.use_kl_in_reward=$use_kl_in_reward \
+ algorithm.kl_ctrl.kl_coef=$kl_coef \
+ data.train_files="$train_files" \
+ data.val_files="$test_files" \
+ data.return_raw_chat=True \
+ data.train_batch_size=$train_batch_size \
+ data.max_prompt_length=$max_prompt_length \
+ data.max_response_length=$max_response_length \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.custom_cls.path=$retool_path \
+ data.custom_cls.name=CustomRLHFDataset \
+ custom_reward_function.path=$retool_path \
+ custom_reward_function.name=compute_score \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.model.path=$model_path \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \
+ actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \
+ actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \
+ actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.actor.optim.lr=$actor_lr \
+ actor_rollout_ref.actor.use_dynamic_bsz=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=$train_sp \
+ actor_rollout_ref.actor.fsdp_config.param_offload=$offload \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=$offload \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=$log_prob_max_token_len_per_gpu \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.mode=async \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \
+ actor_rollout_ref.rollout.multi_turn.enable=True \
+ actor_rollout_ref.rollout.multi_turn.max_user_turns=$max_turns \
+ actor_rollout_ref.rollout.multi_turn.max_assistant_turns=$max_turns \
+ actor_rollout_ref.rollout.multi_turn.tool_config_path=$tool_config_path \
+ actor_rollout_ref.rollout.multi_turn.format=hermes \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=$n_resp_per_prompt \
+ actor_rollout_ref.rollout.val_kwargs.top_p=0.6 \
+ actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \
+ actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name=$project_name \
+ trainer.experiment_name=$experiment_name \
+ trainer.val_before_train=True \
+ trainer.log_val_generations=20 \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir=$default_local_dir \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ trainer.nnodes=$NNODES \
+ trainer.n_gpus_per_node=$n_gpus_training \
+ rollout.nnodes=$NNODES \
+ rollout.n_gpus_per_node=$n_gpus_rollout \
+ rollout.total_rollout_steps=$total_rollout_steps \
+ rollout.total_epochs=10 \
+ rollout.test_freq=$test_freq \
+ async_training.staleness_threshold=$staleness_threshold \
+ async_training.trigger_parameter_sync_step=$trigger_parameter_sync_step \
+ async_training.require_batches=$require_batches \
+ async_training.partial_rollout=$partial_rollout \
+ async_training.use_rollout_log_probs=True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_16_16.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_16_16.sh
new file mode 100644
index 0000000000000000000000000000000000000000..59c83b166b6cfb20baf1dd59536d22dffd9b8b81
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_16_16.sh
@@ -0,0 +1,162 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_16-16'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 28))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=4
+sp_size=4
+fsdp_size=8
+
+# Fully async specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-2}
+NNODES_TRAIN=${NNODES_TRAIN:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*400)))
+test_freq=20
+staleness_threshold=0.1
+trigger_parameter_sync_step=4
+require_batches=4
+partial_rollout=True
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_32_32.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_32_32.sh
new file mode 100644
index 0000000000000000000000000000000000000000..7203652da414bb6cc9bc20a4abf031b905e01802
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_32_32.sh
@@ -0,0 +1,162 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_32-32'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 28))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=4
+sp_size=4
+fsdp_size=8
+
+# Fully async specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-4}
+NNODES_TRAIN=${NNODES_TRAIN:-4}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*400)))
+test_freq=20
+staleness_threshold=0.1
+trigger_parameter_sync_step=4
+require_batches=4
+partial_rollout=True
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_12.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_12.sh
new file mode 100644
index 0000000000000000000000000000000000000000..300cc4551db0b5aac8acaf184312a683b65fdc1f
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_12.sh
@@ -0,0 +1,164 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-fully-async-4-12'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=1
+sp_size=1
+fsdp_size=2
+
+# Fully async specific parameters
+NNODES=${NNODES:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=2
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*100)))
+test_freq=10
+staleness_threshold=0.1
+trigger_parameter_sync_step=4
+require_batches=4
+partial_rollout=True
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.test_freq="${test_freq}" \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_4.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_4.sh
new file mode 100644
index 0000000000000000000000000000000000000000..2dd0adc0ef7063b8e5b7ea134518e2a788c87d77
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_4.sh
@@ -0,0 +1,164 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-fully-async-4-4'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=1
+sp_size=1
+fsdp_size=2
+
+# Fully async specific parameters
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=4
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*100)))
+test_freq=10
+staleness_threshold=0.1
+trigger_parameter_sync_step=4
+require_batches=4
+partial_rollout=True
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=False \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64.sh
new file mode 100644
index 0000000000000000000000000000000000000000..6c8341691a83182d4cf48c678c25548f707c629c
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64.sh
@@ -0,0 +1,162 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_64-64'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 28))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=4
+sp_size=4
+fsdp_size=8
+
+# Fully async specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-8}
+NNODES_TRAIN=${NNODES_TRAIN:-8}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*400)))
+test_freq=20
+staleness_threshold=0.5
+trigger_parameter_sync_step=4
+require_batches=4
+partial_rollout=True
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64_mis.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64_mis.sh
new file mode 100644
index 0000000000000000000000000000000000000000..70237d8725a52bab7bf21707430adcd0f345a6c8
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64_mis.sh
@@ -0,0 +1,173 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_64-64'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 28))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=4
+sp_size=4
+fsdp_size=8
+
+# Fully async specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-8}
+NNODES_TRAIN=${NNODES_TRAIN:-8}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*400)))
+test_freq=20
+staleness_threshold=0.5
+trigger_parameter_sync_step=4
+require_batches=4
+partial_rollout=True
+
+# Rollout Correction
+rollout_is=token
+rollout_is_threshold=2.0
+rollout_rs=seq_mean_k1
+rollout_rs_threshold="0.99_1.001"
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True \
+ async_training.compute_prox_log_prob=True \
+ algorithm.rollout_correction.rollout_is=${rollout_is} \
+ algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} \
+ algorithm.rollout_correction.rollout_rs=${rollout_rs} \
+ algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold}
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_8_8.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_8_8.sh
new file mode 100644
index 0000000000000000000000000000000000000000..ec107948395014db3d7ff62d693d43d51cfa1fdc
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_8_8.sh
@@ -0,0 +1,162 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-fully-async-8-8'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=1
+sp_size=1
+fsdp_size=2
+
+# Fully async specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-1}
+NNODES_TRAIN=${NNODES_TRAIN:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+total_rollout_steps=$(((512*100)))
+test_freq=10
+staleness_threshold=0.1
+trigger_parameter_sync_step=4
+require_batches=4
+partial_rollout=True
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.save_freq=-1 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/geo3k_qwen25vl_7b_megatron_4_4.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/geo3k_qwen25vl_7b_megatron_4_4.sh
new file mode 100644
index 0000000000000000000000000000000000000000..251c0ae840a6d489b3f5daf16c8d0e4577fa49d7
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/geo3k_qwen25vl_7b_megatron_4_4.sh
@@ -0,0 +1,111 @@
+set -x
+ENGINE=${1:-vllm}
+export CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping
+
+
+HF_MODEL_PATH=${HF_MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-VL-7B-Instruct"}
+
+train_path=$HOME/data/geo3k/train.parquet
+test_path=$HOME/data/geo3k/test.parquet
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# Fully async specific parameters
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=4
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=4
+train_prompt_mini_bsz=128
+total_rollout_steps=$(((512*100)))
+test_freq=5
+staleness_threshold=0.1
+trigger_parameter_sync_step=4
+require_batches=2
+partial_rollout=True
+total_epochs=200
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ --config-path=config \
+ --config-name='fully_async_ppo_megatron_trainer.yaml'\
+ algorithm.adv_estimator=grpo \
+ data.train_files="$train_path" \
+ data.val_files="$test_path" \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.max_prompt_length=1024 \
+ data.max_response_length=2048 \
+ actor_rollout_ref.rollout.max_model_len=32768 \
+ actor_rollout_ref.rollout.max_num_batched_tokens=32768 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.gen_batch_size=${gen_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.model.path=$HF_MODEL_PATH \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_decay_steps=51200 \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
+ actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=1 \
+ actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.01 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.actor.use_dynamic_bsz=True \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=5120 \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=5120 \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=5120 \
+ actor_rollout_ref.rollout.name=$ENGINE \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \
+ actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=1 \
+ actor_rollout_ref.ref.megatron.tensor_model_parallel_size=4 \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \
+ actor_rollout_ref.actor.megatron.use_mbridge=True \
+ actor_rollout_ref.actor.megatron.param_offload=True \
+ actor_rollout_ref.actor.megatron.optimizer_offload=True \
+ actor_rollout_ref.actor.megatron.grad_offload=True \
+ actor_rollout_ref.ref.megatron.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_example_geo3k' \
+ trainer.experiment_name='qwen2_5_vl_7b_megatron_async' \
+ trainer.test_freq="${test_freq}" \
+ trainer.total_epochs="${total_epochs}" \
+ trainer.val_before_train=False \
+ trainer.save_freq=-1 \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs="${total_epochs}" \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32.sh
new file mode 100644
index 0000000000000000000000000000000000000000..bb25144481ef792383be9bc1a6e231efe11ee5b7
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32.sh
@@ -0,0 +1,230 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='GRPO-Qwen3-30b-Base-MATH'
+exp_name='GRPO-Qwen3-30b-Base-MATH-megatron-fully-async_96-32'
+
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-30B-A3B-Base"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=True
+kl_loss_coef=0.001
+kl_loss_type=low_var_kl
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length)))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length)))
+offload=True
+train_ppo_micro_batch_size_per_gpu=2
+infer_ppo_micro_batch_size_per_gpu=2
+
+optimizer_offload_fraction=${OFFLOAD_FRACTION:-1.}
+
+COMMON_PP=${COMMON_PP:-1}
+COMMON_VPP=${COMMON_VPP:-null}
+COMMON_CP=${COMMON_CP:-2}
+COMMON_TP=${COMMON_TP:-2}
+COMMON_EP=${COMMON_EP:-8}
+COMMON_ETP=${COMMON_ETP:-1}
+
+TRAIN_TP=${TRAIN_TP:-$COMMON_TP}
+INFER_TP=${INFER_TP:-4}
+
+ACTOR_PP=${ACTOR_PP:-$COMMON_PP}
+ACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP}
+ACTOR_CP=${ACTOR_CP:-$COMMON_CP}
+ACTOR_TP=${ACTOR_TP:-$TRAIN_TP}
+ACTOR_EP=${ACTOR_EP:-$COMMON_EP}
+ACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP}
+ROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP}
+REF_PP=${REF_PP:-$COMMON_PP}
+REF_VPP=${REF_VPP:-$COMMON_VPP}
+REF_CP=${REF_CP:-$COMMON_CP}
+REF_TP=${REF_TP:-$TRAIN_TP}
+REF_EP=${REF_EP:-$COMMON_EP}
+REF_ETP=${REF_ETP:-$COMMON_ETP}
+CRITIC_PP=${CRITIC_PP:-$COMMON_PP}
+CRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP}
+CRITIC_CP=${CRITIC_CP:-$COMMON_CP}
+CRITIC_TP=${CRITIC_TP:-$TRAIN_TP}
+CRITIC_EP=${CRITIC_EP:-$COMMON_EP}
+CRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP}
+RM_PP=${RM_PP:-$COMMON_PP}
+RM_VPP=${RM_VPP:-$COMMON_VPP}
+RM_CP=${RM_CP:-$COMMON_CP}
+RM_TP=${RM_TP:-$TRAIN_TP}
+RM_EP=${RM_EP:-$COMMON_EP}
+RM_ETP=${RM_ETP:-$COMMON_ETP}
+
+# install mbridge
+# pip3 install git+https://github.com/ISEEKYAN/mbridge
+USE_MBRIDGE=True
+USE_DIST_CKPT=False
+
+# Fully async specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-12}
+NNODES_TRAIN=${NNODES_TRAIN:-4}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=128
+total_rollout_steps=$(((512*400)))
+test_freq=20
+staleness_threshold=0.5
+trigger_parameter_sync_step=4
+require_batches=1
+partial_rollout=True
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ --config-path=config \
+ --config-name='fully_async_ppo_megatron_trainer.yaml'\
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ +actor_rollout_ref.model.override_config.model_config.max_position_embeddings=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.model.use_fused_kernels=False \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_ppo_micro_batch_size_per_gpu} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.lr_decay_style='constant' \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.optim.lr_decay_steps=${total_rollout_steps} \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \
+ actor_rollout_ref.actor.megatron.use_mbridge=$USE_MBRIDGE \
+ actor_rollout_ref.actor.megatron.use_dist_checkpointing=$USE_DIST_CKPT \
+ actor_rollout_ref.actor.megatron.param_offload=${offload} \
+ actor_rollout_ref.actor.megatron.grad_offload=${offload} \
+ actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \
+ actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \
+ actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} \
+ actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} \
+ actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} \
+ actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} \
+ actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type="flex" \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${INFER_TP} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.rollout.enforce_eager=True \
+ actor_rollout_ref.rollout.free_cache_engine=True \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \
+ actor_rollout_ref.ref.megatron.param_offload=${offload} \
+ actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${REF_TP} \
+ actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${REF_PP} \
+ actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${REF_VPP} \
+ actor_rollout_ref.ref.megatron.context_parallel_size=${REF_CP} \
+ actor_rollout_ref.ref.megatron.expert_model_parallel_size=${REF_EP} \
+ actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${REF_ETP} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10 \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True \
+
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32_mis.sh b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32_mis.sh
new file mode 100644
index 0000000000000000000000000000000000000000..ed0716e8c24c89dc587397b442b6f2430d79b1a0
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32_mis.sh
@@ -0,0 +1,239 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='GRPO-Qwen3-30b-Base-MATH'
+exp_name='GRPO-Qwen3-30b-Base-MATH-megatron-fully-async_96-32'
+
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-30B-A3B-Base"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=True
+kl_loss_coef=0.001
+kl_loss_type=low_var_kl
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length)))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length)))
+offload=True
+train_ppo_micro_batch_size_per_gpu=2
+infer_ppo_micro_batch_size_per_gpu=2
+
+optimizer_offload_fraction=${OFFLOAD_FRACTION:-1.}
+
+COMMON_PP=${COMMON_PP:-1}
+COMMON_VPP=${COMMON_VPP:-null}
+COMMON_CP=${COMMON_CP:-2}
+COMMON_TP=${COMMON_TP:-2}
+COMMON_EP=${COMMON_EP:-8}
+COMMON_ETP=${COMMON_ETP:-1}
+
+TRAIN_TP=${TRAIN_TP:-$COMMON_TP}
+INFER_TP=${INFER_TP:-4}
+
+ACTOR_PP=${ACTOR_PP:-$COMMON_PP}
+ACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP}
+ACTOR_CP=${ACTOR_CP:-$COMMON_CP}
+ACTOR_TP=${ACTOR_TP:-$TRAIN_TP}
+ACTOR_EP=${ACTOR_EP:-$COMMON_EP}
+ACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP}
+ROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP}
+REF_PP=${REF_PP:-$COMMON_PP}
+REF_VPP=${REF_VPP:-$COMMON_VPP}
+REF_CP=${REF_CP:-$COMMON_CP}
+REF_TP=${REF_TP:-$TRAIN_TP}
+REF_EP=${REF_EP:-$COMMON_EP}
+REF_ETP=${REF_ETP:-$COMMON_ETP}
+CRITIC_PP=${CRITIC_PP:-$COMMON_PP}
+CRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP}
+CRITIC_CP=${CRITIC_CP:-$COMMON_CP}
+CRITIC_TP=${CRITIC_TP:-$TRAIN_TP}
+CRITIC_EP=${CRITIC_EP:-$COMMON_EP}
+CRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP}
+RM_PP=${RM_PP:-$COMMON_PP}
+RM_VPP=${RM_VPP:-$COMMON_VPP}
+RM_CP=${RM_CP:-$COMMON_CP}
+RM_TP=${RM_TP:-$TRAIN_TP}
+RM_EP=${RM_EP:-$COMMON_EP}
+RM_ETP=${RM_ETP:-$COMMON_ETP}
+
+# install mbridge
+# pip3 install git+https://github.com/ISEEKYAN/mbridge
+USE_MBRIDGE=True
+USE_DIST_CKPT=False
+
+# Fully async specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-12}
+NNODES_TRAIN=${NNODES_TRAIN:-4}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=0
+gen_prompt_bsz=1
+n_resp_per_prompt=16
+train_prompt_mini_bsz=128
+total_rollout_steps=$(((512*400)))
+test_freq=20
+staleness_threshold=0.5
+trigger_parameter_sync_step=4
+require_batches=1
+partial_rollout=True
+
+# Rollout Importance Sampling
+
+rollout_is=null
+rollout_rs=seq_mean_k1
+rollout_rs_threshold="0.999_1.001"
+
+python -m verl.experimental.fully_async_policy.fully_async_main \
+ --config-path=config \
+ --config-name='fully_async_ppo_megatron_trainer.yaml'\
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ data.return_raw_chat=${return_raw_chat} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ async_training.compute_prox_log_prob=True \
+ algorithm.rollout_correction.rollout_is=${rollout_is} \
+ algorithm.rollout_correction.rollout_rs=${rollout_rs} \
+ algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ +actor_rollout_ref.model.override_config.model_config.max_position_embeddings=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.model.use_fused_kernels=False \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_ppo_micro_batch_size_per_gpu} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.lr_decay_style='constant' \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.optim.lr_decay_steps=${total_rollout_steps} \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \
+ +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \
+ actor_rollout_ref.actor.megatron.use_mbridge=$USE_MBRIDGE \
+ actor_rollout_ref.actor.megatron.use_dist_checkpointing=$USE_DIST_CKPT \
+ actor_rollout_ref.actor.megatron.param_offload=${offload} \
+ actor_rollout_ref.actor.megatron.grad_offload=${offload} \
+ actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \
+ actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \
+ actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} \
+ actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} \
+ actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} \
+ actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} \
+ actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type="flex" \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \
+ +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${INFER_TP} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=${rollout_name} \
+ actor_rollout_ref.rollout.mode=${rollout_mode} \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.rollout.enforce_eager=True \
+ actor_rollout_ref.rollout.free_cache_engine=True \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \
+ actor_rollout_ref.ref.megatron.param_offload=${offload} \
+ actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${REF_TP} \
+ actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${REF_PP} \
+ actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${REF_VPP} \
+ actor_rollout_ref.ref.megatron.context_parallel_size=${REF_CP} \
+ actor_rollout_ref.ref.megatron.expert_model_parallel_size=${REF_EP} \
+ actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${REF_ETP} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10 \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.total_rollout_steps="${total_rollout_steps}" \
+ rollout.total_epochs=10 \
+ rollout.test_freq="${test_freq}" \
+ async_training.staleness_threshold="${staleness_threshold}" \
+ async_training.trigger_parameter_sync_step="${trigger_parameter_sync_step}" \
+ async_training.require_batches="${require_batches}" \
+ async_training.partial_rollout="${partial_rollout}" \
+ async_training.use_rollout_log_probs=True \
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/runtime_env.yaml b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/runtime_env.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..88467b8c2435de0eb2c7aaf9988798b6dfd8da78
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/shell/runtime_env.yaml
@@ -0,0 +1,4 @@
+env_vars:
+ VLLM_USE_V1: "1"
+ NCCL_DEBUG: "INFO"
+ HYDRA_FULL_ERROR: "1"
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/unittest/simple_streaming_demo.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/unittest/simple_streaming_demo.py
new file mode 100644
index 0000000000000000000000000000000000000000..209c2aae39bf2d7386d7b88085ceffb3deff6433
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/unittest/simple_streaming_demo.py
@@ -0,0 +1,176 @@
+# Copyright 2025 Meituan 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.
+
+import asyncio
+import random
+import time
+
+
+class SimpleStreamingSystem:
+ """Simplified streaming system demonstration"""
+
+ def __init__(self, max_concurrent_tasks: int = 4):
+ self.max_concurrent_tasks = max_concurrent_tasks
+ self.data_queue = asyncio.Queue()
+ self.result_queue = asyncio.Queue()
+ self.consumer_count = 0
+
+ # Data stream coroutine
+ async def data_stream(self):
+ # Add initial data
+ # Prepare test data
+ test_data = [{"id": f"task_{i}", "content": f"data_{i}"} for i in range(8)]
+ await self.add_data_stream(test_data)
+
+ # Simulate subsequent data stream
+ await asyncio.sleep(3)
+ print("\nAdding second batch of data...")
+ extra_data = [{"id": f"extra_{i}", "content": f"extra_data_{i}"} for i in range(5)]
+ await self.add_data_stream(extra_data)
+
+ # Send termination signal
+ await asyncio.sleep(1)
+ await self.data_queue.put("DONE")
+ print("Sending termination signal")
+
+ async def add_data_stream(self, data_list: list[dict]):
+ """Simulate data stream"""
+ print("Starting to add data stream...")
+
+ for i, data_item in enumerate(data_list):
+ await self.data_queue.put(data_item)
+ print(f"Data {data_item['id']} added to pending queue")
+
+ # Simulate interval between data streams
+ if i < len(data_list) - 1: # Don't wait after the last item
+ await asyncio.sleep(0.8)
+
+ print("Initial data stream added successfully")
+
+ async def _process_data_async(self, data_item: dict):
+ """Asynchronously process a single data item"""
+ data_id = data_item["id"]
+ content = data_item["content"]
+
+ # Simulate different processing times (1-3 seconds)
+ processing_time = random.uniform(1, 3)
+
+ print(f" Starting to process {data_id}, estimated time {processing_time:.1f}s")
+
+ # Asynchronously wait for processing completion
+ await asyncio.sleep(processing_time)
+
+ result = {
+ "id": data_id,
+ "processed_content": f"Processed {content}",
+ "processing_time": round(processing_time, 2),
+ "completed_at": time.time(),
+ }
+
+ # Immediately put into result queue
+ await self.result_queue.put(result)
+ print(f" {data_id} processing completed! (took {processing_time:.1f}s) -> Added to result queue")
+
+ async def _submit_worker(self):
+ """Stream submission worker coroutine"""
+ active_tasks = set()
+
+ print("Stream submitter started...")
+
+ while True:
+ # Get data to process
+ data_item = await self.data_queue.get()
+
+ if data_item == "DONE":
+ print("Received termination signal, waiting for remaining tasks to complete...")
+ if active_tasks:
+ await asyncio.gather(*active_tasks, return_exceptions=True)
+ break
+
+ # Check concurrent limit
+ while len(active_tasks) >= self.max_concurrent_tasks:
+ print(f"Reached maximum concurrency {self.max_concurrent_tasks}, waiting for tasks to complete...")
+ done_tasks, active_tasks = await asyncio.wait(active_tasks, return_when=asyncio.FIRST_COMPLETED)
+
+ # Clean up completed tasks
+ for task in done_tasks:
+ try:
+ await task
+ print(f"Task completed {task}")
+ except Exception as e:
+ print(f"Task execution failed: {e}")
+
+ # Immediately submit new task
+ task = asyncio.create_task(self._process_data_async(data_item), name=f"active {data_item}")
+ active_tasks.add(task)
+
+ print(f"Submitted task {data_item['id']}, current concurrency: {len(active_tasks)}")
+
+ async def _consumer_worker(self):
+ """Result consumer coroutine"""
+ print("Consumer started...")
+
+ while True:
+ try:
+ # Get processing result from result queue
+ result = await asyncio.wait_for(self.result_queue.get(), timeout=2.0)
+
+ self.consumer_count += 1
+
+ print(
+ f"Consumed #{self.consumer_count}: {result['id']} "
+ f"(processing time {result['processing_time']}s) - {result['processed_content']}"
+ )
+
+ except asyncio.TimeoutError:
+ print(" Consumer waiting...")
+ await asyncio.sleep(0.5)
+
+ async def run_demo(self):
+ """Run demonstration"""
+ print("=" * 60)
+ print(f"Maximum concurrency: {self.max_concurrent_tasks}")
+ print("=" * 60)
+
+ # Start core coroutines
+ stream_task = asyncio.create_task(self.data_stream())
+ submit_task = asyncio.create_task(self._submit_worker())
+ consumer_task = asyncio.create_task(self._consumer_worker())
+
+ try:
+ # Wait for data stream to complete
+ await stream_task
+ print("Data stream completed")
+
+ # Wait for processing to complete
+ await submit_task
+ print("All tasks processed")
+
+ finally:
+ # Cleanup
+ submit_task.cancel()
+ consumer_task.cancel()
+ await asyncio.gather(submit_task, consumer_task, return_exceptions=True)
+
+ print(f"\nFinal statistics: Consumed {self.consumer_count} results")
+
+
+async def main():
+ """Main function"""
+ system = SimpleStreamingSystem(max_concurrent_tasks=3)
+ await system.run_demo()
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/vllm_rollout/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/vllm_rollout/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9cd3ed5b8e9f967b0e91ce33ffa01d4902e69a38
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/vllm_rollout/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2025 Meituan 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.
diff --git a/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/vllm_rollout/vllm_async_server.py b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/vllm_rollout/vllm_async_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..aaed2f948f13adf2f148e10ae5d6c9a57a4a3081
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/fully_async_policy/vllm_rollout/vllm_async_server.py
@@ -0,0 +1,148 @@
+# Copyright 2025 Meituan 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.
+import asyncio
+import logging
+from typing import Any, Optional, Sequence
+
+import ray
+from ray.actor import ActorHandle
+from vllm import SamplingParams
+from vllm.inputs import TokensPrompt
+from vllm.outputs import RequestOutput
+
+from verl.workers.config import HFModelConfig, RolloutConfig
+from verl.workers.rollout.replica import RolloutMode
+from verl.workers.rollout.vllm_rollout.vllm_async_server import (
+ _qwen2_5_vl_dedup_image_tokens,
+ vLLMHttpServer,
+ vLLMReplica,
+)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(logging.INFO)
+
+
+class vLLMHttpServerForPartial(vLLMHttpServer):
+ def __init__(
+ self,
+ config: RolloutConfig,
+ model_config: HFModelConfig,
+ rollout_mode: RolloutMode,
+ workers: list[ActorHandle],
+ replica_rank: int,
+ node_rank: int,
+ gpus_per_node: int,
+ nnodes: int,
+ ):
+ super().__init__(config, model_config, rollout_mode, workers, replica_rank, node_rank, gpus_per_node, nnodes)
+
+ # for cancel LLMServer
+ self.paused = False
+ self.lock = asyncio.Lock()
+ self.cancel_event: dict[str, asyncio.Event] = {}
+ self.req_output: dict[str, Optional[RequestOutput]] = {}
+
+ async def _generate_step(
+ self,
+ prompt_ids: list[int],
+ sampling_params: dict[str, Any],
+ request_id: str,
+ image_data: Optional[list[Any]] = None,
+ ):
+ max_tokens = self.config.max_model_len - len(prompt_ids)
+ sampling_params["logprobs"] = 1
+ sampling_params.setdefault("repetition_penalty", self.config.get("repetition_penalty", 1.0))
+ sampling_params = SamplingParams(max_tokens=max_tokens, **sampling_params)
+ prompt_ids = _qwen2_5_vl_dedup_image_tokens(prompt_ids, self.model_config.processor)
+ prompt = TokensPrompt(
+ prompt_token_ids=prompt_ids, multi_modal_data={"image": image_data} if image_data else None
+ )
+ generator = self.engine.generate(prompt=prompt, sampling_params=sampling_params, request_id=request_id)
+
+ # Get final response
+ async for output in generator:
+ self.req_output[request_id] = output
+ assert self.req_output[request_id] is not None
+
+ async def generate_for_partial(
+ self,
+ prompt_ids: list[int],
+ sampling_params: dict[str, Any],
+ request_id: str,
+ image_data: Optional[list[Any]] = None,
+ ) -> tuple[list[Any], list[Any], bool] | tuple[Sequence[int], list[float], Any]:
+ async with self.lock:
+ if self.paused:
+ # After cancel, all tasks will return directly and wait for the next submission
+ return [], [], True
+ self.req_output[request_id]: Optional[RequestOutput] = None
+ self.cancel_event[request_id] = asyncio.Event()
+ cancel_handle = asyncio.create_task(self.cancel_event[request_id].wait())
+ generation_handle = asyncio.create_task(
+ self._generate_step(prompt_ids, sampling_params, request_id, image_data)
+ )
+
+ done, pend = await asyncio.wait([generation_handle, cancel_handle], return_when=asyncio.FIRST_COMPLETED)
+
+ for task in done:
+ await task
+
+ for task in pend:
+ task.cancel()
+
+ async with self.lock:
+ if self.req_output[request_id] is None:
+ return [], [], True
+ token_ids = self.req_output[request_id].outputs[0].token_ids
+ log_probs: list[float] = []
+ for i, x in enumerate(self.req_output[request_id].outputs[0].logprobs):
+ # In sampling_params, logprobs is set to 1, which should return 1,
+ # but in practice there are multiple. Take the log_prob corresponding to token_id
+ token_id = self.req_output[request_id].outputs[0].token_ids[i]
+ log_probs.append(x[token_id].logprob)
+ is_cancel = generation_handle not in done
+ self.cancel_event.pop(request_id, None)
+ self.req_output.pop(request_id, None)
+ return token_ids, log_probs, is_cancel
+
+ async def cancel(self):
+ async with self.lock:
+ self.paused = True
+ for request_id in self.cancel_event:
+ self.cancel_event[request_id].set()
+
+ async def resume(self):
+ async with self.lock:
+ self.paused = False
+
+
+class FullyAsyncvLLMReplica(vLLMReplica):
+ def __init__(
+ self,
+ replica_rank: int,
+ config: RolloutConfig,
+ model_config: HFModelConfig,
+ gpus_per_node: int = 8,
+ is_reward_model: bool = False,
+ ):
+ super().__init__(replica_rank, config, model_config, gpus_per_node, is_reward_model)
+ self.server_class = ray.remote(vLLMHttpServerForPartial)
+
+ async def cancel(self):
+ """Cancel each rollout server."""
+ await asyncio.gather(*[server.cancel.remote() for server in self.servers])
+
+ async def resume(self):
+ """Resume each rollout server."""
+ await asyncio.gather(*[server.resume.remote() for server in self.servers])
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/README.md b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..e46c95dfb60bd0eaa534bcae154a70fa9ae78737
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/README.md
@@ -0,0 +1,306 @@
+# Recipe: One Step Off Policy Async Trainer
+
+**Author:** `https://github.com/meituan-search`
+
+Last updated: 07/17/2025.
+
+## Introduction
+
+### Background
+
+The current reinforcement learning training process implemented by verl is synchronous, adhering to the algorithmic
+workflows of established methods like PPO, GRPO, and DAPO. In each step, training samples are generated by the latest
+model, and the model is updated after training completes. While this approach aligns with off-policy reinforcement
+learning and stabilizes RL training, but it suffers from severe efficiency issues.
+Model updates must wait for the longest output in the generation phase to complete.
+During the generation of long-tail samples, GPUs remain idle, resulting in significant underutilization.
+The more severe the long-tail problem in sample generation, the lower the overall training efficiency.
+For example, in DAPO 32B training, the Rollout phase accounts for approximately 70% of the total time,
+and increasing resources does not reduce the Rollout duration.
+
+
+
+> source data: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=nwusertongyuxuan361
+
+### Solution
+
+We have implemented the **One Step Off Async Trainer** to help alleviate this issue. This approach parallelizes the
+generation and training processes, utilizing samples generated in the previous step for current training.
+It also involves appropriately partitioning resources, allocating dedicated resources for generation while automatically
+assigning the remainder to training. By reducing resources allocated to the generation phase, we mitigate GPU idle time
+during long-tail sample generation. Throughout this process, generation and training parameters maintain a one-step off
+policy.
+
+
+
+> reference: [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://arxiv.org/abs/2505.24298)
+> original work: [Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models](https://arxiv.org/abs/2410.18252)
+
+Our core contributions include:
+
+1. **Parallel Generation and Training**:
+ Samples for the next batch are asynchronously generated while the current batch is being trained.
+
+2. **Resource Isolation**:
+ Unlike `hybrid_engine`, this method requires explicit resource allocation for rollout, with remaining resources
+ automatically assigned to training.
+
+3. **NCCL Parameter Synchronization**:
+ Employs NCCL communication primitives for seamless parameter transfer between generation and training modules.
+
+### Experimental Results
+
+- **Machine Configuration**: 2 nodes with 16 H20 GPUs each
+ - Generation: 4 GPUs
+ - Training: 12 GPUs
+- **Model**: Qwen2.5-Math-7B
+- **Rollout Configuration**:
+- **Max Response Length**: FSDP2: 20,480 tokens; Megatron: 8,192 tokens
+- **Algorithm**: DAPO
+- **Rollout Engine**: vLLM
+
+| training mode | engine | step | gen | wait_prev_gen | generate_sequences | old_log_prob | update_actor | total time | acc/best@32/mean | acc/maj@32/mean |
+| ---------------------- | ------------- | ---- | --- | ------------- | ------------------ | ------------ | ------------ | -------------- | ---------------- | --------------- |
+| colocate sync | VLLM+FSDP2 | 749 | 321 | - | 247 | 88 | 286 | 19h18m | 0.5948 | 0.417 |
+| one-step-overlap async | VLLM+FSDP2 | 520 | - | 45 | 458 | 108 | 337 | 15h34m（+23%） | 0.6165 | 0.494 |
+| colocate sync | VLLM+Megatron | 699 | 207 | - | 162 | 119 | 344 | 18h21m | 0.605 | 0.4217 |
+| one-step-overlap async | VLLM+Megatron | 566 | - | 59 | 501 | 120 | 347 | 13h06m (+40%) | 0.6569 | 0.4038 |
+
+- colocate sync: step ≈ gen + old_log_prob + update_actor
+- one-step-overlap async: step ≈ wait_prev_gen + old_log_prob + update_actor
+
+
+
+> source data: https://wandb.ai/hou-zg-meituan/one-step-off-policy?nw=nwuserhouzg
+
+## Implementation
+
+### One Step Off Policy Async Pipeline
+
+Our implemented **One Step Off Policy Async Pipeline** integrates seamlessly into existing training logic at minimal
+cost,
+eliminating the need for additional sample storage management. The core mechanism uses `async_gen_next_batch`
+for asynchronous rollout generation while maintaining continuous operation during epoch transitions
+via `create_continuous_iterator`.
+
+```python
+# iterator generator, simplify one-step integration of the training process
+def _create_continuous_iterator(self):
+ for epoch in range(self.config.trainer.total_epochs):
+ iterator = iter(self.train_dataloader)
+ for batch_dict in iterator:
+ yield epoch, batch_dict
+
+
+# read next batch samples, parameters sync and launch asyn gen_seq
+def _async_gen_next_batch(self, continuous_iterator):
+ # read train_data
+ try:
+ epoch, batch_dict = next(continuous_iterator)
+ except StopIteration:
+ return None
+ batch = DataProto.from_single_dict(batch_dict)
+ gen_batch = batch_pocess(batch)
+ # sync weights from actor to rollout
+ self.sync_rollout_weights()
+ # async generation
+ gen_batch_output = self.rollout_wg.async_generate_sequences(gen_batch)
+ # future encapsulated
+ return GenerationBatchFuture(epoch, batch, gen_batch_output)
+
+
+continuous_iterator = self._create_continuous_iterator()
+# run rollout first to achieve one-step-off
+batch_data_future = self._async_gen_next_batch(continuous_iterator)
+
+while batch_data_future is not None:
+ # wait for the gen_seq result from the previous step
+ batch = batch_data_future.get()
+ # launch the next async call to generate sequences
+ batch_data_future = self._async_gen_next_batch(continuous_iterator)
+
+ # compute advantages
+ batch = critic.compute_values(batch)
+ batch = reference.compute_log_prob(batch)
+ batch = reward.compute_reward(batch)
+ batch = compute_advantages(batch)
+
+ # model update
+ critic_metrics = critic.update_critic(batch)
+ actor_metrics = actor.update_actor(batch)
+```
+
+### Parameter Synchronization
+
+The exciting point is that our nccl based weights updating for rollout model has great performance.
+At most of time, the latency is under 300ms, which is negligible for RLHF.
+
+> **sync_rollout_weights**：The time for synchronizing parameters from actor to rollout is extremely fast and can almost
+> be ignored because it is implemented with nccl.
+
+```python
+class ActorRolloutRefWorker:
+ # actor acquires the meta-info of model parameters for parameter sync
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def get_actor_weights_info(self):
+ params = self._get_actor_params()
+ ret = []
+ for key, tensor in params.items():
+ ret.append((key, tensor.size(), tensor.dtype))
+ self._weights_info = ret
+ return ret
+
+ # rollout sets the meta-info of model parameters for parameter sync
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def set_actor_weights_info(self, weights_info):
+ self._weights_info = weights_info
+
+
+class AsyncRayPPOTrainer(RayPPOTrainer):
+ def init_workers(self):
+
+
+...
+# rollout obtains the meta-info of model parameters from the actor for parameter sync
+weights_info = self.actor_wg.get_actor_weights_info()[0]
+self.rollout_wg.set_actor_weights_info(weights_info)
+
+# Create an actor-rollout communication group for parameter sync
+actor_rollout_workers = self.actor_wg.workers + self.rollout_wg.workers
+collective.create_collective_group(
+ actor_rollout_workers,
+ len(actor_rollout_workers),
+ list(range(0, len(actor_rollout_workers))),
+ backend="nccl",
+ group_name="actor_rollout"
+)
+```
+
+```python
+# drive process call the actor and rollout respectively to sync parameters by nccl
+def sync_rollout_weights(self):
+ self.actor_wg.sync_rollout_weights()
+ ray.get(self.rollout_wg.sync_rollout_weights())
+
+
+# fsdp model parameter sync
+@register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+def sync_rollout_weights(self):
+ params = self._get_actor_params() if self._is_actor else None
+ if self._is_rollout:
+ inference_model = (
+ self.rollout.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model
+ )
+ from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+ patch_vllm_moe_model_weight_loader(inference_model)
+ # Model parameters are broadcast tensor-by-tensor from actor to rollout
+ for key, shape, dtype in self._weights_info:
+ tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device())
+ if self._is_actor:
+ assert key in params
+ origin_data = params[key]
+ if hasattr(origin_data, "full_tensor"):
+ origin_data = origin_data.full_tensor()
+ if torch.distributed.get_rank() == 0:
+ tensor.copy_(origin_data)
+ from ray.util.collective import collective
+
+ collective.broadcast(tensor, src_rank=0, group_name="actor_rollout")
+ if self._is_rollout:
+ inference_model.load_weights([(key, tensor)])
+```
+
+### PPO Correctness
+
+To ensure the correctness of the PPO algorithm, we use rollout log_probs for PPO importance sampling.
+For the related algorithm details, please refer to: https://verl.readthedocs.io/en/latest/algo/rollout_corr_math.html
+The default mode is `bypass_ppo_clip`, but other modification strategies can also be explored.
+
+### AgentLoop
+
+In the current implementation, we no longer provide SPMD model rollout mode.
+Instead, we have switched to AgentLoop mode, which also supports multi-turn tool calling.
+
+## Usage
+
+### FSDP2 Configuration Example
+
+```shell
+python3 -m verl.experimental.one_step_off_policy.async_main_ppo \
+ --config-path=config \
+ --config-name='one_step_off_ppo_trainer.yaml' \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ # actor and rollout are placed separately
+ actor_rollout_ref.hybrid_engine=False \
+ # actor and rollout resource
+ trainer.nnodes=1 \
+ trainer.n_gpus_per_node=6 \
+ rollout.nnodes=1 \
+ rollout.n_gpus_per_node=2
+```
+
+### Megatron Configuration Example
+
+```shell
+python3 -m verl.experimental.one_step_off_policy.async_main_ppo \
+ --config-path=config \
+ --config-name='one_step_off_ppo_megatron_trainer.yaml' \
+ actor_rollout_ref.actor.strategy=megatron \
+ # actor and rollout are placed separately
+ actor_rollout_ref.hybrid_engine=False \
+ # actor and rollout resource
+ trainer.nnodes=1 \
+ trainer.n_gpus_per_node=6 \
+ rollout.nnodes=1 \
+ rollout.n_gpus_per_node=2
+```
+
+### Configuration Guidelines
+
+1. **Card Number Relationships**
+ Maintain either of these relationships for optimal batch distribution:
+
+ - `actor_rollout_ref.rollout.n` should be an integer divisor of:
+ `trainer.n_gpus_per_node * trainer.nnodes`
+ - `actor_rollout_ref.rollout.n * data.train_batch_size` should be evenly divisible by:
+ `trainer.n_gpus_per_node * trainer.nnodes`
+
+ > Rationale: Ensures training samples can be evenly distributed across training GPUs when using partial resources for
+ > generation.
+
+2. **Dynamic Resource Tuning**
+ Adjust `trainer.nnodes` `trainer.n_gpus_per_node` `rollout.nnodes` `rollout.n_gpus_per_node` based on phase
+ durations:
+ - **Ideal state**: Rollout and training phases have comparable durations
+ - **Diagnostic metrics**:
+ - Monitor `wait_prev_gen` duration
+ - Analyze `sequence_length` distribution
+ - **Adjustment strategy**: - High `wait_prev_gen` + uniform sequence lengths → Increase rollout resources - High `wait_prev_gen` + long-tail sequences → Optimize stopping criteria (resource increase won't help)
+ > **wait_prev_gen**：The time consumed waiting for the previous rollout to end (the part that is not fully
+ > overlapped).
+ > **Resource Configuration Strategies:**
+ - **Resource-constrained scenario**: Optimize resource utilization by adjusting GPU allocation ratios,
+ keeping the number of nodes equal to allow training and rollout to share nodes;
+ - Configure `trainer.nnodes = rollout.nnodes` with
+ `trainer.n_gpus_per_node + rollout.n_gpus_per_node = physical_gpus_per_node`. Control rollout resource
+ allocation by adjusting `n_gpus_per_node`.
+ - **Resource-abundant scenario**: Optimize performance by adjusting the number of nodes,
+ keeping the number of GPUs per node equal to enable independent scaling of training and rollout
+ parallelism. - Configure `trainer.n_gpus_per_node = rollout.n_gpus_per_node` and control rollout resource allocation by
+ adjusting `trainer.nnodes` and `rollout.nnodes`to achieve optimal performance.
+ > **Note**: The total number of nodes required by the system is not simply `trainer.nnodes + rollout.nnodes`. The
+ > actual calculation depends on GPU capacity:
+ >
+ > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node <= physical_gpus_per_node`,
+ > the required node count is `max(trainer.nnodes, rollout.nnodes)`
+ > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node > physical_gpus_per_node`,
+ > the required node count is `trainer.nnodes + rollout.nnodes`
+
+## Functional Support
+
+| Category | Support Situation |
+| ------------------ | --------------------------------------------------------------------------------------------------------------- |
+| train engine | FSDP2
Megatron |
+| rollout engine | vLLM
SGLang |
+| AdvantageEstimator | GRPO
GRPO_PASSK
REINFORCE_PLUS_PLUS
RLOO
OPO
REINFORCE_PLUS_PLUS_BASELINE
GPG |
+| Reward | all |
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/agent_loop/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/agent_loop/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a9eb0705e41bd8db7f9e6b706d82b20fa52d0d13
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/agent_loop/__init__.py
@@ -0,0 +1,17 @@
+# Copyright 2025 Meituan 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 .agent_loop import OneStepOffAgentLoopManager
+
+__all__ = [OneStepOffAgentLoopManager]
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/agent_loop/agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/agent_loop/agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..85455d655b2ecba630559086166995717edf2073
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/agent_loop/agent_loop.py
@@ -0,0 +1,64 @@
+# Copyright 2025 Meituan 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.
+import asyncio
+import logging
+import os
+
+import ray
+
+from verl.experimental.agent_loop.agent_loop import AgentLoopManager
+from verl.protocol import DataProto
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class OneStepOffAgentLoopManager(AgentLoopManager):
+ async def generate_sequences_async(self, prompts: DataProto) -> DataProto:
+ """Split input batch and dispatch to agent loop workers (async version).
+
+ Args:
+ prompts (DataProto): Input batch.
+
+ Returns:
+ DataProto: Output batch.
+ """
+
+ chunkes = prompts.chunk(len(self.agent_loop_workers))
+ # Use asyncio.gather with ray.get wrapped in asyncio.to_thread to avoid blocking
+ import asyncio
+
+ outputs = await asyncio.gather(
+ *[
+ asyncio.to_thread(ray.get, worker.generate_sequences.remote(chunk))
+ for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True)
+ ]
+ )
+ output = DataProto.concat(outputs)
+
+ # calculate performance metrics
+ metrics = [output.meta_info.pop("metrics") for output in outputs] # List[List[Dict[str, str]]]
+ timing = self._performance_metrics(metrics, output)
+
+ output.meta_info = {"timing": timing, **outputs[0].meta_info}
+ return output
+
+ async def wake_up(self):
+ await asyncio.gather(*[replica.wake_up() for replica in self.rollout_replicas])
+
+ async def sleep(self):
+ await asyncio.gather(*[replica.sleep() for replica in self.rollout_replicas])
+
+ async def clear_kv_cache(self):
+ await asyncio.gather(*[replica.clear_kv_cache() for replica in self.rollout_replicas])
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/config/one_step_off_ppo_megatron_trainer.yaml b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/config/one_step_off_ppo_megatron_trainer.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..3aea4e4c94d5aa1a5701eeba8c16511fcb0bb496
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/config/one_step_off_ppo_megatron_trainer.yaml
@@ -0,0 +1,28 @@
+hydra:
+ searchpath:
+ - file://verl/trainer/config
+
+defaults:
+ - ppo_megatron_trainer
+ - _self_
+
+# config for the rollout (only for resource isolation)
+rollout:
+ # Number of nodes used in the rollout
+ nnodes: 1
+ # Number of GPUs per node
+ n_gpus_per_node: 8
+
+# To adapt to the current logic of AgentLoopManager
+actor_rollout_ref:
+ rollout:
+ # Must be turned off! Otherwise, Parameter synchronization cannot be performed.
+ free_cache_engine: False
+ # Must be enabled! Otherwise, log_probs cannot be calculated.
+ calculate_log_probs: True
+
+# Only then will the use of log probs be correct.
+# And it can be used in conjunction with other rollout_correction algorithms.
+algorithm:
+ rollout_correction:
+ bypass_mode: True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/config/one_step_off_ppo_trainer.yaml b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/config/one_step_off_ppo_trainer.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..4c4deb485e1ea5918db452d61db4368d1be99494
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/config/one_step_off_ppo_trainer.yaml
@@ -0,0 +1,28 @@
+hydra:
+ searchpath:
+ - file://verl/trainer/config
+
+defaults:
+ - ppo_trainer
+ - _self_
+
+# config for the rollout (only for resource isolation)
+rollout:
+ # Number of nodes used in the rollout
+ nnodes: 1
+ # Number of GPUs per node
+ n_gpus_per_node: 8
+
+# To adapt to the current logic of AgentLoopManager
+actor_rollout_ref:
+ rollout:
+ # Must be turned off! Otherwise, Parameter synchronization cannot be performed.
+ free_cache_engine: False
+ # Must be enabled! Otherwise, log_probs cannot be calculated.
+ calculate_log_probs: True
+
+# Only then will the use of log probs be correct.
+# And it can be used in conjunction with other rollout_correction algorithms.
+algorithm:
+ rollout_correction:
+ bypass_mode: True
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/distributed_utils.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/distributed_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d117fb96f14d306386cd4d956067ef54e1c69ef5
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/distributed_utils.py
@@ -0,0 +1,137 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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.
+
+import ipaddress
+import socket
+from datetime import timedelta
+
+import vllm
+from torch.distributed import TCPStore
+from vllm.distributed.utils import StatelessProcessGroup
+
+from verl.utils.device import is_npu_available
+
+
+@staticmethod
+def create(
+ host: str,
+ port: int,
+ rank: int,
+ world_size: int,
+ data_expiration_seconds: int = 3600,
+ store_timeout: int = 300,
+) -> "StatelessProcessGroup":
+ """A replacement for `torch.distributed.init_process_group` that does not
+ pollute the global state.
+
+ If we have process A and process B called `torch.distributed.init_process_group`
+ to form a group, and then we want to form another group with process A, B, C,
+ D, it is not possible in PyTorch, because process A and process B have already
+ formed a group, and process C and process D cannot join that group. This
+ function is a workaround for this issue.
+
+ `torch.distributed.init_process_group` is a global call, while this function
+ is a stateless call. It will return a `StatelessProcessGroup` object that can be
+ used for exchanging metadata. With this function, process A and process B
+ can call `StatelessProcessGroup.create` to form a group, and then process A, B,
+ C, and D can call `StatelessProcessGroup.create` to form another group.
+
+ Args:
+ host: Host address (IPv4 or IPv6). For IPv6, can be in format like "::1" or "[::1]".
+ port: Port number to bind/listen on.
+ rank: Rank of the current process.
+ world_size: Total number of processes in the group.
+ data_expiration_seconds: Time in seconds before data entries expire (default: 3600).
+ store_timeout: Timeout in seconds for TCPStore connection (default: 300).
+
+ Returns:
+ StatelessProcessGroup: A stateless process group instance.
+ """ # noqa
+ # Detect address family (IPv4 or IPv6)
+ try:
+ # Try to parse as IPv6 first (IPv6 addresses are more specific)
+ ipaddress.IPv6Address(host.strip("[]"))
+ address_family = socket.AF_INET6
+ except (ipaddress.AddressValueError, ValueError):
+ address_family = socket.AF_INET
+
+ launch_server = rank == 0
+ if launch_server:
+ # listen on the specified interface (instead of 0.0.0.0 or ::)
+ listen_socket = socket.socket(address_family, socket.SOCK_STREAM)
+ listen_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+
+ # For IPv6, set IPV6_V6ONLY to only listen on IPv6 (not dual-stack)
+ # This ensures consistent behavior across different systems
+ if address_family == socket.AF_INET6:
+ try:
+ listen_socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
+ except (AttributeError, OSError):
+ # IPV6_V6ONLY might not be available on all systems
+ pass
+
+ # Remove brackets from IPv6 address if present (socket.bind handles it)
+ bind_host = host.strip("[]")
+ listen_socket.bind((bind_host, port))
+ listen_socket.listen()
+ listen_fd = listen_socket.fileno()
+ else:
+ listen_socket = None
+ listen_fd = None
+
+ store = TCPStore(
+ host_name=host,
+ port=port,
+ world_size=world_size,
+ is_master=launch_server,
+ timeout=timedelta(seconds=store_timeout),
+ use_libuv=False, # for now: github.com/pytorch/pytorch/pull/150215
+ master_listen_fd=listen_fd,
+ )
+
+ return StatelessProcessGroup(
+ rank=rank,
+ world_size=world_size,
+ store=store,
+ socket=listen_socket,
+ data_expiration_seconds=data_expiration_seconds,
+ )
+
+
+vllm.distributed.utils.StatelessProcessGroup.create = create
+
+
+def vllm_stateless_init_process_group(master_address, master_port, rank, world_size, device):
+ """
+ vLLM provides `StatelessProcessGroup` to create a process group
+ without considering the global process group in torch.distributed.
+ It is recommended to create `StatelessProcessGroup`, and then initialize
+ the data-plane communication (NCCL) between external (train processes)
+ and vLLM workers.
+ """
+ # NOTE: If it is necessary to support weight synchronization with the sglang backend in the future,
+ # the following can be used:
+ # from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator
+ # from sglang.srt.distributed.utils import statelessprocessgroup
+ if is_npu_available:
+ from vllm_ascend.distributed.device_communicators.pyhccl import (
+ PyHcclCommunicator as PyNcclCommunicator,
+ )
+ else:
+ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+
+ pg = StatelessProcessGroup.create(host=master_address, port=master_port, rank=rank, world_size=world_size)
+ pynccl = PyNcclCommunicator(pg, device=device)
+ return pynccl
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/fsdp_workers.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/fsdp_workers.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab5c2955f7aecf861c77049713ef732e5b7e7f52
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/fsdp_workers.py
@@ -0,0 +1,172 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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.
+
+import logging
+import os
+
+import torch
+import torch.distributed
+from omegaconf import DictConfig
+from ray.util.collective import collective
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+from verl.experimental.one_step_off_policy.distributed_utils import vllm_stateless_init_process_group
+from verl.single_controller.base.decorator import Dispatch, register
+from verl.utils.device import (
+ get_device_name,
+ get_torch_device,
+)
+from verl.utils.fsdp_utils import (
+ fsdp_version,
+ load_fsdp_model_to_gpu,
+ offload_fsdp_model_to_cpu,
+)
+from verl.utils.ray_utils import get_event_loop
+from verl.workers.fsdp_workers import (
+ ActorRolloutRefWorker,
+ AsyncActorRolloutRefWorker,
+ CriticWorker,
+ RewardModelWorker,
+)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+device_name = get_device_name()
+
+__all__ = ["DetachActorWorker", "DetachAsyncRolloutWorker", "CriticWorker", "RewardModelWorker"]
+
+
+class DetachSync(AsyncActorRolloutRefWorker):
+ def _get_actor_params(self):
+ pass
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def create_weight_sync_group(self, master_address, master_port, rank_offset, world_size):
+ rank = torch.distributed.get_rank() + rank_offset
+ self._weight_sync_group = vllm_stateless_init_process_group(
+ master_address,
+ master_port,
+ rank,
+ world_size,
+ get_torch_device().current_device(),
+ )
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def sync_rollout_weights(self):
+ assert (self._is_actor or self._is_rollout) and not self.config.hybrid_engine
+ assert hasattr(self, "_weights_info") and self._weights_info is not None
+
+ if self._is_actor and self._is_offload_param:
+ load_fsdp_model_to_gpu(self.actor_module_fsdp)
+ params = self._get_actor_params() if self._is_actor else None
+
+ rollout_name = self.config.rollout.name
+ if self._is_rollout:
+ if rollout_name == "vllm":
+ from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+
+ inference_model = self.rollout.inference_engine.worker.model_runner.model
+ patch_vllm_moe_model_weight_loader(inference_model)
+ elif rollout_name == "sglang":
+ inference_model = self.rollout._engine
+ else:
+ raise NotImplementedError(f"Unknown rollout name: {rollout_name}")
+ loop = get_event_loop()
+ for key, shape, dtype in self._weights_info:
+ tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device())
+ if self._is_actor:
+ assert key in params
+ origin_data = params[key]
+ if hasattr(origin_data, "full_tensor"):
+ origin_data = origin_data.full_tensor()
+ if torch.distributed.get_rank() == 0:
+ tensor.copy_(origin_data)
+
+ if device_name == "npu":
+ self._weight_sync_group.broadcast(tensor, src=0, stream=get_torch_device().current_stream())
+ else:
+ collective.broadcast(tensor, src_rank=0, group_name="actor_rollout")
+
+ if self._is_rollout:
+ if rollout_name == "vllm":
+ inference_model.load_weights([(key, tensor)])
+ elif rollout_name == "sglang":
+ # first_rank_in_node = self._tp_rank % tp_size_per_node == 0，
+ # Only the first rank within each node (i.e., the local rank is 0) initializes the engine;
+ # engines for other ranks are set to None.
+
+ if inference_model is not None:
+ loop.run_until_complete(self.update_weights(inference_model, [(key, tensor)]))
+
+ if self._is_actor and self._is_offload_param:
+ offload_fsdp_model_to_cpu(self.actor_module_fsdp)
+ get_torch_device().empty_cache()
+
+ async def update_weights(self, inference_engine, params):
+ from sglang.srt.weight_sync.utils import update_weights as sgl_update_weights
+
+ await sgl_update_weights(
+ engine=inference_engine,
+ params_batch=params,
+ device_mesh_key="infer_tp",
+ device_mesh=self.rollout_device_mesh,
+ )
+
+ if self.rollout_device_mesh["infer_tp"].get_local_rank() == 0:
+ await inference_engine.flush_cache()
+
+
+class DetachActorWorker(DetachSync):
+ def _get_actor_params(self):
+ assert self._is_actor
+ params = self.actor_module_fsdp.state_dict()
+ from verl.utils.model import convert_weight_keys
+
+ params = convert_weight_keys(
+ params, getattr(self.actor_module_fsdp, "_fsdp_wrapped_module", self.actor_module_fsdp)
+ )
+ return params
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def get_actor_weights_info(self):
+ assert self._is_actor
+ if hasattr(self, "_weights_info"):
+ return self._weights_info
+ if fsdp_version(self.actor_module_fsdp) == 1:
+ from torch.distributed.fsdp.api import ShardedStateDictConfig, StateDictType
+
+ FSDP.set_state_dict_type(
+ self.actor_module_fsdp,
+ state_dict_type=StateDictType.SHARDED_STATE_DICT,
+ state_dict_config=ShardedStateDictConfig(),
+ )
+ params = self._get_actor_params()
+ ret = []
+ for key, tensor in params.items():
+ ret.append((key, tensor.size(), tensor.dtype))
+ self._weights_info = ret
+ return ret
+
+
+class DetachAsyncRolloutWorker(DetachSync):
+ def __init__(self, config: DictConfig, role: str):
+ print(f"[DetachAsyncRolloutWorker] {DetachAsyncRolloutWorker.__mro__}")
+ ActorRolloutRefWorker.__init__(self, config, role)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def set_actor_weights_info(self, weights_info):
+ assert self._is_rollout
+ self._weights_info = weights_info
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/main_ppo.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/main_ppo.py
new file mode 100644
index 0000000000000000000000000000000000000000..d19c40ffbe263359a7919159608a456f9d21a11f
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/main_ppo.py
@@ -0,0 +1,235 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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.
+"""
+Note that we don't combine the main with ray_trainer as ray_trainer is used by other main.
+"""
+
+import asyncio
+import os
+import socket
+
+import hydra
+import ray
+
+from verl.experimental.one_step_off_policy.ray_trainer import OneStepOffRayTrainer
+from verl.experimental.one_step_off_policy.utils import need_critic
+from verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler
+from verl.trainer.ppo.ray_trainer import ResourcePoolManager
+from verl.trainer.ppo.reward import load_reward_manager
+from verl.trainer.ppo.utils import Role, need_reference_policy
+from verl.utils.config import validate_config
+from verl.utils.device import auto_set_device
+
+
+def create_resource_pool_manager(config, roles: list) -> ResourcePoolManager:
+ """
+ Create resource pool manager
+
+ Args:
+ config: Configuration object
+ roles: List of roles that need to create resource pools
+
+ Returns:
+ ResourcePoolManager: Resource pool manager
+ """
+ resource_pool_spec = {}
+ mapping = {}
+
+ # Actor/Critic resource pool
+ if any(role in roles for role in [Role.Actor, Role.Critic, Role.RefPolicy, Role.RewardModel]):
+ assert config.trainer.n_gpus_per_node > 0, "config.trainer.n_gpus_per_node must be greater than 0"
+ assert config.trainer.nnodes > 0, "config.trainer.nnodes must be greater than 0"
+
+ trainer_pool = [config.trainer.n_gpus_per_node] * config.trainer.nnodes
+ resource_pool_spec["trainer_pool"] = trainer_pool
+
+ # Map training-related roles to the same resource pool
+ for role in [Role.Actor, Role.Critic, Role.RefPolicy, Role.RewardModel]:
+ if role in roles:
+ mapping[role] = "trainer_pool"
+
+ # Rollout resource pool
+ if Role.Rollout in roles:
+ assert config.rollout.n_gpus_per_node > 0, "config.rollout.n_gpus_per_node must be greater than 0"
+ assert config.rollout.nnodes > 0, "config.rollout.nnodes must be greater than 0"
+
+ rollout_pool = [config.rollout.n_gpus_per_node] * config.rollout.nnodes
+ resource_pool_spec["rollout_pool"] = rollout_pool
+ mapping[Role.Rollout] = "rollout_pool"
+
+ return ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
+
+
+def create_role_worker_mapping(config):
+ """
+ Create mapping from roles to worker classes
+
+ Args:
+ config: Configuration object
+
+ Returns:
+ dict: Mapping from roles to worker classes
+ """
+ # Select worker class based on strategy
+ if config.actor_rollout_ref.actor.strategy in ["fsdp", "fsdp2"]:
+ assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
+ from verl.experimental.one_step_off_policy.fsdp_workers import (
+ CriticWorker,
+ DetachActorWorker,
+ DetachAsyncRolloutWorker,
+ )
+ from verl.single_controller.ray import RayWorkerGroup
+
+ ray_worker_group_cls = RayWorkerGroup
+
+ elif config.actor_rollout_ref.actor.strategy == "megatron":
+ assert config.critic.strategy == "megatron"
+ from verl.experimental.one_step_off_policy.megatron_workers import (
+ CriticWorker,
+ DetachActorWorker,
+ DetachAsyncRolloutWorker,
+ )
+ from verl.single_controller.ray import RayWorkerGroup
+
+ ray_worker_group_cls = RayWorkerGroup
+ else:
+ raise NotImplementedError(f"Unsupported strategy: {config.actor_rollout_ref.actor.strategy}")
+
+ role_worker_mapping = {
+ Role.Actor: ray.remote(DetachActorWorker),
+ Role.Rollout: ray.remote(DetachAsyncRolloutWorker),
+ Role.Critic: ray.remote(CriticWorker),
+ }
+
+ if config.reward_model.enable:
+ if config.reward_model.strategy in ["fsdp", "fsdp2"]:
+ from verl.workers.fsdp_workers import RewardModelWorker
+ elif config.reward_model.strategy == "megatron":
+ from verl.workers.megatron_workers import RewardModelWorker
+ else:
+ raise NotImplementedError(f"Unsupported reward model strategy: {config.reward_model.strategy}")
+
+ role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker)
+
+ # Add reference policy (if KL loss or reward is required)
+ if need_reference_policy(config):
+ role_worker_mapping[Role.RefPolicy] = ray.remote(DetachActorWorker)
+
+ return role_worker_mapping, ray_worker_group_cls
+
+
+@ray.remote(num_cpus=10, max_concurrency=100) # please make sure main_task is not scheduled on head
+class OneStepTaskRunner:
+ def run(self, config):
+ # Print the initial configuration. `resolve=True` will evaluate symbolic values.
+ from pprint import pprint
+
+ from omegaconf import OmegaConf
+
+ from verl.utils.fs import copy_to_local
+
+ print(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
+
+ pprint(OmegaConf.to_container(config, resolve=True))
+
+ OmegaConf.resolve(config)
+
+ role_worker_mapping, ray_worker_group_cls = create_role_worker_mapping(config)
+
+ # validate config
+ validate_config(
+ config=config,
+ use_reference_policy=need_reference_policy(config),
+ use_critic=need_critic(config),
+ )
+
+ # Download the checkpoint from HDFS to the local machine.
+ # `use_shm` determines whether to use shared memory, which could lead to faster model loading if turned on
+ local_path = copy_to_local(
+ config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get("use_shm", False)
+ )
+
+ # Instantiate the tokenizer and processor.
+ from verl.utils 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)
+ # Used for multimodal LLM, could be None
+ processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)
+
+ # Load the reward manager for training and validation.
+ reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})
+ )
+ val_reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {})
+ )
+
+ resource_pool_manager = create_resource_pool_manager(config, role_worker_mapping.keys())
+
+ from verl.utils.dataset.rl_dataset import collate_fn
+
+ # Create training and validation datasets.
+ train_dataset = create_rl_dataset(
+ config.data.train_files,
+ config.data,
+ tokenizer,
+ processor,
+ max_samples=config.data.get("train_max_samples", -1),
+ )
+ val_dataset = create_rl_dataset(
+ config.data.val_files, config.data, tokenizer, processor, max_samples=config.data.get("val_max_samples", -1)
+ )
+ train_sampler = create_rl_sampler(config.data, train_dataset)
+
+ # Initialize the PPO trainer.
+ trainer = OneStepOffRayTrainer(
+ config=config,
+ tokenizer=tokenizer,
+ processor=processor,
+ role_worker_mapping=role_worker_mapping,
+ resource_pool_manager=resource_pool_manager,
+ ray_worker_group_cls=ray_worker_group_cls,
+ reward_fn=reward_fn,
+ val_reward_fn=val_reward_fn,
+ train_dataset=train_dataset,
+ val_dataset=val_dataset,
+ collate_fn=collate_fn,
+ train_sampler=train_sampler,
+ device_name=config.trainer.device,
+ )
+ # Initialize the workers of the trainer.
+ trainer.init_workers()
+ # Start the training process.
+ asyncio.run(trainer.fit())
+
+
+@hydra.main(config_path="config", config_name="one_step_off_ppo_trainer", version_base=None)
+def main(config):
+ from time import time
+
+ from verl.trainer.main_ppo import run_ppo
+
+ start_time = time()
+
+ # Automatically set `config.trainer.device = npu` when running on Ascend NPU.
+ auto_set_device(config)
+
+ run_ppo(config, task_runner_class=OneStepTaskRunner)
+ print(f"total time: {time() - start_time:.2f} seconds")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/megatron_workers.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/megatron_workers.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc1e1ceeb22887a158c09dd1fbfcb568af34928f
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/megatron_workers.py
@@ -0,0 +1,177 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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.
+
+import logging
+import os
+
+import torch
+import torch.distributed
+from omegaconf import DictConfig
+from ray.util.collective import collective
+
+from verl.experimental.one_step_off_policy.distributed_utils import vllm_stateless_init_process_group
+from verl.single_controller.base.decorator import Dispatch, register
+from verl.utils.device import (
+ get_device_name,
+ get_torch_device,
+)
+from verl.utils.megatron_utils import load_megatron_model_to_gpu, offload_megatron_model_to_cpu
+from verl.utils.ray_utils import get_event_loop
+from verl.workers.megatron_workers import (
+ ActorRolloutRefWorker,
+ AsyncActorRolloutRefWorker,
+ CriticWorker,
+ RewardModelWorker,
+)
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+device_name = get_device_name()
+
+__all__ = ["DetachActorWorker", "DetachAsyncRolloutWorker", "CriticWorker", "RewardModelWorker"]
+
+
+class DetachSync(AsyncActorRolloutRefWorker):
+ def _get_actor_params(self):
+ pass
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def create_weight_sync_group(self, master_address, master_port, rank_offset, world_size):
+ rank = torch.distributed.get_rank() + rank_offset
+ self._weight_sync_group = vllm_stateless_init_process_group(
+ master_address,
+ master_port,
+ rank,
+ world_size,
+ get_torch_device().current_device(),
+ )
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
+ def sync_rollout_weights(self):
+ assert (self._is_actor or self._is_rollout) and not self.config.hybrid_engine
+ assert hasattr(self, "_weights_info") and self._weights_info is not None
+
+ params_generator = self._get_actor_params_generator() if self._is_actor else None
+
+ if self._is_actor and self._is_offload_param:
+ load_megatron_model_to_gpu(self.actor_module)
+
+ rollout_name = self.config.rollout.name
+ if self._is_rollout:
+ if rollout_name == "vllm":
+ from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+
+ inference_model = self.rollout.inference_engine.worker.model_runner.model
+ patch_vllm_moe_model_weight_loader(inference_model)
+ elif rollout_name == "sglang":
+ inference_model = self.rollout._engine
+ else:
+ raise NotImplementedError(f"Unknown rollout name: {rollout_name}")
+
+ loop = get_event_loop()
+ for key, shape, dtype in self._weights_info:
+ if self._is_actor:
+ weight_key, weight = next(params_generator)
+ assert key == weight_key
+ assert shape == weight.size()
+ assert dtype == weight.dtype
+
+ tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device())
+ if self._is_actor and torch.distributed.get_rank() == 0:
+ tensor.copy_(weight)
+
+ if device_name == "npu":
+ self._weight_sync_group.broadcast(tensor, src=0, stream=get_torch_device().current_stream())
+ else:
+ collective.broadcast(tensor, src_rank=0, group_name="actor_rollout")
+
+ if self._is_rollout:
+ if rollout_name == "vllm":
+ inference_model.load_weights([(key, tensor)])
+ elif rollout_name == "sglang":
+ # first_rank_in_node = self._tp_rank % tp_size_per_node == 0，
+ # Only the first rank within each node (i.e., the local rank is 0) initializes the engine;
+ # engines for other ranks are set to None.
+
+ if inference_model is not None:
+ loop.run_until_complete(self.update_weights(inference_model, [(key, tensor)]))
+
+ if self._is_actor and self._is_offload_param:
+ offload_megatron_model_to_cpu(self.actor_module)
+
+ async def update_weights(self, inference_engine, params):
+ from sglang.srt.weight_sync.utils import update_weights as sgl_update_weights
+
+ await sgl_update_weights(
+ engine=inference_engine,
+ params_batch=params,
+ device_mesh_key="infer_tp",
+ device_mesh=self.rollout_device_mesh,
+ )
+
+ if self.rollout_device_mesh["infer_tp"].get_local_rank() == 0:
+ await inference_engine.flush_cache()
+
+
+class DetachActorWorker(DetachSync):
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def _get_actor_params_generator(self):
+ assert self._is_actor
+ from verl.models.mcore import get_mcore_weight_converter
+ from verl.utils.megatron_utils import per_tensor_generator
+
+ layer_name_mapping = {
+ "qkv_layer_name": "self_attention.linear_qkv.",
+ "gate_proj_layer_name": "linear_fc1.",
+ }
+ weight_converter = get_mcore_weight_converter(self.actor_model_config, self.dtype)
+ generator = per_tensor_generator(
+ self.actor.actor_module,
+ self.actor_model_config,
+ weight_converter,
+ self.tf_config,
+ layer_name_mapping,
+ )
+ return generator
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def get_actor_weights_info(self):
+ assert self._is_actor
+ if hasattr(self, "_weights_info"):
+ return self._weights_info
+ if self._is_offload_param:
+ load_megatron_model_to_gpu(self.actor_module)
+ params_generator = self._get_actor_params_generator()
+ ret = []
+ for key, tensor in params_generator:
+ ret.append((key, tensor.size(), tensor.dtype))
+
+ self._weights_info = ret
+ # Here, we only call this function at the beginning,
+ # and immediately afterwards we call sync_rollout_weights.
+ # So we no longer call offload in this.
+ return ret
+
+
+class DetachAsyncRolloutWorker(DetachSync):
+ def __init__(self, config: DictConfig, role: str):
+ print(f"[DetachAsyncRolloutWorker] {DetachAsyncRolloutWorker.__mro__}")
+ ActorRolloutRefWorker.__init__(self, config, role)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def set_actor_weights_info(self, weights_info):
+ assert self._is_rollout
+ self._weights_info = weights_info
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/ray_trainer.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/ray_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..43da0d5322db81d9ba9ab2b5820aa52859e403c0
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/ray_trainer.py
@@ -0,0 +1,768 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023-2024 SGLang Team
+# Copyright 2025 ModelBest Inc. and/or its affiliates
+# Copyright 2025 Meituan 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.
+"""
+This trainer supports model-agonistic model initialization with huggingface
+"""
+
+import asyncio
+import uuid
+from pprint import pprint
+
+import numpy as np
+import ray
+import torch
+from omegaconf import OmegaConf
+from ray.util.collective import collective
+from torch.utils.data import Dataset, Sampler
+from tqdm import tqdm
+
+from verl import DataProto
+from verl.experimental.dataset.sampler import AbstractCurriculumSampler
+from verl.experimental.one_step_off_policy.utils import need_critic
+from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup
+from verl.single_controller.ray.base import create_colocated_worker_cls
+from verl.trainer.ppo import core_algos
+from verl.trainer.ppo.core_algos import agg_loss
+from verl.trainer.ppo.metric_utils import compute_data_metrics, compute_throughout_metrics, compute_timing_metrics
+from verl.trainer.ppo.ray_trainer import (
+ RayPPOTrainer,
+ ResourcePoolManager,
+ apply_kl_penalty,
+ compute_advantage,
+ compute_response_mask,
+)
+from verl.trainer.ppo.reward import compute_reward, compute_reward_async
+from verl.trainer.ppo.utils import Role, WorkerType, need_reference_policy, need_reward_model
+from verl.utils import omega_conf_to_dataclass
+from verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi
+from verl.utils.debug import marked_timer
+from verl.utils.metric import reduce_metrics
+from verl.utils.tracking import ValidationGenerationsLogger
+
+
+class OneStepOffRayTrainer(RayPPOTrainer):
+ # TODO: support each role have individual ray_worker_group_cls,
+ # i.e., support different backend of different role
+ def __init__(
+ self,
+ config,
+ tokenizer,
+ role_worker_mapping: dict[Role, WorkerType],
+ resource_pool_manager: ResourcePoolManager,
+ ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,
+ processor=None,
+ reward_fn=None,
+ val_reward_fn=None,
+ train_dataset: Dataset | None = None,
+ val_dataset: Dataset | None = None,
+ collate_fn=None,
+ train_sampler: Sampler | None = None,
+ device_name=None,
+ ):
+ """
+ Initialize distributed PPO trainer with Ray backend.
+ Note that this trainer runs on the driver process on a single CPU/GPU node.
+
+ Args:
+ config: Configuration object containing training parameters.
+ tokenizer: Tokenizer used for encoding and decoding text.
+ role_worker_mapping (dict[Role, WorkerType]): Mapping from roles to worker classes.
+ resource_pool_manager (ResourcePoolManager): Manager for Ray resource pools.
+ ray_worker_group_cls (RayWorkerGroup, optional): Class for Ray worker groups. Defaults to RayWorkerGroup.
+ processor: Optional data processor, used for multimodal data
+ reward_fn: Function for computing rewards during training.
+ val_reward_fn: Function for computing rewards during validation.
+ train_dataset (Optional[Dataset], optional): Training dataset. Defaults to None.
+ val_dataset (Optional[Dataset], optional): Validation dataset. Defaults to None.
+ collate_fn: Function to collate data samples into batches.
+ train_sampler (Optional[Sampler], optional): Sampler for the training dataset. Defaults to None.
+ device_name (str, optional): Device name for training (e.g., "cuda", "cpu"). Defaults to "cuda".
+ """
+
+ # Store the tokenizer for text processing
+ self.tokenizer = tokenizer
+ self.processor = processor
+ self.config = config
+ self.reward_fn = reward_fn
+ self.val_reward_fn = val_reward_fn
+
+ self.hybrid_engine = config.actor_rollout_ref.hybrid_engine
+
+ assert not self.hybrid_engine
+
+ self.role_worker_mapping = role_worker_mapping
+ self.resource_pool_manager = resource_pool_manager
+ 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
+ self.device_name = device_name if device_name else self.config.trainer.device
+ self.validation_generations_logger = ValidationGenerationsLogger()
+
+ 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)
+ # if ref_in_actor is True, the reference policy will be actor without lora applied
+ self.ref_in_actor = lora_rank > 0
+
+ # define in-reward KL control
+ # kl loss control currently not suppoorted
+ if config.algorithm.use_kl_in_reward:
+ self.kl_ctrl_in_reward = core_algos.get_kl_controller(config.algorithm.kl_ctrl)
+
+ self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)
+
+ def _validate(self):
+ self.actor_rollout_wg = self.rollout_wg
+ ret = super()._validate()
+ self.actor_rollout_wg = self.actor_wg
+ return ret
+
+ def init_workers(self):
+ """Initialize distributed training workers using Ray backend.
+
+ Creates:
+ 1. Ray resource pools from configuration
+ 2. Worker groups for each role (actor, critic, etc.)
+ """
+ self._init_resource_pools()
+ self._create_worker_classes()
+ self._init_worker_groups()
+ self._init_models()
+ self._init_async_rollout_manager()
+
+ def _init_resource_pools(self):
+ self.resource_pool_manager.create_resource_pool()
+
+ self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
+
+ def _create_worker_classes(self):
+ self._create_actor_rollout_classes()
+ self._create_critic_class()
+ self._create_reference_policy_class()
+ self._create_reward_model_class()
+
+ def _create_actor_rollout_classes(self):
+ for role in [Role.Actor, Role.Rollout]:
+ resource_pool = self.resource_pool_manager.get_resource_pool(role)
+ role_cls = RayClassWithInitArgs(
+ cls=self.role_worker_mapping[role],
+ config=self.config.actor_rollout_ref,
+ role=str(role),
+ )
+ self.resource_pool_to_cls[resource_pool][str(role)] = role_cls
+
+ def _create_critic_class(self):
+ # create critic
+ if self.use_critic:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)
+ critic_cfg = omega_conf_to_dataclass(self.config.critic)
+ critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg)
+ self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls
+
+ def _create_reference_policy_class(self):
+ # create reference policy if needed
+ if self.use_reference_policy:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)
+ ref_policy_cls = RayClassWithInitArgs(
+ self.role_worker_mapping[Role.RefPolicy],
+ config=self.config.actor_rollout_ref,
+ role=str(Role.RefPolicy),
+ # profile_option=self.config.trainer.npu_profile.options,
+ )
+ self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls
+
+ def _create_reward_model_class(self):
+ # create a reward model if reward_fn is None
+ 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
+
+ def _init_worker_groups(self):
+ # 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`.
+ # Instead, directly pass different resource pool to different worker groups.
+ # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.
+ all_wg = {}
+ wg_kwargs = {} # Setting up kwargs for RayWorkerGroup
+ if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None:
+ wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout
+ if OmegaConf.select(self.config.global_profiler, "steps") is not None:
+ wg_kwargs["profile_steps"] = OmegaConf.select(self.config.global_profiler, "steps")
+ # Only require nsight worker options when tool is nsys
+ if OmegaConf.select(self.config.global_profiler, "tool") == "nsys":
+ assert (
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ is not None
+ ), "worker_nsight_options must be set when using nsys with profile_steps"
+ wg_kwargs["worker_nsight_options"] = OmegaConf.to_container(
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ )
+ wg_kwargs["device_name"] = self.device_name
+
+ for resource_pool, class_dict in self.resource_pool_to_cls.items():
+ worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
+ wg_dict = self.ray_worker_group_cls(
+ resource_pool=resource_pool,
+ ray_cls_with_init=worker_dict_cls,
+ **wg_kwargs,
+ )
+ spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
+ all_wg.update(spawn_wg)
+ self.all_wg = all_wg
+
+ def _init_models(self):
+ if self.use_critic:
+ self.critic_wg = self.all_wg[str(Role.Critic)]
+ self.critic_wg.init_model()
+
+ if self.use_reference_policy and not self.ref_in_actor:
+ self.ref_policy_wg = self.all_wg[str(Role.RefPolicy)]
+ self.ref_policy_wg.init_model()
+
+ self.rm_wg = None
+ if self.use_rm:
+ self.rm_wg = self.all_wg[str(Role.RewardModel)]
+ self.rm_wg.init_model()
+
+ self.actor_wg = self.all_wg[str(Role.Actor)]
+ self.rollout_wg = self.all_wg[str(Role.Rollout)]
+ self.actor_wg.init_model()
+ self.rollout_wg.init_model()
+ self.actor_rollout_wg = self.actor_wg
+ weights_info = self.actor_wg.get_actor_weights_info()[0]
+ self.rollout_wg.set_actor_weights_info(weights_info)
+ self._create_weight_sync_group()
+
+ def _create_weight_sync_group(self):
+ from verl.utils.device import get_nccl_backend
+
+ actor_rollout_workers = self.actor_wg.workers + self.rollout_wg.workers
+ n_workers = len(actor_rollout_workers)
+
+ if self.device_name == "npu":
+ master_address = ray.get(self.actor_wg.workers[0]._get_node_ip.remote()).strip("[]")
+ master_port = ray.get(self.actor_wg.workers[0]._get_free_port.remote())
+ self.actor_wg.create_weight_sync_group(
+ master_address,
+ master_port,
+ 0,
+ n_workers,
+ )
+ ray.get(
+ self.rollout_wg.create_weight_sync_group(
+ master_address,
+ master_port,
+ len(self.actor_wg.workers),
+ n_workers,
+ )
+ )
+ else:
+ # Create Ray collective group for fallback communication
+ collective.create_collective_group(
+ actor_rollout_workers,
+ n_workers,
+ list(range(0, n_workers)),
+ backend=get_nccl_backend(),
+ group_name="actor_rollout",
+ )
+
+ def _init_async_rollout_manager(self):
+ # create async rollout manager and request scheduler
+ assert self.config.actor_rollout_ref.rollout.mode == "async"
+ from verl.experimental.one_step_off_policy.agent_loop import OneStepOffAgentLoopManager
+
+ self.async_rollout_mode = True
+
+ if self.config.reward_model.enable and self.config.reward_model.enable_resource_pool:
+ rm_resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
+ else:
+ rm_resource_pool = None
+
+ self.async_rollout_manager = OneStepOffAgentLoopManager(
+ config=self.config, worker_group=self.rollout_wg, rm_resource_pool=rm_resource_pool
+ )
+
+ def sync_rollout_weights(self):
+ self.actor_wg.sync_rollout_weights()
+ ray.get(self.rollout_wg.sync_rollout_weights())
+
+ def _create_continuous_iterator(self):
+ """
+ Create a continuous data iterator across epoch
+ """
+ for epoch in range(self.config.trainer.total_epochs):
+ iterator = iter(self.train_dataloader)
+ for batch_dict in iterator:
+ yield epoch, batch_dict
+
+ async def _async_gen_next_batch(self, continuous_iterator):
+ """
+ Call parameter synchronization and asynchronous sequence generation.
+ """
+ try:
+ epoch, batch_dict = next(continuous_iterator)
+ except StopIteration:
+ return None
+ except Exception as e:
+ print(f"Error in async_gen_next_batch: {e}")
+ return None
+
+ metrics = {}
+ timing_raw = {}
+
+ # Create the initial batch from the data loader
+ batch = DataProto.from_single_dict(batch_dict)
+
+ # add uid to batch
+ batch.non_tensor_batch["uid"] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object)
+
+ gen_batch = self._get_gen_batch(batch)
+
+ # pass global_steps to trace
+ gen_batch.meta_info["global_steps"] = self.global_steps
+ gen_batch_output = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
+
+ # async generation
+ with marked_timer("generate_async", timing_raw, color="purple"):
+ gen_batch_output = await self.async_rollout_manager.generate_sequences_async(gen_batch_output)
+
+ # repeat to align with repeated responses in rollout
+ batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
+ batch = batch.union(gen_batch_output)
+
+ if "response_mask" not in batch.batch.keys():
+ batch.batch["response_mask"] = compute_response_mask(batch)
+ # Balance the number of valid tokens across DP ranks.
+ # NOTE: This usually changes the order of data in the `batch`,
+ # which won't affect the advantage calculation (since it's based on uid),
+ # but might affect the loss calculation (due to the change of mini-batching).
+ if self.config.trainer.balance_batch:
+ self._balance_batch(batch, metrics=metrics)
+
+ # compute global_valid tokens
+ batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
+
+ # Launch individual reward computations as each generation completes
+ future_reward = None
+ if self.config.reward_model.launch_reward_fn_async:
+ # Store the object reference and set up callback
+ future_reward = self._launch_individual_rewards.remote(batch, self.config, self.tokenizer)
+
+ # Return the original, now-modified `batch` and the `future_reward`
+ return metrics, timing_raw, epoch, batch, future_reward
+
+ @staticmethod
+ @ray.remote
+ def _launch_individual_rewards(batch, config, tokenizer):
+ # Get generation results
+ gen_batch_result = batch
+ original_non_tensor_batch = batch.non_tensor_batch
+
+ # Repeat non_tensor_batch to match the number of responses
+ n = config.actor_rollout_ref.rollout.n
+ repeated_non_tensor_batch = {}
+ for key, value in original_non_tensor_batch.items():
+ repeated_non_tensor_batch[key] = np.repeat(value, n, axis=0)
+
+ # Split into individual responses with preserved non_tensor_batch
+ responses_split = []
+ for i in range(len(gen_batch_result)):
+ response_data = gen_batch_result[i : i + 1] # Get single response
+ # Add repeated non_tensor_batch values
+ for key in repeated_non_tensor_batch:
+ response_data.non_tensor_batch[key] = repeated_non_tensor_batch[key][i : i + 1]
+ responses_split.append(response_data)
+
+ # Launch async reward computation
+ reward_futures = [
+ compute_reward_async.remote(response_data, config, tokenizer) for response_data in responses_split
+ ]
+
+ # Wait for results and combine
+ results = ray.get(reward_futures)
+ rewards_list = [r[0] for r in results]
+ extras_list = [r[1] for r in results]
+
+ combined_reward_tensor = torch.cat(rewards_list, dim=0)
+ combined_extras_dict = {}
+ if extras_list and extras_list[0]:
+ for key in extras_list[0].keys():
+ combined_extras_dict[key] = [d[key] for d in extras_list if key in d]
+
+ return combined_reward_tensor, combined_extras_dict
+
+ async def fit(self):
+ """
+ The training loop of PPO.
+ The driver process only need to call the compute functions of the worker group through RPC
+ to construct the PPO dataflow.
+ The light-weight advantage computation is done on the driver process.
+ """
+
+ from omegaconf import OmegaConf
+
+ from verl.utils.tracking import Tracking
+
+ logger = Tracking(
+ project_name=self.config.trainer.project_name,
+ experiment_name=self.config.trainer.experiment_name,
+ default_backend=self.config.trainer.logger,
+ config=OmegaConf.to_container(self.config, resolve=True),
+ )
+
+ self.global_steps = 0
+
+ # load checkpoint before doing anything
+ self._load_checkpoint()
+
+ # after load checkpoint sync rollout weights
+ self.sync_rollout_weights()
+ await self.async_rollout_manager.clear_kv_cache()
+
+ # perform validation before training
+ # currently, we only support validation using the reward_function.
+ if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
+ val_metrics = self._validate()
+ assert val_metrics, f"{val_metrics=}"
+ pprint(f"Initial validation metrics: {val_metrics}")
+ logger.log(data=val_metrics, step=self.global_steps)
+ if self.config.trainer.get("val_only", False):
+ return
+
+ # add tqdm
+ progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress")
+
+ # we start from step 1
+ self.global_steps += 1
+ last_val_metrics = None
+ self.max_steps_duration = 0
+
+ prev_step_profile = False
+ curr_step_profile = (
+ self.global_steps in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+
+ # across epoch iterator
+ continuous_iterator = self._create_continuous_iterator()
+
+ # Start the first asynchronous generation task.
+ batch_data_future = asyncio.create_task(self._async_gen_next_batch(continuous_iterator))
+
+ while batch_data_future is not None:
+ do_profile = (
+ self.global_steps in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ if do_profile:
+ self.actor_wg.start_profile()
+ if not self.hybrid_engine:
+ self.rollout_wg.start_profile()
+ if self.use_reference_policy:
+ self.ref_policy_wg.start_profile()
+ if self.use_critic:
+ self.critic_wg.start_profile()
+ if self.use_rm:
+ self.rm_wg.start_profile()
+
+ metrics = {}
+ timing_raw = {}
+ is_last_step = self.global_steps >= self.total_training_steps
+
+ with marked_timer("start_profile", timing_raw):
+ self._start_profiling(
+ not prev_step_profile and curr_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+
+ with marked_timer("step", timing_raw):
+ # wait for the previous batch
+ with marked_timer("gen", timing_raw, color="red"):
+ _metrics, _timing_raw, epoch, batch, future_reward = await batch_data_future
+ timing_raw.update(batch.meta_info["timing"])
+ timing_raw.update(_timing_raw)
+ metrics.update(_metrics)
+ batch.meta_info.pop("timing", None)
+
+ # sync weights from actor to rollout
+ with marked_timer("sync_rollout_weights", timing_raw, color="purple"):
+ self.sync_rollout_weights()
+ await self.async_rollout_manager.clear_kv_cache()
+
+ # async next generation
+ if not is_last_step:
+ batch_data_future = asyncio.create_task(self._async_gen_next_batch(continuous_iterator))
+ await asyncio.sleep(0)
+
+ with marked_timer("reward", timing_raw, color="yellow"):
+ # compute reward model score
+ if self.use_rm and "rm_scores" not in batch.batch.keys():
+ reward_tensor = self.rm_wg.compute_rm_score(batch)
+ batch = batch.union(reward_tensor)
+
+ if self.config.reward_model.launch_reward_fn_async:
+ future_reward = compute_reward_async.remote(
+ data=batch, config=self.config, tokenizer=self.tokenizer
+ )
+ else:
+ reward_tensor, reward_extra_infos_dict = compute_reward(batch, self.reward_fn)
+
+ # await asyncio.sleep(0) ensures:
+ # Asynchronous tasks can start executing immediately
+ # The event loop can handle other pending coroutines
+ # Prevents computations in a certain phase from blocking the entire asynchronous workflow
+ #
+ # The purpose here is to ensure that after triggering
+ # `self.async_rollout_manager.generate_sequences_async(gen_batch_output)`,
+ # the subsequent relevant logic can proceed in a timely manner
+ await asyncio.sleep(0)
+
+ # Operating Mode Selection:
+ # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ)
+ # - Decoupled mode: Recomputes old_log_probs as proximal anchor (3 policies: π_rollout, π_old, π_θ)
+ # Note: π_old computed once per data batch, serves as stable reference during mini-batch updates
+ rollout_corr_config = self.config.algorithm.get("rollout_correction", None)
+ bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get("bypass_mode", False)
+ if bypass_recomputing_logprobs: # Use `rollout_log_probs`
+ 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,
+ )
+ else: # Recompute old_log_probs
+ with marked_timer("old_log_prob", timing_raw, color="blue"):
+ old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
+ entropys = old_log_prob.batch["entropys"]
+ response_masks = batch.batch["response_mask"]
+ actor_config = self.config.actor_rollout_ref.actor
+ entropy_agg = agg_loss(
+ loss_mat=entropys,
+ loss_mask=response_masks,
+ loss_agg_mode=actor_config.loss_agg_mode,
+ loss_scale_factor=actor_config.loss_scale_factor,
+ )
+ old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()}
+ metrics.update(old_log_prob_metrics)
+ old_log_prob.batch.pop("entropys")
+ batch = batch.union(old_log_prob)
+ if "rollout_log_probs" in batch.batch.keys():
+ # TODO: we may want to add diff of probs too.
+ from verl.utils.debug.metrics import calculate_debug_metrics
+
+ metrics.update(calculate_debug_metrics(batch))
+
+ assert "old_log_probs" in batch.batch, f'"old_log_prob" not in {batch.batch.keys()=}'
+ await asyncio.sleep(0)
+
+ if self.use_reference_policy:
+ # compute reference log_prob
+ with marked_timer(str(Role.RefPolicy), timing_raw, color="olive"):
+ if not self.ref_in_actor:
+ ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
+ else:
+ ref_log_prob = self.actor_rollout_wg.compute_ref_log_prob(batch)
+ batch = batch.union(ref_log_prob)
+ await asyncio.sleep(0)
+
+ # compute values
+ if self.use_critic:
+ with marked_timer("values", timing_raw, color="cyan"):
+ values = self.critic_wg.compute_values(batch)
+ batch = batch.union(values)
+ await asyncio.sleep(0)
+
+ with marked_timer("adv", timing_raw, color="brown"):
+ # we combine with rule-based rm
+ reward_extra_infos_dict: dict[str, list]
+ if self.config.reward_model.launch_reward_fn_async:
+ reward_tensor, reward_extra_infos_dict = ray.get(future_reward)
+ batch.batch["token_level_scores"] = reward_tensor
+
+ if reward_extra_infos_dict:
+ batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})
+
+ # compute rewards. apply_kl_penalty if available
+ if self.config.algorithm.use_kl_in_reward:
+ batch, kl_metrics = apply_kl_penalty(
+ batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
+ )
+ metrics.update(kl_metrics)
+ else:
+ batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
+
+ # Compute rollout correction: IS weights, rejection sampling, and metrics
+ # Only runs in decoupled mode (computes once per batch using stable π_old)
+ # In bypass mode, this is skipped - actor computes metrics from evolving π_θ vs π_rollout
+ if (
+ rollout_corr_config is not None
+ and "rollout_log_probs" in batch.batch
+ and not bypass_recomputing_logprobs # Only in decoupled mode
+ ):
+ from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch
+
+ # Compute IS weights, apply rejection sampling, compute metrics
+ batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config)
+ # IS and off-policy metrics already have rollout_corr/ prefix
+ metrics.update(is_metrics)
+
+ # compute advantages, executed on the driver process
+ norm_adv_by_std_in_grpo = self.config.algorithm.get(
+ "norm_adv_by_std_in_grpo", True
+ ) # GRPO adv normalization factor
+
+ batch = compute_advantage(
+ batch,
+ adv_estimator=self.config.algorithm.adv_estimator,
+ gamma=self.config.algorithm.gamma,
+ lam=self.config.algorithm.lam,
+ num_repeat=self.config.actor_rollout_ref.rollout.n,
+ norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+ config=self.config.algorithm,
+ )
+ await asyncio.sleep(0)
+
+ # update critic
+ if self.use_critic:
+ with marked_timer("update_critic", timing_raw, color="pink"):
+ critic_output = self.critic_wg.update_critic(batch)
+ critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"])
+ metrics.update(critic_output_metrics)
+ await asyncio.sleep(0)
+
+ # implement critic warmup
+ if self.config.trainer.critic_warmup <= self.global_steps:
+ # update actor
+ with marked_timer("update_actor", timing_raw, color="red"):
+ rollout_config = self.config.actor_rollout_ref.rollout
+ batch.meta_info["multi_turn"] = rollout_config.multi_turn.enable
+ # TODO: Make "temperature" single source of truth from generation.
+ batch.meta_info["temperature"] = rollout_config.temperature
+ actor_output = self.actor_rollout_wg.update_actor(batch)
+ actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])
+ metrics.update(actor_output_metrics)
+ await asyncio.sleep(0)
+
+ # Log rollout generations if enabled
+ rollout_data_dir = self.config.trainer.get("rollout_data_dir", None)
+ if rollout_data_dir:
+ self._log_rollout_data(batch, reward_extra_infos_dict, timing_raw, rollout_data_dir)
+
+ await asyncio.sleep(0)
+ # validate
+ if (
+ self.val_reward_fn is not None
+ and self.config.trainer.test_freq > 0
+ and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
+ ):
+ with marked_timer("testing", timing_raw, color="green"):
+ val_metrics: dict = self._validate()
+ if is_last_step:
+ last_val_metrics = val_metrics
+ metrics.update(val_metrics)
+ await asyncio.sleep(0)
+
+ # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.
+ esi_close_to_expiration = should_save_ckpt_esi(
+ max_steps_duration=self.max_steps_duration,
+ redundant_time=self.config.trainer.esi_redundant_time,
+ )
+ # Check if the conditions for saving a checkpoint are met.
+ # The conditions include a mandatory condition (1) and
+ # one of the following optional conditions (2/3/4):
+ # 1. The save frequency is set to a positive value.
+ # 2. It's the last training step.
+ # 3. The current step number is a multiple of the save frequency.
+ # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.
+ if self.config.trainer.save_freq > 0 and (
+ is_last_step or self.global_steps % self.config.trainer.save_freq == 0 or esi_close_to_expiration
+ ):
+ if esi_close_to_expiration:
+ print("Force saving checkpoint: ESI instance expiration approaching.")
+ with marked_timer("save_checkpoint", timing_raw, color="green"):
+ self._save_checkpoint()
+
+ with marked_timer("stop_profile", timing_raw):
+ next_step_profile = (
+ self.global_steps + 1 in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ self._stop_profiling(
+ curr_step_profile and not next_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+ prev_step_profile = curr_step_profile
+ curr_step_profile = next_step_profile
+
+ steps_duration = timing_raw["step"]
+ self.max_steps_duration = max(self.max_steps_duration, steps_duration)
+
+ # training metrics
+ metrics.update(
+ {
+ "training/global_step": self.global_steps,
+ "training/epoch": epoch,
+ }
+ )
+ # collect metrics
+ metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
+ metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
+ # TODO: implement actual tflpo and theoretical tflpo
+ n_gpus = self.resource_pool_manager.get_n_gpus()
+ metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))
+ # Note: mismatch metrics (KL, PPL, etc.) are collected at line 1179 after advantage computation
+
+ # this is experimental and may be changed/removed in the future in favor of a general-purpose one
+ if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):
+ self.train_dataloader.sampler.update(batch=batch)
+
+ # TODO: make a canonical logger that supports various backend
+ logger.log(data=metrics, step=self.global_steps)
+
+ progress_bar.update(1)
+ self.global_steps += 1
+
+ if (
+ hasattr(self.config.actor_rollout_ref.actor, "profiler")
+ and self.config.actor_rollout_ref.actor.profiler.tool == "torch_memory"
+ ):
+ self.actor_rollout_wg.dump_memory_snapshot(
+ tag=f"post_update_step{self.global_steps}", sub_dir=f"step{self.global_steps}"
+ )
+
+ if is_last_step:
+ if hasattr(self.actor_rollout_wg, "async_calls_finalize_fn_exec"):
+ self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=True)
+ pprint(f"Final validation metrics: {last_val_metrics}")
+ progress_bar.close()
+ return
+
+ # this is experimental and may be changed/removed in the future
+ # in favor of a general-purpose data buffer pool
+ if hasattr(self.train_dataset, "on_batch_end"):
+ # The dataset may be changed after each training batch
+ self.train_dataset.on_batch_end(batch=batch)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_4_12.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_4_12.sh
new file mode 100644
index 0000000000000000000000000000000000000000..68b3343f354fafbe8e8d2612aa6053d4998bd2dc
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_4_12.sh
@@ -0,0 +1,139 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-one-step-off-4-12'
+
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+train_prompt_bsz=512
+n_resp_per_prompt=12
+train_prompt_mini_bsz=32
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+NNODES=${NNODES:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=2
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=2
+sp_size=4
+fsdp_size=2
+
+python3 -m verl.experimental.one_step_off_policy.main_ppo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=10 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=100 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10 \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}"
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64.sh
new file mode 100644
index 0000000000000000000000000000000000000000..2db7548980dccd72dc849d96485b051b79d4e593
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64.sh
@@ -0,0 +1,140 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='dapo_qwen2-7B-math_28k_fsdp2_one_step_off_64-64'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 28))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=4
+sp_size=4
+fsdp_size=8
+
+# one stepa specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-8}
+NNODES_TRAIN=${NNODES_TRAIN:-8}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=512
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+test_freq=20
+
+python -m verl.experimental.one_step_off_policy.main_ppo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=20 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=400 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}"
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64_ris.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64_ris.sh
new file mode 100644
index 0000000000000000000000000000000000000000..873979f6eb7bb7269268461f6c4ee8969a99800f
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64_ris.sh
@@ -0,0 +1,155 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='dapo_qwen2-7B-math_28k_fsdp2_one_step_off_64-64-ris'
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+# Algorithm parameters
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+# Response length parameters
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 28))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+# Training parameters
+loss_agg_mode="token-mean"
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=4
+sp_size=4
+fsdp_size=8
+
+# one stepa specific parameters
+NNODES_ROLLOUT=${NNODES_ROLLOUT:-8}
+NNODES_TRAIN=${NNODES_TRAIN:-8}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+train_prompt_bsz=512
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+test_freq=20
+
+# https://github.com/volcengine/verl/blob/main/docs/algo/rollout_corr.md
+# use decoupled_geo_rs
+#algorithm:
+# rollout_correction:
+# rollout_is: null
+# rollout_is_threshold=null
+# rollout_rs: seq_mean_k1
+# rollout_rs_threshold: 0.999_1.001
+# bypass_mode: false # Decoupled mode
+
+python -m verl.experimental.one_step_off_policy.main_ppo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=10 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=400 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES_TRAIN}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ rollout.nnodes="${NNODES_ROLLOUT}" \
+ rollout.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ algorithm.rollout_correction.rollout_is=null \
+ algorithm.rollout_correction.rollout_is_threshold=null \
+ algorithm.rollout_correction.rollout_rs=seq_mean_k1 \
+ algorithm.rollout_correction.rollout_rs_threshold="0.999_1.001" \
+ algorithm.rollout_correction.bypass_mode=false
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_colocate.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_colocate.sh
new file mode 100644
index 0000000000000000000000000000000000000000..617d7a7c8479e68d0804841f4d49c81c96b91d88
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_colocate.sh
@@ -0,0 +1,132 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-colocate'
+
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+train_prompt_bsz=512
+n_resp_per_prompt=12
+train_prompt_mini_bsz=32
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+NNODES=${NNODES:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+offload=True
+gen_tp=2
+sp_size=4
+fsdp_size=2
+
+# reference run wandb: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/runs/ow47vvon?nw=nwusertongyuxuan361
+
+python3 -m verl.trainer.main_ppo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ trainer.nnodes="${NNODES}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=10 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=100 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_4_12.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_4_12.sh
new file mode 100644
index 0000000000000000000000000000000000000000..35ac6af16d966eb133d3b6b740683c327ba5a1c6
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_4_12.sh
@@ -0,0 +1,140 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-sglang-one-step-off-4-12'
+
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+train_prompt_bsz=512
+n_resp_per_prompt=12
+train_prompt_mini_bsz=32
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+NNODES=${NNODES:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=2
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=2
+sp_size=4
+fsdp_size=2
+
+python3 -m verl.experimental.one_step_off_policy.main_ppo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.layered_summon=True \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=sglang \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=10 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=100 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10 \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}"
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_colocate.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_colocate.sh
new file mode 100644
index 0000000000000000000000000000000000000000..694fa13caf0fa6736bd7532c2b8869ca17975edd
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_colocate.sh
@@ -0,0 +1,133 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-sglang-colocate'
+
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+train_prompt_bsz=512
+n_resp_per_prompt=12
+train_prompt_mini_bsz=32
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+NNODES=${NNODES:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+offload=True
+gen_tp=2
+sp_size=4
+fsdp_size=2
+
+# reference run wandb: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/runs/ow47vvon?nw=nwusertongyuxuan361
+
+python3 -m verl.trainer.main_ppo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.grad_clip=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.layered_summon=True \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=sglang \
+ actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
+ trainer.nnodes="${NNODES}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=10 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=100 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_4_12.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_4_12.sh
new file mode 100644
index 0000000000000000000000000000000000000000..0b97d2d1aeb107b8850ec156d947b113a58e409a
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_4_12.sh
@@ -0,0 +1,146 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0527a1-megatron-one-step-off-4-12'
+
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+train_prompt_bsz=512
+n_resp_per_prompt=12
+train_prompt_mini_bsz=32
+
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+NNODES=${NNODES:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=2
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+ref_offload=True
+actor_offload=False
+gen_tp=2
+train_tp=2
+train_pp=2
+
+# TODO: support dynamic_bsz for megatron
+# actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+# actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+# actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+# actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+# actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+# actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+
+python3 -m verl.experimental.one_step_off_policy.main_ppo \
+ --config-path=config \
+ --config-name='one_step_off_ppo_megatron_trainer.yaml' \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=megatron \
+ critic.strategy=megatron \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ actor_rollout_ref.hybrid_engine=False \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.megatron.param_offload=${actor_offload} \
+ actor_rollout_ref.actor.megatron.optimizer_offload=${actor_offload} \
+ actor_rollout_ref.actor.megatron.grad_offload=${actor_offload} \
+ actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \
+ actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.optim.clip_grad=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \
+ actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \
+ actor_rollout_ref.ref.megatron.param_offload=${ref_offload} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=10 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=100 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10 \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}"
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_colocate.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_colocate.sh
new file mode 100644
index 0000000000000000000000000000000000000000..df0c451e845b9d003305731df28059145a7dca87
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_colocate.sh
@@ -0,0 +1,138 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+project_name='DAPO'
+exp_name='DAPO-Qwen2.5-7b-MATH-0519a1-megatron-colocate'
+
+adv_estimator=grpo
+
+use_kl_in_reward=False
+kl_coef=0.0
+use_kl_loss=False
+kl_loss_coef=0.0
+
+clip_ratio_low=0.2
+clip_ratio_high=0.28
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 8))
+enable_overlong_buffer=True
+overlong_buffer_len=$((1024 * 4))
+overlong_penalty_factor=1.0
+
+loss_agg_mode="token-mean"
+
+train_prompt_bsz=512
+n_resp_per_prompt=16
+train_prompt_mini_bsz=32
+
+# Ray
+# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
+# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
+# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/verl/trainer/runtime_env.yaml"}
+NNODES=${NNODES:-2}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
+
+# Algorithm
+temperature=1.0
+top_p=1.0
+top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
+val_top_p=0.7
+
+# Performance Related Parameter
+use_dynamic_bsz=True
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
+offload=True
+gen_tp=2
+train_tp=2
+train_pp=2
+
+# TODO: support dynamic_bsz for megatron
+# actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+# actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+# actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+# actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+# actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+# actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+
+python3 -m verl.trainer.main_ppo \
+ --config-path=config \
+ --config-name='ppo_megatron_trainer.yaml' \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.prompt_key=prompt \
+ data.truncation='left' \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.train_batch_size=${train_prompt_bsz} \
+ actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
+ algorithm.adv_estimator=${adv_estimator} \
+ algorithm.use_kl_in_reward=${use_kl_in_reward} \
+ algorithm.kl_ctrl.kl_coef=${kl_coef} \
+ actor_rollout_ref.actor.strategy=megatron \
+ critic.strategy=megatron \
+ actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
+ actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
+ actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
+ actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
+ actor_rollout_ref.actor.clip_ratio_c=10.0 \
+ +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
+ actor_rollout_ref.actor.optim.weight_decay=0.1 \
+ actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
+ actor_rollout_ref.actor.megatron.param_offload=${offload} \
+ actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \
+ actor_rollout_ref.actor.megatron.grad_offload=${offload} \
+ actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \
+ actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.optim.clip_grad=1.0 \
+ actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+ actor_rollout_ref.rollout.enable_chunked_prefill=True \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.temperature=${temperature} \
+ actor_rollout_ref.rollout.top_p=${top_p} \
+ actor_rollout_ref.rollout.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
+ actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
+ actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
+ actor_rollout_ref.rollout.val_kwargs.do_sample=True \
+ actor_rollout_ref.rollout.val_kwargs.n=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \
+ actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \
+ actor_rollout_ref.ref.megatron.param_offload=${offload} \
+ reward_model.reward_manager=dapo \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
+ +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
+ +reward_model.reward_kwargs.max_resp_len=${max_response_length} \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes="${NNODES}" \
+ trainer.val_before_train=True \
+ trainer.test_freq=10 \
+ trainer.save_freq=-1 \
+ trainer.total_epochs=10 \
+ trainer.total_training_steps=100 \
+ trainer.default_local_dir="${CKPTS_DIR}" \
+ trainer.resume_mode=auto \
+ trainer.log_val_generations=10
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_2_6.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_2_6.sh
new file mode 100644
index 0000000000000000000000000000000000000000..b2dfa578ed701ecd69c0270c35d451f31aaca48b
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_2_6.sh
@@ -0,0 +1,65 @@
+set -x
+
+project_name='GRPO'
+exp_name='GRPO-Qwen3-0.6b-gsm8k-fsdp2-one-step-off-2-6'
+
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-0.6B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/gsm8k/train.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/gsm8k/test.parquet"}
+
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=2
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+
+python3 -m verl.experimental.one_step_off_policy.main_ppo \
+ algorithm.adv_estimator=grpo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.train_batch_size=1152 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=192 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+ actor_rollout_ref.rollout.n=5 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.layered_summon=True \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.val_before_train=True \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.save_freq=-1 \
+ trainer.test_freq=5 \
+ trainer.total_epochs=2 \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" $@
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_2_6.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_2_6.sh
new file mode 100644
index 0000000000000000000000000000000000000000..1f5f72e6bcc8110d51b33ce1f48398b33c13c290
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_2_6.sh
@@ -0,0 +1,65 @@
+set -x
+
+project_name='GRPO'
+exp_name='GRPO-Qwen3-0.6b-gsm8k-fsdp2-sglang-one-step-off-2-6'
+
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen3-0.6B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/gsm8k/train.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/gsm8k/test.parquet"}
+
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=2
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+
+python3 -m verl.experimental.one_step_off_policy.main_ppo \
+ algorithm.adv_estimator=grpo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.train_batch_size=1152 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=192 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
+ actor_rollout_ref.rollout.name=sglang \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+ actor_rollout_ref.rollout.n=5 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.layered_summon=True \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.val_before_train=True \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.save_freq=-1 \
+ trainer.test_freq=5 \
+ trainer.total_epochs=2 \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" $@
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_3b_gsm8k_fsdp2_2_6.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_3b_gsm8k_fsdp2_2_6.sh
new file mode 100644
index 0000000000000000000000000000000000000000..b94a66f588bc435a5e366c7ef0e09f2d9a61fd7f
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_3b_gsm8k_fsdp2_2_6.sh
@@ -0,0 +1,64 @@
+set -x
+
+project_name='GRPO'
+exp_name='GRPO-Qwen3-0.6b-gsm8k-fsdp2-one-step-off-2-6'
+
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen/Qwen2.5-3B-Instruct"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/gsm8k/train.parquet"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/gsm8k/test.parquet"}
+
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+n_gpus_rollout=2
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+python3 -m verl.experimental.one_step_off_policy.main_ppo \
+ algorithm.adv_estimator=grpo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.train_batch_size=1152 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=192 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+ actor_rollout_ref.rollout.n=5 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.layered_summon=True \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.val_before_train=True \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.save_freq=-1 \
+ trainer.test_freq=5 \
+ trainer.total_epochs=2 \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" $@
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_qwen3_8b_gsm8k_fsdp2_8_8_npu.sh b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_qwen3_8b_gsm8k_fsdp2_8_8_npu.sh
new file mode 100644
index 0000000000000000000000000000000000000000..d6f884ad53af4e1f01cdc37736b4adf06db0bb1e
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/shell/grpo_qwen3_8b_gsm8k_fsdp2_8_8_npu.sh
@@ -0,0 +1,93 @@
+# The script has been validated on the Ascend Atlas 800T A3.
+set -x
+
+export HCCL_EXEC_TIMEOUT=60000
+export HCCL_CONNECT_TIMEOUT=7200
+
+project_name='GRPO'
+exp_name='GRPO-Qwen3-8b-gsm8k-fsdp2-one-step-off-8-8-npu'
+
+# Paths
+RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
+MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/Qwen/Qwen3-8B"}
+CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
+TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/BytedTsinghua-SIA/DAPO-Math-17k"}
+TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/BytedTsinghua-SIA/DAPO-Math-17k"}
+
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-16}
+
+n_gpus_rollout=8
+n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))
+
+max_prompt_length=$((1024 * 2))
+max_response_length=$((1024 * 32))
+
+use_dynamic_bsz=True
+sp_size=8
+fsdp_size=8
+actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))
+infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))
+
+python3 -m verl.experimental.one_step_off_policy.main_ppo \
+ algorithm.adv_estimator=grpo \
+ data.train_files="${TRAIN_FILE}" \
+ data.val_files="${TEST_FILE}" \
+ data.train_batch_size=32 \
+ data.max_prompt_length=${max_prompt_length} \
+ data.max_response_length=${max_response_length} \
+ data.filter_overlong_prompts=True \
+ data.filter_overlong_prompts_workers=64 \
+ data.truncation='error' \
+ actor_rollout_ref.actor.strategy=fsdp2 \
+ critic.strategy=fsdp2 \
+ actor_rollout_ref.model.path="${MODEL_PATH}" \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.hybrid_engine=False \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
+ actor_rollout_ref.actor.use_kl_loss=False \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.use_torch_compile=False \
+ actor_rollout_ref.ref.use_torch_compile=False \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=True \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.mode=async \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.enforce_eager=True \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \
+ actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \
+ actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \
+ actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
+ actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
+ algorithm.use_kl_in_reward=False \
+ actor_rollout_ref.nccl_timeout=14400 \
+ trainer.critic_warmup=0 \
+ trainer.val_before_train=False \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name="${project_name}" \
+ trainer.experiment_name="${exp_name}" \
+ trainer.default_local_dir=${CKPTS_DIR} \
+ trainer.save_freq=10 \
+ trainer.test_freq=-1 \
+ trainer.total_epochs=15 \
+ trainer.resume_mode=auto \
+ trainer.nnodes="${NNODES}" \
+ trainer.n_gpus_per_node="${n_gpus_training}" \
+ rollout.nnodes="${NNODES}" \
+ rollout.n_gpus_per_node="${n_gpus_rollout}" $@
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/utils.py b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..1879b0672fa68eda19a1b8e6553f4354b17816fe
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/one_step_off_policy/utils.py
@@ -0,0 +1,38 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 Meituan 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 omegaconf import DictConfig
+
+from verl.trainer.ppo.core_algos import AdvantageEstimator
+
+
+def need_critic(config: DictConfig) -> bool:
+ """Given a config, do we need critic"""
+ if config.algorithm.adv_estimator == AdvantageEstimator.GAE:
+ return True
+ elif config.algorithm.adv_estimator in [
+ AdvantageEstimator.GRPO,
+ AdvantageEstimator.GRPO_PASSK,
+ AdvantageEstimator.REINFORCE_PLUS_PLUS,
+ # AdvantageEstimator.REMAX, # TODO:REMAX advantage estimator is not yet supported in one_step_off_policy
+ AdvantageEstimator.RLOO,
+ AdvantageEstimator.OPO,
+ AdvantageEstimator.REINFORCE_PLUS_PLUS_BASELINE,
+ AdvantageEstimator.GPG,
+ ]:
+ return False
+ else:
+ raise NotImplementedError
diff --git a/code/RL_model/verl/verl_train/verl/experimental/reward_loop/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..03807f0277bffce8e11d1b59bb50d999a8909b96
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2024 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 .reward_loop import RewardLoopManager, RewardLoopWorker
+from .reward_model import RewardModelManager
+
+__all__ = ["RewardModelManager", "RewardLoopWorker", "RewardLoopManager"]
diff --git a/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_loop.py b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..77077dc084e49a646eb676bd9dc28fb3940287ae
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_loop.py
@@ -0,0 +1,321 @@
+# Copyright 2024 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.
+
+import asyncio
+import logging
+import os
+
+import aiohttp
+import numpy as np
+import ray
+import torch
+from omegaconf import DictConfig
+from tensordict import TensorDict
+
+from verl.protocol import DataProto
+from verl.single_controller.ray.base import RayResourcePool
+from verl.trainer.ppo.reward import get_custom_reward_fn
+from verl.utils import hf_tokenizer
+from verl.utils.fs import copy_to_local
+
+from .reward_manager import get_reward_manager_cls
+from .reward_model import RewardModelManager
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+@ray.remote
+class RewardLoopWorker:
+ def __init__(self, config: DictConfig, reward_router_address: str = None):
+ """
+ RewardLoopWork can tackle reward computation:
+ (1) rule-based reward computation
+ (2) reward model-based reward computation (both disrm and genrm)
+ (3) high-flexible user-customized reward function (can access rm by posting requests to reward_model_router)
+
+ Reward Computation Logic:
+ - if user-customized reward function is provided:
+ -> directly use user-customized reward function
+ - if user-customized reward function is not provided:
+ -> rm is not enabled: use default rule-based reward function
+ -> rm is disrm: compute reward score using disrm
+ -> rm is genrm: raise error (user-costomized reward func must be provided)
+
+ Args:
+ config: DictConfig, the config for reward loop worker.
+ reward_router_address: str, the address of reward router.
+ """
+ self.config = config
+ self.reward_router_address = reward_router_address
+ self._init_reward_fn()
+
+ def _init_reward_fn(self):
+ input_tokenizer_local_path = copy_to_local(self.config.actor_rollout_ref.model.path)
+ self.input_tokenizer = hf_tokenizer(input_tokenizer_local_path, trust_remote_code=True)
+ self.reward_model_tokenizer = None
+ if self.config.reward_model.enable:
+ reward_model_tokenizer_local_path = copy_to_local(self.config.reward_model.model.path)
+ self.reward_model_tokenizer = hf_tokenizer(reward_model_tokenizer_local_path, trust_remote_code=True)
+ self.reward_fn = get_custom_reward_fn(self.config)
+
+ # Load reward loop manager class
+ # Support both registry and importlib loading methods
+ reward_loop_source = self.config.reward_model.get("reward_loop_source", "register")
+
+ if reward_loop_source == "register":
+ # Load from registry (default behavior)
+ reward_manager_cls = get_reward_manager_cls(self.config.reward_model.reward_manager)
+ elif reward_loop_source == "importlib":
+ # Load from external module using importlib
+ from verl.utils.import_utils import load_extern_object
+
+ reward_loop_module_path = self.config.reward_model.get("reward_loop_module_path", None)
+ reward_loop_class_name = self.config.reward_model.get("reward_loop_class_name", None)
+
+ assert reward_loop_module_path is not None, (
+ "reward_loop_module_path must be set when reward_loop_source='importlib'"
+ )
+ assert reward_loop_class_name is not None, (
+ "reward_loop_class_name must be set when reward_loop_source='importlib'"
+ )
+
+ reward_manager_cls = load_extern_object(
+ module_path=reward_loop_module_path, object_name=reward_loop_class_name
+ )
+ else:
+ raise ValueError(f"Unknown reward_loop_source: {reward_loop_source}. Must be 'register' or 'importlib'")
+
+ self.reward_loop = reward_manager_cls(
+ self.config, self.input_tokenizer, self.reward_fn, self.reward_router_address, self.reward_model_tokenizer
+ )
+
+ async def compute_score_batch(self, data: DataProto) -> list[dict]:
+ tasks = []
+ for i in range(len(data)):
+ tasks.append(asyncio.create_task(self.compute_score(data[i : i + 1])))
+ outputs = await asyncio.gather(*tasks)
+ return outputs
+
+ async def compute_score(self, data: DataProto) -> dict:
+ assert len(data) == 1, "RewardLoopWorker only support single data item"
+ if self.config.custom_reward_function.path is not None:
+ # directly use user-customized reward function
+ return await self.reward_loop.run_single(data)
+ else:
+ if self.config.reward_model.enable:
+ # we assume the rm is disrm
+ # genrm must set custom_reward_function
+ return await self.compute_score_disrm(data)
+ else:
+ return await self.reward_loop.run_single(data)
+
+ async def _post_request(self, payload: dict, endpoint: str, max_retries: int = 16):
+ url = f"http://{self.reward_router_address}/{endpoint}"
+ last_exception = None
+ for attempt in range(max_retries):
+ try:
+ # It's safer to have a timeout instead of None, which can hang indefinitely.
+ timeout = aiohttp.ClientTimeout(total=None)
+ async with aiohttp.ClientSession(timeout=timeout) as session:
+ async with session.post(url, json=payload) as resp:
+ resp.raise_for_status()
+ return await resp.json()
+ except aiohttp.ClientResponseError as e:
+ # Do not retry on 4xx client errors, but retry on 5xx server errors.
+ if 400 <= e.status < 500:
+ logger.error(f"Request to {url} failed with client error HTTP {e.status}: {e}. Not retrying.")
+ raise
+ last_exception = e
+ logger.warning(
+ f"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed with HTTP {e.status}: {e}. "
+ "Retrying..."
+ )
+ except (asyncio.TimeoutError, aiohttp.ClientConnectorError) as e:
+ last_exception = e
+ logger.warning(f"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed: {e}. Retrying...")
+ except Exception as e:
+ last_exception = e
+ logger.warning(
+ f"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed with unexpected error: {e}. "
+ "Retrying..."
+ )
+
+ if attempt < max_retries - 1:
+ # Using exponential backoff is generally better than a fixed sleep.
+ backoff_seconds = 2**attempt
+ await asyncio.sleep(min(backoff_seconds, 30))
+
+ logger.error(f"Max retries ({max_retries}) reached for request to {url}.")
+ if last_exception:
+ raise last_exception
+
+ async def _preprocess_reward_inputs(self, data: DataProto) -> str:
+ assert len(data) == 1, "RewardLoopWorker only support single data item"
+ data_item = data[0]
+ assert "raw_prompt" in data_item.non_tensor_batch
+
+ # extract raw prompt
+ chat: list = list(data_item.non_tensor_batch["raw_prompt"])
+
+ # extract response
+ response_ids = data_item.batch["responses"]
+ response_length = response_ids.shape[-1]
+ valid_response_length = data_item.batch["attention_mask"][-response_length:].sum()
+ valid_response_ids = response_ids[:valid_response_length]
+
+ # decode
+ rollout_response = self.input_tokenizer.decode(valid_response_ids)
+ # remove bos and eos
+ rollout_response = rollout_response.replace(self.input_tokenizer.eos_token, "")
+
+ chat.append({"role": "assistant", "content": rollout_response})
+
+ rm_prompt = self.reward_model_tokenizer.apply_chat_template(
+ chat,
+ add_generation_prompt=False,
+ tokenize=False,
+ )
+
+ # llama tokenizer will add bos token by default
+ # will be removed in vllm >= 0.11.2, where we can add "add_special_tokens" = False
+ if self.reward_model_tokenizer.bos_token is not None and rm_prompt.startswith(
+ self.reward_model_tokenizer.bos_token
+ ):
+ rm_prompt = rm_prompt[len(self.reward_model_tokenizer.bos_token) :]
+
+ return rm_prompt
+
+ async def compute_score_disrm(self, data: DataProto) -> dict:
+ disrm_prompt = await self._preprocess_reward_inputs(data)
+ engine_name = self.config.reward_model.rollout.name
+ model_name = self.config.reward_model.model.path
+ if engine_name == "vllm":
+ # TODO (dyy): the "activation" has been changed to "use_activation" in vllm 0.11.2
+ payloads = {
+ "model": model_name,
+ "input": disrm_prompt,
+ "activation": False,
+ # "add_special_tokens": False, # vllm >= 0.11.2
+ }
+ output = await self._post_request(payloads, "classify")
+ rm_score = output["data"][-1]["probs"][-1]
+ elif engine_name == "sglang":
+ payloads = {
+ "model": model_name,
+ "input": disrm_prompt,
+ }
+ output = await self._post_request(payloads, "v1/embeddings")
+ rm_score = output["data"][-1]["embedding"][-1]
+ elif engine_name == "trtllm":
+ # TODO: remove this once TRT-LLM switches to TorchSampler
+ raise ValueError("TensorRT-LLM backend does not support reward models currently.")
+
+ payloads = {
+ "model": model_name,
+ "prompt": disrm_prompt,
+ "return_context_logits": True,
+ }
+ output = await self._post_request(payloads, "v1/completions")
+ rm_score = output["choices"][0]["context_logits"]
+ assert isinstance(rm_score, list) and len(rm_score) > 0, (
+ "TensorRT-LLM OpenAI server response for reward score is not in the expected format."
+ )
+
+ rm_score = float(rm_score[0][0])
+ logger.debug(f"rm score: {rm_score}")
+ else:
+ raise NotImplementedError(f"RewardLoopManager does not support {engine_name}")
+
+ return {"reward_score": rm_score}
+
+
+class RewardLoopManager:
+ """
+ RewardLoopManager run in single controller.
+ This class will create reward loop workers and manage them.
+ RewardLoopManager will deprecate fsdp/megatron RewardModelWorker in the future.
+ """
+
+ def __init__(self, config: DictConfig, rm_resource_pool: RayResourcePool = None):
+ self.config = config
+ if self.config.reward_model.enable:
+ self.reward_model_manager = RewardModelManager(config.reward_model, rm_resource_pool)
+ self.reward_router_address = self.reward_model_manager.get_router_address()
+ else:
+ self.reward_model_manager = None
+ self.reward_router_address = None
+
+ self._init_reward_loop_workers()
+
+ def _init_reward_loop_workers(self):
+ self.reward_loop_workers = []
+ num_workers = self.config.reward_model.num_workers
+ node_ids = [node["NodeID"] for node in ray.nodes() if node["Alive"] and node["Resources"].get("CPU", 0) > 0]
+
+ for i in range(num_workers):
+ # Round-robin scheduling over the all nodes
+ node_id = node_ids[i % len(node_ids)]
+ self.reward_loop_workers.append(
+ RewardLoopWorker.options(
+ name=f"reward_loop_worker_{i}",
+ scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+ node_id=node_id,
+ soft=True,
+ ),
+ ).remote(self.config, self.reward_router_address)
+ )
+
+ # this func is used to replace the legacy fsdp/megatron RewardModelWorker.compute_rm_score
+ def compute_rm_score(self, data: DataProto) -> DataProto:
+ if self.reward_model_manager is not None:
+ self.reward_model_manager.wake_up()
+
+ chunks = data.chunk(len(self.reward_loop_workers))
+ outputs = ray.get(
+ [
+ worker.compute_score_batch.remote(chunk)
+ for worker, chunk in zip(self.reward_loop_workers, chunks, strict=True)
+ ]
+ )
+ outputs_flat = [item for sublist in outputs for item in sublist]
+
+ # compute rm score
+ scores = [item["reward_score"] for item in outputs_flat]
+ prompt_length = data.batch["prompts"].size(1)
+ valid_response_length = data.batch["attention_mask"][:, prompt_length:].sum(dim=1)
+ rm_scores = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
+ rm_scores[torch.arange(rm_scores.size(0)), valid_response_length - 1] = torch.tensor(
+ scores, dtype=torch.float32
+ )
+ batch = TensorDict({"rm_scores": rm_scores}, batch_size=len(data))
+
+ reward_extra_infos = [output.get("reward_extra_info", {}) for output in outputs_flat]
+ reward_extra_keys = list(reward_extra_infos[0].keys())
+ non_tensor_batch = {}
+ for key in reward_extra_keys:
+ non_tensor_batch[key] = np.array([info[key] for info in reward_extra_infos])
+
+ if self.reward_model_manager is not None:
+ self.reward_model_manager.sleep()
+
+ return DataProto(
+ batch=batch, non_tensor_batch=non_tensor_batch, meta_info={"reward_extra_keys": reward_extra_keys}
+ )
+
+ def _run_all(self, tasks: list[asyncio.Task]):
+ async def run_all():
+ return await asyncio.gather(*tasks)
+
+ return asyncio.run(run_all())
diff --git a/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_manager/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_manager/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..75a440a2324f36c000eec87278414e829f44221c
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_manager/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 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 .registry import get_reward_manager_cls, register # noqa: I001
+from .dapo import DAPORewardManager
+from .naive import NaiveRewardManager
+from .limited import RateLimitedRewardManager
+from .remote import RemoteRewardManager
+
+__all__ = [
+ "DAPORewardManager",
+ "NaiveRewardManager",
+ "RateLimitedRewardManager",
+ "RemoteRewardManager",
+ "register",
+ "get_reward_manager_cls",
+]
diff --git a/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_model.py b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bc05e1eea142247d6313c221b000ce31dd092f9
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/reward_loop/reward_model.py
@@ -0,0 +1,119 @@
+# Copyright 2024 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.
+
+import asyncio
+import logging
+import os
+
+from verl.single_controller.ray.base import RayResourcePool, split_resource_pool
+from verl.workers.config import HFModelConfig, RewardModelConfig
+from verl.workers.rollout.replica import get_rollout_replica_class
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class RewardModelManager:
+ """Reward model manager."""
+
+ def __init__(
+ self,
+ config: RewardModelConfig,
+ resource_pool: RayResourcePool = None,
+ ):
+ """
+ Initialize the reward model manager.
+
+ Args:
+ config (RewardModelConfig): Reward model configuration.
+ resource_pool (RayResourcePool, optional): Resource pool. Defaults to None.
+ """
+ self.config = config
+ self.resource_pool = resource_pool
+ self._initialize_llm_servers()
+ self._initialize_router()
+ assert self.config.rollout.skip_tokenizer_init is False, "Reward model should not skip tokenizer init."
+ if self.config.rollout.free_cache_engine:
+ self.sleep()
+
+ def _initialize_llm_servers(self):
+ rollout_world_size = self.config.rollout.tensor_model_parallel_size
+ world_size = (
+ self.resource_pool.world_size
+ if self.resource_pool # colocate mode
+ else self.config.n_gpus_per_node * self.config.nnodes # standalone mode
+ )
+ num_replicas = world_size // rollout_world_size
+
+ rollout_replica_class = get_rollout_replica_class(self.config.rollout.name)
+ rollout_config = self.config.rollout
+ model_config = HFModelConfig(
+ path=self.config.model.path,
+ external_lib=self.config.model.external_lib,
+ trust_remote_code=self.config.model.trust_remote_code,
+ )
+ self.tokenizer = model_config.get_processor()
+ self.rollout_replicas = [
+ rollout_replica_class(
+ replica_rank=replica_rank,
+ config=rollout_config,
+ model_config=model_config,
+ gpus_per_node=self.config.n_gpus_per_node,
+ is_reward_model=True,
+ )
+ for replica_rank in range(num_replicas)
+ ]
+ if self.resource_pool:
+ split_resource_pools = split_resource_pool(self.resource_pool, split_size=rollout_world_size)
+ assert len(split_resource_pools) == len(self.rollout_replicas)
+ self._run_all(
+ [
+ server.init_colocated(resource_pool)
+ for server, resource_pool in zip(self.rollout_replicas, split_resource_pools, strict=True)
+ ]
+ )
+ else:
+ self._run_all([server.init_standalone() for server in self.rollout_replicas])
+ self.server_handles = [server._server_handle for server in self.rollout_replicas]
+ self.server_addresses = [server._server_address for server in self.rollout_replicas]
+
+ def _initialize_router(self):
+ worker_urls = [f"http://{server_address}" for server_address in self.server_addresses]
+
+ # TODO (dyy): sglang router is not ready yet.
+ # if self.config.rollout.name == "sglang":
+ # from .router.inner_sglang_router import launch_router_process
+ # else:
+ # from .router.naive_router import launch_router_process
+
+ from .router.naive_router import launch_router_process
+
+ self.router_address, _ = launch_router_process(worker_urls=worker_urls)
+
+ def get_router_address(self):
+ return self.router_address
+
+ def wake_up(self):
+ """Wake up all rollout replica instances."""
+ self._run_all([replica.wake_up() for replica in self.rollout_replicas])
+
+ def sleep(self):
+ """Sleep all rollout replica instances."""
+ self._run_all([replica.sleep() for replica in self.rollout_replicas])
+
+ def _run_all(self, tasks: list[asyncio.Task]):
+ async def run_all():
+ await asyncio.gather(*tasks)
+
+ asyncio.run(run_all())
diff --git a/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/agent_loop.py b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/agent_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..0887b4600e9eecb5e4cabc4417ed5d93fe4fade3
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/agent_loop.py
@@ -0,0 +1,95 @@
+# Copyright 2025 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.
+import asyncio
+
+import numpy as np
+import ray
+from transfer_queue import BatchMeta
+
+import verl.experimental.agent_loop.agent_loop as agent_loop
+
+
+class AgentLoopManager(agent_loop.AgentLoopManager):
+ def generate_sequences(self, prompts: BatchMeta) -> BatchMeta:
+ """Split input batch and dispatch to agent loop workers.
+
+ Args:
+ prompts (BatchMeta): Input batch.
+
+ Returns:
+ BatchMeta: Output batch metadata.
+ """
+
+ if self.reward_model_manager and self.config.reward_model.rollout.free_cache_engine:
+ self.reward_model_manager.wake_up()
+
+ chunkes = prompts.chunk(len(self.agent_loop_workers))
+ outputs = ray.get(
+ [
+ worker.generate_sequences.remote(chunk)
+ for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True)
+ ]
+ )
+ output = BatchMeta.concat(outputs)
+ if self.reward_model_manager and self.config.reward_model.rollout.free_cache_engine:
+ self.reward_model_manager.sleep()
+
+ # calculate performance metrics
+ metrics = [output.extra_info.pop("metrics") for output in outputs] # List[List[Dict[str, str]]]
+ timing = self._performance_metrics(metrics, output)
+
+ output.set_extra_info("timing", timing)
+ return output
+
+ def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: BatchMeta) -> dict[str, float]:
+ timing = {}
+ t_generate_sequences = np.array([metric["generate_sequences"] for chunk in metrics for metric in chunk])
+ t_tool_calls = np.array([metric["tool_calls"] for chunk in metrics for metric in chunk])
+ timing["agent_loop/generate_sequences/min"] = t_generate_sequences.min()
+ timing["agent_loop/generate_sequences/max"] = t_generate_sequences.max()
+ timing["agent_loop/generate_sequences/mean"] = t_generate_sequences.mean()
+ timing["agent_loop/tool_calls/min"] = t_tool_calls.min()
+ timing["agent_loop/tool_calls/max"] = t_tool_calls.max()
+ timing["agent_loop/tool_calls/mean"] = t_tool_calls.mean()
+
+ # TODO (TQ): initialize tq during init when enable TQ switch is stable
+ tq_client = self._create_transferqueue_client()
+ # batch sequence generation is bounded by the slowest sample
+ slowest = np.argmax(t_generate_sequences + t_tool_calls)
+ attention_mask = asyncio.run(tq_client.async_get_data(output[slowest]))["attention_mask"]
+ prompt_length = output.samples[0].fields["prompts"].shape[0]
+ timing["agent_loop/slowest/generate_sequences"] = t_generate_sequences[slowest]
+ timing["agent_loop/slowest/tool_calls"] = t_tool_calls[slowest]
+ timing["agent_loop/slowest/prompt_length"] = attention_mask[:prompt_length].sum().item()
+ timing["agent_loop/slowest/response_length"] = attention_mask[prompt_length:].sum().item()
+
+ return timing
+
+ def create_transferqueue_client_for_workers(self):
+ # TODO (TQ): initialize tq during worker init when enable TQ switch is stable
+ ray.get([worker.create_transferqueue_client.remote() for worker in self.agent_loop_workers])
+
+ def _create_transferqueue_client(self):
+ """Create a client for data system (TransferQueue)."""
+ from verl.single_controller.ray.base import get_random_string
+ from verl.utils.transferqueue_utils import create_transferqueue_client
+
+ client_name = get_random_string(length=6)
+
+ tq_client = create_transferqueue_client(
+ client_id=f"AgentLoopManager_{client_name}",
+ config=self.config.transfer_queue,
+ )
+
+ return tq_client
diff --git a/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/config/transfer_queue_ppo_megatron_trainer.yaml b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/config/transfer_queue_ppo_megatron_trainer.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..37b19b45708fc4a09ab3c15bffb92af3f4a50d07
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/config/transfer_queue_ppo_megatron_trainer.yaml
@@ -0,0 +1,14 @@
+hydra:
+ searchpath:
+ - file://verl/trainer/config
+
+defaults:
+ - ppo_megatron_trainer
+ - _self_
+
+# config for TransferQueue
+transfer_queue:
+ enable: True
+ num_global_batch: 1
+ storage_backend: AsyncSimpleStorageManager
+ num_data_storage_units: 8
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/config/transfer_queue_ppo_trainer.yaml b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/config/transfer_queue_ppo_trainer.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..7a5f57ddd4f12bf4fd385cbdde9c53c626d02a0d
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/config/transfer_queue_ppo_trainer.yaml
@@ -0,0 +1,14 @@
+hydra:
+ searchpath:
+ - file://verl/trainer/config
+
+defaults:
+ - ppo_trainer
+ - _self_
+
+# config for TransferQueue
+transfer_queue:
+ enable: True
+ num_global_batch: 1
+ storage_backend: AsyncSimpleStorageManager
+ num_data_storage_units: 8
diff --git a/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/main_ppo.py b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/main_ppo.py
new file mode 100644
index 0000000000000000000000000000000000000000..a29ff2cf86235e7dea758409043f5af18ea2166e
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/main_ppo.py
@@ -0,0 +1,203 @@
+# Copyright 2024 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.
+"""
+Note that we don't combine the main with ray_trainer as ray_trainer is used by other main.
+"""
+
+import os
+import socket
+
+import hydra
+import ray
+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.ppo.reward import load_reward_manager
+from verl.trainer.ppo.utils import need_critic, need_reference_policy
+from verl.utils.config import validate_config
+from verl.utils.device import auto_set_device, is_cuda_available
+
+from .ray_trainer import RayPPOTrainer
+
+
+@hydra.main(config_path="config", config_name="ppo_trainer", version_base=None)
+def main(config):
+ """Main entry point for PPO training with Hydra configuration management.
+
+ Args:
+ config_dict: Hydra configuration dictionary containing training parameters.
+ """
+ # Automatically set `config.trainer.device = npu` when running on Ascend NPU.
+ auto_set_device(config)
+
+ run_ppo(config)
+
+
+# Define a function to run the PPO-like training process
+def run_ppo(config, task_runner_class=None) -> None:
+ """Initialize Ray cluster and run distributed PPO training process.
+
+ Args:
+ config: Training configuration object containing all necessary parameters
+ for distributed PPO training including Ray initialization settings,
+ model paths, and training hyperparameters.
+ task_runner_class: For recipe to change TaskRunner.
+ """
+ # Check if Ray is not initialized
+ if not ray.is_initialized():
+ # Initialize Ray with a local cluster configuration
+ # Set environment variables in the runtime environment to control tokenizer parallelism,
+ # NCCL debug level, VLLM logging level, and allow runtime LoRA updating
+ # `num_cpus` specifies the number of CPU cores Ray can use, obtained from the configuration
+ default_runtime_env = get_ppo_ray_runtime_env()
+ ray_init_kwargs = config.ray_kwargs.get("ray_init", {})
+ runtime_env_kwargs = ray_init_kwargs.get("runtime_env", {})
+
+ if config.transfer_queue.enable:
+ # Add runtime environment variables for transfer queue
+ runtime_env_vars = runtime_env_kwargs.get("env_vars", {})
+ runtime_env_vars["TRANSFER_QUEUE_ENABLE"] = "1"
+ runtime_env_kwargs["env_vars"] = runtime_env_vars
+
+ runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)
+ ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, "runtime_env": runtime_env})
+ print(f"ray init kwargs: {ray_init_kwargs}")
+ ray.init(**OmegaConf.to_container(ray_init_kwargs))
+
+ if task_runner_class is None:
+ task_runner_class = ray.remote(num_cpus=1)(TaskRunner) # please make sure main_task is not scheduled on head
+
+ # Create a remote instance of the TaskRunner class, and
+ # Execute the `run` method of the TaskRunner instance remotely and wait for it to complete
+ if (
+ is_cuda_available
+ and config.global_profiler.tool == "nsys"
+ and config.global_profiler.get("steps") is not None
+ and len(config.global_profiler.get("steps", [])) > 0
+ ):
+ from verl.utils.import_utils import is_nvtx_available
+
+ assert is_nvtx_available(), "nvtx is not available in CUDA platform. Please 'pip3 install nvtx'"
+ nsight_options = OmegaConf.to_container(
+ config.global_profiler.global_tool_config.nsys.controller_nsight_options
+ )
+ runner = task_runner_class.options(runtime_env={"nsight": nsight_options}).remote()
+ else:
+ runner = task_runner_class.remote()
+ ray.get(runner.run.remote(config))
+
+ # [Optional] get the path of the timeline trace file from the configuration, default to None
+ # This file is used for performance analysis
+ timeline_json_file = config.ray_kwargs.get("timeline_json_file", None)
+ if timeline_json_file:
+ ray.timeline(filename=timeline_json_file)
+
+
+class TaskRunner(MainTaskRunner):
+ def run(self, config):
+ """Execute the main PPO training workflow.
+
+ This method sets up the distributed training environment, initializes
+ workers, datasets, and reward functions, then starts the training process.
+
+ Args:
+ config: Training configuration object containing all parameters needed
+ for setting up and running the PPO training process.
+ """
+ # Print the initial configuration. `resolve=True` will evaluate symbolic values.
+ from pprint import pprint
+
+ from verl.utils.fs import copy_to_local
+
+ print(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}")
+ pprint(OmegaConf.to_container(config, resolve=True))
+ OmegaConf.resolve(config)
+
+ actor_rollout_cls, ray_worker_group_cls = self.add_actor_rollout_worker(config)
+ self.add_critic_worker(config)
+
+ # We should adopt a multi-source reward function here:
+ # - for rule-based rm, we directly call a reward score
+ # - for model-based rm, we call a model
+ # - for code related prompt, we send to a sandbox if there are test cases
+ # finally, we combine all the rewards together
+ # The reward type depends on the tag of the data
+ self.add_reward_model_worker(config)
+
+ # Add a reference policy worker if KL loss or KL reward is used.
+ self.add_ref_policy_worker(config, actor_rollout_cls)
+
+ # validate config
+ validate_config(
+ config=config,
+ use_reference_policy=need_reference_policy(config),
+ use_critic=need_critic(config),
+ )
+
+ # Download the checkpoint from HDFS to the local machine.
+ # `use_shm` determines whether to use shared memory, which could lead to faster model loading if turned on
+ local_path = copy_to_local(
+ config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get("use_shm", False)
+ )
+
+ # Instantiate the tokenizer and processor.
+ from verl.utils 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)
+ # Used for multimodal LLM, could be None
+ processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)
+
+ # Load the reward manager for training and validation.
+ reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})
+ )
+ val_reward_fn = load_reward_manager(
+ config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {})
+ )
+
+ resource_pool_manager = self.init_resource_pool_mgr(config)
+
+ from verl.utils.dataset.rl_dataset import collate_fn
+
+ # Create training and validation datasets.
+ train_dataset = create_rl_dataset(config.data.train_files, config.data, tokenizer, processor, is_train=True)
+ val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor, is_train=False)
+ train_sampler = create_rl_sampler(config.data, train_dataset)
+
+ # Initialize the PPO trainer.
+ trainer = RayPPOTrainer(
+ config=config,
+ tokenizer=tokenizer,
+ processor=processor,
+ role_worker_mapping=self.role_worker_mapping,
+ resource_pool_manager=resource_pool_manager,
+ ray_worker_group_cls=ray_worker_group_cls,
+ reward_fn=reward_fn,
+ val_reward_fn=val_reward_fn,
+ train_dataset=train_dataset,
+ val_dataset=val_dataset,
+ collate_fn=collate_fn,
+ train_sampler=train_sampler,
+ )
+ # Initialize the workers of the trainer.
+ trainer.init_workers()
+ # Start the training process.
+ trainer.fit()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/ray_trainer.py b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/ray_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f2be802b0fe62f216aafe2bc39f908f8ce75df7
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/ray_trainer.py
@@ -0,0 +1,1660 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023-2024 SGLang Team
+# Copyright 2025 ModelBest Inc. 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.
+"""
+PPO Trainer with Ray-based single controller.
+This trainer supports model-agonistic model initialization with huggingface
+"""
+
+import json
+import logging
+import math
+import os
+import uuid
+from collections import defaultdict
+from pprint import pprint
+from typing import Any, Optional
+
+import numpy as np
+import ray
+import tensordict
+import torch
+from omegaconf import OmegaConf, open_dict
+from packaging.version import parse as parse_version
+from tensordict import TensorDict
+from torch.utils.data import Dataset, Sampler
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm import tqdm
+from transfer_queue import (
+ BatchMeta,
+ SimpleStorageUnit,
+ TransferQueueController,
+ get_placement_group,
+ process_zmq_server_info,
+)
+
+from verl import DataProto
+from verl.checkpoint_engine import CheckpointEngineManager
+from verl.experimental.dataset.sampler import AbstractCurriculumSampler
+from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup, ResourcePoolManager
+from verl.single_controller.ray.base import create_colocated_worker_cls
+from verl.trainer.config import AlgoConfig
+from verl.trainer.ppo import core_algos
+from verl.trainer.ppo.core_algos import AdvantageEstimator, agg_loss
+from verl.trainer.ppo.metric_utils import (
+ compute_data_metrics,
+ compute_throughout_metrics,
+ compute_timing_metrics,
+ process_validation_metrics,
+)
+from verl.trainer.ppo.reward import compute_reward, compute_reward_async
+from verl.trainer.ppo.utils import Role, WorkerType, need_critic, need_reference_policy, need_reward_model
+from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi
+from verl.utils.config import omega_conf_to_dataclass
+from verl.utils.debug import marked_timer
+from verl.utils.metric import reduce_metrics
+from verl.utils.rollout_skip import RolloutSkip
+from verl.utils.seqlen_balancing import calculate_workload, get_seqlen_balanced_partitions, log_seqlen_unbalance
+from verl.utils.torch_functional import masked_mean
+from verl.utils.tracking import ValidationGenerationsLogger
+from verl.utils.transferqueue_utils import create_transferqueue_client, get_transferqueue_client, tqbridge
+
+
+@tqbridge(put_data=False)
+def compute_reward_decorated(data, reward_fn):
+ return compute_reward(data, reward_fn)
+
+
+@tqbridge(put_data=False)
+def compute_reward_async_decorated(data, reward_fn):
+ return compute_reward_async.remote(data, reward_fn)
+
+
+@tqbridge(put_data=False)
+def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl"):
+ """Apply KL penalty to the token-level rewards.
+
+ This function computes the KL divergence between the reference policy and current policy,
+ then applies a penalty to the token-level rewards based on this divergence.
+
+ Args:
+ data (DataProto): The data containing batched model outputs and inputs.
+ kl_ctrl (core_algos.AdaptiveKLController): Controller for adaptive KL penalty.
+ kl_penalty (str, optional): Type of KL penalty to apply. Defaults to "kl".
+
+ Returns:
+ tuple: A tuple containing:
+ - The updated data with token-level rewards adjusted by KL penalty
+ - A dictionary of metrics related to the KL penalty
+ """
+ response_mask = data.batch["response_mask"]
+ token_level_scores = data.batch["token_level_scores"]
+ batch_size = data.batch.batch_size[0]
+
+ # compute kl between ref_policy and current policy
+ # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference model has been enabled.
+ kld = core_algos.kl_penalty(
+ data.batch["old_log_probs"], data.batch["ref_log_prob"], kl_penalty=kl_penalty
+ ) # (batch_size, response_length)
+ kld = kld * response_mask
+ beta = kl_ctrl.value
+
+ token_level_rewards = token_level_scores - beta * kld
+
+ current_kl = masked_mean(kld, mask=response_mask, axis=-1) # average over sequence
+ current_kl = torch.mean(current_kl, dim=0).item()
+
+ # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837
+ kl_ctrl.update(current_kl=current_kl, n_steps=batch_size)
+
+ metrics = {"actor/reward_kl_penalty": current_kl, "actor/reward_kl_penalty_coeff": beta}
+
+ return token_level_rewards, metrics
+
+
+def compute_response_mask(batch_meta: BatchMeta, tq_client):
+ """Compute the attention mask for the response part of the sequence.
+
+ This function extracts the portion of the attention mask that corresponds to the model's response,
+ which is used for masking computations that should only apply to response tokens.
+
+ Args:
+ batch_meta (BatchMeta): The data containing batched model outputs and inputs.
+
+ Returns:
+ BatchMeta: The BatchMeta of attention mask for the response tokens.
+ """
+ data = tq_client.get_data(batch_meta)
+
+ responses = data["responses"]
+ response_length = responses.size(1)
+ attention_mask = data["attention_mask"]
+ response_mask = attention_mask[:, -response_length:]
+ output = TensorDict({"response_mask": response_mask}, batch_size=response_mask.size(0))
+
+ batch_meta = tq_client.put(data=output, metadata=batch_meta)
+
+ return batch_meta
+
+
+@tqbridge(put_data=False)
+def compute_advantage(
+ data: DataProto,
+ adv_estimator: AdvantageEstimator,
+ gamma: float = 1.0,
+ lam: float = 1.0,
+ num_repeat: int = 1,
+ norm_adv_by_std_in_grpo: bool = True,
+ config: Optional[AlgoConfig] = None,
+) -> tuple[Any, Any]:
+ """Compute advantage estimates for policy optimization.
+
+ This function computes advantage estimates using various estimators like GAE, GRPO, REINFORCE++, etc.
+ The advantage estimates are used to guide policy optimization in RL algorithms.
+
+ Args:
+ data (DataProto): The data containing batched model outputs and inputs.
+ adv_estimator (AdvantageEstimator): The advantage estimator to use (e.g., GAE, GRPO, REINFORCE++).
+ gamma (float, optional): Discount factor for future rewards. Defaults to 1.0.
+ lam (float, optional): Lambda parameter for GAE. Defaults to 1.0.
+ num_repeat (int, optional): Number of times to repeat the computation. Defaults to 1.
+ norm_adv_by_std_in_grpo (bool, optional): Whether to normalize advantages by standard deviation in
+ GRPO. Defaults to True.
+ config (dict, optional): Configuration dictionary for algorithm settings. Defaults to None.
+
+ Returns:
+ tuple: A tuple containing:
+ - advantages: The computed advantage estimates.
+ - returns: The computed returns.
+ """
+ # prepare response group
+ if adv_estimator == AdvantageEstimator.GAE:
+ # Compute advantages and returns using Generalized Advantage Estimation (GAE)
+ advantages, returns = core_algos.compute_gae_advantage_return(
+ token_level_rewards=data.batch["token_level_rewards"],
+ values=data.batch["values"],
+ response_mask=data.batch["response_mask"],
+ gamma=gamma,
+ lam=lam,
+ )
+ # TODO (TQ): adapt core_algos.compute_pf_ppo_reweight_data function to support transfer queue
+ if config.get("use_pf_ppo", False):
+ data = core_algos.compute_pf_ppo_reweight_data(
+ data,
+ config.pf_ppo.get("reweight_method"),
+ config.pf_ppo.get("weight_pow"),
+ )
+ elif adv_estimator == AdvantageEstimator.GRPO:
+ # Initialize the mask for GRPO calculation
+ grpo_calculation_mask = data.batch["response_mask"]
+ # Call compute_grpo_outcome_advantage with parameters matching its definition
+ advantages, returns = core_algos.compute_grpo_outcome_advantage(
+ token_level_rewards=data.batch["token_level_rewards"],
+ response_mask=grpo_calculation_mask,
+ index=data.non_tensor_batch["uid"],
+ norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+ )
+ else:
+ # handle all other adv estimator type other than GAE and GRPO
+ adv_estimator_fn = core_algos.get_adv_estimator_fn(adv_estimator)
+ adv_kwargs = {
+ "token_level_rewards": data.batch["token_level_rewards"],
+ "response_mask": data.batch["response_mask"],
+ "config": config,
+ }
+ if "uid" in data.non_tensor_batch: # optional
+ adv_kwargs["index"] = data.non_tensor_batch["uid"]
+ if "reward_baselines" in data.batch: # optional
+ adv_kwargs["reward_baselines"] = data.batch["reward_baselines"]
+
+ # calculate advantage estimator
+ advantages, returns = adv_estimator_fn(**adv_kwargs)
+ return advantages, returns
+
+
+@tqbridge(put_data=False)
+def compute_data_metrics_decorated(batch, use_critic: bool = True):
+ return compute_data_metrics(batch, use_critic)
+
+
+@tqbridge(put_data=False)
+def compute_timing_metrics_decorated(batch, timing_raw: dict[str, float]) -> dict[str, Any]:
+ return compute_timing_metrics(batch, timing_raw)
+
+
+@tqbridge(put_data=False)
+def compute_throughout_metrics_decorated(batch, timing_raw: dict[str, float], n_gpus: int) -> dict[str, Any]:
+ return compute_throughout_metrics(batch, timing_raw, n_gpus)
+
+
+@tqbridge(put_data=False)
+def calculate_debug_metrics_decorated(data):
+ from verl.utils.debug.metrics import calculate_debug_metrics
+
+ return calculate_debug_metrics(data)
+
+
+@tqbridge(put_data=False)
+def compute_val_reward_decorated(reward_fn, data, return_dict):
+ return reward_fn(data, return_dict)
+
+
+class RayPPOTrainer:
+ """Distributed PPO trainer using Ray for scalable reinforcement learning.
+
+ This trainer orchestrates distributed PPO training across multiple nodes and GPUs,
+ managing actor rollouts, critic training, and reward computation with Ray backend.
+ Supports various model architectures including FSDP, Megatron, vLLM, and SGLang integration.
+ """
+
+ # TODO: support each role have individual ray_worker_group_cls,
+ # i.e., support different backend of different role
+ def __init__(
+ self,
+ config,
+ tokenizer,
+ role_worker_mapping: dict[Role, WorkerType],
+ resource_pool_manager: ResourcePoolManager,
+ ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup,
+ processor=None,
+ reward_fn=None,
+ val_reward_fn=None,
+ train_dataset: Optional[Dataset] = None,
+ val_dataset: Optional[Dataset] = None,
+ collate_fn=None,
+ train_sampler: Optional[Sampler] = None,
+ device_name=None,
+ ):
+ """
+ Initialize distributed PPO trainer with Ray backend.
+ Note that this trainer runs on the driver process on a single CPU/GPU node.
+
+ Args:
+ config: Configuration object containing training parameters.
+ tokenizer: Tokenizer used for encoding and decoding text.
+ role_worker_mapping (dict[Role, WorkerType]): Mapping from roles to worker classes.
+ resource_pool_manager (ResourcePoolManager): Manager for Ray resource pools.
+ ray_worker_group_cls (RayWorkerGroup, optional): Class for Ray worker groups. Defaults to RayWorkerGroup.
+ processor: Optional data processor, used for multimodal data
+ reward_fn: Function for computing rewards during training.
+ val_reward_fn: Function for computing rewards during validation.
+ train_dataset (Optional[Dataset], optional): Training dataset. Defaults to None.
+ val_dataset (Optional[Dataset], optional): Validation dataset. Defaults to None.
+ collate_fn: Function to collate data samples into batches.
+ train_sampler (Optional[Sampler], optional): Sampler for the training dataset. Defaults to None.
+ device_name (str, optional): Device name for training (e.g., "cuda", "cpu"). Defaults to None.
+ """
+
+ # Store the tokenizer for text processing
+ self.tokenizer = tokenizer
+ self.processor = processor
+ self.config = config
+ self.reward_fn = reward_fn
+ self.val_reward_fn = val_reward_fn
+
+ self.hybrid_engine = config.actor_rollout_ref.hybrid_engine
+ assert self.hybrid_engine, "Currently, only support hybrid engine"
+
+ if self.hybrid_engine:
+ assert Role.ActorRollout in role_worker_mapping, f"{role_worker_mapping.keys()=}"
+
+ self.role_worker_mapping = role_worker_mapping
+ self.resource_pool_manager = resource_pool_manager
+ 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
+ self.device_name = device_name if device_name else self.config.trainer.device
+ self.validation_generations_logger = ValidationGenerationsLogger(
+ project_name=self.config.trainer.project_name,
+ experiment_name=self.config.trainer.experiment_name,
+ )
+
+ 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)
+ # if ref_in_actor is True, the reference policy will be actor without lora applied
+ self.ref_in_actor = lora_rank > 0
+
+ # define in-reward KL control
+ # kl loss control currently not suppoorted
+ if self.config.algorithm.use_kl_in_reward:
+ self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl)
+
+ self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)
+
+ self.tq_client = self._initialize_transferqueue()
+
+ def _initialize_transferqueue(self):
+ # 1. initialize TransferQueueStorage
+ if self.config.transfer_queue.storage_backend == "AsyncSimpleStorageManager":
+ train_data_size = (
+ self.config.data.train_batch_size
+ * self.config.transfer_queue.num_global_batch
+ * self.config.actor_rollout_ref.rollout.n
+ )
+ val_data_size = self.val_dataset_size * self.config.actor_rollout_ref.rollout.val_kwargs.n
+
+ total_storage_size = train_data_size + val_data_size
+ self.data_system_storage_units = {}
+ storage_placement_group = get_placement_group(
+ self.config.transfer_queue.num_data_storage_units, num_cpus_per_actor=1
+ )
+ for storage_unit_rank in range(self.config.transfer_queue.num_data_storage_units):
+ storage_node = SimpleStorageUnit.options(
+ placement_group=storage_placement_group, placement_group_bundle_index=storage_unit_rank
+ ).remote(
+ storage_unit_size=math.ceil(total_storage_size / self.config.transfer_queue.num_data_storage_units)
+ )
+ self.data_system_storage_units[storage_unit_rank] = storage_node
+ logging.info(f"SimpleStorageUnit #{storage_unit_rank} has been created.")
+ else:
+ raise NotImplementedError("Currently only support AsyncSimpleStorageManager backend in TransferQueue")
+
+ # 2. Initialize TransferQueueController (single controller only)
+
+ # Sampler usage instructions:
+ # For GRPO grouped sampling, you can initialize the controller with GRPOGroupNSampler:
+ # Option 1: Pass sampler class (will be instantiated automatically)
+ # self.data_system_controller = TransferQueueController.remote(sampler=GRPOGroupNSampler)
+
+ # Option 2: Pass sampler instance (if you need custom configuration)
+ # grpo_sampler = GRPOGroupNSampler()
+ # self.data_system_controller = TransferQueueController.remote(sampler=grpo_sampler)
+
+ # Then use sampling_config in get_meta calls:
+ # sampling_config={"n_samples_per_prompt": 4}
+ self.data_system_controller = TransferQueueController.remote()
+ logging.info("TransferQueueController has been created.")
+
+ # 3. register controller & storage and prepare necessary information
+ self.data_system_controller_info = process_zmq_server_info(self.data_system_controller)
+ if self.config.transfer_queue.storage_backend == "AsyncSimpleStorageManager":
+ self.data_system_storage_unit_infos = process_zmq_server_info(self.data_system_storage_units)
+
+ # Note: Need to generate a new DictConfig with allow_objects=True to preserve ZMQServerInfo instances
+ # (which contain socket connection details). Without this flag, OmegaConf would flatten these objects to dicts,
+ # breaking the transfer queue client initialization.
+ tq_config = OmegaConf.create({"transfer_queue": {}}, flags={"allow_objects": True})
+ tq_config.transfer_queue.controller_info = self.data_system_controller_info
+
+ if self.config.transfer_queue.storage_backend == "AsyncSimpleStorageManager":
+ tq_config.transfer_queue.storage_unit_infos = self.data_system_storage_unit_infos
+
+ self.config = OmegaConf.merge(tq_config, self.config)
+
+ # 4. create client
+ create_transferqueue_client(client_id="Trainer", config=self.config.transfer_queue, sync=True)
+ tq_client = get_transferqueue_client()
+ return tq_client
+
+ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampler: Optional[Sampler]):
+ """
+ Creates the train and validation dataloaders.
+ """
+ # TODO: we have to make sure the batch size is divisible by the dp size
+ from verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler
+
+ if train_dataset is None:
+ train_dataset = create_rl_dataset(
+ self.config.data.train_files, self.config.data, self.tokenizer, self.processor
+ )
+ if val_dataset is None:
+ val_dataset = create_rl_dataset(
+ self.config.data.val_files, self.config.data, self.tokenizer, self.processor
+ )
+ self.train_dataset, self.val_dataset = train_dataset, val_dataset
+
+ self.val_dataset_size = len(val_dataset)
+
+ if train_sampler is None:
+ train_sampler = create_rl_sampler(self.config.data, self.train_dataset)
+ if collate_fn is None:
+ from verl.utils.dataset.rl_dataset import collate_fn as default_collate_fn
+
+ collate_fn = default_collate_fn
+
+ num_workers = self.config.data["dataloader_num_workers"]
+
+ self.train_dataloader = StatefulDataLoader(
+ dataset=self.train_dataset,
+ batch_size=self.config.data.get("gen_batch_size", self.config.data.train_batch_size),
+ num_workers=num_workers,
+ drop_last=True,
+ collate_fn=collate_fn,
+ sampler=train_sampler,
+ )
+
+ val_batch_size = self.config.data.val_batch_size # Prefer config value if set
+ if val_batch_size is None:
+ val_batch_size = len(self.val_dataset)
+ self.val_batch_size = val_batch_size
+
+ self.val_dataloader = StatefulDataLoader(
+ dataset=self.val_dataset,
+ batch_size=val_batch_size,
+ num_workers=num_workers,
+ shuffle=self.config.data.get("validation_shuffle", True),
+ drop_last=False,
+ collate_fn=collate_fn,
+ )
+
+ assert len(self.train_dataloader) >= 1, "Train dataloader is empty!"
+ assert len(self.val_dataloader) >= 1, "Validation dataloader is empty!"
+
+ print(
+ f"Size of train dataloader: {len(self.train_dataloader)}, Size of val dataloader: "
+ f"{len(self.val_dataloader)}"
+ )
+
+ total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs
+
+ if self.config.trainer.total_training_steps is not None:
+ total_training_steps = self.config.trainer.total_training_steps
+
+ self.total_training_steps = total_training_steps
+ print(f"Total training steps: {self.total_training_steps}")
+
+ try:
+ OmegaConf.set_struct(self.config, True)
+ with open_dict(self.config):
+ if OmegaConf.select(self.config, "actor_rollout_ref.actor.optim"):
+ self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps
+ if OmegaConf.select(self.config, "critic.optim"):
+ self.config.critic.optim.total_training_steps = total_training_steps
+ except Exception as e:
+ print(f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}")
+
+ def _dump_generations(self, inputs, outputs, gts, scores, reward_extra_infos_dict, dump_path):
+ """Dump rollout/validation samples as JSONL."""
+ os.makedirs(dump_path, exist_ok=True)
+ filename = os.path.join(dump_path, f"{self.global_steps}.jsonl")
+
+ n = len(inputs)
+ base_data = {
+ "input": inputs,
+ "output": outputs,
+ "gts": gts,
+ "score": scores,
+ "step": [self.global_steps] * n,
+ }
+
+ for k, v in reward_extra_infos_dict.items():
+ if len(v) == n:
+ base_data[k] = v
+
+ lines = []
+ for i in range(n):
+ entry = {k: v[i] for k, v in base_data.items()}
+ lines.append(json.dumps(entry, ensure_ascii=False))
+
+ with open(filename, "w") as f:
+ f.write("\n".join(lines) + "\n")
+
+ print(f"Dumped generations to {filename}")
+
+ def _log_rollout_data(
+ self, log_rollout_meta: BatchMeta, reward_extra_infos_dict: dict, timing_raw: dict, rollout_data_dir: str
+ ):
+ """
+ Log rollout data to disk.
+
+ Args:
+ log_rollout_meta (BatchMeta): The batch_meta of rollout data
+ reward_extra_infos_dict (dict): Additional reward information to log
+ timing_raw (dict): Timing information for profiling
+ rollout_data_dir (str): Directory path to save the rollout data
+ """
+ with marked_timer("dump_rollout_generations", timing_raw, color="green"):
+ data = self.tq_client.get_data(log_rollout_meta)
+
+ inputs = self.tokenizer.batch_decode(data["prompts"], skip_special_tokens=True)
+ outputs = self.tokenizer.batch_decode(data["responses"], skip_special_tokens=True)
+ scores = data["token_level_scores"].sum(-1).cpu().tolist()
+ sample_gts = [item.get("ground_truth", None) for item in data.get("reward_model", {})]
+
+ reward_extra_infos_to_dump = reward_extra_infos_dict.copy()
+ if "request_id" in log_rollout_meta.field_names:
+ reward_extra_infos_dict.setdefault(
+ "request_id",
+ data["request_id"].tolist(),
+ )
+
+ self._dump_generations(
+ inputs=inputs,
+ outputs=outputs,
+ gts=sample_gts,
+ scores=scores,
+ reward_extra_infos_dict=reward_extra_infos_to_dump,
+ dump_path=rollout_data_dir,
+ )
+
+ def _maybe_log_val_generations(self, inputs, outputs, scores):
+ """Log a table of validation samples to the configured logger (wandb or swanlab)"""
+
+ generations_to_log = self.config.trainer.log_val_generations
+
+ if generations_to_log == 0:
+ return
+
+ import numpy as np
+
+ # Create tuples of (input, output, score) and sort by input text
+ samples = list(zip(inputs, outputs, scores, strict=True))
+ samples.sort(key=lambda x: x[0]) # Sort by input text
+
+ # Use fixed random seed for deterministic shuffling
+ rng = np.random.RandomState(42)
+ rng.shuffle(samples)
+
+ # Take first N samples after shuffling
+ samples = samples[:generations_to_log]
+
+ # Log to each configured logger
+ self.validation_generations_logger.log(self.config.trainer.logger, samples, self.global_steps)
+
+ def _get_gen_batch(self, batch: DataProto) -> DataProto:
+ reward_model_keys = set({"data_source", "reward_model", "extra_info", "uid"}) & batch.non_tensor_batch.keys()
+
+ # pop those keys for generation
+ batch_keys_to_pop = []
+ non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys()) - reward_model_keys
+ gen_batch = batch.pop(
+ batch_keys=batch_keys_to_pop,
+ non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop),
+ )
+
+ # For agent loop, we need reward model keys to compute score.
+ if self.async_rollout_mode:
+ gen_batch.non_tensor_batch.update(batch.non_tensor_batch)
+
+ return gen_batch
+
+ def _validate(self):
+ data_source_lst = []
+ reward_extra_infos_dict: dict[str, list] = defaultdict(list)
+
+ # Lists to collect samples for the table
+ sample_inputs = []
+ sample_outputs = []
+ sample_gts = []
+ sample_scores = []
+ sample_turns = []
+ sample_uids = []
+
+ for test_data in self.val_dataloader:
+ if "uid" not in test_data.keys():
+ test_data["uid"] = np.array(
+ [str(uuid.uuid4()) for _ in range(len(test_data["raw_prompt"]))], dtype=object
+ )
+
+ # repeat test data
+ repeated_test_data = self.repeat_dict(
+ test_data, repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True
+ )
+
+ test_batch: TensorDict = self.dict_to_tensordict(repeated_test_data)
+
+ # we only do validation on rule-based rm
+ if self.config.reward_model.enable and test_batch[0]["reward_model"]["style"] == "model":
+ return {}
+
+ batch_meta = self.tq_client.put(data=test_batch, partition_id=f"val_{self.global_steps - 1}")
+
+ batch_meta.update_extra_info(
+ {
+ "eos_token_id": self.tokenizer.eos_token_id,
+ "pad_token_id": self.tokenizer.pad_token_id,
+ "recompute_log_prob": False,
+ "do_sample": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample,
+ "validate": True,
+ "global_steps": self.global_steps,
+ }
+ )
+ print(f"batch_meta extra_info: {batch_meta.extra_info}")
+
+ # TODO: (TQ) Support padding and unpadding to make DataProto divisible by dp_size with TransferQueue
+ if not self.async_rollout_mode:
+ test_output_gen_meta = self.actor_rollout_wg.generate_sequences(batch_meta)
+ else:
+ test_output_gen_meta = self.async_rollout_manager.generate_sequences(batch_meta)
+
+ batch_meta = batch_meta.union(test_output_gen_meta)
+
+ print("validation generation end")
+
+ # Store generated outputs
+ test_response_meta = batch_meta.select_fields(["prompts", "responses", "uid", "reward_model"])
+ data = self.tq_client.get_data(test_response_meta)
+ output_ids = data["responses"]
+ output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids]
+ sample_outputs.extend(output_texts)
+
+ # TODO: Can we keep special tokens except for padding tokens?
+ input_ids = data["prompts"]
+ input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]
+ sample_inputs.extend(input_texts)
+ sample_uids.extend(data["uid"])
+
+ ground_truths = [item.get("ground_truth", None) for item in data.get("reward_model", {})]
+ sample_gts.extend(ground_truths)
+
+ # evaluate using reward_function
+ if self.val_reward_fn is None:
+ raise ValueError("val_reward_fn must be provided for validation.")
+
+ compute_reward_fields = [
+ "responses",
+ "prompts",
+ "attention_mask",
+ "reward_model",
+ "data_source",
+ ]
+ if "rm_scores" in batch_meta.field_names:
+ compute_reward_fields = ["rm_scores", *set(batch_meta.extra_info["reward_extra_keys"])]
+
+ val_reward_meta = batch_meta.select_fields(compute_reward_fields)
+ result = compute_val_reward_decorated(self.val_reward_fn, val_reward_meta, return_dict=True)
+ reward_tensor = result["reward_tensor"]
+ scores = reward_tensor.sum(-1).cpu().tolist()
+ sample_scores.extend(scores)
+
+ reward_extra_infos_dict["reward"].extend(scores)
+ print(f"len reward_extra_infos_dict['reward']: {len(reward_extra_infos_dict['reward'])}")
+ if "reward_extra_info" in result:
+ for key, lst in result["reward_extra_info"].items():
+ reward_extra_infos_dict[key].extend(lst)
+ print(f"len reward_extra_infos_dict['{key}']: {len(reward_extra_infos_dict[key])}")
+
+ # collect num_turns of each prompt
+ if "__num_turns__" in batch_meta.field_names:
+ data = self.tq_client.get_data(batch_meta.select_fields(["__num_turns__"]))
+ sample_turns.append(data["__num_turns__"])
+
+ data_source = ["unknown"] * reward_tensor.shape[0]
+ if "data_source" in batch_meta.field_names:
+ data_source_meta = batch_meta.select_fields(["data_source"])
+ data = self.tq_client.get_data(data_source_meta)
+ data_source = data["data_source"]
+
+ data_source_lst.append(data_source)
+
+ self.tq_client.clear_samples(batch_meta)
+
+ self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores)
+
+ # dump generations
+ val_data_dir = self.config.trainer.get("validation_data_dir", None)
+ if val_data_dir:
+ self._dump_generations(
+ inputs=sample_inputs,
+ outputs=sample_outputs,
+ gts=sample_gts,
+ scores=sample_scores,
+ reward_extra_infos_dict=reward_extra_infos_dict,
+ dump_path=val_data_dir,
+ )
+
+ for key_info, lst in reward_extra_infos_dict.items():
+ assert len(lst) == 0 or len(lst) == len(sample_scores), f"{key_info}: {len(lst)=}, {len(sample_scores)=}"
+
+ data_sources = np.concatenate(data_source_lst, axis=0)
+
+ data_src2var2metric2val = process_validation_metrics(data_sources, sample_uids, reward_extra_infos_dict)
+ metric_dict = {}
+ for data_source, var2metric2val in data_src2var2metric2val.items():
+ core_var = "acc" if "acc" in var2metric2val else "reward"
+ for var_name, metric2val in var2metric2val.items():
+ n_max = max([int(name.split("@")[-1].split("/")[0]) for name in metric2val.keys()])
+ for metric_name, metric_val in metric2val.items():
+ if (
+ (var_name == core_var)
+ and any(metric_name.startswith(pfx) for pfx in ["mean", "maj", "best"])
+ and (f"@{n_max}" in metric_name)
+ ):
+ metric_sec = "val-core"
+ else:
+ metric_sec = "val-aux"
+ pfx = f"{metric_sec}/{data_source}/{var_name}/{metric_name}"
+ metric_dict[pfx] = metric_val
+
+ if len(sample_turns) > 0:
+ sample_turns = np.concatenate(sample_turns)
+ metric_dict["val-aux/num_turns/min"] = sample_turns.min()
+ metric_dict["val-aux/num_turns/max"] = sample_turns.max()
+ metric_dict["val-aux/num_turns/mean"] = sample_turns.mean()
+
+ return metric_dict
+
+ def init_workers(self):
+ """Initialize distributed training workers using Ray backend.
+
+ Creates:
+ 1. Ray resource pools from configuration
+ 2. Worker groups for each role (actor, critic, etc.)
+ """
+ self.resource_pool_manager.create_resource_pool()
+
+ self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
+
+ # create actor and rollout
+ if self.hybrid_engine:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.ActorRollout)
+ actor_rollout_cls = RayClassWithInitArgs(
+ cls=self.role_worker_mapping[Role.ActorRollout],
+ config=self.config.actor_rollout_ref,
+ role="actor_rollout",
+ )
+ self.resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls
+ else:
+ raise NotImplementedError
+
+ # create critic
+ if self.use_critic:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)
+ critic_cfg = omega_conf_to_dataclass(self.config.critic)
+ critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg)
+ self.resource_pool_to_cls[resource_pool]["critic"] = critic_cls
+
+ # create reference policy if needed
+ if self.use_reference_policy:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)
+ ref_policy_cls = RayClassWithInitArgs(
+ self.role_worker_mapping[Role.RefPolicy],
+ config=self.config.actor_rollout_ref,
+ role="ref",
+ )
+ self.resource_pool_to_cls[resource_pool]["ref"] = ref_policy_cls
+
+ # create a reward model if reward_fn is None
+ 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]["rm"] = rm_cls
+
+ # 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`.
+ # Instead, directly pass different resource pool to different worker groups.
+ # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.
+ all_wg = {}
+ wg_kwargs = {} # Setting up kwargs for RayWorkerGroup
+ if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None:
+ wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout
+ if OmegaConf.select(self.config.global_profiler, "steps") is not None:
+ wg_kwargs["profile_steps"] = OmegaConf.select(self.config.global_profiler, "steps")
+ # Only require nsight worker options when tool is nsys
+ if OmegaConf.select(self.config.global_profiler, "tool") == "nsys":
+ assert (
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ is not None
+ ), "worker_nsight_options must be set when using nsys with profile_steps"
+ wg_kwargs["worker_nsight_options"] = OmegaConf.to_container(
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ )
+ wg_kwargs["device_name"] = self.device_name
+
+ for resource_pool, class_dict in self.resource_pool_to_cls.items():
+ worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
+ wg_dict = self.ray_worker_group_cls(
+ resource_pool=resource_pool,
+ ray_cls_with_init=worker_dict_cls,
+ **wg_kwargs,
+ )
+ spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
+ all_wg.update(spawn_wg)
+
+ if self.use_critic:
+ self.critic_wg = all_wg["critic"]
+ self.critic_wg.init_model()
+
+ if self.use_reference_policy and not self.ref_in_actor:
+ self.ref_policy_wg = all_wg["ref"]
+ self.ref_policy_wg.init_model()
+
+ self.rm_wg = None
+ if self.use_rm:
+ self.rm_wg = all_wg["rm"]
+ 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["actor_rollout"]
+ self.actor_rollout_wg.init_model()
+
+ # set transferqueue server info for each worker
+ for _, wg in all_wg.items():
+ wg.create_transferqueue_client(self.config)
+
+ # create async rollout manager and request scheduler
+ self.async_rollout_mode = False
+ if self.config.actor_rollout_ref.rollout.mode == "async":
+ from .agent_loop import AgentLoopManager
+
+ self.async_rollout_mode = True
+ if self.config.reward_model.enable and self.config.reward_model.enable_resource_pool:
+ rm_resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
+ else:
+ rm_resource_pool = None
+
+ self.async_rollout_manager = AgentLoopManager(
+ config=self.config,
+ worker_group=self.actor_rollout_wg,
+ rm_resource_pool=rm_resource_pool,
+ )
+
+ self.checkpoint_manager = CheckpointEngineManager(
+ backend=self.config.actor_rollout_ref.rollout.checkpoint_engine.backend,
+ trainer=self.actor_rollout_wg,
+ replicas=self.async_rollout_manager.rollout_replicas,
+ )
+
+ # sleep all replicas to load checkpoint
+ self.checkpoint_manager.sleep_replicas()
+
+ # TODO (TQ): initialize tq during worker init when enable TQ switch is stable
+ self.async_rollout_manager.create_transferqueue_client_for_workers()
+
+ def _save_checkpoint(self):
+ from verl.utils.fs import local_mkdir_safe
+
+ # path: given_path + `/global_step_{global_steps}` + `/actor`
+ local_global_step_folder = os.path.join(
+ self.config.trainer.default_local_dir, f"global_step_{self.global_steps}"
+ )
+
+ print(f"local_global_step_folder: {local_global_step_folder}")
+ actor_local_path = os.path.join(local_global_step_folder, "actor")
+
+ actor_remote_path = (
+ None
+ if self.config.trainer.default_hdfs_dir is None
+ else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "actor")
+ )
+
+ remove_previous_ckpt_in_save = self.config.trainer.get("remove_previous_ckpt_in_save", False)
+ if remove_previous_ckpt_in_save:
+ print(
+ "Warning: remove_previous_ckpt_in_save is deprecated,"
+ + " set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead"
+ )
+ max_actor_ckpt_to_keep = (
+ self.config.trainer.get("max_actor_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1
+ )
+ max_critic_ckpt_to_keep = (
+ self.config.trainer.get("max_critic_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1
+ )
+
+ self.actor_rollout_wg.save_checkpoint(
+ actor_local_path, actor_remote_path, self.global_steps, max_ckpt_to_keep=max_actor_ckpt_to_keep
+ )
+
+ if self.use_critic:
+ critic_local_path = os.path.join(local_global_step_folder, "critic")
+ critic_remote_path = (
+ None
+ if self.config.trainer.default_hdfs_dir is None
+ else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "critic")
+ )
+ self.critic_wg.save_checkpoint(
+ critic_local_path, critic_remote_path, self.global_steps, max_ckpt_to_keep=max_critic_ckpt_to_keep
+ )
+
+ # save dataloader
+ local_mkdir_safe(local_global_step_folder)
+ dataloader_local_path = os.path.join(local_global_step_folder, "data.pt")
+ dataloader_state_dict = self.train_dataloader.state_dict()
+ torch.save(dataloader_state_dict, dataloader_local_path)
+
+ # latest checkpointed iteration tracker (for atomic usage)
+ local_latest_checkpointed_iteration = os.path.join(
+ self.config.trainer.default_local_dir, "latest_checkpointed_iteration.txt"
+ )
+ with open(local_latest_checkpointed_iteration, "w") as f:
+ f.write(str(self.global_steps))
+
+ def _load_checkpoint(self):
+ if self.config.trainer.resume_mode == "disable":
+ return 0
+
+ # load from hdfs
+ if self.config.trainer.default_hdfs_dir is not None:
+ raise NotImplementedError("load from hdfs is not implemented yet")
+ else:
+ checkpoint_folder = self.config.trainer.default_local_dir # TODO: check path
+ if not os.path.isabs(checkpoint_folder):
+ working_dir = os.getcwd()
+ checkpoint_folder = os.path.join(working_dir, checkpoint_folder)
+ global_step_folder = find_latest_ckpt_path(checkpoint_folder) # None if no latest
+
+ # find global_step_folder
+ if self.config.trainer.resume_mode == "auto":
+ if global_step_folder is None:
+ print("Training from scratch")
+ return 0
+ else:
+ if self.config.trainer.resume_mode == "resume_path":
+ assert isinstance(self.config.trainer.resume_from_path, str), "resume ckpt must be str type"
+ assert "global_step_" in self.config.trainer.resume_from_path, (
+ "resume ckpt must specify the global_steps"
+ )
+ global_step_folder = self.config.trainer.resume_from_path
+ if not os.path.isabs(global_step_folder):
+ working_dir = os.getcwd()
+ global_step_folder = os.path.join(working_dir, global_step_folder)
+ print(f"Load from checkpoint folder: {global_step_folder}")
+ # set global step
+ self.global_steps = int(global_step_folder.split("global_step_")[-1])
+
+ print(f"Setting global step to {self.global_steps}")
+ print(f"Resuming from {global_step_folder}")
+
+ actor_path = os.path.join(global_step_folder, "actor")
+ critic_path = os.path.join(global_step_folder, "critic")
+ # load actor
+ self.actor_rollout_wg.load_checkpoint(
+ actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load
+ )
+ # load critic
+ if self.use_critic:
+ self.critic_wg.load_checkpoint(
+ critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load
+ )
+
+ # load dataloader,
+ # TODO: from remote not implemented yet
+ dataloader_local_path = os.path.join(global_step_folder, "data.pt")
+ if os.path.exists(dataloader_local_path):
+ dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False)
+ self.train_dataloader.load_state_dict(dataloader_state_dict)
+ else:
+ print(f"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch")
+
+ def _start_profiling(self, do_profile: bool) -> None:
+ """Start profiling for all worker groups if profiling is enabled."""
+ if do_profile:
+ self.actor_rollout_wg.start_profile(role="e2e", profile_step=self.global_steps)
+ if self.use_reference_policy:
+ self.ref_policy_wg.start_profile(profile_step=self.global_steps)
+ if self.use_critic:
+ self.critic_wg.start_profile(profile_step=self.global_steps)
+ if self.use_rm:
+ self.rm_wg.start_profile(profile_step=self.global_steps)
+
+ def _stop_profiling(self, do_profile: bool) -> None:
+ """Stop profiling for all worker groups if profiling is enabled."""
+ if do_profile:
+ self.actor_rollout_wg.stop_profile()
+ if self.use_reference_policy:
+ self.ref_policy_wg.stop_profile()
+ if self.use_critic:
+ self.critic_wg.stop_profile()
+ if self.use_rm:
+ self.rm_wg.stop_profile()
+
+ def _balance_batch(
+ self, batch: BatchMeta, tq_client, metrics, logging_prefix="global_seqlen", keep_minibatch=False
+ ):
+ """Reorder the batchmeta on single controller such that each dp rank gets similar total tokens"""
+ data = tq_client.get_data(batch)
+
+ attention_mask = data["attention_mask"]
+ batch_size = attention_mask.shape[0]
+ global_seqlen_lst = data["attention_mask"].view(batch_size, -1).sum(-1) # (train_batch_size,)
+ global_seqlen_lst = calculate_workload(global_seqlen_lst)
+ world_size = self.actor_rollout_wg.world_size
+ if keep_minibatch:
+ # Decouple the DP balancing and mini-batching.
+ minibatch_size = self.config.actor_rollout_ref.actor.get("ppo_mini_batch_size", None)
+ if minibatch_size is None:
+ raise ValueError("'ppo_mini_batch_size' must be set in actor config when 'keep_minibatch' is True.")
+ minibatch_num = len(global_seqlen_lst) // minibatch_size
+ global_partition_lst = [[] for _ in range(world_size)]
+ for i in range(minibatch_num):
+ rearrange_minibatch_lst = get_seqlen_balanced_partitions(
+ global_seqlen_lst[i * minibatch_size : (i + 1) * minibatch_size],
+ k_partitions=world_size,
+ equal_size=True,
+ )
+ for j, part in enumerate(rearrange_minibatch_lst):
+ global_partition_lst[j].extend([x + minibatch_size * i for x in part])
+ else:
+ global_partition_lst = get_seqlen_balanced_partitions(
+ global_seqlen_lst, k_partitions=world_size, equal_size=True
+ )
+ # Place smaller micro-batches at both ends to reduce the bubbles in pipeline parallel.
+ for idx, partition in enumerate(global_partition_lst):
+ partition.sort(key=lambda x: (global_seqlen_lst[x], x))
+ ordered_partition = partition[::2] + partition[1::2][::-1]
+ global_partition_lst[idx] = ordered_partition
+ # reorder based on index. The data will be automatically equally partitioned by dispatch function
+ global_idx = [j for partition in global_partition_lst for j in partition]
+ global_balance_stats = log_seqlen_unbalance(
+ seqlen_list=global_seqlen_lst, partitions=global_partition_lst, prefix=logging_prefix
+ )
+ metrics.update(global_balance_stats)
+ return global_idx
+
+ @classmethod
+ def repeat_dict(
+ cls, batch_dict: dict[str, torch.Tensor | np.ndarray], repeat_times=2, interleave=True
+ ) -> dict[str, torch.Tensor | np.ndarray]:
+ """
+ Repeat the batch dict a specified number of times.
+
+ Args:
+ repeat_times (int): Number of times to repeat the data.
+ interleave (bool): Whether to interleave the repeated data.
+
+ Returns:
+ dict: A new dict with repeated data.
+ """
+ if repeat_times == 1:
+ return batch_dict
+
+ repeated_batch_dict = {}
+ if batch_dict:
+ if interleave:
+ # Interleave the data
+ for key, val in batch_dict.items():
+ if isinstance(val, torch.Tensor):
+ repeated_batch_dict[key] = val.repeat_interleave(repeat_times, dim=0)
+ elif isinstance(val, np.ndarray):
+ repeated_batch_dict[key] = np.repeat(val, repeat_times, axis=0)
+ else:
+ raise ValueError(f"Unsupported type in data {type(val)}")
+ else:
+ # Stack the data
+ for key, val in batch_dict.items():
+ if isinstance(val, torch.Tensor):
+ repeated_batch_dict[key] = (
+ val.unsqueeze(0).expand(repeat_times, *val.shape).reshape(-1, *val.shape[1:])
+ )
+ elif isinstance(val, np.ndarray):
+ repeated_batch_dict[key] = np.tile(val, (repeat_times,) + (1,) * (val.ndim - 1))
+ else:
+ raise ValueError(f"Unsupported type in data {type(val)}")
+ return repeated_batch_dict
+
+ @classmethod
+ def dict_to_tensordict(cls, data: dict[str, torch.Tensor | np.ndarray]) -> TensorDict:
+ """
+ Create a TensorDict from a dict of tensors and non_tensors.
+ Note that this requires tensordict version at least 0.10
+ """
+ assert parse_version(tensordict.__version__) >= parse_version("0.10"), (
+ "Storing non-tensor data in TensorDict at least requires tensordict version 0.10"
+ )
+ tensors_batch = {}
+ batch_size = None
+
+ for key, val in data.items():
+ if isinstance(val, torch.Tensor | np.ndarray):
+ tensors_batch[key] = val
+ else:
+ raise ValueError(f"Unsupported type in data {type(val)}")
+
+ if batch_size is None:
+ batch_size = len(val)
+ else:
+ assert len(val) == batch_size
+
+ if batch_size is None:
+ batch_size = []
+ else:
+ batch_size = [batch_size]
+
+ return TensorDict(tensors_batch, batch_size=batch_size)
+
+ def fit(self):
+ """
+ The training loop of PPO.
+ The driver process only need to call the compute functions of the worker group through RPC
+ to construct the PPO dataflow.
+ The light-weight advantage computation is done on the driver process.
+ """
+ from omegaconf import OmegaConf
+
+ from verl.utils.tracking import Tracking
+
+ logger = Tracking(
+ project_name=self.config.trainer.project_name,
+ experiment_name=self.config.trainer.experiment_name,
+ default_backend=self.config.trainer.logger,
+ config=OmegaConf.to_container(self.config, resolve=True),
+ )
+
+ self.global_steps = 0
+
+ # load checkpoint and update weights before doing anything
+ self._load_checkpoint()
+ self.checkpoint_manager.update_weights()
+
+ # perform validation before training
+ # currently, we only support validation using the reward_function.
+ if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
+ val_metrics = self._validate()
+ assert val_metrics, f"{val_metrics=}"
+ pprint(f"Initial validation metrics: {val_metrics}")
+ logger.log(data=val_metrics, step=self.global_steps)
+ if self.config.trainer.get("val_only", False):
+ return
+
+ if self.config.actor_rollout_ref.rollout.get("skip_rollout", False):
+ rollout_skip = RolloutSkip(self.config, self.actor_rollout_wg)
+ rollout_skip.wrap_generate_sequences()
+
+ # add tqdm
+ progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress")
+
+ # we start from step 1
+ self.global_steps += 1
+ last_val_metrics = None
+ self.max_steps_duration = 0
+
+ prev_step_profile = False
+ curr_step_profile = (
+ self.global_steps in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ next_step_profile = False
+
+ for epoch in range(self.config.trainer.total_epochs):
+ for batch_dict in self.train_dataloader:
+ metrics = {}
+ timing_raw = {}
+ base_get_meta_kwargs = dict(
+ batch_size=self.config.data.train_batch_size * self.config.actor_rollout_ref.rollout.n,
+ partition_id=f"train_{self.global_steps - 1}", # self.global_steps starts from 1
+ )
+
+ with marked_timer("start_profile", timing_raw):
+ self._start_profiling(
+ not prev_step_profile and curr_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+
+ # add uid to batch
+ batch_dict["uid"] = np.array(
+ [str(uuid.uuid4()) for _ in range(len(batch_dict["raw_prompt"]))], dtype=object
+ )
+ # When n > 1, repeat input data before putting to data system, simulating DataProto repeat.
+ repeated_batch_dict = self.repeat_dict(
+ batch_dict, repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True
+ )
+ batch: TensorDict = self.dict_to_tensordict(repeated_batch_dict)
+ gen_meta = self.tq_client.put(data=batch, partition_id=f"train_{self.global_steps - 1}")
+
+ # pass global_steps to trace
+ gen_meta.set_extra_info("global_steps", self.global_steps)
+
+ is_last_step = self.global_steps >= self.total_training_steps
+
+ with marked_timer("step", timing_raw):
+ # generate a batch
+ with marked_timer("gen", timing_raw, color="red"):
+ if not self.async_rollout_mode:
+ gen_output_meta = self.actor_rollout_wg.generate_sequences(gen_meta)
+ else:
+ gen_output_meta = self.async_rollout_manager.generate_sequences(gen_meta)
+ self.checkpoint_manager.sleep_replicas()
+ timing_raw.update(gen_output_meta.extra_info["timing"])
+ gen_output_meta.extra_info.pop("timing", None)
+
+ # TODO (TQ): support transfer queue
+ # if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
+ # if self.reward_fn is None:
+ # raise ValueError("A reward_fn is required for REMAX advantage estimation.")
+ #
+ # with marked_timer("gen_max", timing_raw, color="purple"):
+ # gen_baseline_meta = deepcopy(gen_meta)
+ # gen_baseline_meta.extra_info["do_sample"] = False
+ # if not self.async_rollout_mode:
+ # gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_meta)
+ # else:
+ # gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_meta)
+ # batch = batch.union(gen_baseline_output)
+ # reward_baseline_tensor = self.reward_fn(batch)
+ # reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
+ #
+ # batch.pop(batch_keys=list(gen_baseline_output.batch.keys()))
+ #
+ # batch.batch["reward_baselines"] = reward_baseline_tensor
+ #
+ # del gen_baseline_batch, gen_baseline_output
+
+ batch_meta: BatchMeta = gen_meta.union(gen_output_meta)
+
+ if "response_mask" not in batch_meta.field_names:
+ response_mask_meta = self.tq_client.get_meta(
+ data_fields=["responses", "attention_mask"],
+ task_name="compute_response_mask",
+ **base_get_meta_kwargs,
+ )
+ response_mask_output_meta = compute_response_mask(response_mask_meta, self.tq_client)
+ batch_meta = batch_meta.union(response_mask_output_meta)
+
+ # Balance the number of valid tokens across DP ranks.
+ # NOTE: This usually changes the order of data in the `batch`,
+ # which won't affect the advantage calculation (since it's based on uid),
+ # but might affect the loss calculation (due to the change of mini-batching).
+ # TODO: Decouple the DP balancing and mini-batching.
+
+ attention_mask_meta = batch_meta.select_fields(["attention_mask"])
+ balanced_idx = None
+ if self.config.trainer.balance_batch:
+ balanced_idx = self._balance_batch(attention_mask_meta, self.tq_client, metrics=metrics)
+ batch_meta.reorder(balanced_idx)
+
+ # compute global_valid tokens
+ data = self.tq_client.get_data(attention_mask_meta)
+ batch_meta.extra_info["global_token_num"] = torch.sum(data["attention_mask"], dim=-1).tolist()
+
+ with marked_timer("reward", timing_raw, color="yellow"):
+ # compute reward model score
+ if self.use_rm and "rm_scores" not in batch_meta.field_names:
+ reward_meta = self.rm_wg.compute_rm_score(batch_meta)
+ batch_meta = batch_meta.union(reward_meta)
+
+ compute_reward_fields = [
+ "responses",
+ "prompts",
+ "attention_mask",
+ "reward_model",
+ "data_source",
+ ]
+ if "rm_scores" in batch_meta.field_names:
+ compute_reward_fields.extend(
+ ["rm_scores", *set(batch_meta.extra_info["reward_extra_keys"])]
+ )
+
+ compute_reward_meta = batch_meta.select_fields(compute_reward_fields)
+
+ if self.config.reward_model.launch_reward_fn_async:
+ future_reward = compute_reward_async_decorated(
+ data=compute_reward_meta,
+ reward_fn=self.reward_fn,
+ )
+ else:
+ reward_tensor, reward_extra_infos_dict = compute_reward_decorated(
+ compute_reward_meta, self.reward_fn
+ )
+ batch_meta = batch_meta.union(compute_reward_meta)
+
+ # recompute old_log_probs
+ with marked_timer("old_log_prob", timing_raw, color="blue"):
+ old_log_prob_meta_fields = [
+ "input_ids",
+ "attention_mask",
+ "position_ids",
+ "prompts",
+ "responses",
+ "response_mask",
+ "data_source",
+ "reward_model",
+ "extra_info",
+ "uid",
+ "index",
+ "tools_kwargs",
+ "interaction_kwargs",
+ "ability",
+ ]
+ old_log_prob_meta = batch_meta.select_fields(old_log_prob_meta_fields)
+ old_log_prob_output_meta = self.actor_rollout_wg.compute_log_prob(old_log_prob_meta)
+ batch_meta = batch_meta.union(old_log_prob_output_meta)
+
+ data = self.tq_client.get_data(old_log_prob_output_meta)
+ entropys = data["entropys"]
+ response_masks = data["response_mask"]
+ actor_config = self.config.actor_rollout_ref.actor
+ entropy_agg = agg_loss(
+ loss_mat=entropys,
+ loss_mask=response_masks,
+ loss_agg_mode=actor_config.loss_agg_mode,
+ loss_scale_factor=actor_config.loss_scale_factor,
+ )
+ old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()}
+ metrics.update(old_log_prob_metrics)
+
+ if "rollout_log_probs" in batch_meta.field_names:
+ # TODO: we may want to add diff of probs too.
+ calculate_debug_metrics_fields = ["rollout_log_probs", "old_log_probs", "responses"]
+
+ if "response_mask" in batch_meta.field_names:
+ calculate_debug_metrics_fields.append("response_mask")
+ if "attention_mask" in batch_meta.field_names:
+ calculate_debug_metrics_fields.append("attention_mask")
+
+ calculate_debug_metrics_meta = batch_meta.select_fields(calculate_debug_metrics_fields)
+ metrics.update(calculate_debug_metrics_decorated(calculate_debug_metrics_meta))
+
+ if self.use_reference_policy:
+ # compute reference log_prob
+ ref_log_prob_fields = [
+ "input_ids",
+ "attention_mask",
+ "position_ids",
+ "prompts",
+ "responses",
+ "response_mask",
+ "old_log_probs",
+ "data_source",
+ "reward_model",
+ "extra_info",
+ "uid",
+ "index",
+ "tools_kwargs",
+ "interaction_kwargs",
+ "ability",
+ ]
+ ref_log_prob_meta = batch_meta.select_fields(ref_log_prob_fields)
+
+ with marked_timer("ref", timing_raw, color="olive"):
+ if not self.ref_in_actor:
+ ref_log_prob_output_meta = self.ref_policy_wg.compute_ref_log_prob(ref_log_prob_meta)
+ else:
+ ref_log_prob_output_meta = self.actor_rollout_wg.compute_ref_log_prob(ref_log_prob_meta)
+ batch_meta = batch_meta.union(ref_log_prob_output_meta)
+
+ # compute values
+ if self.use_critic:
+ with marked_timer("values", timing_raw, color="cyan"):
+ values_meta = self.critic_wg.compute_values(batch_meta)
+ batch_meta = batch_meta.union(values_meta)
+
+ with marked_timer("adv", timing_raw, color="brown"):
+ # we combine with rule-based rm
+ reward_extra_infos_dict: dict[str, list]
+ if self.config.reward_model.launch_reward_fn_async:
+ reward_tensor, reward_extra_infos_dict = ray.get(future_reward)
+ reward_td = TensorDict({"token_level_scores": reward_tensor}, batch_size=reward_tensor.size(0))
+ batch_meta = self.tq_client.put(data=reward_td, metadata=batch_meta)
+
+ if reward_extra_infos_dict:
+ reward_extra_infos_dict_new = {k: np.array(v) for k, v in reward_extra_infos_dict.items()}
+ reward_extra_infos_td = self.dict_to_tensordict(reward_extra_infos_dict_new)
+ batch_meta = self.tq_client.put(data=reward_extra_infos_td, metadata=batch_meta)
+
+ # compute rewards. apply_kl_penalty if available
+ if self.config.algorithm.use_kl_in_reward:
+ apply_kl_penalty_fields = [
+ "response_mask",
+ "token_level_scores",
+ "old_log_probs",
+ "ref_log_prob",
+ ]
+
+ apply_kl_penalty_meta = batch_meta.select_fields(apply_kl_penalty_fields)
+
+ token_level_rewards, kl_metrics = apply_kl_penalty(
+ apply_kl_penalty_meta,
+ kl_ctrl=self.kl_ctrl_in_reward,
+ kl_penalty=self.config.algorithm.kl_penalty,
+ )
+ token_level_rewards_td = TensorDict(
+ {"token_level_rewards": token_level_rewards}, batch_size=token_level_rewards.size(0)
+ )
+ apply_kl_penalty_meta = self.tq_client.put(
+ data=token_level_rewards_td, metadata=apply_kl_penalty_meta
+ )
+
+ metrics.update(kl_metrics)
+ batch_meta = batch_meta.union(apply_kl_penalty_meta)
+ else:
+ token_level_scores_meta = batch_meta.select_fields(["token_level_scores"])
+
+ data = self.tq_client.get_data(token_level_scores_meta)
+ token_level_rewards_td = TensorDict(
+ {"token_level_rewards": data["token_level_scores"]},
+ batch_size=data["token_level_scores"].size(0),
+ )
+ token_level_scores_meta = self.tq_client.put(
+ data=token_level_rewards_td, metadata=token_level_scores_meta
+ )
+ batch_meta = batch_meta.union(token_level_scores_meta)
+
+ # compute advantages, executed on the driver process
+
+ norm_adv_by_std_in_grpo = self.config.algorithm.get(
+ "norm_adv_by_std_in_grpo", True
+ ) # GRPO adv normalization factor
+
+ assert "response_mask" in batch_meta.field_names, (
+ f"`response_mask` must be in batch_meta {batch_meta.field_names} for advantage computation"
+ )
+ compute_advantage_fields = [
+ "response_mask",
+ "token_level_rewards",
+ ]
+ if self.config.algorithm.adv_estimator == AdvantageEstimator.GAE:
+ compute_advantage_fields.append("values")
+ elif self.config.algorithm.adv_estimator == AdvantageEstimator.GRPO:
+ compute_advantage_fields.append("uid")
+ else:
+ if "uid" in batch_meta.field_names:
+ compute_advantage_fields.append("uid")
+ if "reward_baselines" in batch_meta.field_names:
+ compute_advantage_fields.append("reward_baselines")
+
+ compute_advantage_meta = batch_meta.select_fields(compute_advantage_fields)
+
+ advantages, returns = compute_advantage(
+ compute_advantage_meta,
+ adv_estimator=self.config.algorithm.adv_estimator,
+ gamma=self.config.algorithm.gamma,
+ lam=self.config.algorithm.lam,
+ num_repeat=self.config.actor_rollout_ref.rollout.n,
+ norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+ config=self.config.algorithm,
+ )
+
+ advantages_td = TensorDict(
+ {"advantages": advantages, "returns": returns}, batch_size=advantages.size(0)
+ )
+ compute_advantage_meta = self.tq_client.put(data=advantages_td, metadata=compute_advantage_meta)
+ batch_meta = batch_meta.union(compute_advantage_meta)
+
+ # update critic
+ if self.use_critic:
+ with marked_timer("update_critic", timing_raw, color="pink"):
+ critic_output_meta = self.critic_wg.update_critic(batch_meta)
+ batch_meta = batch_meta.union(critic_output_meta)
+ critic_output_metrics = reduce_metrics(critic_output_meta.extra_info["metrics"])
+ metrics.update(critic_output_metrics)
+
+ # implement critic warmup
+ if self.config.trainer.critic_warmup <= self.global_steps:
+ # update actor
+ with marked_timer("update_actor", timing_raw, color="red"):
+ batch_meta.extra_info["multi_turn"] = (
+ self.config.actor_rollout_ref.rollout.multi_turn.enable
+ )
+
+ update_actor_fields = [
+ "input_ids",
+ "attention_mask",
+ "position_ids",
+ "prompts",
+ "responses",
+ "response_mask",
+ "old_log_probs",
+ "ref_log_prob",
+ "advantages",
+ "returns",
+ "token_level_rewards",
+ "token_level_scores",
+ "data_source",
+ "reward_model",
+ "extra_info",
+ "uid",
+ "index",
+ "tools_kwargs",
+ "interaction_kwargs",
+ "ability",
+ ]
+ update_actor_meta = batch_meta.select_fields(update_actor_fields)
+
+ update_actor_meta.set_extra_info(
+ "global_token_num", batch_meta.get_extra_info("global_token_num")
+ )
+ update_actor_meta.set_extra_info("temperature", batch_meta.get_extra_info("temperature"))
+
+ actor_output_meta = self.actor_rollout_wg.update_actor(update_actor_meta)
+ batch_meta = batch_meta.union(actor_output_meta)
+
+ # update weights from trainer to rollout
+ with marked_timer("update_weights", timing_raw, color="red"):
+ self.checkpoint_manager.update_weights()
+
+ actor_output_metrics = reduce_metrics(actor_output_meta.extra_info["metrics"])
+ metrics.update(actor_output_metrics)
+
+ # Log rollout generations if enabled
+ rollout_data_dir = self.config.trainer.get("rollout_data_dir", None)
+ if rollout_data_dir:
+ log_rollout_fields = ["prompts", "responses", "token_level_scores", "reward_model"]
+ if "request_id" in batch_meta.field_names:
+ log_rollout_fields.append("request_id")
+ log_rollout_meta = batch_meta.select_fields(log_rollout_fields)
+ self._log_rollout_data(log_rollout_meta, reward_extra_infos_dict, timing_raw, rollout_data_dir)
+
+ # TODO: validate
+ if (
+ self.val_reward_fn is not None
+ and self.config.trainer.test_freq > 0
+ and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
+ ):
+ with marked_timer("testing", timing_raw, color="green"):
+ val_metrics: dict = self._validate()
+ if is_last_step:
+ last_val_metrics = val_metrics
+ metrics.update(val_metrics)
+
+ # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.
+ esi_close_to_expiration = should_save_ckpt_esi(
+ max_steps_duration=self.max_steps_duration,
+ redundant_time=self.config.trainer.esi_redundant_time,
+ )
+ # Check if the conditions for saving a checkpoint are met.
+ # The conditions include a mandatory condition (1) and
+ # one of the following optional conditions (2/3/4):
+ # 1. The save frequency is set to a positive value.
+ # 2. It's the last training step.
+ # 3. The current step number is a multiple of the save frequency.
+ # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.
+ if self.config.trainer.save_freq > 0 and (
+ is_last_step or self.global_steps % self.config.trainer.save_freq == 0 or esi_close_to_expiration
+ ):
+ if esi_close_to_expiration:
+ print("Force saving checkpoint: ESI instance expiration approaching.")
+ with marked_timer("save_checkpoint", timing_raw, color="green"):
+ self._save_checkpoint()
+
+ with marked_timer("stop_profile", timing_raw):
+ next_step_profile = (
+ self.global_steps + 1 in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ self._stop_profiling(
+ curr_step_profile and not next_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+ prev_step_profile = curr_step_profile
+ curr_step_profile = next_step_profile
+
+ steps_duration = timing_raw["step"]
+ self.max_steps_duration = max(self.max_steps_duration, steps_duration)
+
+ # training metrics
+ metrics.update(
+ {
+ "training/global_step": self.global_steps,
+ "training/epoch": epoch,
+ }
+ )
+ # collect metrics
+ compute_data_metrics_fields = [
+ "token_level_rewards",
+ "token_level_scores",
+ "advantages",
+ "returns",
+ "responses",
+ "attention_mask",
+ "response_mask",
+ ]
+ if "__num_turns__" in batch_meta.field_names:
+ compute_data_metrics_fields.append("__num_turns__")
+ if "tool_call_counts" in batch_meta.field_names:
+ compute_data_metrics_fields.append("tool_call_counts")
+ compute_data_metrics_meta = batch_meta.select_fields(compute_data_metrics_fields)
+ compute_data_metrics_meta.reorder(balanced_idx)
+ metrics.update(
+ compute_data_metrics_decorated(batch=compute_data_metrics_meta, use_critic=self.use_critic)
+ )
+
+ compute_timing_metrics_fields = ["responses", "attention_mask"]
+ compute_timing_metrics_meta = batch_meta.select_fields(compute_timing_metrics_fields)
+ compute_timing_metrics_meta.reorder(balanced_idx)
+ metrics.update(
+ compute_timing_metrics_decorated(batch=compute_timing_metrics_meta, timing_raw=timing_raw)
+ )
+
+ compute_throughout_metrics_meta = BatchMeta(
+ samples=[],
+ extra_info={"global_token_num": batch_meta.get_extra_info("global_token_num")},
+ )
+ # TODO: implement actual tflpo and theoretical tflpo
+ n_gpus = self.resource_pool_manager.get_n_gpus()
+ metrics.update(
+ compute_throughout_metrics_decorated(
+ batch=compute_throughout_metrics_meta, timing_raw=timing_raw, n_gpus=n_gpus
+ )
+ )
+
+ # this is experimental and may be changed/removed in the future in favor of a general-purpose one
+ if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):
+ # TODO (TQ) :support transfer queue
+ self.train_dataloader.sampler.update(batch=batch)
+
+ self.tq_client.clear_samples(batch_meta)
+ # TODO: make a canonical logger that supports various backend
+ logger.log(data=metrics, step=self.global_steps)
+
+ progress_bar.update(1)
+ self.global_steps += 1
+
+ if (
+ hasattr(self.config.actor_rollout_ref.actor, "profiler")
+ and self.config.actor_rollout_ref.actor.profiler.tool == "torch_memory"
+ ):
+ self.actor_rollout_wg.dump_memory_snapshot(
+ tag=f"post_update_step{self.global_steps}", sub_dir=f"step{self.global_steps}"
+ )
+
+ if is_last_step:
+ pprint(f"Final validation metrics: {last_val_metrics}")
+ progress_bar.close()
+ return
+
+ # this is experimental and may be changed/removed in the future
+ # in favor of a general-purpose data buffer pool
+ if hasattr(self.train_dataset, "on_batch_end"):
+ # The dataset may be changed after each training batch
+ # TODO (TQ): support transfer queue
+ self.train_dataset.on_batch_end(batch=batch)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/run_qwen3-8b_transferqueue.sh b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/run_qwen3-8b_transferqueue.sh
new file mode 100644
index 0000000000000000000000000000000000000000..bd6d09e32d7be2199e8332bc63bb1296d53b3b27
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/transfer_queue/run_qwen3-8b_transferqueue.sh
@@ -0,0 +1,70 @@
+set -x
+
+MODEL_PATH="/workspace/models/Qwen3-8B"
+TRAIN_FILE="/workspace/datasets/preprocessed/gsm8k/train.parquet"
+TEST_FILE="/workspace/datasets/preprocessed/gsm8k/test.parquet"
+
+log_dir="./logs"
+mkdir -p ${log_dir}
+timestamp=$(date +"%Y%m%d%H%M%S")
+log_file="${log_dir}/qwen3-8b_tq_${timestamp}.log"
+
+# You may try to enable zero-copy serialization for TransferQueue when using SimpleStorageUnit backend.
+export TQ_ZERO_COPY_SERIALIZATION=False
+
+rollout_mode="async"
+rollout_name="vllm" # sglang or vllm
+if [ "$rollout_mode" = "async" ]; then
+ export VLLM_USE_V1=1
+ return_raw_chat="True"
+fi
+
+# You may also refer to tests/special_e2e/run_transferqueue.sh for more demo scripts
+
+python3 -m verl.experimental.transfer_queue.main_ppo \
+ --config-name='transfer_queue_ppo_trainer' \
+ algorithm.adv_estimator=grpo \
+ data.train_files=${TRAIN_FILE} \
+ data.val_files=${TEST_FILE} \
+ data.return_raw_chat=$return_raw_chat \
+ data.train_batch_size=128 \
+ data.max_prompt_length=2048 \
+ data.max_response_length=8192 \
+ data.filter_overlong_prompts_workers=128 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ actor_rollout_ref.model.path=${MODEL_PATH} \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=True \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
+ actor_rollout_ref.rollout.max_num_batched_tokens=10240 \
+ actor_rollout_ref.rollout.name=$rollout_name \
+ actor_rollout_ref.rollout.mode=$rollout_mode \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=5 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger=console \
+ trainer.project_name='verl_grpo_example_gsm8k' \
+ trainer.experiment_name='qwen3_8b_function_rm' \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes=1 \
+ trainer.save_freq=-1 \
+ trainer.test_freq=1000 \
+ trainer.total_epochs=15 \
+ trainer.total_training_steps=2 \
+ trainer.val_before_train=False \
+ 2>&1 | tee "$log_file"
+echo "Finished, log is saved in: $log_file"
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/README.md b/code/RL_model/verl/verl_train/verl/experimental/vla/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..5797ea8560da7694636eb0d8ca7d02c24f784986
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/README.md
@@ -0,0 +1,67 @@
+# [WIP] Experimental VLA RL Support
+
+This recipe introduces experimental support for training SimpleVLA-OFT, a VLA model.
+
+A key challenge in VLA RL training, which differs from standard LLM RL training, is that the environment/simulation phase has a higher computational overhead than the generation phase. To achieve high efficiency, RL in this context requires an effective environment scheduling mechanism in addition to verl's existing efficient training and inference scheduling. The goal is to reduce the inefficiency caused by the environment and the model's generation process waiting on each other.
+
+The core computational model of this PR is inspired by the pipeline parallelism design from RLinf. It aims to overlap the environment's execution time with the model's generation time, thereby maximizing environment utilization.
+
+This PR also proposes a future direction: creating a unified `Env` class. This class would encapsulate functionalities like tool calling, MCP, etc., under a single interface. The environment would manage its state internally, allowing the agent to communicate simply by calling `step(action)` to submit an action and receive an observation.
+
+Currently, this code is located independently within the `recipes` folder. Much of the design is tightly coupled with the SimpleVLA model and the Libero environment, serving as an initial version for demonstration and discussion.
+
+## Supported Simulators
+
+| Simulator | Env Name | Difference | Benchmark data source |
+| --- | --- | --- | --- |
+| Mujoco | LiberoEnv | 1. init task from init_states in Libero dataset
2. each env can have different tasks | https://github.com/Lifelong-Robot-Learning/LIBERO |
+| IsaacSim | IsaacEnv | 1. init task from random states, which has more variety than init_states in dataset
2. each sim process must using the same task for its envs | https://huggingface.co/datasets/china-sae-robotics/IsaacLabPlayGround_Dataset |
+
+## Hardware Requirements
+
+* Simulator GPU: NVIDIA L20 or L40 with 48GB memory and RT Cores
+
+Notes:
+1. Mujoco can failback to CPU mode with degraded performance if no RT Cores is available
+2. IsaacSim only support GPU with RT Cores
+3. RTX GPU will be supported in the future release with remote deployment feature, but it can not work with colocated mode because of the limitation of GPU memory capacity.
+
+## Docker image
+
+The Isaac Lab support for libero dataset depends on RobotLearningLab project from The Isaac Lab Project Developers team. The project is in the process of being public available and is currently build in this image with BSD-3-Clause license.
+
+`recipe/vla/run_simpleVLA_libero_grpo.sh` is the example of training SimpleVLA-OFT with this image:
+
+`vemlp-cn-shanghai.cr.volces.com/preset-images/verl_vla:preview_vla_0.1`
+
+## Disaggregation Mode for Train-Rollout / Simulation
+
+Disaggregate Train-Rollout workers and Simulation workers into different nodes.
+
+To enable disaggregation mode for Train-Rollout nodes and Simulation nodes, we need to establish ray connection before running verl.
+* On Train-Rollout node (default main node):
+```shell
+ray start --head --dashboard-host=0.0.0.0 --resources='{"train_rollout": 1}'
+```
+* On Simulation node:
+```shell
+ray start --address=':6379' --resources='{"sim": 1}'
+```
+
+Then run verl on main node **only**. See `run_simpleVLA_isaac_disagg.sh` for example.
+- `env.disagg_sim.enable=True` enable disagg mode
+- `trainer.n_env_gpus_per_node` GPUs for simulaton per node
+- `trainer.n_rollout_gpus_per_node` GPUs for train-rollout node
+- `env.disagg_sim.nnodes` sim node num
+- `trainer.nnodes` train-rollout node num
+
+*Tips: you can run the following command on the sim node to check whether sim workers are scheduled up*
+```shell
+python -c "import ray; ray.init(address=\":6379\"); print(ray._private.state.available_resources_per_node())"
+```
+*If you see output pattern like "'train_rollout': 0.9992" and "'sim': 0.9992", the sim workers are scheduled up successfully*
+*The actual value depends on your GPUs per node, usually <1 - 1e-4 * num_gpus>*
+
+**References:**
+* [https://github.com/PRIME-RL/SimpleVLA-RL](https://github.com/PRIME-RL/SimpleVLA-RL)
+* [https://github.com/RLinf/RLinf](https://github.com/RLinf/RLinf)
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/config/rob_ppo_trainer.yaml b/code/RL_model/verl/verl_train/verl/experimental/vla/config/rob_ppo_trainer.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..8ad4c7dd7c26637444673efaa761b76e008cdc37
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/config/rob_ppo_trainer.yaml
@@ -0,0 +1,138 @@
+# the rob_ppo config will override default ppo_trainer.yaml
+
+hydra:
+ searchpath:
+ - file://verl/trainer/config
+
+defaults:
+ - ppo_trainer
+ - _self_
+
+env:
+ rollout:
+ pipeline_stage_num: 2
+ actor:
+ model:
+ num_action_chunks: 8
+ action_dim: 7
+ train:
+ simulator_type: libero
+ max_episode_steps: 512
+ reward_coef: 1.0
+ only_eval: False
+ video_cfg:
+ save_video: True
+ video_base_dir: /tmp/videos
+ num_envs: 16
+ seed: 42
+ task_suite_name: libero_10
+ init_params:
+ camera_depths: False
+ camera_heights: 256
+ camera_widths: 256
+ camera_names:
+ - agentview
+ - robot0_eye_in_hand
+
+ # Profile the env worker
+ profiler:
+
+ # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs
+ _target_: verl.utils.profiler.ProfilerConfig
+
+ # Profiling tool to use
+ # options: nsys, npu, torch, torch_memory
+ # Defaults to global_profiler.tool if set
+ tool: ${oc.select:global_profiler.tool,null}
+
+ # Whether to enable profiling for env worker
+ enable: False
+
+ # Whether to profile all ranks
+ all_ranks: False
+
+ # List of ranks to profile (empty means no specific ranks)
+ ranks: []
+
+ # Path to save profiling results
+ # Defaults to global_profiler.save_path if set
+ save_path: ${oc.select:global_profiler.save_path,null}
+
+ # Tool-specific configurations
+ tool_config:
+
+ # nsys tool config
+ nsys:
+
+ # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs
+ _target_: verl.utils.profiler.config.NsightToolConfig
+
+ # True for each task has its own database, False for all tasks in one training step share one database.
+ discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}
+
+ # npu config
+ npu:
+
+ # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs
+ _target_: verl.utils.profiler.config.NPUToolConfig
+
+ # Contents to profile, can be empty
+ # options: npu, cpu, memory, shapes, module, stack
+ contents: []
+
+ # Collection level, optional values: level_none, level0, level1, level2.
+ level: "level1"
+
+ # Whether to automatically parse the data.
+ analysis: True
+
+ # True for each task has its own database, False for all tasks in one training step share one database.
+ discrete: False
+
+ # torch profiler config
+ torch:
+
+ # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs
+ _target_: verl.utils.profiler.config.TorchProfilerToolConfig
+
+ # Contents to profile, can be empty
+ # options: cuda, cpu, memory, shapes, stack
+ contents: []
+
+ # True for each task has its own database, False for all tasks in one training step share one database.
+ discrete: False
+
+
+ # torch memory profiler config
+ torch_memory:
+
+ # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs
+ _target_: verl.utils.profiler.config.TorchMemoryToolConfig
+
+ # Maximum number of memory allocation entries to track
+ trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}
+
+ # Stack trace depth for memory allocations
+ stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}
+ disagg_sim:
+ enable: False
+ nnodes: 1
+
+
+actor_rollout_ref:
+ actor:
+ num_images_in_input: 1
+ traj_mini_batch_size: 16
+ fsdp_config:
+ wrap_policy:
+ transformer_layer_cls_to_wrap:
+ - PrismaticProjector
+ - LlamaDecoderLayer
+ min_num_params: 0
+ param_offload: False
+ optimizer_offload: False
+ forward_prefetch: True
+ fsdp_size: -1
+ rollout:
+ mode: async_envloop
+ prompt_length: 512
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/dp_rob.py b/code/RL_model/verl/verl_train/verl/experimental/vla/dp_rob.py
new file mode 100644
index 0000000000000000000000000000000000000000..1830aa81a3f0ca3b0946d85b00dde9115cab1630
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/dp_rob.py
@@ -0,0 +1,323 @@
+# Copyright 2024 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.
+"""
+Single Process Actor
+"""
+
+import logging
+
+import torch
+from tensordict.base import TensorDictBase
+from torch import nn
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+import verl.utils.torch_functional as verl_F
+from verl.protocol import DataProto
+from verl.trainer.ppo import core_algos
+from verl.utils.device import get_device_id, get_device_name
+from verl.utils.py_functional import append_to_dict
+from verl.utils.seqlen_balancing import prepare_dynamic_batch, restore_dynamic_batch
+from verl.utils.torch_functional import logprobs_from_logits
+from verl.workers.actor import BasePPOActor
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["RobDataParallelPPOActor"]
+
+
+class RobDataParallelPPOActor(BasePPOActor):
+ def __init__(
+ self,
+ config,
+ actor_module: nn.Module,
+ actor_optimizer: torch.optim.Optimizer = None,
+ ):
+ """When optimizer is None, it is Reference Policy"""
+ super().__init__(config)
+ self.actor_module = actor_module
+ self.actor_optimizer = actor_optimizer
+ self.use_remove_padding = self.config.get("use_remove_padding", False)
+ logger.info(f"Actor use_remove_padding={self.use_remove_padding}")
+ logger.info(f"PRM use dynamic bsz={self.config.get('use_dynamic_bsz', False)}")
+ self.ulysses_sequence_parallel_size = self.config.ulysses_sequence_parallel_size
+ self.use_ulysses_sp = False # self.ulysses_sequence_parallel_size > 1
+ self.compute_entropy_from_logits = torch.compile(verl_F.entropy_from_logits, dynamic=True)
+
+ def process_tensor(self, tensor, pad_id):
+ mask = tensor != pad_id
+ if not torch.all(mask == mask[0:1], dim=1).all():
+ raise ValueError("Padding error!")
+ base_mask = mask[0]
+ valid_len = base_mask.sum().item()
+ return tensor[:, base_mask], valid_len
+
+ def generate_traj_mask(self, end_step, traj_len):
+ """
+ Args:
+ end_step: (batch_size,),
+ traj_len:
+ Returns:
+ mask: (batch_size, traj_len),
+ """
+ steps = torch.arange(traj_len, device=end_step.device) # (traj_len,)
+ steps_expanded = steps.unsqueeze(0).expand(end_step.size(0), -1)
+ mask = steps_expanded < end_step.unsqueeze(1) # (batch_size, traj_len)
+ return mask
+
+ def apply_mask_with_grad_control(self, log_probs, entropy, mask):
+ """
+ Args:
+ log_probs: (batch_size, 7*8)
+ entropy: (batch_size, 7*8)
+ # mask: (batch_size, 8)
+ mask: (batch_size, 7*8)
+ Returns:
+ log_probs_masked:
+ entropy_masked:
+ """
+
+ mask = mask.to(log_probs.device)
+ log_probs_masked = torch.where(mask, log_probs, torch.zeros_like(log_probs, requires_grad=False))
+ entropy_masked = torch.where(mask, entropy, torch.zeros_like(entropy, requires_grad=False))
+ return log_probs_masked, entropy_masked
+
+ def _forward_micro_batch(self, micro_batch, temperature) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ micro_batch:
+
+ Returns:
+ entropy: # (bs, response_len)
+ log_probs: # (bs, response_len)
+ """
+
+ with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):
+ input_ids = micro_batch["input_ids"]
+ attention_mask = micro_batch["attention_mask"]
+ pixel_values = micro_batch["pixel_values"]
+ responses = micro_batch["responses"]
+
+ input_ids_unpad, _ = self.process_tensor(input_ids, self.pad_token_id)
+ attention_mask_unpad, _ = self.process_tensor(attention_mask, 0)
+
+ logits = self.actor_module(
+ input_ids=input_ids_unpad,
+ attention_mask=attention_mask_unpad,
+ pixel_values=pixel_values,
+ ) # prevent model thinks we are generating
+
+ assert self.actor_module.vocab_size == 32000
+ start_index = self.actor_module.vocab_size - 256
+ logits = logits[..., -256 - 64 : -64] # Shape: [batch_size, seq_len, 256]
+ responses = responses - start_index
+ # assert (0<=responses<=255).all()
+
+ logits = logits.div(temperature)
+
+ log_probs = logprobs_from_logits(logits, responses.to(logits.device))
+ entropy = verl_F.entropy_from_logits(logits) # (bsz, response_length)
+
+ # assert len(log_probs.shape) == 2 and len(entropy.shape) == 2
+
+ # TODO(caiyunke.astra): check here
+
+ mask = micro_batch["response_mask"]
+ log_probs, entropy = self.apply_mask_with_grad_control(log_probs, entropy, mask)
+
+ return entropy, log_probs
+
+ def _forward_micro_batch_update(
+ self, input_ids, attention_mask, pixel_values, responses, temperature
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):
+ input_ids_unpad, _ = self.process_tensor(input_ids, self.pad_token_id)
+ attention_mask_unpad, _ = self.process_tensor(attention_mask, 0)
+
+ logits = self.actor_module(
+ input_ids=input_ids_unpad,
+ attention_mask=attention_mask_unpad,
+ pixel_values=pixel_values,
+ )
+
+ assert logits.requires_grad
+
+ assert self.actor_module.vocab_size == 32000
+ start_index = self.actor_module.vocab_size - 256
+ logits = logits[..., -256 - 64 : -64] # Shape: [batch_size, seq_len, 256]
+ responses = responses - start_index
+
+ logits = logits.div(temperature)
+
+ log_probs = logprobs_from_logits(logits, responses)
+ entropy = verl_F.entropy_from_logits(logits) # (bsz, response_length)
+ return entropy, log_probs
+
+ def _optimizer_step(self):
+ assert self.config.grad_clip is not None
+
+ if isinstance(self.actor_module, FSDP):
+ grad_norm = self.actor_module.clip_grad_norm_(max_norm=self.config.grad_clip)
+ else:
+ grad_norm = torch.nn.utils.clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip)
+ self.actor_optimizer.step()
+ return grad_norm
+
+ def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:
+ """Compute the log probability of the responses given input_ids, attention_mask and position_ids
+
+ Args:
+ data (DataProto): a DataProto containing keys
+
+ ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the
+ concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``.
+
+ ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64.
+
+ ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64.
+
+ ``responses``: tensor of shape [batch_size, response_length]. torch.int64.
+
+ Returns:
+ torch.Tensor: the log_prob tensor
+ """
+ self.actor_module.eval()
+
+ micro_batch_size = data.meta_info["micro_batch_size"] # 256
+ temperature = data.meta_info[
+ "temperature"
+ ] # temperature must be in the data.meta_info to avoid slient error # 1
+ use_dynamic_bsz = data.meta_info["use_dynamic_bsz"] # trues
+ self.pad_token_id = data.meta_info["pad_token_id"]
+
+ select_keys = ["responses", "input_ids", "attention_mask", "pixel_values", "response_mask"]
+ data = data.select(batch_keys=select_keys).batch
+
+ if use_dynamic_bsz:
+ max_token_len = data.meta_info["max_token_len"] * self.ulysses_sequence_parallel_size
+ micro_batches, batch_idx_list = prepare_dynamic_batch(data, max_token_len=max_token_len)
+ else:
+ micro_batches = data.split(micro_batch_size)
+
+ log_probs_lst = []
+ entropy_lst = []
+ for micro_batch in micro_batches:
+ with torch.no_grad():
+ entropy, log_probs = self._forward_micro_batch(micro_batch, temperature=temperature)
+ log_probs_lst.append(log_probs)
+ if calculate_entropy:
+ entropy_lst.append(entropy)
+ log_probs = torch.concat(log_probs_lst, dim=0)
+ entropys = None
+ if calculate_entropy:
+ entropys = torch.concat(entropy_lst, dim=0)
+
+ if use_dynamic_bsz:
+ log_probs = restore_dynamic_batch(log_probs, batch_idx_list)
+ if calculate_entropy:
+ entropys = restore_dynamic_batch(entropys, batch_idx_list)
+
+ return log_probs, entropys
+
+ def update_policy(self, data: DataProto):
+ self.actor_module.train()
+
+ assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size_per_gpu == 0
+ self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu
+ temperature = data.meta_info["temperature"] # temperature must be in the data.meta_info to avoid slient error
+
+ select_keys = [
+ "responses",
+ "response_mask",
+ "input_ids",
+ "attention_mask",
+ "pixel_values",
+ "old_log_probs",
+ "advantages",
+ ]
+ batch = data.select(batch_keys=select_keys).batch
+ self.pad_token_id = data.meta_info["pad_token_id"]
+ # TODO(caiyunke.astra): check here
+ # assert self.config.ppo_micro_batch_size_per_gpu == 1
+
+ # Split to make minibatch iterator for updating the actor
+ # See PPO paper for details. https://arxiv.org/abs/1707.06347
+ mini_batches = batch.split(self.config.ppo_mini_batch_size)
+ metrics = {}
+ for batch_idx, mini_batch in enumerate(mini_batches):
+ if self.config.use_dynamic_bsz:
+ max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size
+ micro_batches, _ = prepare_dynamic_batch(mini_batch, max_token_len=max_token_len)
+ else:
+ self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu
+ micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu)
+
+ self.actor_optimizer.zero_grad()
+
+ for _, micro_batch in enumerate[DataProto | TensorDictBase](micro_batches):
+ micro_batch = micro_batch.to(get_device_id()) # actor device is cpu when using offload
+ responses = micro_batch["responses"]
+
+ response_mask = micro_batch["response_mask"] # (batch_size, traj_len)
+
+ old_log_prob = micro_batch["old_log_probs"]
+ advantages = micro_batch["advantages"]
+
+ # clip_ratio = self.config.clip_ratio
+ clip_ratio_high = self.config.clip_ratio_high
+ clip_ratio_low = self.config.clip_ratio_low
+
+ input_ids = micro_batch["input_ids"]
+ attention_mask = micro_batch["attention_mask"]
+ pixel_values = micro_batch["pixel_values"]
+ responses = micro_batch["responses"]
+
+ loss_info = {
+ "actor/pg_loss": 0,
+ "actor/pg_clipfrac": 0,
+ "actor/ppo_kl": 0,
+ "actor/pg_clipfrac_lower": 0,
+ }
+
+ _, log_prob = self._forward_micro_batch_update(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ pixel_values=pixel_values,
+ responses=responses,
+ temperature=temperature,
+ )
+
+ pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = core_algos.compute_policy_loss(
+ old_log_prob=old_log_prob,
+ log_prob=log_prob,
+ advantages=advantages,
+ response_mask=response_mask,
+ cliprange_high=clip_ratio_high,
+ cliprange_low=clip_ratio_low,
+ )
+ loss = pg_loss / self.gradient_accumulation
+
+ loss.backward()
+
+ loss_info["actor/pg_loss"] = loss_info["actor/pg_loss"] + pg_loss.detach().item()
+ loss_info["actor/pg_clipfrac"] = loss_info["actor/pg_clipfrac"] + pg_clipfrac.detach().item()
+ loss_info["actor/ppo_kl"] = loss_info["actor/ppo_kl"] + ppo_kl.detach().item()
+ loss_info["actor/pg_clipfrac_lower"] = (
+ loss_info["actor/pg_clipfrac_lower"] + pg_clipfrac_lower.detach().item()
+ )
+ append_to_dict(metrics, loss_info)
+
+ grad_norm = self._optimizer_step()
+ mini_batch_metrics = {"actor/grad_norm": grad_norm.detach().item()}
+ append_to_dict(metrics, mini_batch_metrics)
+ self.actor_optimizer.zero_grad()
+ return metrics
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/env_loop.py b/code/RL_model/verl/verl_train/verl/experimental/vla/env_loop.py
new file mode 100644
index 0000000000000000000000000000000000000000..699e62441dcdbce97e8d87c2bfc43fd81ca1920d
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/env_loop.py
@@ -0,0 +1,199 @@
+# Copyright 2025 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.
+
+import asyncio
+import logging
+import os
+
+import numpy as np
+import torch
+from omegaconf import DictConfig
+
+from verl import DataProto
+from verl.single_controller.ray import RayWorkerGroup
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class EnvLoop:
+ """An env loop manages interactions between models and vectorized environments. It's designed for computationally
+ intensive environments, such as robotics simulators."""
+
+ def __init__(self, env_wg: RayWorkerGroup, rollout_wg: RayWorkerGroup, config: DictConfig):
+ """
+ Initialize the EnvLoop.
+
+ Args:
+ env_wg (RayWorkerGroup): Environment worker group.
+ rollout_wg (RayWorkerGroup): Rollout worker group for model inference.
+ config (DictConfig): YAML config.
+ """
+ self.env_wg = env_wg
+ self.rollout_wg = rollout_wg
+ self.config = config
+ # Extract relevant configuration
+ self.max_interactions = config.env.train.max_episode_steps // config.env.actor.model.num_action_chunks
+ self.stage_num = config.env.rollout.pipeline_stage_num
+ self.num_envs_per_worker = config.env.train.num_envs
+ self.action_dim = config.env.actor.model.action_dim
+ self.num_action_chunks = config.env.actor.model.num_action_chunks
+ # Derived properties
+ self.total_envs = self.env_wg.world_size * self.num_envs_per_worker
+ if self.total_envs % self.stage_num != 0:
+ raise ValueError(f"Total envs ({self.total_envs}) must be divisible by stage_num ({self.stage_num})")
+ self.envs_per_stage = self.total_envs // self.stage_num
+
+ self.env_wg.init_worker()
+ self.env_wg.init_simulator()
+
+ def generate_sequences(self, prompts: DataProto, reset_future: asyncio.Future) -> DataProto:
+ """Split input batch and dispatch to env loop workers.
+
+ Args:
+ prompts (DataProto): Input batch.
+
+ Returns:
+ DataProto: Output batch.
+ """
+
+ reset_results = reset_future.get()
+
+ loop = asyncio.get_event_loop()
+ self.rollout_wg.switch_to_rollout()
+ output = loop.run_until_complete(self.run(prompts, reset_results))
+ self.rollout_wg.switch_to_train()
+ # TODO(caiyunke.astra): add timing metrics
+ return output
+
+ async def run(self, prompts: DataProto, reset_results: DataProto) -> DataProto:
+ """
+ Run the environment interaction loop.
+ This method orchestrates a pipelined process:
+ 1. Resets environments to specified initial states.
+ 2. In a loop, it gets actions from the rollout workers and applies them to the environments.
+ 3. Collects all trajectory data (observations, actions, rewards, dones).
+ 4. Formats and returns the collected trajectories as a single batch.
+ Args:
+ prompts (DataProto): Contains initial state IDs and other settings.
+ - 'non_tensor_batch.state_ids': A numpy array of state IDs to reset envs.
+ Returns:
+ DataProto: A batch containing the complete trajectories.
+ """
+ initial_state_ids = prompts.non_tensor_batch["state_ids"]
+
+ staged_obs = self._restructure_obs_data(reset_results)
+ # --- Pipeline state ---
+ trajectories = {i: [] for i in range(self.stage_num)} # To store (obs, action, rew, done) tuples
+ rollout_futures = {}
+ # is_complete = torch.zeros((self.total_envs,), dtype=torch.bool)
+
+ for stage_id in range(self.stage_num):
+ # trajectories[stage_id].append({'obs': staged_obs[stage_id]})
+ trajectories[stage_id].append({})
+ vla_input = staged_obs[stage_id]
+ vla_input.meta_info = prompts.meta_info # Pass along rollout config
+ rollout_futures[stage_id] = self.rollout_wg.generate_sequences(vla_input)
+
+ async def _stage_loop(stage_id: int):
+ for step_idx in range(self.max_interactions):
+ action_result: DataProto = await asyncio.to_thread(rollout_futures[stage_id].get)
+
+ trajectories[stage_id][-1]["action"] = action_result
+ action_data = DataProto.from_dict(
+ non_tensors={"actions": action_result.batch["action"].cpu().numpy()},
+ meta_info={"stage_id": stage_id},
+ )
+
+ env_ref = self.env_wg.env_interact_step(action_data)
+ env_result: DataProto = await asyncio.to_thread(env_ref.get)
+
+ trajectories[stage_id][-1]["rew"] = env_result.batch["rews"]
+ trajectories[stage_id][-1]["done"] = env_result.batch["terminations"]
+
+ next_obs = DataProto(
+ batch=env_result.batch.select("full_image", "state"),
+ non_tensor_batch={"task_descriptions": env_result.non_tensor_batch["task_descriptions"]},
+ )
+
+ if step_idx < self.max_interactions - 1:
+ trajectories[stage_id].append({})
+ vla_input = next_obs
+ vla_input.meta_info = prompts.meta_info
+ rollout_futures[stage_id] = self.rollout_wg.generate_sequences(vla_input)
+
+ await asyncio.gather(*[asyncio.create_task(_stage_loop(sid)) for sid in range(self.stage_num)])
+ self.env_wg.finish_rollout()
+
+ return self._collate_trajectories(trajectories, initial_state_ids, meta_info=prompts.meta_info)
+
+ def _restructure_obs_data(self, data_proto: DataProto) -> list[DataProto]:
+ """Reshapes flat observation data from env_wg into a list of per-stage DataProto objects."""
+ # env_wg returns a flat batch ordered by [worker0_stage0, worker0_stage1, ...,
+ # worker1_stage0, worker1_stage1, ...]
+ # First, un-flatten by worker, then by stage
+
+ num_workers = self.env_wg.world_size
+
+ staged_data = [[] for _ in range(self.stage_num)]
+ chunks = data_proto.chunk(num_workers)
+ for worker_chunk in chunks:
+ stage_chunks = worker_chunk.chunk(self.stage_num)
+ for stage_id, data in enumerate(stage_chunks):
+ staged_data[stage_id].append(data)
+
+ # Concatenate data from all workers for each stage
+ return [DataProto.concat(data_list) for data_list in staged_data]
+
+ def _collate_trajectories(self, trajectories: dict, initial_state_ids: np.ndarray, meta_info) -> DataProto:
+ """
+ Collates the collected trajectory data into the final batch format.
+ """
+ flat_trajs = [{} for _ in range(len(trajectories[0]))]
+ for stage_id in range(self.stage_num):
+ for step_idx, step_data in enumerate(trajectories[stage_id]):
+ if not flat_trajs[step_idx]: # if dict is empty
+ flat_trajs[step_idx] = step_data
+ else:
+ # Concatenate DataProto objects
+ for key, value in step_data.items():
+ if isinstance(value, DataProto):
+ flat_trajs[step_idx][key] = DataProto.concat([flat_trajs[step_idx][key], value])
+ elif isinstance(value, torch.Tensor):
+ flat_trajs[step_idx][key] = torch.cat([flat_trajs[step_idx][key], value], dim=0)
+
+ all_pixel_values = [step["action"].batch["pixel_values"] for step in flat_trajs]
+ all_responses = [step["action"].batch["responses"] for step in flat_trajs]
+ all_input_ids = [step["action"].batch["input_ids"] for step in flat_trajs]
+ all_attn_masks = [step["action"].batch["attention_mask"] for step in flat_trajs]
+ all_actions = [step["action"].batch["action"] for step in flat_trajs]
+ all_dones = [step["done"] for step in flat_trajs]
+
+ pixel_values = torch.stack(all_pixel_values, dim=1)
+ responses = torch.stack(all_responses, dim=1)
+ input_ids = torch.stack(all_input_ids, dim=1)
+ attention_mask = torch.stack(all_attn_masks, dim=1)
+ actions = torch.stack(all_actions, dim=1)
+ complete = torch.stack(all_dones, dim=1).squeeze(-1) # Shape [bs, steps]
+ batch_dict = {
+ "pixel_values": pixel_values,
+ "responses": responses,
+ "input_ids": input_ids,
+ "attention_mask": attention_mask,
+ "complete": complete,
+ "action": actions,
+ "env_state_id": torch.from_numpy(initial_state_ids.astype(int)),
+ }
+
+ return DataProto.from_single_dict(batch_dict, meta_info=meta_info)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2171666b035340542d81212af41b2ca1f96fab69
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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.
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/action_utils.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/action_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d361de0814a4e91e679cbebe8f33fe60eaf1dbd8
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/action_utils.py
@@ -0,0 +1,303 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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.
+
+import os
+from io import BytesIO
+from typing import Any, Optional
+
+import imageio
+import numpy as np
+import torch
+import torchvision.transforms.functional as F
+from PIL import Image, ImageDraw, ImageFont
+
+
+def prepare_actions_simplevla(
+ raw_chunk_actions,
+) -> torch.Tensor:
+ from verl.experimental.vla.envs.libero_env.utils import invert_gripper_action, normalize_gripper_action
+
+ normalized_action = normalize_gripper_action(raw_chunk_actions, binarize=True)
+ inverted_action = invert_gripper_action(normalized_action)
+ return inverted_action
+
+
+def prepare_actions(
+ simulator_type,
+ raw_chunk_actions,
+ num_action_chunks,
+ action_dim,
+ action_scale: float = 1.0,
+ policy: str = "widowx_bridge",
+) -> torch.Tensor:
+ # TODO: prepare_actions according to simulator_type
+ chunk_actions = prepare_actions_simplevla(
+ raw_chunk_actions=raw_chunk_actions,
+ )
+
+ return chunk_actions
+
+
+def to_tensor(array: dict | torch.Tensor | np.ndarray | list | Any, device: str = "cpu") -> dict | torch.Tensor:
+ """
+ Copied from ManiSkill!
+ Maps any given sequence to a torch tensor on the CPU/GPU. If physx gpu
+ is not enabled then we use CPU, otherwise GPU, unless specified
+ by the device argument
+
+ Args:
+ array: The data to map to a tensor
+ device: The device to put the tensor on. By default this is None
+ and to_tensor will put the device on the GPU if physx is enabled
+ and CPU otherwise
+
+ """
+ if isinstance(array, (dict)):
+ return {k: to_tensor(v, device=device) for k, v in array.items()}
+ elif isinstance(array, torch.Tensor):
+ ret = array.to(device)
+ elif isinstance(array, np.ndarray):
+ if array.dtype == np.uint16:
+ array = array.astype(np.int32)
+ elif array.dtype == np.uint32:
+ array = array.astype(np.int64)
+ ret = torch.tensor(array).to(device)
+ else:
+ if isinstance(array, list) and isinstance(array[0], np.ndarray):
+ array = np.array(array)
+ ret = torch.tensor(array, device=device)
+ if ret.dtype == torch.float64:
+ ret = ret.to(torch.float32)
+ return ret
+
+
+def tile_images(images: list[np.ndarray | torch.Tensor], nrows: int = 1) -> np.ndarray | torch.Tensor:
+ """
+ Copied from maniskill https://github.com/haosulab/ManiSkill
+ Tile multiple images to a single image comprised of nrows and an
+ appropriate number of columns to fit all the images.
+ The images can also be batched (e.g. of shape (B, H, W, C)), but
+ give images must all have the same batch size.
+
+ if nrows is 1, images can be of different sizes. If nrows > 1,
+ they must all be the same size.
+ """
+ # Sort images in descending order of vertical height
+ batched = False
+ if len(images[0].shape) == 4:
+ batched = True
+ if nrows == 1:
+ images = sorted(images, key=lambda x: x.shape[0 + batched], reverse=True)
+
+ columns: list[list[np.ndarray | torch.Tensor]] = []
+ if batched:
+ max_h = images[0].shape[1] * nrows
+ cur_h = 0
+ cur_w = images[0].shape[2]
+ else:
+ max_h = images[0].shape[0] * nrows
+ cur_h = 0
+ cur_w = images[0].shape[1]
+
+ # Arrange images in columns from left to right
+ column = []
+ for im in images:
+ if cur_h + im.shape[0 + batched] <= max_h and cur_w == im.shape[1 + batched]:
+ column.append(im)
+ cur_h += im.shape[0 + batched]
+ else:
+ columns.append(column)
+ column = [im]
+ cur_h, cur_w = im.shape[0 + batched : 2 + batched]
+ columns.append(column)
+
+ # Tile columns
+ total_width = sum(x[0].shape[1 + batched] for x in columns)
+
+ is_torch = False
+ if torch is not None:
+ is_torch = isinstance(images[0], torch.Tensor)
+
+ output_shape = (max_h, total_width, 3)
+ if batched:
+ output_shape = (images[0].shape[0], max_h, total_width, 3)
+ if is_torch:
+ output_image = torch.zeros(output_shape, dtype=images[0].dtype)
+ else:
+ output_image = np.zeros(output_shape, dtype=images[0].dtype)
+ cur_x = 0
+ for column in columns:
+ cur_w = column[0].shape[1 + batched]
+ next_x = cur_x + cur_w
+ if is_torch:
+ column_image = torch.concatenate(column, dim=0 + batched)
+ else:
+ column_image = np.concatenate(column, axis=0 + batched)
+ cur_h = column_image.shape[0 + batched]
+ output_image[..., :cur_h, cur_x:next_x, :] = column_image
+ cur_x = next_x
+ return output_image
+
+
+def put_text_on_image(image: np.ndarray, lines: list[str], max_width: int = 200) -> np.ndarray:
+ """
+ Put text lines on an image with automatic line wrapping.
+
+ Args:
+ image: Input image as numpy array
+ lines: List of text lines to add
+ max_width: Maximum width for text wrapping
+ """
+ assert image.dtype == np.uint8, image.dtype
+ image = image.copy()
+ image = Image.fromarray(image)
+ draw = ImageDraw.Draw(image)
+ font = ImageFont.load_default(size=20)
+
+ new_lines = []
+ for line in lines:
+ words = line.split()
+ current_line = []
+
+ for word in words:
+ test_line = " ".join(current_line + [word])
+ test_width = font.getlength(test_line)
+
+ if test_width <= max_width:
+ current_line.append(word)
+ else:
+ new_lines.append(" ".join(current_line))
+ current_line = [word]
+ if current_line:
+ new_lines.append(" ".join(current_line))
+
+ y = -10
+ for line in new_lines:
+ bbox = draw.textbbox((0, 0), text=line)
+ textheight = bbox[3] - bbox[1]
+ y += textheight + 10
+ x = 10
+ draw.text((x, y), text=line, fill=(0, 0, 0))
+ return np.array(image)
+
+
+def put_info_on_image(
+ image: np.ndarray,
+ info: dict[str, float],
+ extras: Optional[list[str]] = None,
+ overlay: bool = True,
+) -> np.ndarray:
+ """
+ Put information dictionary and extra lines on an image.
+
+ Args:
+ image: Input image
+ info: Dictionary of key-value pairs to display
+ extras: Additional text lines to display
+ overlay: Whether to overlay text on image
+ """
+ lines = [f"{k}: {v:.3f}" if isinstance(v, float) else f"{k}: {v}" for k, v in info.items()]
+ if extras is not None:
+ lines.extend(extras)
+ return put_text_on_image(image, lines)
+
+
+def list_of_dict_to_dict_of_list(
+ list_of_dict: list[dict[str, Any]],
+) -> dict[str, list[Any]]:
+ """
+ Convert a list of dictionaries to a dictionary of lists.
+
+ Args:
+ list_of_dict: List of dictionaries with same keys
+
+ Returns:
+ Dictionary where each key maps to a list of values
+ """
+ if len(list_of_dict) == 0:
+ return {}
+ keys = list_of_dict[0].keys()
+ output = {key: [] for key in keys}
+ for data in list_of_dict:
+ for key, item in data.items():
+ assert key in output
+ output[key].append(item)
+ return output
+
+
+def save_rollout_video(rollout_images: list[np.ndarray], output_dir: str, video_name: str, fps: int = 30) -> None:
+ """
+ Saves an MP4 replay of an episode.
+
+ Args:
+ rollout_images: List of images from the episode
+ output_dir: Directory to save the video
+ video_name: Name of the output video file
+ fps: Frames per second for the video
+ """
+ os.makedirs(output_dir, exist_ok=True)
+ mp4_path = os.path.join(output_dir, f"{video_name}.mp4")
+ video_writer = imageio.get_writer(mp4_path, fps=fps)
+ for img in rollout_images:
+ video_writer.append_data(img)
+ video_writer.close()
+
+
+def resize_image(img: np.ndarray, resize_size: tuple[int, int]) -> np.ndarray:
+ """
+ Takes numpy array corresponding to a single image and returns resized image as numpy array.
+
+ Args:
+ img: Input image as numpy array
+ resize_size: Target size for resizing
+
+ Returns:
+ Resized image as numpy array
+ """
+
+ assert isinstance(resize_size, tuple), "resize_size must be a tuple"
+ assert isinstance(img, np.ndarray), "img must be a numpy array"
+
+ # Convert numpy array to PIL Image
+ pil_img = Image.fromarray(img)
+
+ # Encode as JPEG, as done in RLDS dataset builder
+ buffer = BytesIO()
+ pil_img.save(buffer, format="JPEG")
+ buffer.seek(0)
+
+ # Immediately decode back
+ img = Image.open(buffer)
+
+ img = img.resize(resize_size, Image.Resampling.LANCZOS)
+ img = np.array(img)
+ img = np.clip(np.round(img), 0, 255).astype(np.uint8)
+
+ return img
+
+
+def center_crop_image(image: Image.Image) -> Image.Image:
+ crop_scale = 0.9
+ orig_w, orig_h = image.size
+ image_tensor = F.to_tensor(image)
+ crop_h = int(orig_h * crop_scale)
+ crop_w = int(orig_w * crop_scale)
+ image_tensor = F.center_crop(image_tensor, (crop_h, crop_w))
+ image_tensor = F.resize(image_tensor, (orig_h, orig_w))
+ final_image = F.to_pil_image(image_tensor)
+
+ final_image = final_image.convert("RGB")
+ return final_image
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/isaac_env/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/isaac_env/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..31a9171f0262536d9ccac845e33441336dd670b0
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/isaac_env/__init__.py
@@ -0,0 +1,17 @@
+# Copyright 2024 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
+#
+# https://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 .isaac_env import IsaacEnv
+
+__all__ = ["IsaacEnv"]
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/isaac_env/isaac_env.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/isaac_env/isaac_env.py
new file mode 100644
index 0000000000000000000000000000000000000000..665b2eaaecc3bf8343cacec0d7c74423e15ee1a8
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/isaac_env/isaac_env.py
@@ -0,0 +1,325 @@
+# Copyright 2025 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
+#
+# https://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.
+
+import logging
+import os
+from typing import Optional
+
+import gymnasium as gym
+import numpy as np
+import omni
+import torch
+
+from verl.experimental.vla.envs.action_utils import (
+ put_info_on_image,
+ save_rollout_video,
+ tile_images,
+ to_tensor,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class IsaacEnv(gym.Env):
+ def __init__(self, cfg, rank, world_size):
+ self.rank = rank
+ self.cfg = cfg
+ self.world_size = world_size
+ self.seed = self.cfg.seed + rank
+ self.num_envs = self.cfg.num_envs
+ self.action_dim = self.cfg.get("action_dim", 7)
+ self.device = self.cfg.get("device", "cuda:0")
+
+ self._generator = np.random.default_rng(seed=self.seed)
+
+ self.task_suite_name = self.cfg.task_suite_name
+
+ self.env = None
+ self.prev_step_reward = np.zeros(self.num_envs)
+ self.use_rel_reward = False
+
+ self._init_metrics()
+ self._elapsed_steps = np.zeros(self.num_envs, dtype=np.int32)
+ self.max_episode_steps = cfg.max_episode_steps
+ self.video_cfg = cfg.video_cfg
+
+ self.render_images = []
+ self.video_cnt = 0
+ self.camera_name = cfg.init_params.camera_names
+
+ # sys env must be set before import isaaclab
+ from isaaclab.app import AppLauncher
+
+ launch_args = {"headless": True, "enable_cameras": True}
+ app_launcher = AppLauncher(**launch_args)
+ self.app = app_launcher.app
+ # force franka registration
+ import isaaclab_playground.tasks.manipulation.libero.config.franka # noqa
+
+ def _init_env(self, task_id=0):
+ """Initializes the Isaac Sim environment."""
+
+ self.task_name = self.cfg.get("task_name")
+ self.task_id = task_id
+ # FIXME since isaac use env to set task id, all env have to use the same task id
+ if self.task_suite_name.startswith("libero"):
+ os.environ["LIBERO_TASK_SUITE"] = self.task_suite_name
+ os.environ["LIBERO_TASK_ID"] = str(task_id)
+ os.environ["LIBERO_OSC_TYPE"] = "pose_rel"
+
+ if not self.task_name:
+ self.task_name = "Isaac-Libero-Franka-OscPose-v0"
+
+ from isaaclab_tasks.utils import parse_env_cfg
+
+ self.env_cfg = parse_env_cfg(self.task_name, num_envs=self.num_envs)
+ self.env_cfg.env_name = self.cfg.get("env_name", str(self.task_id))
+ self.env_cfg.sim.device = self.device
+ self.env_cfg.sim.physx.enable_ccd = True
+ self.env_cfg.terminations.time_out = None
+ self.env_cfg.observations.policy.concatenate_terms = False
+
+ # create environment from loaded config
+ if self.env:
+ self.env.close()
+ omni.usd.get_context().new_stage()
+ self.env = gym.make(self.task_name, cfg=self.env_cfg).unwrapped
+
+ if self.cfg.video_cfg.save_video:
+ video_dir = os.path.join(self.cfg.video_cfg.video_base_dir, f"rank_{self.rank}")
+ os.makedirs(video_dir, exist_ok=True)
+
+ self.action_space = self.env.action_space
+ self.observation_space = self.env.observation_space
+
+ # TODO support other task suite
+ if self.task_suite_name.startswith("libero"):
+ self.task_descriptions = self.env.cfg.libero_config.task_info["language_instruction"]
+ assert self.env_cfg.osc_type == "pose_rel", (
+ f"Only pose_rel osc type is supported for libero. Received: {self.env_cfg.osc_type}"
+ )
+ else:
+ raise ValueError(f"Task suite {self.task_suite_name} is not supported.")
+ logger.info("Isaac Sim environment initialized")
+
+ def _init_metrics(self):
+ self.success_once = np.zeros(self.num_envs, dtype=bool)
+ self.returns = np.zeros(self.num_envs)
+
+ def _reset_metrics(self, env_idx=None):
+ if env_idx is not None:
+ mask = np.zeros(self.num_envs, dtype=bool)
+ mask[env_idx] = True
+ self.prev_step_reward[mask] = 0.0
+ self.success_once[mask] = False
+ self.returns[mask] = 0
+ self._elapsed_steps[env_idx] = 0
+ else:
+ self.prev_step_reward[:] = 0
+ self.success_once[:] = False
+ self.returns[:] = 0.0
+ self._elapsed_steps[:] = 0
+
+ def _record_metrics(self, step_reward, terminations, infos):
+ episode_info = {}
+ self.returns += step_reward
+ # Ensure terminations is a numpy array before the bitwise OR
+ if isinstance(terminations, torch.Tensor):
+ terminations = terminations.cpu().numpy()
+ self.success_once = self.success_once | terminations
+ episode_info["success_once"] = self.success_once.copy()
+ episode_info["return"] = self.returns.copy()
+ episode_info["episode_len"] = self.elapsed_steps.copy()
+ if any(self.elapsed_steps > 0):
+ episode_info["reward"] = episode_info["return"] / self.elapsed_steps
+ else:
+ episode_info["reward"] = 0
+ infos["episode"] = to_tensor(episode_info)
+ return infos
+
+ def reset(self, env_idx: Optional[int | list[int] | np.ndarray] = None, options: Optional[dict] = None):
+ if env_idx is None:
+ env_idx = np.arange(self.num_envs)
+
+ raw_obs, infos = self.env.reset()
+
+ obs = self._wrap_obs(raw_obs)
+
+ self._reset_metrics(env_idx)
+
+ return obs, infos
+
+ def step(self, actions=None):
+ if actions is None:
+ # isaac should start with reset_envs_to_initial_state
+ # do nothing for None
+ return (None, None, None, None, None)
+
+ truncations = self.elapsed_steps >= self.max_episode_steps
+ # _actions = torch.zeros(self.action_space.shape)
+
+ if isinstance(actions, np.ndarray):
+ actions = torch.from_numpy(actions)
+
+ self._elapsed_steps += 1
+ raw_obs, _reward, terminations, _, infos = self.env.step(actions)
+ self.last_obs = raw_obs
+ self.last_infos = infos
+
+ obs = self._wrap_obs(raw_obs)
+
+ step_reward = self._calc_step_reward(_reward.cpu().numpy())
+
+ if self.video_cfg.save_video:
+ plot_infos = {
+ "rewards": step_reward,
+ "terminations": terminations,
+ "task": self.task_descriptions,
+ }
+ self.add_new_frames(obs, plot_infos)
+
+ infos = self._record_metrics(step_reward, terminations, infos)
+
+ return (
+ obs,
+ to_tensor(step_reward),
+ to_tensor(terminations),
+ to_tensor(truncations),
+ infos,
+ )
+
+ def chunk_step(self, chunk_actions):
+ # chunk_actions: [num_envs, chunk_step, action_dim]
+ chunk_size = chunk_actions.shape[1]
+
+ chunk_rewards = []
+
+ raw_chunk_terminations = []
+ raw_chunk_truncations = []
+ for i in range(chunk_size):
+ actions = chunk_actions[:, i]
+ extracted_obs, step_reward, terminations, truncations, infos = self.step(actions)
+
+ chunk_rewards.append(step_reward)
+ raw_chunk_terminations.append(terminations)
+ raw_chunk_truncations.append(truncations)
+
+ chunk_rewards = torch.stack(chunk_rewards, dim=1) # [num_envs, chunk_steps]
+ raw_chunk_terminations = torch.stack(raw_chunk_terminations, dim=1) # [num_envs, chunk_steps]
+ raw_chunk_truncations = torch.stack(raw_chunk_truncations, dim=1) # [num_envs, chunk_steps]
+
+ chunk_terminations = raw_chunk_terminations.clone()
+ chunk_truncations = raw_chunk_truncations.clone()
+ return (
+ extracted_obs,
+ chunk_rewards,
+ chunk_terminations,
+ chunk_truncations,
+ infos,
+ )
+
+ def _calc_step_reward(self, reward):
+ if self.use_rel_reward:
+ reward_diff = reward - self.prev_step_reward
+ self.prev_step_reward = reward
+ return reward_diff
+ else:
+ return reward
+
+ def _wrap_obs(self, raw_obs):
+ images_and_states = self._extract_image_and_state(raw_obs)
+
+ obs = {
+ "images_and_states": to_tensor(images_and_states),
+ "task_descriptions": [self.task_descriptions] * self.num_envs,
+ }
+ return obs
+
+ def _extract_image_and_state(self, obs):
+ # TODO support multiple camera
+ camera_name = self.camera_name[0]
+ for key in self.env.unwrapped.scene.keys():
+ if key.startswith(camera_name):
+ cam = self.env.unwrapped.scene[key]
+ break
+ assert cam is not None, f"camera {camera_name} not found in scene"
+
+ rgb = cam.data.output["rgb"]
+
+ full_image = rgb.cpu().numpy()
+ return {
+ "full_image": full_image,
+ "state": np.concatenate(
+ [
+ obs["policy"]["eef_pose"].cpu(),
+ # quat2axisangle(obs["robot0_eef_quat"]), # isaac do not return robot0_eef_quat
+ # obs["policy"]["gripper_pos"].cpu(),
+ ],
+ axis=-1,
+ ),
+ }
+
+ def add_new_frames(self, obs, plot_infos):
+ images = []
+ for env_id, img in enumerate(obs["images_and_states"]["full_image"]):
+ info_item = {k: v if np.size(v) == 1 else v[env_id] for k, v in plot_infos.items()}
+ img = put_info_on_image(img.cpu().numpy(), info_item)
+ images.append(img)
+ full_image = tile_images(images, nrows=int(np.sqrt(self.num_envs)))
+ self.render_images.append(full_image)
+
+ def flush_video(self, video_sub_dir: Optional[str] = None):
+ output_dir = os.path.join(self.video_cfg.video_base_dir, f"rank_{self.rank}")
+ if video_sub_dir is not None:
+ output_dir = os.path.join(output_dir, f"{video_sub_dir}")
+ save_rollout_video(
+ self.render_images,
+ output_dir=output_dir,
+ video_name=f"{self.video_cnt}",
+ )
+ self.video_cnt += 1
+ self.render_images = []
+
+ def close(self):
+ if self.env is not None:
+ self.env.close()
+ self.app.close()
+
+ def load_state(self, state_buffer: bytes):
+ self.env.load_state(state_buffer)
+
+ def get_state(self):
+ return None
+
+ def reset_envs_to_state_ids(self, state_ids_list, task_ids_list):
+ logger.info(f"IsaacEnv reset_envs_to_state_ids task_ids_list: {task_ids_list}")
+ assert len(set(task_ids_list)) == 1, "Isaac env only support single task"
+
+ self._init_env(task_ids_list[0])
+
+ # In Isaac, reset to random status in groups to have more test coverage
+ # TODO support reset in group with options = {"group": len(set(state_ids_list))}
+ raw_obs, infos = self.env.reset()
+ env_idx = np.arange(self.num_envs)
+ self._reset_metrics(env_idx)
+
+ self.elapsed_steps = np.zeros(self.num_envs, dtype=np.int32)
+
+ # stablize the environment
+ for _ in range(10):
+ zero_actions = torch.zeros((self.num_envs, self.action_dim), device=self.device)
+ raw_obs, _, _, _, infos = self.env.step(zero_actions)
+
+ obs = self._wrap_obs(raw_obs)
+ return obs, infos
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2171666b035340542d81212af41b2ca1f96fab69
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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.
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/libero_env.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/libero_env.py
new file mode 100644
index 0000000000000000000000000000000000000000..2fc78444e0acb01df9f9e6b6d6f3bfcaaeb74770
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/libero_env.py
@@ -0,0 +1,413 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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.
+
+import logging
+import os
+from typing import Optional
+
+import gymnasium as gym
+import numpy as np
+import torch
+from libero.libero import get_libero_path
+from libero.libero.benchmark import Benchmark, get_benchmark
+from libero.libero.envs import OffScreenRenderEnv
+from omegaconf.omegaconf import OmegaConf
+
+from verl.experimental.vla.envs.action_utils import (
+ list_of_dict_to_dict_of_list,
+ put_info_on_image,
+ save_rollout_video,
+ tile_images,
+ to_tensor,
+)
+from verl.experimental.vla.envs.libero_env.utils import (
+ get_libero_image,
+)
+from verl.experimental.vla.envs.libero_env.venv import ReconfigureSubprocEnv
+
+logger = logging.getLogger(__name__)
+
+
+def patched_get_task_init_states(self, i):
+ init_states_path = os.path.join(
+ get_libero_path("init_states"),
+ self.tasks[i].problem_folder,
+ self.tasks[i].init_states_file,
+ )
+ init_states = torch.load(init_states_path, weights_only=False)
+ return init_states
+
+
+Benchmark.get_task_init_states = patched_get_task_init_states
+
+
+class LiberoEnv(gym.Env):
+ def __init__(self, cfg, rank, world_size):
+ self.rank = rank
+ self.cfg = cfg
+ self.world_size = world_size
+ self.seed = self.cfg.seed + rank
+ self.num_envs = self.cfg.num_envs
+
+ self.ignore_terminations = False
+
+ self._generator = np.random.default_rng(seed=self.seed)
+ self._generator_ordered = np.random.default_rng(seed=0)
+ self.start_idx = 0
+
+ self.task_suite: Benchmark = get_benchmark(cfg.task_suite_name)()
+
+ self._compute_total_num_group_envs()
+ self.reset_state_ids_all = self.get_reset_state_ids_all()
+ self.reset_state_ids = self._get_ordered_reset_state_ids(self.num_envs)
+ self._init_task_and_trial_ids()
+ self._init_env()
+
+ self.prev_step_reward = np.zeros(self.num_envs)
+ self.use_rel_reward = False
+
+ self._init_metrics()
+ self._elapsed_steps = np.zeros(self.num_envs, dtype=np.int32)
+
+ self.video_cfg = cfg.video_cfg
+ self.video_cnt = 0
+ self.render_images = []
+
+ @property
+ def elapsed_steps(self):
+ return self._elapsed_steps
+
+ def get_all_state_ids(self):
+ """Returns all possible state IDs from the entire benchmark."""
+ return np.arange(self.total_num_group_envs) # (total_num_states,)
+
+ def _init_env(self):
+ env_fns = self.get_env_fns()
+ self.env = ReconfigureSubprocEnv(env_fns)
+
+ def get_env_fns(self):
+ env_fn_params = self.get_env_fn_params()
+ env_fns = []
+ for env_fn_param in env_fn_params:
+
+ def env_fn(param=env_fn_param):
+ seed = param.pop("seed")
+ env = OffScreenRenderEnv(**param)
+ env.seed(seed)
+ return env
+
+ env_fns.append(env_fn)
+ return env_fns
+
+ def get_env_fn_params(self, env_idx=None):
+ env_fn_params = []
+ base_env_args = OmegaConf.to_container(self.cfg.init_params, resolve=True)
+
+ task_descriptions = []
+ if env_idx is None:
+ env_idx = np.arange(self.cfg.num_envs)
+ for env_id in range(self.cfg.num_envs):
+ if env_id not in env_idx:
+ task_descriptions.append(self.task_descriptions[env_id])
+ continue
+ task = self.task_suite.get_task(self.task_ids[env_id])
+ task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file)
+ env_fn_params.append(
+ {
+ **base_env_args,
+ "bddl_file_name": task_bddl_file,
+ "seed": self.seed,
+ }
+ )
+ task_descriptions.append(task.language)
+ self.task_descriptions = task_descriptions
+ return env_fn_params
+
+ def _compute_total_num_group_envs(self):
+ self.total_num_group_envs = 0
+ self.trial_id_bins = []
+ for task_id in range(self.task_suite.get_num_tasks()):
+ task_num_trials = len(self.task_suite.get_task_init_states(task_id))
+ self.trial_id_bins.append(task_num_trials)
+
+ self.total_num_group_envs += task_num_trials
+
+ self.cumsum_trial_id_bins = np.cumsum(self.trial_id_bins)
+
+ def _init_task_and_trial_ids(self):
+ self.task_ids, self.trial_ids = self._get_task_and_trial_ids_from_reset_state_ids(self.reset_state_ids)
+
+ def _get_random_reset_state_ids(self, num_reset_states):
+ reset_state_ids = self._generator.integers(low=0, high=self.total_num_group_envs, size=(num_reset_states,))
+ return reset_state_ids
+
+ def get_reset_state_ids_all(self):
+ reset_state_ids = np.arange(self.total_num_group_envs)
+ valid_size = len(reset_state_ids) - (len(reset_state_ids) % self.world_size)
+ if not self.cfg.only_eval:
+ self._generator_ordered.shuffle(reset_state_ids)
+ reset_state_ids = reset_state_ids[:valid_size]
+ reset_state_ids = reset_state_ids.reshape(self.world_size, -1)
+ return reset_state_ids
+
+ def _get_ordered_reset_state_ids(self, num_reset_states):
+ reset_state_ids = self.reset_state_ids_all[self.rank][self.start_idx : self.start_idx + num_reset_states]
+ self.start_idx = self.start_idx + num_reset_states
+ if self.start_idx >= len(self.reset_state_ids_all[0]):
+ self.reset_state_ids_all = self.get_reset_state_ids_all()
+ self.start_idx = 0
+ return reset_state_ids
+
+ def _get_task_and_trial_ids_from_reset_state_ids(self, reset_state_ids):
+ task_ids = []
+ trial_ids = []
+ # get task id and trial id from reset state ids
+ for reset_state_id in reset_state_ids:
+ start_pivot = 0
+ for task_id, end_pivot in enumerate(self.cumsum_trial_id_bins):
+ if reset_state_id < end_pivot and reset_state_id >= start_pivot:
+ task_ids.append(task_id)
+ trial_ids.append(reset_state_id - start_pivot)
+ break
+ start_pivot = end_pivot
+ logger.debug(
+ "get task and trial id",
+ self.cumsum_trial_id_bins,
+ reset_state_ids,
+ task_ids,
+ trial_ids,
+ )
+ return np.array(task_ids), np.array(trial_ids)
+
+ def _get_reset_states(self, env_idx):
+ if env_idx is None:
+ env_idx = np.arange(self.num_envs)
+ init_state = [
+ self.task_suite.get_task_init_states(self.task_ids[env_id])[self.trial_ids[env_id]] for env_id in env_idx
+ ]
+ return init_state
+
+ def _init_metrics(self):
+ self.success_once = np.zeros(self.num_envs, dtype=bool)
+ self.fail_once = np.zeros(self.num_envs, dtype=bool)
+ self.returns = np.zeros(self.num_envs)
+
+ def _reset_metrics(self, env_idx=None):
+ if env_idx is not None:
+ mask = np.zeros(self.num_envs, dtype=bool)
+ mask[env_idx] = True
+ self.prev_step_reward[mask] = 0.0
+ self.success_once[mask] = False
+ self.fail_once[mask] = False
+ self.returns[mask] = 0
+ self._elapsed_steps[env_idx] = 0
+ else:
+ self.prev_step_reward[:] = 0
+ self.success_once[:] = False
+ self.fail_once[:] = False
+ self.returns[:] = 0.0
+ self._elapsed_steps[:] = 0
+
+ def _record_metrics(self, step_reward, terminations, infos):
+ episode_info = {}
+ self.returns += step_reward
+ self.success_once = self.success_once | terminations
+ episode_info["success_once"] = self.success_once.copy()
+ episode_info["return"] = self.returns.copy()
+ episode_info["episode_len"] = self.elapsed_steps.copy()
+ episode_info["reward"] = episode_info["return"] / episode_info["episode_len"]
+ infos["episode"] = to_tensor(episode_info)
+ return infos
+
+ def _extract_image_and_state(self, obs):
+ return {
+ "full_image": get_libero_image(obs),
+ "state": np.concatenate(
+ [
+ obs["robot0_eef_pos"],
+ # quat2axisangle(obs["robot0_eef_quat"]),
+ # obs["robot0_gripper_qpos"],
+ ]
+ ),
+ }
+
+ def _wrap_obs(self, obs_list):
+ images_and_states_list = []
+ for obs in obs_list:
+ images_and_states = self._extract_image_and_state(obs)
+ images_and_states_list.append(images_and_states)
+
+ obs = {
+ "images_and_states": to_tensor(list_of_dict_to_dict_of_list(images_and_states_list)),
+ "task_descriptions": self.task_descriptions,
+ }
+ return obs
+
+ def _reconfigure(self, reset_state_ids, env_idx):
+ reconfig_env_idx = []
+ task_ids, trial_ids = self._get_task_and_trial_ids_from_reset_state_ids(reset_state_ids)
+ for j, env_id in enumerate(env_idx):
+ if self.task_ids[env_id] != task_ids[j]:
+ reconfig_env_idx.append(env_id)
+ self.task_ids[env_id] = task_ids[j]
+ self.trial_ids[env_id] = trial_ids[j]
+ if reconfig_env_idx:
+ env_fn_params = self.get_env_fn_params(reconfig_env_idx)
+ self.env.reconfigure_env_fns(env_fn_params, reconfig_env_idx)
+
+ self.env.seed([0] * len(env_idx))
+ self.env.reset(id=env_idx)
+ init_state = self._get_reset_states(env_idx=env_idx)
+ self.env.set_init_state(init_state=init_state, id=env_idx)
+
+ def reset(
+ self,
+ env_idx: Optional[int | list[int] | np.ndarray] = None,
+ reset_state_ids=None,
+ options: Optional[dict] = None,
+ ):
+ if env_idx is None:
+ env_idx = np.arange(self.num_envs)
+
+ if reset_state_ids is None:
+ num_reset_states = len(env_idx)
+ reset_state_ids = self._get_random_reset_state_ids(num_reset_states)
+
+ self._reconfigure(reset_state_ids, env_idx)
+
+ for _ in range(10):
+ zero_actions = np.zeros((self.num_envs, 7))
+ raw_obs, _reward, terminations, info_lists = self.env.step(zero_actions)
+
+ obs = self._wrap_obs(raw_obs)
+ if env_idx is not None:
+ self._reset_metrics(env_idx)
+ else:
+ self._reset_metrics()
+ infos = {}
+ return obs, infos
+
+ def step(self, actions=None):
+ if actions is None:
+ obs, infos = self.reset(reset_state_ids=self.reset_state_ids)
+ terminations = np.zeros(self.num_envs, dtype=bool)
+ truncations = np.zeros(self.num_envs, dtype=bool)
+
+ return obs, None, to_tensor(terminations), to_tensor(truncations), infos
+
+ if isinstance(actions, torch.Tensor):
+ actions = actions.detach().cpu().numpy()
+
+ self._elapsed_steps += 1
+ raw_obs, _reward, terminations, info_lists = self.env.step(actions)
+ infos = list_of_dict_to_dict_of_list(info_lists)
+ truncations = self.elapsed_steps >= self.cfg.max_episode_steps
+
+ obs = self._wrap_obs(raw_obs)
+ step_reward = self._calc_step_reward(terminations)
+
+ if self.video_cfg.save_video:
+ plot_infos = {
+ "rewards": step_reward,
+ "terminations": terminations,
+ "task": self.task_descriptions,
+ }
+ self.add_new_frames(raw_obs, plot_infos)
+
+ infos = self._record_metrics(step_reward, terminations, infos)
+
+ return (
+ obs,
+ to_tensor(step_reward),
+ to_tensor(terminations),
+ to_tensor(truncations),
+ infos,
+ )
+
+ def chunk_step(self, chunk_actions):
+ # chunk_actions: [num_envs, chunk_step, action_dim]
+ chunk_size = chunk_actions.shape[1]
+
+ chunk_rewards = []
+
+ raw_chunk_terminations = []
+ raw_chunk_truncations = []
+ for i in range(chunk_size):
+ actions = chunk_actions[:, i]
+ extracted_obs, step_reward, terminations, truncations, infos = self.step(actions)
+
+ chunk_rewards.append(step_reward)
+ raw_chunk_terminations.append(terminations)
+ raw_chunk_truncations.append(truncations)
+
+ chunk_rewards = torch.stack(chunk_rewards, dim=1) # [num_envs, chunk_steps]
+ raw_chunk_terminations = torch.stack(raw_chunk_terminations, dim=1) # [num_envs, chunk_steps]
+ raw_chunk_truncations = torch.stack(raw_chunk_truncations, dim=1) # [num_envs, chunk_steps]
+
+ chunk_terminations = raw_chunk_terminations.clone()
+ chunk_truncations = raw_chunk_truncations.clone()
+ return (
+ extracted_obs,
+ chunk_rewards,
+ chunk_terminations,
+ chunk_truncations,
+ infos,
+ )
+
+ def _calc_step_reward(self, terminations):
+ reward = self.cfg.reward_coef * terminations
+ reward_diff = reward - self.prev_step_reward
+ self.prev_step_reward = reward
+
+ if self.use_rel_reward:
+ return reward_diff
+ else:
+ return reward
+
+ def add_new_frames(self, raw_obs, plot_infos):
+ images = []
+ for env_id, raw_single_obs in enumerate(raw_obs):
+ info_item = {k: v if np.size(v) == 1 else v[env_id] for k, v in plot_infos.items()}
+ img = raw_single_obs["agentview_image"][::-1, ::-1]
+ img = put_info_on_image(img, info_item)
+ images.append(img)
+ full_image = tile_images(images, nrows=int(np.sqrt(self.num_envs)))
+ self.render_images.append(full_image)
+
+ def flush_video(self, video_sub_dir: Optional[str] = None):
+ output_dir = os.path.join(self.video_cfg.video_base_dir, f"rank_{self.rank}")
+ if video_sub_dir is not None:
+ output_dir = os.path.join(output_dir, f"{video_sub_dir}")
+ save_rollout_video(
+ self.render_images,
+ output_dir=output_dir,
+ video_name=f"{self.video_cnt}",
+ )
+ self.video_cnt += 1
+ self.render_images = []
+
+ def reset_envs_to_state_ids(self, state_ids_list, task_ids_list):
+ """Reset environments to specified state IDs.
+
+ Args:
+ state_ids_list: List of state IDs to reset environments to
+ """
+ env_idx = np.arange(len(state_ids_list))
+ obs, infos = self.reset(env_idx=env_idx, reset_state_ids=state_ids_list)
+ return obs, infos
+
+ def load_state(self, state_buffer: bytes):
+ self.env.load_state(state_buffer)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/utils.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..870741e6e3f1f68753101ed921499429adb3e2b4
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/utils.py
@@ -0,0 +1,138 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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.
+
+"""Utils for evaluating policies in LIBERO simulation environments."""
+
+import math
+
+import numpy as np
+
+from verl.experimental.vla.envs.action_utils import resize_image
+
+
+def get_libero_image(obs: dict[str, np.ndarray]) -> np.ndarray:
+ """
+ Extracts image from observations and preprocesses it.
+
+ Args:
+ obs: Observation dictionary from LIBERO environment
+
+ Returns:
+ Preprocessed image as numpy array
+ """
+ img = obs["agentview_image"]
+ img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing
+ return img
+
+
+def get_libero_wrist_image(obs: dict[str, np.ndarray], resize_size: int | tuple[int, int]) -> np.ndarray:
+ """
+ Extracts wrist camera image from observations and preprocesses it.
+
+ Args:
+ obs: Observation dictionary from LIBERO environment
+ resize_size: Target size for resizing
+
+ Returns:
+ Preprocessed wrist camera image as numpy array
+ """
+ assert isinstance(resize_size, int) or isinstance(resize_size, tuple)
+ if isinstance(resize_size, int):
+ resize_size = (resize_size, resize_size)
+ img = obs["robot0_eye_in_hand_image"]
+ img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing
+ img = resize_image(img, resize_size)
+ return img
+
+
+def quat2axisangle(quat: np.ndarray) -> np.ndarray:
+ """
+ Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
+
+ Converts quaternion to axis-angle format.
+ Returns a unit vector direction scaled by its angle in radians.
+
+ Args:
+ quat (np.array): (x,y,z,w) vec4 float angles
+
+ Returns:
+ np.array: (ax,ay,az) axis-angle exponential coordinates
+ """
+ # clip quaternion
+ if quat[3] > 1.0:
+ quat[3] = 1.0
+ elif quat[3] < -1.0:
+ quat[3] = -1.0
+
+ den = np.sqrt(1.0 - quat[3] * quat[3])
+ if math.isclose(den, 0.0):
+ # This is (close to) a zero degree rotation, immediately return
+ return np.zeros(3)
+
+ return (quat[:3] * 2.0 * math.acos(quat[3])) / den
+
+
+def normalize_gripper_action(action: np.ndarray, binarize: bool = True) -> np.ndarray:
+ """
+ Normalize gripper action from [0,1] to [-1,+1] range.
+
+ This is necessary for some environments because the dataset wrapper
+ standardizes gripper actions to [0,1]. Note that unlike the other action
+ dimensions, the gripper action is not normalized to [-1,+1] by default.
+
+ Normalization formula: y = 2 * (x - orig_low) / (orig_high - orig_low) - 1
+
+ Args:
+ action: Action array with gripper action in the last dimension
+ binarize: Whether to binarize gripper action to -1 or +1
+
+ Returns:
+ np.ndarray: Action array with normalized gripper action
+ """
+ # Create a copy to avoid modifying the original
+ normalized_action = action.copy()
+
+ # Normalize the last action dimension to [-1,+1]
+ orig_low, orig_high = 0.0, 1.0
+ normalized_action[..., -1] = 2 * (normalized_action[..., -1] - orig_low) / (orig_high - orig_low) - 1
+
+ if binarize:
+ # Binarize to -1 or +1
+ normalized_action[..., -1] = np.sign(normalized_action[..., -1])
+
+ return normalized_action
+
+
+def invert_gripper_action(action: np.ndarray) -> np.ndarray:
+ """
+ Flip the sign of the gripper action (last dimension of action vector).
+
+ This is necessary for environments where -1 = open, +1 = close, since
+ the RLDS dataloader aligns gripper actions such that 0 = close, 1 = open.
+
+ Args:
+ action: Action array with gripper action in the last dimension
+
+ Returns:
+ np.ndarray: Action array with inverted gripper action
+ """
+ # Create a copy to avoid modifying the original
+ inverted_action = action.copy()
+
+ # Invert the gripper action
+ inverted_action[..., -1] = inverted_action[..., -1] * -1.0
+
+ return inverted_action
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/venv.py b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/venv.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f9a835e43068666c66ed6e6af1636d40d4f4737
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/envs/libero_env/venv.py
@@ -0,0 +1,162 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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 multiprocessing import Pipe, connection
+from multiprocessing.context import Process
+from typing import Any, Callable, Optional
+
+import gymnasium as gym
+import numpy as np
+from libero.libero.envs import OffScreenRenderEnv
+from libero.libero.envs.venv import (
+ BaseVectorEnv,
+ CloudpickleWrapper,
+ EnvWorker,
+ ShArray,
+ SubprocEnvWorker,
+ SubprocVectorEnv,
+ _setup_buf,
+)
+
+
+def _worker(
+ parent: connection.Connection,
+ p: connection.Connection,
+ env_fn_wrapper: CloudpickleWrapper,
+ obs_bufs: Optional[dict | tuple | ShArray] = None,
+) -> None:
+ def _encode_obs(obs: dict | tuple | np.ndarray, buffer: dict | tuple | ShArray) -> None:
+ if isinstance(obs, np.ndarray) and isinstance(buffer, ShArray):
+ buffer.save(obs)
+ elif isinstance(obs, tuple) and isinstance(buffer, tuple):
+ for o, b in zip(obs, buffer, strict=False):
+ _encode_obs(o, b)
+ elif isinstance(obs, dict) and isinstance(buffer, dict):
+ for k in obs.keys():
+ _encode_obs(obs[k], buffer[k])
+ return None
+
+ parent.close()
+ env = env_fn_wrapper.data()
+ try:
+ while True:
+ try:
+ cmd, data = p.recv()
+ except EOFError: # the pipe has been closed
+ p.close()
+ break
+ if cmd == "step":
+ env_return = env.step(data)
+ if obs_bufs is not None:
+ _encode_obs(env_return[0], obs_bufs)
+ env_return = (None, *env_return[1:])
+ p.send(env_return)
+ elif cmd == "reset":
+ retval = env.reset(**data)
+ reset_returns_info = (
+ isinstance(retval, (tuple | list)) and len(retval) == 2 and isinstance(retval[1], dict)
+ )
+ if reset_returns_info:
+ obs, info = retval
+ else:
+ obs = retval
+ if obs_bufs is not None:
+ _encode_obs(obs, obs_bufs)
+ obs = None
+ if reset_returns_info:
+ p.send((obs, info))
+ else:
+ p.send(obs)
+ elif cmd == "close":
+ p.send(env.close())
+ p.close()
+ break
+ elif cmd == "render":
+ p.send(env.render(**data) if hasattr(env, "render") else None)
+ elif cmd == "seed":
+ if hasattr(env, "seed"):
+ p.send(env.seed(data))
+ else:
+ env.reset(seed=data)
+ p.send(None)
+ elif cmd == "getattr":
+ p.send(getattr(env, data) if hasattr(env, data) else None)
+ elif cmd == "setattr":
+ setattr(env.unwrapped, data["key"], data["value"])
+ elif cmd == "check_success":
+ p.send(env.check_success())
+ elif cmd == "get_segmentation_of_interest":
+ p.send(env.get_segmentation_of_interest(data))
+ elif cmd == "get_sim_state":
+ p.send(env.get_sim_state())
+ elif cmd == "set_init_state":
+ obs = env.set_init_state(data)
+ p.send(obs)
+ elif cmd == "reconfigure":
+ env.close()
+ seed = data.pop("seed")
+ env = OffScreenRenderEnv(**data)
+ env.seed(seed)
+ p.send(None)
+ else:
+ p.close()
+ raise NotImplementedError
+ except KeyboardInterrupt:
+ p.close()
+
+
+class ReconfigureSubprocEnvWorker(SubprocEnvWorker):
+ def __init__(self, env_fn: Callable[[], gym.Env], share_memory: bool = False):
+ self.parent_remote, self.child_remote = Pipe()
+ self.share_memory = share_memory
+ self.buffer: Optional[dict | tuple | ShArray] = None
+ if self.share_memory:
+ dummy = env_fn()
+ obs_space = dummy.observation_space
+ dummy.close()
+ del dummy
+ self.buffer = _setup_buf(obs_space)
+ args = (
+ self.parent_remote,
+ self.child_remote,
+ CloudpickleWrapper(env_fn),
+ self.buffer,
+ )
+ self.process = Process(target=_worker, args=args, daemon=True)
+ self.process.start()
+ self.child_remote.close()
+ EnvWorker.__init__(self, env_fn)
+
+ def reconfigure_env_fn(self, env_fn_param):
+ self.parent_remote.send(["reconfigure", env_fn_param])
+ return self.parent_remote.recv()
+
+
+class ReconfigureSubprocEnv(SubprocVectorEnv):
+ def __init__(self, env_fns: list[Callable[[], gym.Env]], **kwargs: Any) -> None:
+ def worker_fn(fn: Callable[[], gym.Env]) -> ReconfigureSubprocEnvWorker:
+ return ReconfigureSubprocEnvWorker(fn, share_memory=False)
+
+ BaseVectorEnv.__init__(self, env_fns, worker_fn, **kwargs)
+
+ def reconfigure_env_fns(self, env_fns, id=None):
+ self._assert_is_not_closed()
+ id = self._wrap_id(id)
+ if self.is_async:
+ self._assert_id(id)
+
+ for j, i in enumerate(id):
+ self.workers[i].reconfigure_env_fn(env_fns[j])
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/fsdp_workers.py b/code/RL_model/verl/verl_train/verl/experimental/vla/fsdp_workers.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f2d463e5239d151b265f9e79a691fc17a43a338
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/fsdp_workers.py
@@ -0,0 +1,259 @@
+# Copyright 2024 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.
+"""
+The main entry point to run the PPO algorithm
+"""
+
+import asyncio
+import contextlib
+import logging
+import os
+
+import torch
+import torch.distributed
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp._unshard_param_utils import _get_module_fsdp_state, _unshard_params_for_summon
+from torch.distributed.fsdp.api import FullStateDictConfig, ShardedStateDictConfig, StateDictType
+
+from verl import DataProto
+from verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register
+from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager
+from verl.utils.config import omega_conf_to_dataclass
+from verl.utils.device import get_device_id, get_device_name, get_torch_device, set_expandable_segments
+from verl.utils.flops_counter import FlopsCounter
+from verl.utils.fsdp_utils import fsdp_version
+from verl.utils.import_utils import import_external_libs
+from verl.utils.memory_utils import aggressive_empty_cache
+from verl.utils.profiler import DistProfiler, log_gpu_memory_usage, simple_timer
+from verl.utils.profiler.performance import reduce_timing, topk_reduce_ratio_min_max
+from verl.workers.config import HFModelConfig
+from verl.workers.fsdp_workers import ActorRolloutRefWorker
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+device_name = get_device_name()
+
+
+class RobActorRolloutRefWorker(ActorRolloutRefWorker):
+ """
+ This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy
+ or a hybrid engine based on the config.rollout
+ """
+
+ fsdp_unshard_exit_stack = contextlib.ExitStack()
+
+ def _build_rollout(self, trust_remote_code=False):
+ from verl.experimental.vla.naive_rollout_rob import NaiveRolloutRob
+
+ self.base_sync_done = False
+ world_size = torch.distributed.get_world_size()
+ dp = world_size
+ infer_tp = self.config.rollout.tensor_model_parallel_size
+ rollout_device_mesh = init_device_mesh(
+ device_name, mesh_shape=(dp, infer_tp), mesh_dim_names=["dp", "infer_tp"]
+ )
+ # 3. init trainer and rollout random states
+ self.torch_random_states = get_torch_device().get_rng_state()
+ gen_dp_rank = rollout_device_mesh["dp"].get_local_rank()
+ get_torch_device().manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states
+ self.gen_random_states = get_torch_device().get_rng_state()
+ get_torch_device().set_rng_state(self.torch_random_states)
+
+ if torch.distributed.get_world_size() == 1 and fsdp_version(self.actor_module_fsdp) == 1:
+ FSDP.set_state_dict_type(
+ self.actor_module_fsdp,
+ state_dict_type=StateDictType.FULL_STATE_DICT,
+ state_dict_config=FullStateDictConfig(),
+ )
+ elif fsdp_version(self.actor_module_fsdp) == 1:
+ FSDP.set_state_dict_type(
+ self.actor_module_fsdp,
+ state_dict_type=StateDictType.SHARDED_STATE_DICT,
+ state_dict_config=ShardedStateDictConfig(),
+ )
+ else:
+ raise NotImplementedError(f"Unsupported fsdp version {fsdp_version(self.actor_module_fsdp)}")
+
+ self._register_dispatch_collect_info("rollout", dp_rank=self.rank, is_collect=True)
+ self.rollout = NaiveRolloutRob(module=self.actor_module_fsdp, model_config=self.config.model)
+
+ model_config: HFModelConfig = omega_conf_to_dataclass(self.config.model, dataclass_type=HFModelConfig)
+ self.model_config = model_config
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def switch_to_rollout(self):
+ loop = asyncio.get_event_loop()
+ loop.run_until_complete(self.rollout_mode())
+ log_gpu_memory_usage("After switch to rollout mode", logger=logger)
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def switch_to_train(self):
+ loop = asyncio.get_event_loop()
+ loop.run_until_complete(self.trainer_mode())
+ log_gpu_memory_usage("After switch to trainer mode", logger=logger)
+
+ async def rollout_mode(self):
+ """Context switch hybridengine to rollout mode."""
+ aggressive_empty_cache(force_sync=True)
+ fsdp_unshard_exit_stack = contextlib.ExitStack()
+ optional_state = _get_module_fsdp_state(self.actor_module_fsdp)
+ if optional_state is None:
+ self.fsdp_unshard_exit_stack = fsdp_unshard_exit_stack
+ states_and_modules = ([optional_state], [self.actor_module_fsdp])
+
+ self.base_sync_done = True
+ # important: need to manually set the random states of each tp to be identical.
+ self.torch_random_states = get_torch_device().get_rng_state()
+ get_torch_device().set_rng_state(self.gen_random_states)
+ for state, fsdp_module in zip(*states_and_modules, strict=False):
+ fsdp_unshard_exit_stack.enter_context(
+ _unshard_params_for_summon(
+ module=fsdp_module,
+ state=state,
+ writeback=False,
+ rank0_only=False,
+ offload_to_cpu=False,
+ with_grads=False,
+ )
+ )
+
+ self.fsdp_unshard_exit_stack = fsdp_unshard_exit_stack
+ logger.info("rollout mode")
+
+ async def trainer_mode(self):
+ """Context switch hybridengine to trainer mode."""
+
+ 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)
+ if self.fsdp_unshard_exit_stack is not None:
+ self.fsdp_unshard_exit_stack.close()
+ self.fsdp_unshard_exit_stack = None
+ logger.info("trainer mode")
+
+ @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="rollout"), blocking=False)
+ @DistProfiler.annotate(color="red", role="rollout_generate")
+ def generate_sequences(self, prompts: DataProto):
+ # Support all hardwares
+ assert self._is_rollout
+ prompts = prompts.to(get_device_id())
+
+ meta_info = {
+ "eos_token_id": self.model_config.generation_config.eos_token_id
+ if self.model_config.generation_config is not None
+ else self.model_config.tokenizer.eos_token_id,
+ "pad_token_id": self.model_config.generation_config.pad_token_id
+ if self.model_config.generation_config is not None
+ else self.model_config.tokenizer.pad_token_id,
+ }
+ prompts.meta_info.update(meta_info)
+
+ timing_generate = {}
+
+ with simple_timer("generate_sequences", timing_generate):
+ output = self.rollout.generate_sequences(prompts=prompts)
+
+ timing_generate_topk_ratio, timing_generate_min, timing_generate_max = topk_reduce_ratio_min_max(
+ timing_generate["generate_sequences"]
+ )
+ timing_generate = reduce_timing(timing_generate)
+ timing_generate.update(
+ {
+ "generation_timing/max": timing_generate_max,
+ "generation_timing/min": timing_generate_min,
+ "generation_timing/topk_ratio": timing_generate_topk_ratio,
+ }
+ )
+ output.meta_info["metrics"] = timing_generate
+ output = output.to("cpu")
+
+ # clear kv cache
+ get_torch_device().empty_cache()
+ return output
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def init_model(self):
+ from verl.experimental.vla.dp_rob import RobDataParallelPPOActor
+
+ # This is used to import external_lib into the huggingface systems
+ import_external_libs(self.config.model.get("external_lib", None))
+
+ from omegaconf import OmegaConf
+
+ override_model_config = OmegaConf.to_container(self.config.model.get("override_config", OmegaConf.create()))
+ from transformers import AutoConfig, AutoImageProcessor, AutoModelForVision2Seq, AutoProcessor
+
+ from verl.experimental.vla.models.openvla_oft.configuration_prismatic import OpenVLAConfig
+ from verl.experimental.vla.models.openvla_oft.modeling_prismatic import OpenVLAForActionPrediction
+ from verl.experimental.vla.models.openvla_oft.processing_prismatic import (
+ PrismaticImageProcessor,
+ PrismaticProcessor,
+ )
+
+ AutoConfig.register("openvla", OpenVLAConfig)
+ AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor)
+ AutoProcessor.register(OpenVLAConfig, PrismaticProcessor)
+ AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction)
+ if self._is_actor or self._is_rollout:
+ # we need the model for actor and rollout
+ if self._is_actor:
+ optim_config = self.config.actor.optim
+ fsdp_config = self.config.actor.fsdp_config
+ else:
+ optim_config = None
+ fsdp_config = OmegaConf.create()
+ self.actor_module_fsdp, self.actor_optimizer, self.actor_lr_scheduler, self.actor_model_config = (
+ self._build_model_optimizer(
+ model_path=self.config.model.path,
+ fsdp_config=fsdp_config,
+ optim_config=optim_config,
+ override_model_config=override_model_config,
+ enable_gradient_checkpointing=self.config.model.get("enable_gradient_checkpointing", False),
+ trust_remote_code=self.config.model.get("trust_remote_code", False),
+ )
+ )
+
+ if fsdp_version(self.actor_module_fsdp) == 1:
+ # get the original unwrapped module
+ self.actor_module = self.actor_module_fsdp._fsdp_wrapped_module
+
+ if self._is_actor:
+ OmegaConf.set_struct(self.config.actor, True)
+ self.actor = RobDataParallelPPOActor(
+ config=self.config.actor, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer
+ )
+
+ if self._is_rollout:
+ self._build_rollout(trust_remote_code=self.config.model.get("trust_remote_code", False))
+
+ if self._is_actor:
+ self.flops_counter = FlopsCounter(self.actor_model_config)
+ self.checkpoint_manager = FSDPCheckpointManager(
+ model=self.actor_module_fsdp,
+ optimizer=self.actor.actor_optimizer,
+ lr_scheduler=self.actor_lr_scheduler,
+ processing_class=self.processor if self.processor is not None else self.tokenizer,
+ checkpoint_config=self.config.actor.checkpoint,
+ )
+
+ torch.distributed.barrier()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/main_ppo.py b/code/RL_model/verl/verl_train/verl/experimental/vla/main_ppo.py
new file mode 100644
index 0000000000000000000000000000000000000000..633d0a08e9dfccb11f52c03d3971f77885685ee9
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/main_ppo.py
@@ -0,0 +1,171 @@
+# Copyright 2024 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.
+
+
+import logging
+
+import datasets
+import hydra
+import ray
+import torch
+from omegaconf import OmegaConf
+
+from verl import DataProto
+from verl.trainer.constants_ppo import get_ppo_ray_runtime_env
+from verl.trainer.ppo.ray_trainer import ResourcePoolManager
+from verl.trainer.ppo.utils import Role
+from verl.utils.device import is_cuda_available
+
+from .rob_ray_trainer import RobRayPPOTrainer
+
+logger = logging.getLogger(__name__)
+
+
+def calculate_reward(data: DataProto, return_dict: bool = False) -> torch.Tensor:
+ complete_tensor = data.batch["complete"]
+ batch_size, num_steps = complete_tensor.shape[:2]
+ traj_has_complete = torch.any(complete_tensor, dim=(1, 2)) # shape: [batch_size]
+ reward_per_traj = traj_has_complete.float()
+ reward_per_step = reward_per_traj.unsqueeze(1).expand(batch_size, num_steps)
+ if return_dict:
+ return {"reward_tensor": reward_per_step}
+ else:
+ return reward_per_step
+
+
+@hydra.main(config_path="config", config_name="rob_ppo_trainer", version_base=None)
+def main(config):
+ if not ray.is_initialized():
+ default_runtime_env = get_ppo_ray_runtime_env()
+ ray_init_kwargs = config.ray_kwargs.get("ray_init", {})
+ runtime_env_kwargs = ray_init_kwargs.get("runtime_env", {})
+ runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)
+ ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, "runtime_env": runtime_env})
+ logger.info(f"ray init kwargs: {ray_init_kwargs}")
+ ray.init(**OmegaConf.to_container(ray_init_kwargs))
+
+ # Apply controller nsight profiling if configured
+ if (
+ is_cuda_available
+ and config.global_profiler.tool == "nsys"
+ and config.global_profiler.get("steps") is not None
+ and len(config.global_profiler.get("steps", [])) > 0
+ ):
+ from verl.utils.import_utils import is_nvtx_available
+
+ assert is_nvtx_available(), "nvtx is not available in CUDA platform. Please 'pip3 install nvtx'"
+ nsight_options = OmegaConf.to_container(
+ config.global_profiler.global_tool_config.nsys.controller_nsight_options
+ )
+ main_task_with_options = main_task.options(runtime_env={"nsight": nsight_options})
+ ray.get(main_task_with_options.remote(config))
+ else:
+ ray.get(main_task.remote(config))
+
+ # [Optional] get the path of the timeline trace file from the configuration, default to None
+ # This file is used for performance analysis
+ timeline_json_file = config.ray_kwargs.get("timeline_json_file", None)
+ if timeline_json_file:
+ ray.timeline(filename=timeline_json_file)
+
+
+@ray.remote
+def main_task(config):
+ # print initial config
+ from pprint import pprint
+
+ from omegaconf import OmegaConf
+
+ from verl.utils.fs import copy_local_path_from_hdfs
+
+ pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values
+ OmegaConf.resolve(config)
+
+ # download the checkpoint from hdfs
+ local_path = copy_local_path_from_hdfs(config.actor_rollout_ref.model.path)
+
+ # instantiate tokenizer
+ from verl.utils import hf_tokenizer
+
+ tokenizer = hf_tokenizer(local_path)
+
+ # define worker classes
+ if config.actor_rollout_ref.actor.strategy in ["fsdp", "fsdp2"]:
+ assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
+ from verl.experimental.vla.workers.env.env_worker import EnvWorker
+ from verl.single_controller.ray import RayWorkerGroup
+
+ from .fsdp_workers import RobActorRolloutRefWorker
+
+ ray_worker_group_cls = RayWorkerGroup
+
+ else:
+ raise NotImplementedError
+
+ role_worker_mapping = {
+ # Role.Critic: ray.remote(RobActorRolloutRefWorker),
+ Role.ActorRollout: ray.remote(RobActorRolloutRefWorker),
+ # Role.RefPolicy: ray.remote(RobActorRolloutRefWorker),
+ Role.Env: ray.remote(EnvWorker),
+ }
+
+ train_rollout_pool_id = "train_rollout_pool"
+
+ num_nodes_actor_rollout = config.trainer.nnodes
+ train_rollout_gpu_num = config.trainer.n_rollout_gpus_per_node
+ env_gpu_num = config.trainer.n_env_gpus_per_node
+ if config.env.disagg_sim.enable:
+ # disaggregated sim and actor rollout
+ num_nodes_sim = config.env.disagg_sim.nnodes
+ else:
+ # colocated sim and actor rollout
+ num_nodes_sim = config.trainer.nnodes
+
+ resource_pool_spec = {
+ train_rollout_pool_id: [train_rollout_gpu_num] * num_nodes_actor_rollout,
+ "env_gpu_pool": [env_gpu_num] * num_nodes_sim,
+ }
+ mapping = {
+ Role.ActorRollout: train_rollout_pool_id,
+ # Role.Critic: global_pool_id,
+ # Role.RefPolicy: global_pool_id,
+ Role.Env: "env_gpu_pool",
+ }
+
+ reward_fn = calculate_reward
+ val_reward_fn = calculate_reward
+
+ resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
+
+ # Create training and validation datasets.
+ train_dataset = datasets.load_dataset("parquet", data_files=config.data.train_files)["train"]
+ val_dataset = datasets.load_dataset("parquet", data_files=config.data.val_files)["train"]
+
+ trainer = RobRayPPOTrainer(
+ config=config,
+ tokenizer=tokenizer,
+ role_worker_mapping=role_worker_mapping,
+ resource_pool_manager=resource_pool_manager,
+ ray_worker_group_cls=ray_worker_group_cls,
+ reward_fn=reward_fn,
+ val_reward_fn=val_reward_fn,
+ train_dataset=train_dataset,
+ val_dataset=val_dataset,
+ )
+ trainer.init_workers()
+ trainer.fit()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/__init__.py b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6f6b91498b5e61982e3c382964f0a26dd4188bd
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+# https://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.
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/configuration_prismatic.py b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/configuration_prismatic.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e4bb27d05c4ec64c5913dba68f95febf8658675
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/configuration_prismatic.py
@@ -0,0 +1,156 @@
+# Copyright 2025 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
+#
+# https://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 https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/
+# form https://huggingface.co/Haozhan72/Openvla-oft-SFT-libero10-trajall/blob/main/
+"""
+configuration_prismatic.py
+
+HuggingFace-style configuration definition for Prismatic VLMs, inheriting from `transformers.PretrainedConfig`.
+Default configuration specifies `siglip-224px+7b`.
+"""
+
+from typing import Any, Optional
+
+from transformers import PretrainedConfig
+from transformers.models.auto import CONFIG_MAPPING
+
+# === Utilities for Mapping Prismatic names to HF names ===
+# fmt: off
+VISION_BACKBONE_TO_RESOLUTION: dict[str, list[int]] = {
+ "clip-vit-l": [224], "siglip-vit-so400m": [224], "dinov2-vit-l": [224], "in1k-vit-l": [224],
+
+ "clip-vit-l-336px": [336],
+ "siglip-vit-so400m-384px": [384],
+
+ "dinoclip-vit-l-336px": [336, 336],
+ "dinosiglip-vit-so-224px": [224, 224],
+ "dinosiglip-vit-so-384px": [384, 384],
+}
+VISION_BACKBONE_TO_TIMM_ID: dict[str, list[str]] = {
+ "clip-vit-l": ["vit_large_patch14_clip_224.openai"],
+ "clip-vit-l-336px": ["vit_large_patch14_clip_336.openai"],
+
+ "dinov2-vit-l": ["vit_large_patch14_reg4_dinov2.lvd142m"],
+ "in1k-vit-l": ["vit_large_patch16_224.augreg_in21k_ft_in1k"],
+
+ "siglip-vit-so400m": ["vit_so400m_patch14_siglip_224"],
+ "siglip-vit-so400m-384px": ["vit_so400m_patch14_siglip_384"],
+
+ "dinoclip-vit-l-336px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_large_patch14_clip_336.openai"],
+ "dinosiglip-vit-so-224px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_224"],
+ "dinosiglip-vit-so-384px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_384"],
+}
+TIMM_OVERRIDE_ACT_LAYER: dict[str, list[Optional[str]]] = {
+ "clip-vit-l": ["quick_gelu"], "clip-vit-l-336px": ["quick_gelu"],
+ "dinov2-vit-l": [None], "in1k-vit-l": [None],
+ "siglip-vit-so400m": [None], "siglip-vit-so400m-384px": [None],
+ "dinoclip-vit-l-336px": [None, "quick_gelu"],
+ "dinosiglip-vit-so-224px": [None, None], "dinosiglip-vit-so-384px": [None, None]
+}
+
+LLM_BACKBONE_TO_HF_PATH = {
+ "llama2-7b-pure": "meta-llama/Llama-2-7b-hf", "llama2-13b-pure": "meta-llama/Llama-2-13b-hf",
+ "llama2-7b-chat": "meta-llama/Llama-2-7b-chat-hf", "llama2-13b-chat": "meta-llama/Llama-2-13b-chat-hf",
+
+ "vicuna-v15-7b": "lmsys/vicuna-7b-v1.5", "vicuna-v15-13b": "lmsys/vicuna-13b-v1.5",
+
+ "mistral-v0.1-7b-pure": "mistralai/Mistral-7B-v0.1",
+ "mistral-v0.1-7b-instruct": "mistralai/Mistral-7B-Instruct-v0.1",
+
+ "phi-2-3b": "microsoft/phi-2",
+}
+LLM_BACKBONE_TO_HF_METACLASS = {
+ "llama2-7b-pure": "llama", "llama2-13b-pure": "llama", "llama2-7b-chat": "llama", "llama2-13b-chat": "llama",
+ "vicuna-v15-7b": "llama", "vicuna-v15-13b": "llama",
+
+ "mistral-v0.1-7b-pure": "mistral", "mistral-v0.1-7b-instruct": "mistral",
+
+ "phi-2-3b": "phi",
+}
+
+VALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys())
+VALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH)
+# fmt: on
+
+
+class PrismaticConfig(PretrainedConfig):
+ model_type: str = "prismatic"
+ is_composition: bool = False
+
+ def __init__(
+ self,
+ vision_backbone_id: str = "siglip-vit-so400m",
+ llm_backbone_id: str = "vicuna-v15-7b",
+ arch_specifier: str = "no-align+gelu-mlp",
+ use_fused_vision_backbone: Optional[bool] = None,
+ image_resize_strategy: str = "letterbox",
+ text_config: Optional[dict[str, Any]] = None,
+ llm_max_length: int = 2048,
+ pad_token_id: int = 32000,
+ pad_to_multiple_of: int = 64,
+ output_projector_states: bool = False,
+ **kwargs: str,
+ ) -> None:
+ if vision_backbone_id not in VALID_VISION_BACKBONES:
+ raise ValueError(f"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }")
+
+ if llm_backbone_id not in VALID_LLM_BACKBONES:
+ raise ValueError(f"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }")
+
+ # Set Prismatic Configuration Fields
+ self.vision_backbone_id = vision_backbone_id
+ self.llm_backbone_id = llm_backbone_id
+ self.arch_specifier = arch_specifier
+ self.output_projector_states = output_projector_states
+
+ # [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing
+ self.use_fused_vision_backbone = (
+ use_fused_vision_backbone
+ if use_fused_vision_backbone is not None
+ else any(self.vision_backbone_id.startswith(v) for v in ["dinoclip", "dinosiglip"])
+ )
+
+ self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id]
+ self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id]
+ self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id]
+ self.image_resize_strategy = image_resize_strategy
+
+ self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id]
+ self.llm_max_length = llm_max_length
+ self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of
+
+ # [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming!
+ self.text_config = (
+ CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config)
+ if text_config is not None
+ else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]()
+ )
+
+ # Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well...
+ super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+
+class OpenVLAConfig(PrismaticConfig):
+ model_type: str = "openvla"
+
+ def __init__(
+ self,
+ norm_stats: Optional[dict[str, dict[str, dict[str, dict[str, list[float]]]]]] = None,
+ n_action_bins: int = 256,
+ **kwargs: str,
+ ) -> None:
+ self.norm_stats, self.n_action_bins = norm_stats, n_action_bins
+
+ super().__init__(**kwargs)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/constants.py b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6d6b3bce671b47b8977657f3c9fdef6a5635f98
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/constants.py
@@ -0,0 +1,104 @@
+# Copyright 2025 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
+#
+# https://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 https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/
+
+
+"""
+Important constants for VLA training and evaluation.
+
+Attempts to automatically identify the correct constants to set based on the Python command used to launch
+training or evaluation. If it is unclear, defaults to using the LIBERO simulation benchmark constants.
+"""
+
+import sys
+from enum import Enum
+
+# Llama 2 token constants
+IGNORE_INDEX = -100
+ACTION_TOKEN_BEGIN_IDX = 31743
+STOP_INDEX = 2 # ''
+
+
+# Defines supported normalization schemes for action and proprioceptive state.
+class NormalizationType(str, Enum):
+ # fmt: off
+ NORMAL = "normal" # Normalize to Mean = 0, Stdev = 1
+ BOUNDS = "bounds" # Normalize to Interval = [-1, 1]
+ BOUNDS_Q99 = "bounds_q99" # Normalize [quantile_01, ..., quantile_99] --> [-1, ..., 1]
+ # fmt: on
+
+
+# Define constants for each robot platform
+LIBERO_CONSTANTS = {
+ "NUM_ACTIONS_CHUNK": 8,
+ "ACTION_DIM": 7,
+ "PROPRIO_DIM": 8,
+ "ACTION_PROPRIO_NORMALIZATION_TYPE": NormalizationType.BOUNDS_Q99,
+}
+
+ALOHA_CONSTANTS = {
+ "NUM_ACTIONS_CHUNK": 25,
+ "ACTION_DIM": 14,
+ "PROPRIO_DIM": 14,
+ "ACTION_PROPRIO_NORMALIZATION_TYPE": NormalizationType.BOUNDS,
+}
+
+BRIDGE_CONSTANTS = {
+ "NUM_ACTIONS_CHUNK": 5,
+ "ACTION_DIM": 7,
+ "PROPRIO_DIM": 7,
+ "ACTION_PROPRIO_NORMALIZATION_TYPE": NormalizationType.BOUNDS_Q99,
+}
+
+
+# Function to detect robot platform from command line arguments
+def detect_robot_platform():
+ cmd_args = " ".join(sys.argv).lower()
+
+ if "libero" in cmd_args:
+ return "LIBERO"
+ elif "aloha" in cmd_args:
+ return "ALOHA"
+ elif "bridge" in cmd_args:
+ return "BRIDGE"
+ else:
+ # Default to LIBERO if unclear
+ return "LIBERO"
+
+
+# Determine which robot platform to use
+ROBOT_PLATFORM = detect_robot_platform()
+
+# Set the appropriate constants based on the detected platform
+if ROBOT_PLATFORM == "LIBERO":
+ constants = LIBERO_CONSTANTS
+elif ROBOT_PLATFORM == "ALOHA":
+ constants = ALOHA_CONSTANTS
+elif ROBOT_PLATFORM == "BRIDGE":
+ constants = BRIDGE_CONSTANTS
+
+# Assign constants to global variables
+NUM_ACTIONS_CHUNK = constants["NUM_ACTIONS_CHUNK"]
+ACTION_DIM = constants["ACTION_DIM"]
+PROPRIO_DIM = constants["PROPRIO_DIM"]
+ACTION_PROPRIO_NORMALIZATION_TYPE = constants["ACTION_PROPRIO_NORMALIZATION_TYPE"]
+
+# Print which robot platform constants are being used (for debugging)
+print(f"Using {ROBOT_PLATFORM} constants:")
+print(f" NUM_ACTIONS_CHUNK = {NUM_ACTIONS_CHUNK}")
+print(f" ACTION_DIM = {ACTION_DIM}")
+print(f" PROPRIO_DIM = {PROPRIO_DIM}")
+print(f" ACTION_PROPRIO_NORMALIZATION_TYPE = {ACTION_PROPRIO_NORMALIZATION_TYPE}")
+print("If needed, manually set the correct constants in `prismatic/vla/constants.py`!")
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/modeling_prismatic.py b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/modeling_prismatic.py
new file mode 100644
index 0000000000000000000000000000000000000000..a52b56acce048d521bc2a5c8334ed16f969b369d
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/modeling_prismatic.py
@@ -0,0 +1,2000 @@
+# Copyright 2025 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
+#
+# https://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 https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/
+# form https://huggingface.co/Haozhan72/Openvla-oft-SFT-libero10-trajall/blob/main/
+
+"""
+modeling_prismatic.py
+
+Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
+Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
+but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
+"""
+
+import logging
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, ClassVar, Optional
+
+import numpy as np
+import timm
+import tokenizers
+import torch
+import torch.nn as nn
+import transformers
+from timm.models.vision_transformer import LayerScale
+from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
+from .constants import (
+ ACTION_DIM,
+ ACTION_PROPRIO_NORMALIZATION_TYPE,
+ ACTION_TOKEN_BEGIN_IDX,
+ IGNORE_INDEX,
+ NUM_ACTIONS_CHUNK,
+ STOP_INDEX,
+ NormalizationType,
+)
+from .train_utils import (
+ get_current_action_mask,
+ get_next_actions_mask,
+)
+
+# Set up logger
+logger = logging.getLogger(__name__)
+
+
+# === Utility Functions for Monkey-Patching ===
+def unpack_tuple(fn: Callable[[Any], tuple[Any]]) -> Callable[[Any], Any]:
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
+ result = fn(*args, **kwargs)
+ return result[0] if isinstance(result, tuple) else result
+
+ return wrapper
+
+
+# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
+# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
+# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
+def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
+ return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
+
+
+def ls_apply_patch(ls_module: LayerScale):
+ ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
+ ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
+ del ls_module.gamma
+
+
+# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
+class PrismaticVisionBackbone(nn.Module):
+ """
+ Vision backbone for Prismatic models that handles image feature extraction.
+
+ Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
+ For fused backbones, features from both models are concatenated along the feature dimension.
+ """
+
+ def __init__(
+ self,
+ use_fused_vision_backbone: bool,
+ image_sizes: list[int],
+ timm_model_ids: list[str],
+ timm_override_act_layers: list[Optional[str]],
+ ) -> None:
+ """
+ Initialize the vision backbone.
+
+ Args:
+ use_fused_vision_backbone: Whether to use two backbones and fuse their features
+ image_sizes: List of image sizes for each backbone
+ timm_model_ids: List of TIMM model IDs to use for each backbone
+ timm_override_act_layers: List of activation layer overrides for each backbone
+ """
+ super().__init__()
+ self.use_fused_vision_backbone = use_fused_vision_backbone
+ self.num_images_in_input = 1 # Default value, can be overridden later
+
+ # Validate number of (fused) vision backbones
+ if len(timm_model_ids) > 2:
+ raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
+
+ # Create primary featurizer
+ self.featurizer = self._create_featurizer(
+ model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
+ )
+ self.embed_dim = self.featurizer.embed_dim
+
+ # Create secondary featurizer if using fused backbone
+ if self.use_fused_vision_backbone:
+ self.fused_featurizer = self._create_featurizer(
+ model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
+ )
+ self.embed_dim += self.fused_featurizer.embed_dim
+
+ # Patch LayerScale modules for HF compatibility
+ self._patch_layer_scales()
+
+ def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
+ """
+ Create a TIMM-based featurizer model with appropriate configurations.
+
+ Args:
+ model_id: The TIMM model ID to load
+ img_size: Input image size for the model
+ act_layer: Override for the activation layer type
+
+ Returns:
+ A configured featurizer model
+ """
+ featurizer = timm.create_model(
+ model_id,
+ pretrained=False,
+ num_classes=0,
+ img_size=img_size,
+ act_layer=act_layer,
+ )
+
+ # Monkey-patch the forward function to extract the second-to-last layer features
+ num_blocks = len(featurizer.blocks)
+ featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
+
+ return featurizer
+
+ def _patch_layer_scales(self) -> None:
+ """
+ Patch all LayerScale modules to be compatible with HF's parameter naming.
+
+ HF Transformers overwrites parameters with names containing 'gamma',
+ so we need to rename and modify the forward method.
+ """
+ # Patch primary featurizer
+ for module in self.featurizer.modules():
+ if isinstance(module, LayerScale):
+ ls_apply_patch(module)
+
+ # Patch secondary featurizer if it exists
+ if self.use_fused_vision_backbone:
+ for module in self.fused_featurizer.modules():
+ if isinstance(module, LayerScale):
+ ls_apply_patch(module)
+
+ def get_num_patches(self) -> int:
+ """
+ Returns the number of vision patches output by the vision backbone.
+
+ Returns:
+ Number of patches per image
+ """
+ return self.featurizer.patch_embed.num_patches
+
+ def get_num_images_in_input(self) -> int:
+ """
+ Returns the number of input images for the vision backbone.
+
+ Returns:
+ Number of images expected in the input
+ """
+ return self.num_images_in_input
+
+ def set_num_images_in_input(self, num_images_in_input: int) -> None:
+ """
+ Sets the number of input images for the vision backbone.
+
+ Args:
+ num_images_in_input: Number of images to expect in the input
+ """
+ self.num_images_in_input = num_images_in_input
+
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ """
+ Implements the forward pass for the vision backbone.
+
+ If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
+ (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
+
+ Args:
+ pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
+ """
+ if self.num_images_in_input == 1:
+ if not self.use_fused_vision_backbone:
+ return self.featurizer(pixel_values)
+
+ # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
+ img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
+ patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
+
+ return torch.cat([patches, patches_fused], dim=2)
+
+ else:
+ assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
+
+ # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
+ images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
+
+ # Process each image and collect patches
+ all_patches = []
+ for img in images:
+ # Split each image further into two stacks of channels (each with 3 channels)
+ img_regular, img_fused = torch.split(img, [3, 3], dim=1)
+
+ # Get patches from both SigLIP and DINOv2 vision transformers
+ patches = self.featurizer(img_regular)
+ patches_fused = self.fused_featurizer(img_fused)
+
+ # Concatenate SigLIP and DINOv2 patches along the hidden dimension
+ combined_patches = torch.cat([patches, patches_fused], dim=2)
+ all_patches.append(combined_patches)
+
+ # Concatenate all patches along the patch dimension
+ return torch.cat(all_patches, dim=1)
+
+
+# === Prismatic Projector (nn.Module) Definitions ===
+class PrismaticProjector(nn.Module):
+ def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
+ super().__init__()
+ self.use_fused_vision_backbone = use_fused_vision_backbone
+ self.vision_dim, self.llm_dim = vision_dim, llm_dim
+
+ # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
+ if not self.use_fused_vision_backbone:
+ self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
+ self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+ self.act_fn1 = nn.GELU()
+ else:
+ initial_projection_dim = 4 * vision_dim
+ self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
+ self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
+ self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+ self.act_fn1 = nn.GELU()
+ self.act_fn2 = nn.GELU()
+
+ def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+ if not self.use_fused_vision_backbone:
+ projected_features = self.fc1(img_patches)
+ projected_features = self.act_fn1(projected_features)
+ projected_features = self.fc2(projected_features)
+ else:
+ projected_features = self.fc1(img_patches)
+ projected_features = self.act_fn1(projected_features)
+ projected_features = self.fc2(projected_features)
+ projected_features = self.act_fn2(projected_features)
+ projected_features = self.fc3(projected_features)
+
+ return projected_features
+
+
+# === Main HF Class Definitions ===
+@dataclass
+class PrismaticCausalLMOutputWithPast(ModelOutput):
+ """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
+
+ loss: Optional[torch.FloatTensor] = None
+ logits: torch.FloatTensor = None
+ past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None
+ hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+ attentions: Optional[tuple[torch.FloatTensor]] = None
+
+ # Additions for VLMs
+ projector_features: Optional[torch.FloatTensor] = None
+
+
+class PrismaticPreTrainedModel(PreTrainedModel):
+ config_class: PretrainedConfig = PrismaticConfig
+ base_model_prefix: str = "model"
+ supports_gradient_checkpointing: bool = True
+
+ _no_split_modules: ClassVar[list[str]] = ["PrismaticProjector"]
+ _skip_keys_device_placement: str = "past_key_values"
+ _supports_flash_attn_2: bool = True
+
+ def _init_weights(self, module: nn.Module) -> None:
+ # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
+ # => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
+ # https://github.com/TRI-ML/prismatic-vlms
+ std = (
+ self.config.initializer_range
+ if hasattr(self.config, "initializer_range")
+ else self.config.text_config.initializer_range
+ )
+
+ if hasattr(module, "class_embedding"):
+ module.class_embedding.data.normal_(mean=0.0, std=std)
+
+ if isinstance(module, (nn.Linear | nn.Conv2d)):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+ @property
+ def _supports_sdpa(self) -> bool:
+ """Check LLM supports SDPA Attention"""
+ return self.language_model._supports_sdpa
+
+
+class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
+ def __init__(self, config: PrismaticConfig) -> None:
+ super().__init__(config)
+
+ # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
+ if config.use_fused_vision_backbone is None:
+ raise ValueError("Missing config field `use_fused_vision_backbone`")
+
+ if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
+ raise NotImplementedError(
+ "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
+ "if you urgently need support for latest TIMM versions."
+ )
+
+ if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
+ logger.warning(
+ f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
+ f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
+ f"there might be inference-time regressions due to dependency changes. If in doubt, please"
+ f"use the above versions."
+ )
+
+ # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
+ self.vision_backbone = PrismaticVisionBackbone(
+ config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
+ )
+
+ # Create Multimodal Projector
+ self.projector = PrismaticProjector(
+ config.use_fused_vision_backbone,
+ vision_dim=self.vision_backbone.embed_dim,
+ llm_dim=config.text_config.hidden_size,
+ )
+
+ # Instantiate LLM Backbone
+ self.language_model = AutoModelForCausalLM.from_config(
+ config.text_config, attn_implementation=config._attn_implementation
+ )
+ self.vocab_size = config.text_config.vocab_size
+ self.pad_token_id = config.pad_token_id
+ self.llm_dim = config.text_config.hidden_size
+
+ # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
+ self.post_init()
+
+ # === `PreTrainedModel` Boilerplate ===
+ def get_input_embeddings(self) -> nn.Module:
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value: nn.Module) -> None:
+ self.language_model.set_input_embeddings(value)
+
+ def get_output_embeddings(self) -> nn.Module:
+ return self.language_model.get_output_embeddings()
+
+ def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+ self.language_model.set_output_embeddings(new_embeddings)
+
+ def get_decoder(self) -> nn.Module:
+ return self.language_model.get_decoder()
+
+ def set_decoder(self, decoder: nn.Module) -> None:
+ self.language_model.set_decoder(decoder)
+
+ def tie_weights(self) -> None:
+ self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
+
+ def resize_token_embeddings(
+ self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
+ ) -> nn.Embedding:
+ updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+
+ # Update config/instance variables
+ self.config.text_config.vocab_size = updated_embeddings.num_embeddings
+ self.vocab_size = updated_embeddings.num_embeddings
+
+ return updated_embeddings
+
+ def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
+ """
+ Replace embeddings in input_embeddings at positions where all_actions_mask is True
+ with embeddings from noisy_action_features, using vectorized operations.
+
+ Args:
+ input_embeddings: Tensor of shape (B, S, D)
+ all_actions_mask: Boolean tensor of shape (B, S)
+ noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
+
+ Returns:
+ Modified input_embeddings tensor
+ """
+ # Clone input to avoid modifying the original tensor
+ new_input_embeddings = input_embeddings.clone()
+
+ # Create a tensor with the same shape of input_embeddings to hold the noisy action features
+ repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
+
+ # Create batch indices for splicing
+ batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
+ batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
+
+ # Get indices where mask is True for each sample
+ masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
+
+ # Move the noisy action features into their correct positions
+ repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
+
+ # Combine original input embeddings and noisy action embeddings using the mask
+ new_input_embeddings = torch.where(
+ all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
+ )
+
+ return new_input_embeddings
+
+ def _process_action_masks(self, labels):
+ """Helper to get action masks from labels"""
+ current_action_mask = get_current_action_mask(labels)
+ next_actions_mask = get_next_actions_mask(labels)
+ all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
+ return all_actions_mask
+
+ def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False):
+ """Process vision features with optional FiLM conditioning"""
+ if use_film:
+ # FiLM: Infuse language inputs into visual features
+ patch_features = self.vision_backbone(pixel_values, language_embeddings) # (bsz, 256 * num_images, D)
+ else:
+ patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
+
+ # Project patch embeddings into language embedding space
+ return self.projector(patch_features)
+
+ def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
+ """Process proprioceptive features and append to vision features"""
+ if proprio_projector is not None and proprio is not None:
+ # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
+ # proprio: (bsz, proprio_dim) or (propro_dim,)
+ proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
+ proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
+ proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
+ # For simplicity, just append proprio token to the end of projected vision patch tokens
+ return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
+ return projected_patch_embeddings
+
+ def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
+ """Build multimodal embeddings and attention mask"""
+ # Update attention mask
+ projected_patch_attention_mask = None
+ if attention_mask is not None:
+ projected_patch_attention_mask = torch.full(
+ (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+ fill_value=True,
+ dtype=attention_mask.dtype,
+ device=attention_mask.device,
+ )
+
+ # Build multimodal embeddings & attention mask; insert embeddings after token (1:)
+ multimodal_embeddings = torch.cat(
+ [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+ )
+
+ multimodal_attention_mask = None
+ if attention_mask is not None:
+ multimodal_attention_mask = torch.cat(
+ [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+ )
+
+ return multimodal_embeddings, multimodal_attention_mask
+
+ def _build_multimodal_labels(self, labels, projected_patch_embeddings):
+ """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
+ if labels is not None:
+ projected_patch_labels = torch.full(
+ (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+ fill_value=IGNORE_INDEX,
+ dtype=labels.dtype,
+ device=labels.device,
+ )
+ return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)
+ return None
+
+ # === Core Prismatic VLM `forward()` Logic ===
+ # def forward(
+ # self,
+ # input_ids: Optional[torch.LongTensor] = None,
+ # attention_mask: Optional[torch.Tensor] = None,
+ # pixel_values: Optional[torch.FloatTensor] = None,
+ # labels: Optional[torch.LongTensor] = None,
+ # inputs_embeds: Optional[torch.FloatTensor] = None,
+ # past_key_values: Optional[List[torch.FloatTensor]] = None,
+ # use_cache: Optional[bool] = None,
+ # output_attentions: Optional[bool] = None,
+ # output_hidden_states: Optional[bool] = None,
+ # output_projector_features: Optional[bool] = None,
+ # return_dict: Optional[bool] = None,
+ # proprio=None,
+ # proprio_projector=None,
+ # noisy_actions=None,
+ # noisy_action_projector=None,
+ # diffusion_timestep_embeddings=None,
+ # use_film: bool = False,
+ # ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+ # """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+ # output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ # output_hidden_states = (
+ # output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ # )
+ # output_projector_features = output_projector_features if output_projector_features is not None else False
+ # return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ # # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+ # use_cache = use_cache and not self.training
+
+ # # Instantiate Placeholder for Projector Features
+ # projected_patch_embeddings = None
+
+ # # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+ # if input_ids.shape[1] == 1:
+ # assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+ # assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+ # assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+ # language_model_output = self.language_model(
+ # input_ids=input_ids,
+ # attention_mask=None,
+ # position_ids=None,
+ # past_key_values=past_key_values,
+ # inputs_embeds=None,
+ # labels=None,
+ # use_cache=use_cache,
+ # output_attentions=output_attentions,
+ # output_hidden_states=output_hidden_states,
+ # return_dict=return_dict,
+ # )
+
+ # # === Handle Unimodal Forward ===
+ # elif pixel_values is None:
+ # assert (input_ids is not None) and (inputs_embeds is None), \
+ # "Missing `input_ids` in language-only forward!"
+ # assert past_key_values is None, \
+ # "Unexpected key `past_key_values` provided during language-only forward!"
+
+ # language_model_output = self.language_model(
+ # input_ids=input_ids,
+ # attention_mask=attention_mask,
+ # position_ids=None,
+ # past_key_values=None,
+ # inputs_embeds=None,
+ # labels=labels,
+ # use_cache=use_cache,
+ # output_attentions=output_attentions,
+ # output_hidden_states=output_hidden_states,
+ # return_dict=return_dict,
+ # )
+
+ # # === Handle Multimodal Forward ===
+ # elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+ # assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+ # #test
+ #
+ # #test end
+
+ # # Get input embeddings (from language model embeddings)
+ # input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
+
+ # # Extract action masks
+ # all_actions_mask = self._process_action_masks(labels)
+
+ # # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+ # language_embeddings = input_embeddings[~all_actions_mask].reshape(
+ # input_embeddings.shape[0], -1, input_embeddings.shape[2]
+ # ) # (B, lang_seq_len, llm_dim)
+
+ # # Get visual features
+ # projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+ # # Add proprioceptive state if provided
+ # projected_patch_embeddings = self._process_proprio_features(
+ # projected_patch_embeddings, proprio, proprio_projector
+ # )
+
+ # # [Diffusion] Add diffusion timestep embedding if provided
+ # if diffusion_timestep_embeddings is not None:
+ # # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
+ # projected_patch_embeddings = torch.cat(
+ # (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+ # )
+
+ # # Process action embeddings
+ # if noisy_actions is not None:
+ # # Get mask corresponding to all action tokens
+ # all_actions_mask = self._process_action_masks(labels)
+
+ # # Reshape noisy actions into individual action tokens
+ # # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
+ # B = noisy_actions.shape[0]
+ # noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
+
+ # # Project noisy action tokens into language model embedding space
+ # noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
+
+ # # Replace embeddings of the action tokens with noisy action embeddings
+ # input_embeddings = self._replace_input_embeddings(
+ # input_embeddings, all_actions_mask, noisy_action_features
+ # )
+ # else:
+ # # Replace the embeddings of the action tokens with zeros
+ # # (Later on, the positional embeddings will be added to them)
+ # all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
+ # input_embeddings = input_embeddings * ~all_actions_mask
+
+ # # Build multimodal embeddings & attention mask
+ # multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+ # input_embeddings, projected_patch_embeddings, attention_mask
+ # )
+
+ # # Build labels for multimodal sequence if needed
+ # multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+ # # Dispatch to language model
+ # language_model_output = self.language_model(
+ # input_ids=None,
+ # attention_mask=multimodal_attention_mask,
+ # position_ids=None,
+ # past_key_values=None,
+ # inputs_embeds=multimodal_embeddings,
+ # labels=multimodal_labels,
+ # use_cache=use_cache,
+ # output_attentions=output_attentions,
+ # output_hidden_states=output_hidden_states,
+ # return_dict=return_dict,
+ # )
+
+ # # === Otherwise =>> Assume Invalid! ===
+ # elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+ # raise ValueError("Non-homogenous batch of (text, image) input \
+ # -- forward() does not support mixed batches!")
+
+ # else:
+ # raise ValueError(
+ # "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+ # f"=> `input_ids` = {input_ids is not None}\n"
+ # f"=> `attention_mask` = {attention_mask is not None}\n"
+ # f"=> `pixel_values` = {pixel_values is not None}\n"
+ # f"=> `labels` = {labels is not None}\n"
+ # f"=> `input_embeds` = {inputs_embeds is not None}\n"
+ # f"=> `past_key_values` = {past_key_values is not None}\n"
+ # f"=> `use_cache` = {use_cache}"
+ # )
+
+ # # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+ # if not return_dict:
+ # if output_projector_features and (projected_patch_embeddings is not None):
+ # return *language_model_output, projected_patch_embeddings
+
+ # return language_model_output
+
+ # return PrismaticCausalLMOutputWithPast(
+ # loss=language_model_output.loss,
+ # logits=language_model_output.logits,
+ # past_key_values=language_model_output.past_key_values,
+ # hidden_states=language_model_output.hidden_states,
+ # attentions=language_model_output.attentions,
+ # projector_features=projected_patch_embeddings,
+ # )
+
+ # === GenerationMixin Methods ===
+ def prepare_inputs_for_generation(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ past_key_values: Optional[list[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ pixel_values: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs: str,
+ ) -> dict[str, torch.Tensor]:
+ """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
+ if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
+ (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
+ ):
+ raise ValueError("Generation with batch size > 1 is not currently supported!")
+
+ # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
+ if past_key_values is not None:
+ input_ids = input_ids[:, -1:]
+
+ # If `input_embeds` are passed, we only want to use them in the 1st generation step
+ if inputs_embeds is not None and past_key_values is None:
+ model_inputs = {"input_embeds": inputs_embeds}
+ else:
+ model_inputs = {"input_ids": input_ids}
+
+ # Make sure `pixel_values` are preserved in `model_inputs`
+ model_inputs.update(
+ {
+ "attention_mask": attention_mask,
+ "pixel_values": pixel_values,
+ "past_key_values": past_key_values,
+ "use_cache": kwargs.get("use_cache"),
+ }
+ )
+
+ return model_inputs
+
+ # Defer to Language Model (all handle this differently, with different return types)
+ def _reorder_cache(self, *args, **kwargs) -> Any:
+ return self.language_model._reorder_cache(*args, **kwargs)
+
+ def _prepare_input_for_action_prediction_verl(self, input_ids, attention_mask):
+ """Prepares input for action prediction by adding necessary tokens"""
+ # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+ placeholder_action_token_ids = (
+ torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+ )
+ input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+ # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+ stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+ input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+ # Extend the attention mask to fit the new shape of input
+ # Note: Only batch size == 1 supported right now
+ mask_extension = (
+ torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+ .to(attention_mask.device)
+ .to(attention_mask.dtype)
+ )
+ attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+ return input_ids, attention_mask
+
+ def _prepare_labels_for_action_prediction_verl(self, labels, input_ids):
+ """Creates labels tensor for action prediction if not provided"""
+ # Extend labels tensor with fake action labels
+ ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+ labels_extension = (
+ torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+ * ARBITRARY_ACTION_TOKEN_IDX
+ )
+ labels = torch.cat([labels, labels_extension], dim=-1)
+
+ # Replace last label token with stop token
+ labels[:, -1] = STOP_INDEX
+
+ return labels
+
+ def _verl_discrete_compute_logits(
+ self,
+ input_embeddings,
+ all_actions_mask,
+ projected_patch_embeddings,
+ attention_mask,
+ labels,
+ NUM_PATCHES,
+ NUM_PROMPT_TOKENS,
+ action_head=None,
+ ): # contintue!!!!!
+ """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+ # Zero out action token embeddings
+ all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
+ input_embeddings = input_embeddings * ~all_actions_mask
+
+ # Build multimodal embeddings and attention mask
+ multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+ input_embeddings, projected_patch_embeddings, attention_mask
+ )
+
+ # Forward pass through language model
+ language_model_output = self.language_model(
+ input_ids=None,
+ attention_mask=multimodal_attention_mask,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=multimodal_embeddings,
+ labels=None,
+ use_cache=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+
+ # Extract hidden states for action tokens
+ # last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
+ # actions_hidden_states = last_hidden_states[
+ # :,
+ # NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ # :,
+ # ] # (B, act_chunk_len, D)
+
+ # Handle different prediction methods
+ # if action_head is not None:
+ # # L1 regression prediction
+ # normalized_actions = action_head.predict_action(actions_hidden_states)
+ # normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+ # normalized_actions = normalized_actions.float().cpu().detach().numpy()
+ # else:
+ # Discrete token-based prediction
+
+ compute_logits = language_model_output.logits[
+ :,
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ ]
+
+ return compute_logits
+
+ # def forward(
+ # self,
+ # input_ids: Optional[torch.LongTensor] = None,
+ # unnorm_key: Optional[str] = None,
+ # proprio=None,
+ # proprio_projector=None,
+ # action_head=None,
+ # noisy_action_projector=None,
+ # use_film: bool = False,
+ # **kwargs: str,
+ # ) :
+ # """Predict actions from input sequence, with options for different prediction methods.
+
+ # Args:
+ # input_ids: Input token ids
+ # unnorm_key: Key for unnormalization statistics
+ # proprio: Proprioceptive features
+ # proprio_projector: Projector for proprioceptive features
+ # action_head: Optional head for L1 regression or diffusion-based prediction
+ # noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+ # use_film: Whether to use FiLM conditioning
+ # **kwargs: Additional arguments including pixel_values and attention_mask
+
+ # Returns:
+ # Tuple of (unnormalized_actions, action_hidden_states)
+ # """
+ # # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+ # # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+ # # if not torch.all(input_ids[:, -1] == 29871):
+ # # input_ids = torch.cat(
+ # # (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+ # # )
+ # #print("!!!!!!!!!!!!!!Entering forward!!!!!!!!!!")
+ # pixel_values = kwargs["pixel_values"]
+ # attention_mask = kwargs["attention_mask"]
+
+ # # Create fake labels tensor (needed for action mask)
+ # labels = input_ids.clone()
+ # labels[:] = IGNORE_INDEX
+
+ # # Get number of tokens in prompt (excluding the start token)
+ # NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
+
+ # # Prepare inputs by adding necessary tokens
+ # #input_ids, attention_mask = self._prepare_input_for_action_prediction_verl(input_ids, attention_mask)
+
+ # #test
+ # placeholder_action_token_ids = (
+ # torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+ # )
+ # input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+ # # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+ # stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+ # input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+ # # Extend the attention mask to fit the new shape of input
+ # # Note: Only batch size == 1 supported right now
+ # mask_extension = (
+ # torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+ # .to(attention_mask.device)
+ # .to(attention_mask.dtype)
+ # )
+ # attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+ # #return input_ids, attention_mask
+
+ # #test end
+
+ # # Update labels tensor for action mask computation later
+ # #labels = self._prepare_labels_for_action_prediction_verl(labels, input_ids)
+ # #test
+
+ # ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+ # labels_extension = (
+ # torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+ # * ARBITRARY_ACTION_TOKEN_IDX
+ # )
+ # labels = torch.cat([labels, labels_extension], dim=-1)
+
+ # # Replace last label token with stop token
+ # labels[:, -1] = STOP_INDEX
+
+ # #return labels
+
+ # #test ed
+
+ # # Get input embeddings and action masks
+
+ # input_embeddings = self.get_input_embeddings()(input_ids)
+
+ # #all_actions_mask = self._process_action_masks(labels)
+ # #test
+ # #current_action_mask = get_current_action_mask(labels)
+ # newline_positions = labels != IGNORE_INDEX
+
+ # # Calculate cumulative sum to identify regions between newlines
+ # cumsum = torch.cumsum(newline_positions, dim=1)
+
+ # # Create the mask
+ # mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)
+
+ # # Extract the action part only
+ # action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX
+ # current_action_mask = action_tokens_only_mask * mask
+
+ # #next_actions_mask = get_next_actions_mask(labels)
+ # newline_positions = labels != IGNORE_INDEX
+
+ # # Calculate cumulative sum to identify regions between newlines
+ # cumsum = torch.cumsum(newline_positions, dim=1)
+
+ # # Create the mask
+ # mask = cumsum > ACTION_DIM
+
+ # # Extract the action part only
+ # action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX
+ # next_actions_mask = action_tokens_only_mask * mask
+
+ # all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
+
+ # #test end
+
+ # # Extract language embeddings
+ # language_embeddings = input_embeddings[~all_actions_mask].reshape(
+ # input_embeddings.shape[0], -1, input_embeddings.shape[2]
+ # )
+
+ # # Process vision features
+ # #projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+ # #test
+ # if use_film:
+ # # FiLM: Infuse language inputs into visual features
+ # raise ValueError
+ # patch_features = self.vision_backbone(pixel_values, language_embeddings) # (bsz, 256 * num_images, D)
+ # else:
+ # patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
+
+ # projected_patch_embeddings = self.projector(patch_features)
+ # #test end
+
+ # # Add proprioceptive features if provided
+ # use_proprio = proprio_projector is not None and proprio is not None
+ # if use_proprio:
+ # proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device,
+ # dtype=projected_patch_embeddings.dtype)
+ # projected_patch_embeddings = self._process_proprio_features(
+ # projected_patch_embeddings, proprio, proprio_projector
+ # )
+
+ # # Use diffusion if provided, otherwise use regression or discrete prediction
+ # use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+ # # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+ # NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+ # if use_proprio:
+ # NUM_PATCHES += 1
+ # if use_diffusion:
+ # NUM_PATCHES += 1
+
+ # if use_diffusion:
+ # raise ValueError
+ # # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion
+ # noise = torch.randn(
+ # size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype
+ # )
+
+ # # Run diffusion-based prediction
+ # normalized_actions, actions_hidden_states = self._run_diffusion_prediction(
+ # input_embeddings,
+ # all_actions_mask,
+ # noise,
+ # action_head,
+ # projected_patch_embeddings,
+ # labels,
+ # attention_mask,
+ # NUM_PATCHES,
+ # NUM_PROMPT_TOKENS,
+ # noisy_action_projector,
+ # )
+ # else:
+ # # Run regression or discrete token-based prediction
+ # # compute_logits = self._verl_discrete_compute_logits(
+ # # input_embeddings,
+ # # all_actions_mask,
+ # # projected_patch_embeddings,
+ # # attention_mask,
+ # # labels,
+ # # NUM_PATCHES,
+ # # NUM_PROMPT_TOKENS,
+ # # action_head,
+ # # )
+
+ # #test
+
+ # all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
+ # input_embeddings = input_embeddings * ~all_actions_mask
+
+ # # Build multimodal embeddings and attention mask
+ # # multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+ # # input_embeddings, projected_patch_embeddings, attention_mask
+ # # )
+ # #test
+
+ # projected_patch_attention_mask = None
+ # if attention_mask is not None:
+ # projected_patch_attention_mask = torch.full(
+ # (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+ # fill_value=True,
+ # dtype=attention_mask.dtype,
+ # device=attention_mask.device,
+ # )
+
+ # # Build multimodal embeddings & attention mask; insert embeddings after token (1:)
+ # multimodal_embeddings = torch.cat(
+ # [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+ # )
+
+ # multimodal_attention_mask = None
+ # if attention_mask is not None:
+ # multimodal_attention_mask = torch.cat(
+ # [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+ # )
+
+ # #return multimodal_embeddings, multimodal_attention_mask
+
+ # #test end
+
+ # # Forward pass through language model
+ # language_model_output = self.language_model(
+ # input_ids=None,
+ # attention_mask=multimodal_attention_mask,
+ # position_ids=None,
+ # past_key_values=None,
+ # inputs_embeds=multimodal_embeddings,
+ # labels=None,
+ # use_cache=None,
+ # output_attentions=False,
+ # output_hidden_states=False,
+ # return_dict=True,
+ # )
+
+ # compute_logits = language_model_output.logits[
+ # :,
+ # NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + \
+ # NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ # ]
+
+ # #test end
+
+ # return compute_logits
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ pixel_values=None,
+ attention_mask=None,
+ # labels=None,
+ proprio=None,
+ proprio_projector=None,
+ action_head=None,
+ noisy_action_projector=None,
+ use_film: bool = False,
+ **kwargs: str,
+ ):
+ """Predict actions from input sequence, with options for different prediction methods.
+
+ Args:
+ input_ids: Input token ids
+ unnorm_key: Key for unnormalization statistics
+ proprio: Proprioceptive features
+ proprio_projector: Projector for proprioceptive features
+ action_head: Optional head for L1 regression or diffusion-based prediction
+ noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+ use_film: Whether to use FiLM conditioning
+ **kwargs: Additional arguments including pixel_values and attention_mask
+
+ Returns:
+ Tuple of (unnormalized_actions, action_hidden_states)
+ """
+ # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+ # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+ # if not torch.all(input_ids[:, -1] == 29871):
+ # input_ids = torch.cat(
+ # (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+ # )
+
+ # pixel_values = kwargs["pixel_values"]
+ # attention_mask = kwargs["attention_mask"]
+
+ # Create fake labels tensor (needed for action mask)
+ labels = input_ids.clone()
+ labels[:] = IGNORE_INDEX
+
+ # # Get number of tokens in prompt (excluding the start token)
+ NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
+
+ # # Prepare inputs by adding necessary tokens
+ # #input_ids, attention_mask = self._prepare_input_for_action_prediction_verl(input_ids, attention_mask)
+
+ # #test
+ placeholder_action_token_ids = (
+ torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+ )
+ input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+ # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+ stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+ input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+ # Extend the attention mask to fit the new shape of input
+ # Note: Only batch size == 1 supported right now
+ mask_extension = (
+ torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+ .to(attention_mask.device)
+ .to(attention_mask.dtype)
+ )
+ attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+ ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+ labels_extension = (
+ torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+ * ARBITRARY_ACTION_TOKEN_IDX
+ )
+ labels = torch.cat([labels, labels_extension], dim=-1)
+
+ # # Replace last label token with stop token
+ labels[:, -1] = STOP_INDEX
+
+ # Get input embeddings and action masks
+
+ # NUM_PROMPT_TOKENS = kwargs["num_prompt_tokens"]
+
+ input_embeddings = self.get_input_embeddings()(input_ids)
+
+ # all_actions_mask = self._process_action_masks(labels)
+ # test
+ # current_action_mask = get_current_action_mask(labels)
+ newline_positions = labels != IGNORE_INDEX
+
+ # Calculate cumulative sum to identify regions between newlines
+ cumsum = torch.cumsum(newline_positions, dim=1)
+
+ # Create the mask
+ mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)
+
+ # Extract the action part only
+ action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX
+ current_action_mask = action_tokens_only_mask * mask
+
+ # next_actions_mask = get_next_actions_mask(labels)
+ newline_positions = labels != IGNORE_INDEX
+
+ # Calculate cumulative sum to identify regions between newlines
+ cumsum = torch.cumsum(newline_positions, dim=1)
+
+ # Create the mask
+ mask = cumsum > ACTION_DIM
+
+ # Extract the action part only
+ action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX
+ next_actions_mask = action_tokens_only_mask * mask
+
+ all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
+
+ # test end
+
+ # Extract language embeddings
+ language_embeddings = input_embeddings[~all_actions_mask].reshape(
+ input_embeddings.shape[0], -1, input_embeddings.shape[2]
+ )
+
+ # Process vision features
+ # projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+ # test
+ if use_film:
+ # FiLM: Infuse language inputs into visual features
+ raise ValueError
+ patch_features = self.vision_backbone(pixel_values, language_embeddings) # (bsz, 256 * num_images, D)
+ else:
+ patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
+
+ projected_patch_embeddings = self.projector(patch_features)
+ # test end
+
+ # Add proprioceptive features if provided
+ use_proprio = proprio_projector is not None and proprio is not None
+ if use_proprio:
+ proprio = torch.Tensor(proprio).to(
+ projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype
+ )
+ projected_patch_embeddings = self._process_proprio_features(
+ projected_patch_embeddings, proprio, proprio_projector
+ )
+
+ # Use diffusion if provided, otherwise use regression or discrete prediction
+ use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+ # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+ NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+ if use_proprio:
+ NUM_PATCHES += 1
+ if use_diffusion:
+ NUM_PATCHES += 1
+
+ if use_diffusion:
+ raise ValueError
+ # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion
+ noise = torch.randn(
+ size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype
+ )
+
+ # Run diffusion-based prediction
+ normalized_actions, actions_hidden_states = self._run_diffusion_prediction(
+ input_embeddings,
+ all_actions_mask,
+ noise,
+ action_head,
+ projected_patch_embeddings,
+ labels,
+ attention_mask,
+ NUM_PATCHES,
+ NUM_PROMPT_TOKENS,
+ noisy_action_projector,
+ )
+ else:
+ # Run regression or discrete token-based prediction
+ # compute_logits = self._verl_discrete_compute_logits(
+ # input_embeddings,
+ # all_actions_mask,
+ # projected_patch_embeddings,
+ # attention_mask,
+ # labels,
+ # NUM_PATCHES,
+ # NUM_PROMPT_TOKENS,
+ # action_head,
+ # )
+
+ # test
+
+ all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
+ input_embeddings = input_embeddings * ~all_actions_mask
+
+ # Build multimodal embeddings and attention mask
+ # multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+ # input_embeddings, projected_patch_embeddings, attention_mask
+ # )
+ # test
+
+ projected_patch_attention_mask = None
+ if attention_mask is not None:
+ projected_patch_attention_mask = torch.full(
+ (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+ fill_value=True,
+ dtype=attention_mask.dtype,
+ device=attention_mask.device,
+ )
+
+ # Build multimodal embeddings & attention mask; insert embeddings after token (1:)
+ multimodal_embeddings = torch.cat(
+ [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+ )
+
+ multimodal_attention_mask = None
+ if attention_mask is not None:
+ multimodal_attention_mask = torch.cat(
+ [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+ )
+
+ # return multimodal_embeddings, multimodal_attention_mask
+
+ # test end
+
+ # Forward pass through language model
+ language_model_output = self.language_model(
+ input_ids=None,
+ attention_mask=multimodal_attention_mask,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=multimodal_embeddings,
+ labels=None,
+ use_cache=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+
+ compute_logits = language_model_output.logits[
+ :,
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ ]
+
+ # test end
+
+ return compute_logits
+
+
+class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
+ config_class: PretrainedConfig = OpenVLAConfig
+ _supports_sdpa = True
+
+ def __init__(self, config: OpenVLAConfig) -> None:
+ super().__init__(config)
+ self.norm_stats = config.norm_stats
+
+ # Compute action bins
+ self.bins = np.linspace(-1, 1, config.n_action_bins)
+ self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+ # Compute vocab size for de-tokenization -- revert added "multiple of"
+ self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
+
+ def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
+ """Prepares input for action prediction by adding necessary tokens"""
+ # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+ placeholder_action_token_ids = (
+ torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+ )
+ input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+ # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+ stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+ input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+ # Extend the attention mask to fit the new shape of input
+ # Note: Only batch size == 1 supported right now
+ mask_extension = (
+ torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+ .to(attention_mask.device)
+ .to(attention_mask.dtype)
+ )
+ attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+ return input_ids, attention_mask
+
+ def _prepare_labels_for_action_prediction(self, labels, input_ids):
+ """Creates labels tensor for action prediction if not provided"""
+ # Extend labels tensor with fake action labels
+ ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+ labels_extension = (
+ torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+ * ARBITRARY_ACTION_TOKEN_IDX
+ )
+ labels = torch.cat([labels, labels_extension], dim=-1)
+
+ # Replace last label token with stop token
+ labels[:, -1] = STOP_INDEX
+
+ return labels
+
+ def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
+ """Unnormalize actions using dataset statistics"""
+ action_norm_stats = self.get_action_stats(unnorm_key)
+
+ if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+ mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+ action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+ elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+ mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+ action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+ else:
+ raise ValueError("Unsupported action/proprio normalization type detected!")
+
+ actions = np.where(
+ mask,
+ 0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
+ normalized_actions,
+ )
+
+ return actions
+
+ def _run_diffusion_prediction(
+ self,
+ input_embeddings,
+ all_actions_mask,
+ noise,
+ action_head,
+ projected_patch_embeddings,
+ labels,
+ attention_mask,
+ NUM_PATCHES,
+ NUM_PROMPT_TOKENS,
+ noisy_action_projector,
+ ):
+ """Run diffusion-based action prediction"""
+ # Set diffusion timestep values
+ action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
+ # Clone embedding for reuse in each timestep
+ orig_projected_patch_embeddings = projected_patch_embeddings.clone()
+ curr_noisy_actions = noise
+
+ # Reverse diffusion: Iteratively denoise to generate action prediction
+ for t in action_head.noise_scheduler.timesteps:
+ # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
+ # embedding, and diffusion timestep embedding)
+ timesteps = torch.Tensor([t]).to(labels.device)
+ diffusion_timestep_embeddings = (
+ action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
+ ) # (B, llm_dim)
+ diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
+
+ # [Diffusion] Replace the embeddings of the action tokens with noisy actions
+ # (Later on, the positional embeddings will be added to them)
+
+ # For simplicity, append diffusion timestep embedding to the end of projected vision tokens
+ projected_patch_embeddings = torch.cat(
+ (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+ )
+
+ # Reshape and project noisy actions into language embedding space
+ B = curr_noisy_actions.shape[0]
+ orig_curr_noisy_actions_shape = curr_noisy_actions.shape
+ curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
+ noisy_action_features = noisy_action_projector(curr_noisy_actions)
+ curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
+
+ # Replace action token embeddings with noisy action embeddings
+ input_embeddings = self._replace_input_embeddings(
+ input_embeddings.clone(), all_actions_mask, noisy_action_features
+ )
+
+ # Build multimodal embeddings and attention mask
+ multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+ input_embeddings, projected_patch_embeddings, attention_mask
+ )
+
+ # Forward pass through language model
+ language_model_output = self.language_model(
+ input_ids=None,
+ attention_mask=multimodal_attention_mask,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=multimodal_embeddings,
+ labels=None,
+ use_cache=None,
+ output_attentions=False,
+ output_hidden_states=True,
+ return_dict=True,
+ )
+
+ # Extract hidden states for action portion of response
+ last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
+ actions_hidden_states = last_hidden_states[
+ :,
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ :,
+ ] # (B, act_chunk_len, D)
+
+ # Predict noise and update noisy actions: x_t -> x_{t-1}
+ noise_pred = action_head.predict_noise(actions_hidden_states)
+ curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
+
+ curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+ # Return final actions
+ return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
+
+ def _regression_or_discrete_prediction(
+ self,
+ input_embeddings,
+ all_actions_mask,
+ projected_patch_embeddings,
+ attention_mask,
+ labels,
+ NUM_PATCHES,
+ NUM_PROMPT_TOKENS,
+ action_head=None,
+ ):
+ """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+ # Zero out action token embeddings
+ all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
+ input_embeddings = input_embeddings * ~all_actions_mask
+
+ # Build multimodal embeddings and attention mask
+ multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+ input_embeddings, projected_patch_embeddings, attention_mask
+ )
+
+ # Forward pass through language model
+ language_model_output = self.language_model(
+ input_ids=None,
+ attention_mask=multimodal_attention_mask,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=multimodal_embeddings,
+ labels=None,
+ use_cache=None,
+ output_attentions=False,
+ output_hidden_states=True,
+ return_dict=True,
+ )
+
+ # Extract hidden states for action tokens
+ last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
+ actions_hidden_states = last_hidden_states[
+ :,
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ :,
+ ] # (B, act_chunk_len, D)
+
+ # Handle different prediction methods
+ if action_head is not None:
+ # L1 regression prediction
+ normalized_actions = action_head.predict_action(actions_hidden_states)
+ normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+ normalized_actions = normalized_actions.float().cpu().detach().numpy()
+ else:
+ # Discrete token-based prediction
+ predicted_action_token_ids = (
+ language_model_output.logits[
+ :,
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ ]
+ .argmax(dim=2)
+ .cpu()
+ .numpy()
+ )
+ discretized_actions = self.vocab_size - predicted_action_token_ids
+ discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
+ normalized_actions = self.bin_centers[discretized_actions]
+ normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+ return normalized_actions, actions_hidden_states
+
+ def _verl_discrete_prediction(
+ self,
+ input_embeddings,
+ all_actions_mask,
+ projected_patch_embeddings,
+ attention_mask,
+ labels,
+ NUM_PATCHES,
+ NUM_PROMPT_TOKENS,
+ action_head=None,
+ do_sample=True,
+ temperature=1,
+ ):
+ """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+ # Zero out action token embeddings
+ all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
+ input_embeddings = input_embeddings * ~all_actions_mask
+
+ # Build multimodal embeddings and attention mask
+ multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+ input_embeddings, projected_patch_embeddings, attention_mask
+ )
+
+ # Forward pass through language model
+ language_model_output = self.language_model(
+ input_ids=None,
+ attention_mask=multimodal_attention_mask,
+ position_ids=None,
+ past_key_values=None,
+ inputs_embeds=multimodal_embeddings,
+ labels=None,
+ use_cache=None,
+ output_attentions=False,
+ output_hidden_states=False,
+ return_dict=True,
+ )
+
+ # Extract hidden states for action tokens
+ # last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
+ # actions_hidden_states = last_hidden_states[
+ # :,
+ # NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ # :,
+ # ] # (B, act_chunk_len, D)
+
+ # Handle different prediction methods
+ # if action_head is not None:
+ # # L1 regression prediction
+ # normalized_actions = action_head.predict_action(actions_hidden_states)
+ # normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+ # normalized_actions = normalized_actions.float().cpu().detach().numpy()
+ # else:
+ # Discrete token-based prediction
+
+ # test
+ # NUM_PROMPT_TOKENS = NUM_PROMPT_TOKENS + NUM_PATCHES
+ # j = torch.arange(language_model_output.logits.shape[1], device=NUM_PROMPT_TOKENS.device)
+ # start = NUM_PROMPT_TOKENS.unsqueeze(1)
+ # end = start + ACTION_DIM * NUM_ACTIONS_CHUNK
+ # mask_2d = (j >= start) & (j < end)
+ # mask = mask_2d.unsqueeze(-1)
+ # actions_masks = mask.expand_as(language_model_output.logits)
+
+ NUM_PROMPT_TOKENS = NUM_PROMPT_TOKENS + NUM_PATCHES
+ batch_size = language_model_output.logits.shape[0]
+ device = language_model_output.logits.device
+
+ start_indices = NUM_PROMPT_TOKENS.unsqueeze(1) # [batch_size, 1]
+ position_offsets = torch.arange(ACTION_DIM * NUM_ACTIONS_CHUNK, device=device).unsqueeze(0) # [1, seq_length]
+ seq_indices = start_indices + position_offsets # [batch_size, ACTION_DIM*NUM_ACTIONS_CHUNK]
+ # test end
+ # test add
+ # print("language_model_output",language_model_output.logits.shape[-1])
+ # print("self.vocab_size",self.vocab_size) 32000
+ # topk_values, topk_indices = torch.topk(language_model_output.logits, k=256, dim=-1)
+ # print(topk_indices)
+ # assert language_model_output.logits.shape[-1] == self.vocab_size
+ # test add
+ if not do_sample:
+ # org
+ # reponse_ids = language_model_output.logits[
+ # :,
+ # NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES +\
+ # NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ # ].argmax(dim=2)
+ # reponse_ids = language_model_output.logits[actions_masks].argmax(dim=2)
+ # org end
+
+ # padding
+ # reponse_ids = language_model_output.logits[
+ # torch.arange(batch_size, device=device).unsqueeze(-1),
+ # seq_indices,
+ # :
+ # ].argmax(dim=2)
+ # padding end
+
+ # padding + only get last 256 token
+ reponse_ids_logits = language_model_output.logits[
+ torch.arange(batch_size, device=device).unsqueeze(-1), seq_indices, :
+ ]
+ start_index = self.vocab_size - 256
+ response_last256 = reponse_ids_logits[..., -256 - 64 : -64] # Shape: [batch_size, seq_len, 256]
+ last256_argmax = response_last256.argmax(dim=-1) # Shape: [batch_size, seq_len]
+ reponse_ids = last256_argmax + start_index # Shape: [batch_size, seq_len]
+ # padding + only get last 256 token end
+
+ predicted_action_token_ids = reponse_ids.cpu().numpy()
+
+ else:
+ assert temperature > 0
+ # org
+ # action_logits = language_model_output.logits[
+ # :,
+ # NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + \
+ # NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+ # ]
+ # action_logits = language_model_output.logits[actions_masks]
+ # org end
+
+ action_logits = language_model_output.logits[
+ torch.arange(batch_size, device=device).unsqueeze(-1), seq_indices, :
+ ]
+ # padding
+ # scaled_logits = action_logits / temperature
+ # probs = torch.softmax(scaled_logits, dim=-1)
+ # probs_flat = probs.reshape(-1, probs.shape[-1]) # (B*act_chunk_len, vocab_size)
+ # sampled_indices_flat = torch.multinomial(probs_flat, num_samples=1) # (B*act_chunk_len, 1)
+ # reponse_ids = sampled_indices_flat.view(action_logits.shape[0], -1)
+ # padding end
+
+ # padding + only get last 256 token
+ action_logits_last256 = action_logits[..., -256 - 64 : -64]
+ scaled_logits = action_logits_last256 / temperature
+ probs = torch.softmax(scaled_logits, dim=-1)
+ assert probs.shape[-1] == 256
+ probs_flat = probs.reshape(-1, probs.shape[-1])
+ sampled_indices_flat = torch.multinomial(probs_flat, num_samples=1)
+ original_ids_flat = sampled_indices_flat + (self.vocab_size - 256)
+ reponse_ids = original_ids_flat.view(action_logits.shape[0], -1)
+ # padding + only get last 256 token end
+
+ predicted_action_token_ids = reponse_ids.cpu().numpy()
+
+ discretized_actions = self.vocab_size - predicted_action_token_ids
+ discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
+ normalized_actions = self.bin_centers[discretized_actions]
+ # normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+ normalized_actions = normalized_actions.reshape(-1, ACTION_DIM)
+
+ return normalized_actions, reponse_ids
+ # return normalized_actions, actions_hidden_states
+
+ def predict_action(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ unnorm_key: Optional[str] = None,
+ proprio=None,
+ proprio_projector=None,
+ action_head=None,
+ noisy_action_projector=None,
+ use_film: bool = False,
+ **kwargs: str,
+ ) -> np.ndarray:
+ """Predict actions from input sequence, with options for different prediction methods.
+
+ Args:
+ input_ids: Input token ids
+ unnorm_key: Key for unnormalization statistics
+ proprio: Proprioceptive features
+ proprio_projector: Projector for proprioceptive features
+ action_head: Optional head for L1 regression or diffusion-based prediction
+ noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+ use_film: Whether to use FiLM conditioning
+ **kwargs: Additional arguments including pixel_values and attention_mask
+
+ Returns:
+ Tuple of (unnormalized_actions, action_hidden_states)
+ """
+ # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+ # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+ if not torch.all(input_ids[:, -1] == 29871):
+ input_ids = torch.cat(
+ (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+ )
+
+ pixel_values = kwargs["pixel_values"]
+ attention_mask = kwargs["attention_mask"]
+
+ # Create fake labels tensor (needed for action mask)
+ labels = input_ids.clone()
+ labels[:] = IGNORE_INDEX
+
+ # Get number of tokens in prompt (excluding the start token)
+ NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
+
+ # Prepare inputs by adding necessary tokens
+ input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+ # Update labels tensor for action mask computation later
+ labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+
+ # Get input embeddings and action masks
+ input_embeddings = self.get_input_embeddings()(input_ids)
+ all_actions_mask = self._process_action_masks(labels)
+
+ # Extract language embeddings
+ language_embeddings = input_embeddings[~all_actions_mask].reshape(
+ input_embeddings.shape[0], -1, input_embeddings.shape[2]
+ )
+
+ # Process vision features
+ projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+ # Add proprioceptive features if provided
+ use_proprio = proprio_projector is not None and proprio is not None
+ if use_proprio:
+ proprio = torch.Tensor(proprio).to(
+ projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype
+ )
+ projected_patch_embeddings = self._process_proprio_features(
+ projected_patch_embeddings, proprio, proprio_projector
+ )
+
+ # Use diffusion if provided, otherwise use regression or discrete prediction
+ use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+ # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+ NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+ if use_proprio:
+ NUM_PATCHES += 1
+ if use_diffusion:
+ NUM_PATCHES += 1
+
+ if use_diffusion:
+ # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion
+ noise = torch.randn(
+ size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype
+ )
+
+ # Run diffusion-based prediction
+ normalized_actions, actions_hidden_states = self._run_diffusion_prediction(
+ input_embeddings,
+ all_actions_mask,
+ noise,
+ action_head,
+ projected_patch_embeddings,
+ labels,
+ attention_mask,
+ NUM_PATCHES,
+ NUM_PROMPT_TOKENS,
+ noisy_action_projector,
+ )
+ else:
+ # Run regression or discrete token-based prediction
+ normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
+ input_embeddings,
+ all_actions_mask,
+ projected_patch_embeddings,
+ attention_mask,
+ labels,
+ NUM_PATCHES,
+ NUM_PROMPT_TOKENS,
+ action_head,
+ )
+
+ # Unnormalize predicted actions
+ actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+ return actions, actions_hidden_states
+
+ def generate_action_verl(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ unnorm_key: Optional[str] = None,
+ proprio=None,
+ proprio_projector=None,
+ action_head=None,
+ noisy_action_projector=None,
+ use_film: bool = False,
+ **kwargs: str,
+ ) -> np.ndarray:
+ """Predict actions from input sequence, with options for different prediction methods.
+
+ Args:
+ input_ids: Input token ids
+ unnorm_key: Key for unnormalization statistics
+ proprio: Proprioceptive features
+ proprio_projector: Projector for proprioceptive features
+ action_head: Optional head for L1 regression or diffusion-based prediction
+ noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+ use_film: Whether to use FiLM conditioning
+ **kwargs: Additional arguments including pixel_values and attention_mask
+
+ Returns:
+ Tuple of (unnormalized_actions, action_hidden_states)
+ """
+ # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+ # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+ # if not torch.all(input_ids[:, -1] == 29871):
+ # input_ids = torch.cat(
+ # (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+ # )
+
+ pixel_values = kwargs["pixel_values"]
+ attention_mask = kwargs["attention_mask"]
+ do_sample = kwargs["do_sample"]
+ temperature = kwargs["temperature"]
+
+ # Create fake labels tensor (needed for action mask)
+ labels = input_ids.clone()
+ labels[:] = IGNORE_INDEX
+
+ # Get number of tokens in prompt (excluding the start token)
+ # NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
+ # test
+ padding_idx = kwargs["padding_idx"]
+ num_prompt_tokens = input_ids.ne(padding_idx).sum(dim=1) - 1
+ # test end
+
+ # Prepare inputs by adding necessary tokens
+ input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+ # Update labels tensor for action mask computation later
+ labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+
+ # here to convert padding from before to last
+ # test
+ padding_mask = input_ids.ne(padding_idx)
+ assert torch.all(padding_mask == attention_mask.ne(0))
+ # print("in predict_action padding_mask:", padding_mask)
+ padding_mask = padding_mask.int()
+ sorted_indices = torch.argsort(padding_mask, dim=1, descending=True, stable=True)
+ input_ids = torch.gather(input_ids, 1, sorted_indices)
+ attention_mask = torch.gather(attention_mask, 1, sorted_indices)
+ labels = torch.gather(labels, 1, sorted_indices)
+ assert not use_film
+ # test end
+
+ # Get input embeddings and action masks
+ input_embeddings = self.get_input_embeddings()(input_ids)
+ all_actions_mask = self._process_action_masks(labels)
+
+ # Extract language embeddings
+ language_embeddings = input_embeddings[~all_actions_mask].reshape(
+ input_embeddings.shape[0], -1, input_embeddings.shape[2]
+ )
+
+ # Process vision features
+ projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+ # Add proprioceptive features if provided
+ use_proprio = proprio_projector is not None and proprio is not None
+ if use_proprio:
+ proprio = torch.Tensor(proprio).to(
+ projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype
+ )
+ projected_patch_embeddings = self._process_proprio_features(
+ projected_patch_embeddings, proprio, proprio_projector
+ )
+
+ # Use diffusion if provided, otherwise use regression or discrete prediction
+ use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+ # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+ NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+ if use_proprio:
+ NUM_PATCHES += 1
+ if use_diffusion:
+ NUM_PATCHES += 1
+
+ if use_diffusion:
+ raise ValueError
+ # # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion
+ # noise = torch.randn(
+ # size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype
+ # )
+
+ # # Run diffusion-based prediction
+ # normalized_actions, actions_hidden_states = self._run_diffusion_prediction(
+ # input_embeddings,
+ # all_actions_mask,
+ # noise,
+ # action_head,
+ # projected_patch_embeddings,
+ # labels,
+ # attention_mask,
+ # NUM_PATCHES,
+ # NUM_PROMPT_TOKENS,
+ # noisy_action_projector,
+ # )
+ else:
+ # Run regression or discrete token-based prediction
+ normalized_actions, reponse_ids = self._verl_discrete_prediction(
+ input_embeddings,
+ all_actions_mask,
+ projected_patch_embeddings,
+ attention_mask,
+ labels,
+ NUM_PATCHES,
+ num_prompt_tokens,
+ action_head,
+ do_sample=do_sample,
+ temperature=temperature,
+ )
+
+ # Unnormalize predicted actions
+ actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+ # verl add!
+ actions = actions.reshape(-1, NUM_ACTIONS_CHUNK, ACTION_DIM)
+ #
+ return actions, reponse_ids
+
+ @staticmethod
+ def _check_unnorm_key(norm_stats: dict[str, dict[str, Any]], unnorm_key: Optional[str]) -> str:
+ """Validate and resolve the unnormalization key for action statistics"""
+ if unnorm_key is None:
+ assert len(norm_stats) == 1, (
+ f"Your model was trained on more than one dataset, "
+ f"please pass a `unnorm_key` from the following options to choose the statistics "
+ f"used for un-normalizing actions: {norm_stats.keys()}"
+ )
+ unnorm_key = next(iter(norm_stats.keys()))
+
+ assert unnorm_key in norm_stats, (
+ f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
+ f"please choose from: {norm_stats.keys()}"
+ )
+ return unnorm_key
+
+ def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
+ """Get the dimensionality of the policy's action space."""
+ unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+ return len(self.norm_stats[unnorm_key]["action"]["min"])
+
+ def get_action_stats(self, unnorm_key: Optional[str] = None) -> dict[str, Any]:
+ """Get all the logged statistics for the given dataset."""
+ unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+ return self.norm_stats[unnorm_key]["action"]
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/processing_prismatic.py b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/processing_prismatic.py
new file mode 100644
index 0000000000000000000000000000000000000000..4249a5f051001e3bdbc8a9ba77d8481818badd00
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/processing_prismatic.py
@@ -0,0 +1,269 @@
+# Copyright 2025 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
+#
+# https://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 https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/
+# form https://huggingface.co/Haozhan72/Openvla-oft-SFT-libero10-trajall/blob/main/
+
+"""
+processing_prismatic.py
+
+HuggingFace-style preprocessor definitions for Prismatic VLMs, inheriting from `ProcessorMixin`. Default configuration
+specifies `siglip-224px+7b`.
+"""
+
+from typing import Any, ClassVar, Optional
+
+import timm.data
+import torch
+import torchvision.transforms.functional as TVF
+from PIL import Image
+from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
+from transformers import PreTrainedTokenizerBase
+from transformers.image_processing_utils import BatchFeature, ImageProcessingMixin
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType
+
+
+# === Image Processing ===
+def letterbox_pad_transform(image: Image.Image, padding_fill_value: tuple[int, int, int]) -> Image.Image:
+ """Given a PIL.Image, pad to square by adding a symmetric border around the height/width."""
+ (w, h), max_wh = image.size, max(image.size)
+ horizontal_pad, vertical_pad = int((max_wh - w) / 2), int((max_wh - h) / 2)
+ padding = (horizontal_pad, vertical_pad, horizontal_pad, vertical_pad)
+
+ return TVF.pad(image, padding, fill=padding_fill_value, padding_mode="constant")
+
+
+class PrismaticImageProcessor(ImageProcessingMixin):
+ model_input_names: ClassVar[list[str]] = ["pixel_values"]
+
+ def __init__(
+ self,
+ use_fused_vision_backbone: bool = False,
+ image_resize_strategy: str = "letterbox",
+ input_sizes: Optional[list[tuple[int, int, int]]] = None,
+ interpolations: Optional[list[str]] = None,
+ means: Optional[list[tuple[float, float, float]]] = None,
+ stds: Optional[list[tuple[float, float, float]]] = None,
+ **kwargs: str,
+ ) -> None:
+ """
+ Initialize a PrismaticImageProcessor as a wrapper around a torchvision transform; this transform will be
+ created by TIMM, and edited to follow our custom `image_resize_strategy` logic.
+ @param use_fused_vision_backbone: Boolean indicating single or fused (dual) vision backbone
+ @param image_resize_strategy: Prismatic image resize strategy in < resize-naive | resize-crop | letterbox >
+ @param input_size: [TIMM :: `data_cfg`] Input image size as tuple (channels, width, height)
+ @param interpolation: [TIMM :: `data_cfg`] Interpolation as string (default: "bicubic")
+ @param mean: [TIMM :: `data_cfg`] Normalization mean as float tuple (or two-tuple if `fused_backbone`)
+ @param std: [TIMM :: `data_cfg`] Normalization std as float tuple (or two-tuple if `fused_backbone`)
+ """
+ self.use_fused_vision_backbone = use_fused_vision_backbone
+ self.image_resize_strategy = image_resize_strategy
+
+ # Handle `None` default values
+ input_sizes = [(3, 224, 224)] if input_sizes is None else input_sizes
+ means = [(0.5, 0.5, 0.5)] if means is None else means
+ stds = [(0.5, 0.5, 0.5)] if stds is None else stds
+
+ # TIMM `data_cfg` Parameters
+ self.input_sizes, self.interpolations, self.means, self.stds = input_sizes, interpolations, means, stds
+
+ # Grab torchvision transforms via TIMM =>> need to parse for specific "functional" transform values!
+ self.tvf_resize_params, self.tvf_crop_params, self.tvf_normalize_params = [], [], []
+ self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
+
+ for idx in range(len(input_sizes)):
+ transform = timm.data.create_transform(
+ input_size=self.input_sizes[idx],
+ interpolation=self.interpolations[idx],
+ mean=self.means[idx],
+ std=self.stds[idx],
+ crop_pct=1.0, # Set to 1.0 to ignore cropping (initial Resize sets `input_size`)
+ crop_mode="center", # Default crop mode -- no-op when `crop_pct == 1.0`
+ is_training=False, # No image augmentations when loading the transform!
+ )
+
+ # [Validation] Ensure appropriate transform structure, expected sizes
+ if not (
+ isinstance(transform, Compose)
+ and (len(transform.transforms) == 4)
+ and isinstance(transform.transforms[0], Resize)
+ and isinstance(transform.transforms[1], CenterCrop)
+ and isinstance(transform.transforms[2], ToTensor)
+ and isinstance(transform.transforms[3], Normalize)
+ and (transform.transforms[0].size == self.input_sizes[idx][-1])
+ and (transform.transforms[1].size == self.input_sizes[idx][-2:])
+ ):
+ raise ValueError(f"Unexpected TIMM image transformation structure/sizes: `{transform}`")
+
+ # HF Image Processors *must* be JSON-serializable; as such, cannot have torchvision. as an attribute.
+ # => Instead, we're going to parse the transform and call "torchvision.transforms.functional" (`tvf`)
+ resize_t, crop_t, norm_t = transform.transforms[0], transform.transforms[1], transform.transforms[3]
+ self.tvf_resize_params.append(
+ {
+ "size": resize_t.size,
+ "interpolation": TVF.pil_modes_mapping[resize_t.interpolation],
+ "max_size": None,
+ "antialias": True,
+ }
+ )
+ self.tvf_crop_params.append({"output_size": crop_t.size})
+ self.tvf_normalize_params.append(
+ {
+ "mean": norm_t.mean.float().numpy().tolist(),
+ "std": norm_t.std.float().numpy().tolist(),
+ "inplace": False,
+ }
+ )
+ self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
+
+ # Handle Prismatic `image_resize_strategy`
+ if self.image_resize_strategy == "resize-naive":
+ self.tvf_resize_params[idx]["size"] = (resize_t.size, resize_t.size)
+ elif self.image_resize_strategy == "letterbox":
+ self.tvf_do_letterbox, self.tvf_letterbox_fill = True, tuple([int(x * 255) for x in self.means[idx]])
+ elif self.image_resize_strategy == "resize-crop":
+ pass
+ else:
+ raise ValueError(f"Image resize strategy `{self.image_resize_strategy}` is not supported!")
+
+ # Dispatch **kwargs to super()
+ super().__init__(**kwargs)
+
+ def apply_transform(self, img: Image.Image) -> torch.Tensor:
+ """Apply `functional` variant of TIMM's Transform = Compose([Resize -> CenterCrop -> ToTensor -> Normalize])"""
+ if self.tvf_do_letterbox:
+ img = letterbox_pad_transform(img, self.tvf_letterbox_fill)
+
+ # [Contract] Fused Backbones expect "channel-stacked" inputs; we'll unpack on the model side!
+ imgs_t = []
+ for idx in range(len(self.input_sizes)):
+ img_idx = TVF.resize(img, **self.tvf_resize_params[idx])
+ img_idx = TVF.center_crop(img_idx, **self.tvf_crop_params[idx])
+ img_idx_t = TVF.to_tensor(img_idx)
+ img_idx_t = TVF.normalize(img_idx_t, **self.tvf_normalize_params[idx])
+ imgs_t.append(img_idx_t)
+
+ # [Contract] `imgs_t` is a list of Tensors of shape [3, input_size, input_size]; stack along dim = 0
+ img_t = torch.vstack(imgs_t)
+
+ return img_t
+
+ def preprocess(
+ self,
+ images: Image.Image | list[Image.Image],
+ return_tensors: Optional[str | TensorType] = None,
+ **_: str,
+ ) -> BatchFeature:
+ """
+ Preprocess an image (or batch of images); note that unlike the `transformers :: BaseImageProcessor` we
+ explicitly only handle PIL.Image.Image instances for simplicity.
+ @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
+ @param return_tensors: BatchFeature default Tensor format (e.g., "pt" for torch); if None, returns np.ndarray
+ @return: Instance of `transformers :: BatchFeature` with a single key "pixel_values"
+ """
+ if not isinstance(images, list):
+ images = [images]
+
+ # Apply `self.img_transform` to each image (will return list of torch.Tensors); stack into "batched" Tensor
+ pixel_values = torch.stack([self.apply_transform(img.convert("RGB")) for img in images])
+
+ # Return BatchFeature =>> note that for compatibility, constructor expects Dict[str, np.ndarray], so we convert
+ return BatchFeature(data={"pixel_values": pixel_values.float().numpy()}, tensor_type=return_tensors)
+
+ def __call__(self, images: Image.Image | list[Image.Image], **kwargs) -> BatchFeature:
+ return self.preprocess(images, **kwargs)
+
+
+# === PrismaticProcessor =>> Wraps both ImageProcessor and Tokenizer ===
+# =>> https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/processing_llava.py
+class PrismaticProcessor(ProcessorMixin):
+ attributes: ClassVar[list[str]] = ["image_processor", "tokenizer"]
+ image_processor_class: str = "AutoImageProcessor"
+ tokenizer_class: str = "AutoTokenizer"
+
+ def __init__(
+ self,
+ image_processor: Optional[ImageProcessingMixin] = None,
+ tokenizer: Optional[PreTrainedTokenizerBase] = None,
+ ) -> None:
+ super().__init__(image_processor, tokenizer)
+
+ def __call__(
+ self,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput],
+ images: Image.Image | list[Image.Image],
+ padding: bool | str | PaddingStrategy = False,
+ truncation: Optional[bool | str | TruncationStrategy] = None,
+ max_length: Optional[int] = None,
+ return_tensors: Optional[str | TensorType] = TensorType.PYTORCH,
+ ) -> BatchFeature:
+ """
+ Preprocess a given (batch) of text/images for a Prismatic VLM; forwards text to the underlying LLM's tokenizer,
+ forwards images to PrismaticImageProcessor.
+ @param text: The (batch) of text to encode; must be a string or list of strings.
+ @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
+ @param padding: Sequence padding strategy (if multiple specified) in < True = "longest" | "max_length" | False >
+ @param truncation: Truncation strategy for the output sequences; requires `max_length` to be specified
+ @param max_length: Maximum length (in tokens) to truncate
+ @param return_tensors: Type of return tensors (usually "pt" or TensorType.PYTORCH)
+ @return: BatchFeature with keys for `input_ids`, `attention_mask` and `pixel_values`.
+ """
+ pixel_values = self.image_processor(images, return_tensors=return_tensors)["pixel_values"]
+ text_inputs = self.tokenizer(
+ text, return_tensors=return_tensors, padding=padding, truncation=truncation, max_length=max_length
+ )
+
+ # [Validate] Need same number of images and text inputs!
+ if pixel_values.shape[0] != text_inputs.input_ids.shape[0]:
+ raise ValueError("Batch is malformed; expected same number of images and text inputs!")
+
+ return BatchFeature(data={**text_inputs, "pixel_values": pixel_values})
+
+ # === Tokenizer Dispatch Utilities =>> check `PreTrainedTokenizerBase` for documentation ===
+ def batch_decode(
+ self,
+ sequences: list[int] | list[list[int]] | torch.Tensor | Any, # `Any` = np.ndarray | tf.Tensor
+ skip_special_tokens: bool = False,
+ clean_up_tokenization_spaces: Optional[bool] = None,
+ **kwargs: str,
+ ) -> list[str]:
+ return self.tokenizer.batch_decode(
+ sequences=sequences,
+ skip_special_tokens=skip_special_tokens,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **kwargs,
+ )
+
+ def decode(
+ self,
+ token_ids: int | list[int] | torch.Tensor | Any, # `Any` = np.ndarray | tf.Tensor
+ skip_special_tokens: bool = False,
+ clean_up_tokenization_spaces: Optional[bool] = None,
+ **kwargs: str,
+ ) -> str:
+ return self.tokenizer.decode(
+ token_ids=token_ids,
+ skip_special_tokens=skip_special_tokens,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **kwargs,
+ )
+
+ @property
+ def model_input_names(self) -> list[str]:
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+
+ return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/train_utils.py b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/train_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..74f0935db8b227ba84b38f1a107982060aa69478
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/models/openvla_oft/train_utils.py
@@ -0,0 +1,72 @@
+# Copyright 2025 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
+#
+# https://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 https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/
+
+"""Utils for training/fine-tuning scripts."""
+
+import torch
+
+from .constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX
+
+
+def get_current_action_mask(token_ids):
+ # Create a tensor marking positions of IGNORE_INDEX
+ newline_positions = token_ids != IGNORE_INDEX
+
+ # Calculate cumulative sum to identify regions between newlines
+ cumsum = torch.cumsum(newline_positions, dim=1)
+
+ # Create the mask
+ mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)
+
+ # Extract the action part only
+ action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX
+ mask = action_tokens_only_mask * mask
+
+ return mask
+
+
+def get_next_actions_mask(token_ids):
+ # Create a tensor marking positions of IGNORE_INDEX
+ newline_positions = token_ids != IGNORE_INDEX
+
+ # Calculate cumulative sum to identify regions between newlines
+ cumsum = torch.cumsum(newline_positions, dim=1)
+
+ # Create the mask
+ mask = cumsum > ACTION_DIM
+
+ # Extract the action part only
+ action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX
+ mask = action_tokens_only_mask * mask
+
+ return mask
+
+
+def compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask):
+ correct_preds = (predicted_token_ids == ground_truth_token_ids) & mask
+ accuracy = correct_preds.sum().float() / mask.sum().float()
+ return accuracy
+
+
+def compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask):
+ pred_continuous_actions = torch.tensor(
+ action_tokenizer.decode_token_ids_to_actions(predicted_token_ids[mask].cpu().numpy())
+ )
+ true_continuous_actions = torch.tensor(
+ action_tokenizer.decode_token_ids_to_actions(ground_truth_token_ids[mask].cpu().numpy())
+ )
+ l1_loss = torch.nn.functional.l1_loss(pred_continuous_actions, true_continuous_actions)
+ return l1_loss
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/naive_rollout_rob.py b/code/RL_model/verl/verl_train/verl/experimental/vla/naive_rollout_rob.py
new file mode 100644
index 0000000000000000000000000000000000000000..e4b56f4c7697ac1a2f7008a6e401660ef49a7fe8
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/naive_rollout_rob.py
@@ -0,0 +1,226 @@
+# Copyright 2025 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.
+"""
+In single GPU rollout, the sequences are generated directly by sampling from the model.
+The output will contain
+1. output_ids
+2. attention_masks (left padding)
+3. eos_masks
+4. log_probs
+"""
+
+import json
+import logging
+import os
+
+import torch
+from PIL import Image
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.nn.utils.rnn import pad_sequence
+
+from verl import DataProto
+from verl.experimental.vla.envs.action_utils import center_crop_image, resize_image
+from verl.experimental.vla.models.openvla_oft.modeling_prismatic import OpenVLAForActionPrediction
+from verl.experimental.vla.models.openvla_oft.processing_prismatic import PrismaticProcessor
+from verl.utils.device import get_device_id, get_device_name, get_torch_device
+from verl.workers.rollout.base import BaseRollout
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = ["NaiveRolloutRob"]
+
+
+def pad_sequence_to_length(tensors, max_seq_len, pad_token_id, left_pad=False):
+ """
+ pad a 2D tensors (e.g. responses, logprobs) in the last dim to max_seq_length.
+ input shape: [bs, seq_length]
+ output shape: [bs, max_seq_length]
+ (0, max_seq_len - tensors.shape[-1]) means right pad to max_seq_length and no left pad
+ """
+ if tensors.shape[-1] >= max_seq_len:
+ return tensors
+ pad_tuple = (max_seq_len - tensors.shape[-1], 0) if left_pad else (0, max_seq_len - tensors.shape[-1])
+ return torch.nn.functional.pad(tensors, pad_tuple, "constant", pad_token_id)
+
+
+def process_input(task_descriptions, images_and_states, processor):
+ batchdata = {"input_ids": [], "attention_mask": [], "pixel_values": []}
+
+ for i in range(len(task_descriptions)):
+ task_description = task_descriptions[i]
+ image = resize_image(images_and_states["full_image"][i].cpu().numpy(), (224, 224))
+ image = Image.fromarray(image).convert("RGB")
+ image = center_crop_image(image)
+ prompt = f"In: What action should the robot take to {task_description.lower()}?\nOut:"
+ batch_feature = processor(prompt, image)
+
+ input_ids = batch_feature["input_ids"]
+ attention_mask = batch_feature["attention_mask"]
+ pixel_values = batch_feature["pixel_values"]
+
+ if not torch.all(input_ids[:, -1] == 29871):
+ input_ids = torch.cat(
+ (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+ )
+ attention_mask = torch.cat(
+ (attention_mask, torch.unsqueeze(torch.Tensor([True]).bool(), dim=0).to(attention_mask.device)), dim=1
+ )
+
+ batchdata["input_ids"].append(input_ids)
+ batchdata["attention_mask"].append(attention_mask)
+ batchdata["pixel_values"].append(pixel_values)
+
+ device = get_device_id()
+
+ batchdata["input_ids"] = [x.transpose(0, 1) for x in batchdata["input_ids"]]
+ batchdata["attention_mask"] = [x.transpose(0, 1) for x in batchdata["attention_mask"]]
+ batchdata["input_ids"] = (
+ pad_sequence(batchdata["input_ids"], batch_first=True, padding_value=processor.tokenizer.pad_token_id)
+ .squeeze(-1)
+ .to(device)
+ )
+ batchdata["attention_mask"] = (
+ pad_sequence(batchdata["attention_mask"], batch_first=True, padding_value=0).squeeze(-1).to(device)
+ )
+
+ padding_mask = batchdata["input_ids"].ne(processor.tokenizer.pad_token_id)
+ assert torch.all(padding_mask == batchdata["attention_mask"].ne(0))
+ padding_mask = ~padding_mask
+ padding_mask = padding_mask.int()
+ sorted_indices = torch.argsort(padding_mask, dim=1, descending=True, stable=True)
+ batchdata["input_ids"] = torch.gather(batchdata["input_ids"], 1, sorted_indices)
+ batchdata["attention_mask"] = torch.gather(batchdata["attention_mask"], 1, sorted_indices)
+
+ batchdata["pixel_values"] = torch.cat(batchdata["pixel_values"], dim=0).to(device)
+ assert torch.all(batchdata["attention_mask"].ne(0) == batchdata["input_ids"].ne(processor.tokenizer.pad_token_id))
+
+ return batchdata
+
+
+class NaiveRolloutRob(BaseRollout):
+ def __init__(
+ self,
+ model_config: dict,
+ module: torch.nn.Module = None,
+ ):
+ self.model_config = model_config
+ if module is not None:
+ self.module = module
+ else:
+ self.module = OpenVLAForActionPrediction.from_pretrained(model_config["path"], trust_remote_code=True)
+ self.module.vision_backbone.set_num_images_in_input(1)
+ self.processor = PrismaticProcessor.from_pretrained(model_config["path"], trust_remote_code=True)
+ dataset_statistics_path = os.path.join(model_config["path"], "dataset_statistics.json")
+ if os.path.isfile(dataset_statistics_path):
+ with open(dataset_statistics_path) as f:
+ norm_stats = json.load(f)
+ if isinstance(self.module, FSDP):
+ self.module.module.norm_stats = norm_stats
+ else:
+ self.module.norm_stats = norm_stats
+ self.module.eval()
+
+ @torch.no_grad()
+ def _generate_one_step(self, prompts: dict, do_sample, temperature, max_prompt_length):
+ idx = prompts["input_ids"] # (bs, prompt_length)
+ attention_mask = prompts["attention_mask"] # left-padded attention_mask
+ pixel_values = prompts["pixel_values"]
+
+ with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):
+ actions, response = self.module.generate_action_verl(
+ input_ids=idx,
+ pixel_values=pixel_values,
+ attention_mask=attention_mask,
+ padding_idx=self.processor.tokenizer.pad_token_id,
+ do_sample=do_sample,
+ unnorm_key="libero_10_no_noops",
+ temperature=temperature,
+ )
+
+ assert self.processor.tokenizer.pad_token_id is not None
+
+ assert idx.ndim == 2
+ idx = pad_sequence_to_length(
+ idx, max_seq_len=max_prompt_length, pad_token_id=self.processor.tokenizer.pad_token_id, left_pad=True
+ )
+
+ assert attention_mask.ndim == 2
+ attention_mask = pad_sequence_to_length(
+ attention_mask, max_seq_len=max_prompt_length, pad_token_id=0, left_pad=True
+ )
+
+ device_type = get_device_name()
+ assert idx.device.type == device_type
+ assert response.device.type == device_type
+ assert attention_mask.device.type == device_type
+ assert pixel_values.device.type == device_type
+ batch = {
+ "responses": response,
+ "input_ids": idx,
+ "attention_mask": attention_mask,
+ "pixel_values": pixel_values,
+ "action": actions,
+ }
+
+ return batch
+
+ # @conditional_profiler(name="generate_sequences", path="traces/rollout", max_steps=5)
+ @torch.no_grad()
+ def generate_sequences(self, prompts: DataProto) -> DataProto:
+ """Generate sequences"""
+ # make sampling args can be overriden by inputs
+ do_sample = prompts.meta_info["do_sample"]
+ temperature = prompts.meta_info["temperature"]
+ max_prompt_length = prompts.meta_info["prompt_length"]
+ # TODO: split into micro-batches
+ task_descriptions = prompts.non_tensor_batch["task_descriptions"]
+ images_and_states = {"full_image": prompts.batch["full_image"]}
+ vla_input = process_input(task_descriptions, images_and_states, self.processor)
+
+ vla_output = self._generate_one_step(vla_input, do_sample, temperature, max_prompt_length)
+ # batch = TensorDict(vla_output)
+ batch = DataProto.from_dict(tensors=vla_output)
+ return batch
+
+ async def update_weights(self, weights_iterator, **kwargs):
+ prefix = "_fsdp_wrapped_module."
+ target_state_dict = self.module.state_dict()
+ loaded_tensors_count = 0
+ for name, param in weights_iterator:
+ cleaned_name = name.replace(prefix, "")
+ if cleaned_name in target_state_dict:
+ target_tensor = target_state_dict[cleaned_name]
+ try:
+ target_tensor.copy_(param, non_blocking=True)
+ loaded_tensors_count += 1
+ except Exception as e:
+ logger.warning(f"Warning: Failed to copy tensor '{cleaned_name}'. Error: {e}")
+ else:
+ logger.warning(f"Warning: Failed to copy tensor '{cleaned_name}'. Model has no such key.")
+ logger.info(f"Rollout model weights updated. Loaded {loaded_tensors_count} tensors one by one.")
+
+ async def release(self):
+ if self.module.device.type == get_device_name():
+ logger.info("Releasing rollout model to CPU.")
+ self.module.cpu()
+ self.device = torch.device("cpu")
+ get_torch_device().empty_cache()
+
+ async def resume(self, **kwargs):
+ if self.module.device.type == "cpu":
+ target_device = get_device_name()
+ logger.info(f"Resuming rollout model to device: {target_device}.")
+ self.module.to(target_device)
+ self.device = torch.device(target_device)
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/prepare_libero_dataset.py b/code/RL_model/verl/verl_train/verl/experimental/vla/prepare_libero_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..9c2ce204811926437067b098978cd1e5c0ef4853
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/prepare_libero_dataset.py
@@ -0,0 +1,154 @@
+# Copyright 2024 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.
+"""
+Preprocess the Geometry3k dataset to parquet format
+"""
+
+import argparse
+import os
+import random
+
+import numpy as np
+import torch
+from datasets import Dataset
+from libero.libero import get_libero_path
+from libero.libero.benchmark import Benchmark, get_benchmark
+
+
+def patched_get_task_init_states(self, i):
+ init_states_path = os.path.join(
+ get_libero_path("init_states"),
+ self.tasks[i].problem_folder,
+ self.tasks[i].init_states_file,
+ )
+ init_states = torch.load(init_states_path, weights_only=False)
+ return init_states
+
+
+Benchmark.get_task_init_states = patched_get_task_init_states
+
+
+def compute_total_num_group_envs(task_suite: Benchmark):
+ total_num_group_envs = 0
+ trial_id_bins = []
+ for task_id in range(task_suite.get_num_tasks()):
+ task_num_trials = len(task_suite.get_task_init_states(task_id))
+ trial_id_bins.append(task_num_trials)
+
+ total_num_group_envs += task_num_trials
+
+ cumsum_trial_id_bins = np.cumsum(trial_id_bins)
+ return total_num_group_envs, cumsum_trial_id_bins
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--task_suite_name", default="libero_10")
+ parser.add_argument(
+ "--local_save_dir", default="~/data/libero_rl", help="The save directory for the preprocessed dataset."
+ )
+ args = parser.parse_args()
+ random.seed(42)
+ np.random.seed(42)
+ task_suite = get_benchmark("libero_10")()
+ total_num_group_envs, cumsum_trial_id_bins = compute_total_num_group_envs(task_suite)
+ print(f"Total number of group envs: {total_num_group_envs}")
+ print(f"Cumsum trial id bins: {cumsum_trial_id_bins}")
+
+ # Total number of group envs: 500
+ # Cumsum trial id bins: [ 50 100 150 200 250 300 350 400 450 500]
+ def get_state_ids_for_task(task_id):
+ start_id = 0 if task_id == 0 else cumsum_trial_id_bins[task_id - 1]
+ end_id = cumsum_trial_id_bins[task_id]
+ return list(range(start_id, end_id))
+
+ all_task_ids = list(range(task_suite.get_num_tasks()))
+ train_task_ids = sorted(random.sample(all_task_ids, 9))
+ ood_test_task_id = list(set[int](all_task_ids) - set(train_task_ids))[0] # for OOD test
+
+ print("\n[Data Split Plan]")
+ print(f"Training Task IDs: {train_task_ids}")
+ print(f"OOD Test Task ID: {ood_test_task_id}")
+ train_metadata = []
+ test_metadata = []
+ for task_id in train_task_ids:
+ all_trials = get_state_ids_for_task(task_id)
+ random.shuffle(all_trials)
+ selected_train_trials = all_trials[:40]
+ for state_id in selected_train_trials:
+ train_metadata.append({"task_id": task_id, "state_id": state_id, "data_source": "train"})
+
+ # ID
+ for task_id in train_task_ids:
+ all_trials = get_state_ids_for_task(task_id)
+ random.shuffle(all_trials)
+ selected_id_test_trials = all_trials[40:]
+ for state_id in selected_id_test_trials[:10]:
+ test_metadata.append({"task_id": task_id, "state_id": state_id, "data_source": "test_in_distribution"})
+
+ # OOD
+ ood_all_trials = get_state_ids_for_task(ood_test_task_id)
+ random.shuffle(ood_all_trials)
+ selected_ood_trials = ood_all_trials[:20]
+ for state_id in selected_ood_trials:
+ test_metadata.append(
+ {"task_id": ood_test_task_id, "state_id": state_id, "data_source": "test_out_of_distribution"}
+ )
+ print(f"Generated {len(train_metadata)} training samples.")
+ print(f"Generated {len(test_metadata)} testing samples.")
+ print("-" * 20)
+ train_ds_meta = Dataset.from_list(train_metadata)
+ test_ds_meta = Dataset.from_list(test_metadata)
+
+ def map_and_process(example, idx):
+ task_id = example["task_id"]
+ state_id = example["state_id"]
+ data_source = example["data_source"]
+ split = "train" if data_source == "train" else "test"
+ task = task_suite.get_task(task_id)
+ # demonstration = task.get_demonstration(state_id)
+
+ data = {
+ "data_source": data_source,
+ "prompt": task.language,
+ "state_ids": state_id,
+ "task_ids": task_id,
+ "ability": "robot",
+ "extra_info": {
+ "split": split,
+ "state_ids": state_id,
+ "index": idx,
+ "task": task,
+ "task_ids": task_id,
+ },
+ }
+ return data
+
+ print("Mapping and processing training dataset...")
+ train_dataset = train_ds_meta.map(map_and_process, with_indices=True, num_proc=8)
+ print("Mapping and processing test dataset...")
+ test_dataset = test_ds_meta.map(map_and_process, with_indices=True, num_proc=8)
+ local_save_dir = os.path.expanduser(args.local_save_dir)
+ os.makedirs(local_save_dir, exist_ok=True)
+ print(f"Saving training dataset to {os.path.join(local_save_dir, 'train.parquet')}")
+ train_dataset.to_parquet(os.path.join(local_save_dir, "train.parquet"))
+ print(f"Saving test dataset to {os.path.join(local_save_dir, 'test.parquet')}")
+ test_dataset.to_parquet(os.path.join(local_save_dir, "test.parquet"))
+ print("\nDataset generation complete!")
+
+ print("\n--- Verification ---")
+ print("Train dataset data sources:", train_dataset.unique("data_source"))
+ print("Test dataset data sources:", test_dataset.unique("data_source"))
+ print("Train dataset length:", len(train_dataset))
+ print("Test dataset length:", len(test_dataset))
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/requirements_vla.txt b/code/RL_model/verl/verl_train/verl/experimental/vla/requirements_vla.txt
new file mode 100644
index 0000000000000000000000000000000000000000..11705d307c455bc7a2dfb66cabed89f5c4a7246e
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/requirements_vla.txt
@@ -0,0 +1,19 @@
+# libero
+timm<1.0.0
+imageio
+draccus==0.8.0
+einops
+huggingface_hub
+json-numpy
+jsonlines
+matplotlib
+rich
+sentencepiece==0.1.99
+# dlimp @ git+https://github.com/moojink/dlimp_openvla
+diffusers==0.30.3
+imageio
+uvicorn
+fastapi
+json-numpy
+wandb==0.19.11
+protobuf==3.20.3
\ No newline at end of file
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/rob_ray_trainer.py b/code/RL_model/verl/verl_train/verl/experimental/vla/rob_ray_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2046d6c4844d82e08863f04b95961775c82cb50
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/rob_ray_trainer.py
@@ -0,0 +1,669 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023-2024 SGLang Team
+# Copyright 2025 ModelBest Inc. 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.
+"""
+PPO Trainer with Ray-based single controller.
+This trainer supports model-agonistic model initialization with huggingface
+"""
+
+import asyncio
+import uuid
+from collections import defaultdict
+from pprint import pprint
+
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from tqdm import tqdm
+
+from verl import DataProto
+from verl.experimental.dataset.sampler import AbstractCurriculumSampler
+from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
+from verl.single_controller.ray import RayClassWithInitArgs
+from verl.single_controller.ray.base import create_colocated_worker_cls
+from verl.trainer.ppo.core_algos import agg_loss
+from verl.trainer.ppo.metric_utils import (
+ compute_data_metrics,
+ compute_throughout_metrics,
+ compute_timing_metrics,
+ process_validation_metrics,
+)
+from verl.trainer.ppo.ray_trainer import RayPPOTrainer, apply_kl_penalty, compute_advantage
+from verl.trainer.ppo.reward import compute_reward
+from verl.trainer.ppo.utils import Role
+from verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi
+from verl.utils.debug import marked_timer
+from verl.utils.metric import reduce_metrics
+
+
+def compute_response_mask(data: DataProto) -> torch.Tensor:
+ """Compute the attention mask for the response part of the sequence.
+
+ This function extracts the portion of the attention mask that corresponds to the model's response,
+ which is used for masking computations that should only apply to response tokens.
+
+ Args:
+ data (DataProto): The data containing batched model outputs and inputs.
+
+ Returns:
+ torch.Tensor: The attention mask for the response tokens.
+ """
+ complete = data.batch["complete"] # shape: [batch_size, num_steps, chunk_size]
+
+ complete_traj = complete.view(complete.shape[0], -1) # # shape: [batch_size, num_steps * chunk_size]
+ batch_size, action_steps = complete_traj.shape
+
+ step_indices = torch.arange(action_steps, device=complete.device).unsqueeze(0).expand(batch_size, -1)
+
+ first_true_idx_approx = torch.argmax(complete_traj.long(), dim=1)
+
+ has_any_true = complete_traj.any(dim=1)
+
+ final_first_true_idx = torch.where(
+ has_any_true, first_true_idx_approx, torch.tensor(action_steps - 1, device=complete.device)
+ )
+
+ mask_traj = step_indices <= final_first_true_idx.unsqueeze(1)
+
+ mask = mask_traj.view(complete.shape) # shape: [batch_size, num_steps, chunk_size]
+ mask = mask.repeat_interleave(7, dim=-1) # eapand to action dim
+ return mask
+
+
+def flatten_trajectories(data: DataProto) -> DataProto:
+ batch_size, num_steps = data.batch["action"].shape[:2]
+ new_batch_fields = {}
+ for key, tensor in data.batch.items():
+ if len(tensor.shape) >= 2 and tensor.shape[0] == batch_size and tensor.shape[1] == num_steps:
+ # (B, S, H, W) -> (B*S, H, W)
+ new_shape = (batch_size * num_steps, *tensor.shape[2:])
+ new_batch_fields[key] = tensor.reshape(new_shape)
+ elif len(tensor.shape) == 1 and tensor.shape[0] == batch_size:
+ # [e1, e2] -> [e1, e1, ..., e2, e2, ...] (S times each)
+ new_batch_fields[key] = tensor.repeat_interleave(num_steps)
+ else:
+ new_batch_fields[key] = tensor
+ new_data = DataProto.from_dict(tensors=new_batch_fields, meta_info=data.meta_info)
+ return new_data
+
+
+# def filter_by_acc(data: DataProto, accuracy_lower_bound, accuracy_upper_bound) -> torch.Tensor:
+
+
+class RobRayPPOTrainer(RayPPOTrainer):
+ """Distributed PPO trainer using Ray for scalable reinforcement learning.
+
+ This trainer orchestrates distributed PPO training across multiple nodes and GPUs,
+ managing actor rollouts, critic training, and reward computation with Ray backend.
+ Supports various model architectures including FSDP, Megatron, vLLM, and SGLang integration.
+ """
+
+ def _start_profiling(self, do_profile: bool) -> None:
+ """Start profiling for all worker groups including env workers."""
+ super()._start_profiling(do_profile)
+ if do_profile and hasattr(self, "env_wg"):
+ self.env_wg.start_profile(role="env", profile_step=self.global_steps)
+
+ def _stop_profiling(self, do_profile: bool) -> None:
+ """Stop profiling for all worker groups including env workers."""
+ super()._stop_profiling(do_profile)
+ if do_profile and hasattr(self, "env_wg"):
+ self.env_wg.stop_profile()
+
+ def init_workers(self):
+ self.resource_pool_manager.create_resource_pool()
+
+ if self.config.env.disagg_sim.enable:
+ # pin EnvWorker to Simulator GPU nodes
+ self.resource_pool_manager.get_resource_pool(Role.Env).accelerator_type = "sim"
+ self.resource_pool_manager.get_resource_pool(Role.ActorRollout).accelerator_type = "train_rollout"
+
+ self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.ActorRollout)
+ actor_rollout_cls = RayClassWithInitArgs(
+ cls=self.role_worker_mapping[Role.ActorRollout],
+ config=self.config.actor_rollout_ref,
+ role="actor_rollout",
+ )
+ self.resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls
+
+ assert Role.Env in self.role_worker_mapping
+ if Role.Env in self.role_worker_mapping:
+ resource_pool = self.resource_pool_manager.get_resource_pool(Role.Env)
+ env_cls = RayClassWithInitArgs(self.role_worker_mapping[Role.Env], config=self.config.env)
+ self.resource_pool_to_cls[resource_pool]["env"] = env_cls
+
+ # 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`.
+ # Instead, directly pass different resource pool to different worker groups.
+ # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.
+ all_wg = {}
+ wg_kwargs = {} # Setting up kwargs for RayWorkerGroup
+ if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None:
+ wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout
+ if OmegaConf.select(self.config.global_profiler, "steps") is not None:
+ wg_kwargs["profile_steps"] = OmegaConf.select(self.config.global_profiler, "steps")
+ # Only require nsight worker options when tool is nsys
+ if OmegaConf.select(self.config.global_profiler, "tool") == "nsys":
+ assert (
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ is not None
+ ), "worker_nsight_options must be set when using nsys with profile_steps"
+ wg_kwargs["worker_nsight_options"] = OmegaConf.to_container(
+ OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options")
+ )
+ wg_kwargs["device_name"] = self.device_name
+
+ for resource_pool, class_dict in self.resource_pool_to_cls.items():
+ worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
+ wg_dict = self.ray_worker_group_cls(
+ resource_pool=resource_pool,
+ ray_cls_with_init=worker_dict_cls,
+ **wg_kwargs,
+ )
+ spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())
+ all_wg.update(spawn_wg)
+
+ # 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["actor_rollout"]
+ self.actor_rollout_wg.init_model()
+ self.env_wg = all_wg["env"]
+
+ # create async rollout manager and request scheduler
+ self.async_rollout_mode = False
+ if self.config.actor_rollout_ref.rollout.mode == "async_envloop":
+ from verl.experimental.vla.env_loop import EnvLoop
+
+ self.async_rollout_mode = True
+ self.async_rollout_manager = EnvLoop(
+ config=self.config, rollout_wg=self.actor_rollout_wg, env_wg=self.env_wg
+ )
+
+ def _get_gen_batch(self, batch: DataProto) -> DataProto:
+ # pop those keys for generation
+ batch_keys_to_pop = []
+ non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys())
+ gen_batch = batch.pop(
+ batch_keys=batch_keys_to_pop,
+ non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop),
+ )
+
+ return gen_batch
+
+ def _reset_envs(self, gen_batch: DataProto) -> asyncio.Future:
+ initial_state_ids = gen_batch.non_tensor_batch["state_ids"]
+ task_ids = gen_batch.non_tensor_batch["task_ids"]
+ reset_prompts = DataProto.from_dict(non_tensors={"state_ids": initial_state_ids, "task_ids": task_ids})
+ reset_future = self.env_wg.reset_envs_to_state_ids(reset_prompts)
+ return reset_future
+
+ def fit(self):
+ """
+ The training loop of PPO.
+ The driver process only need to call the compute functions of the worker group through RPC
+ to construct the PPO dataflow.
+ The light-weight advantage computation is done on the driver process.
+ """
+ from omegaconf import OmegaConf
+
+ from verl.utils.tracking import Tracking
+
+ logger = Tracking(
+ project_name=self.config.trainer.project_name,
+ experiment_name=self.config.trainer.experiment_name,
+ default_backend=self.config.trainer.logger,
+ config=OmegaConf.to_container(self.config, resolve=True),
+ )
+
+ self.global_steps = 0
+
+ # load checkpoint before doing anything
+ self._load_checkpoint()
+
+ # perform validation before training
+ # currently, we only support validation using the reward_function.
+ if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
+ val_metrics = self._validate()
+ assert val_metrics, f"{val_metrics=}"
+ pprint(f"Initial validation metrics: {val_metrics}")
+ logger.log(data=val_metrics, step=self.global_steps)
+ if self.config.trainer.get("val_only", False):
+ return
+
+ # add tqdm
+ progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress")
+
+ # we start from step 1
+ self.global_steps += 1
+ last_val_metrics = None
+ self.max_steps_duration = 0
+
+ prev_step_profile = False
+ curr_step_profile = (
+ self.global_steps in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ next_step_profile = False
+
+ for epoch in range(self.config.trainer.total_epochs):
+ train_iter = iter(self.train_dataloader)
+ next_batch_dict = next(train_iter)
+ need_validate = False
+ dataloader_len = len(self.train_dataloader)
+ print(f"Starting epoch {epoch}, dataloader length: {dataloader_len}")
+ for step_idx in range(dataloader_len):
+ batch_dict = next_batch_dict
+ try:
+ next_batch_dict = next(train_iter)
+ except StopIteration:
+ next_batch_dict = None
+
+ metrics = {}
+ timing_raw = {}
+
+ with marked_timer("start_profile", timing_raw):
+ self._start_profiling(
+ not prev_step_profile and curr_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+
+ batch: DataProto = DataProto.from_single_dict(batch_dict)
+
+ # add uid to batch
+ batch.non_tensor_batch["uid"] = np.array([str(uuid.uuid4()) for _ in range(len(batch))], dtype=object)
+
+ gen_batch = self._get_gen_batch(batch)
+
+ # pass global_steps to trace
+ gen_batch.meta_info["global_steps"] = self.global_steps
+ # pass generation config to gen_batch
+ gen_batch.meta_info["do_sample"] = True
+ gen_batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature
+ gen_batch.meta_info["prompt_length"] = self.config.actor_rollout_ref.rollout.prompt_length
+ gen_batch.meta_info["eos_token_id"] = self.tokenizer.eos_token_id
+ gen_batch.meta_info["n_samples"] = self.config.actor_rollout_ref.rollout.n
+ gen_batch.meta_info["pad_token_id"] = self.tokenizer.pad_token_id
+
+ gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
+
+ is_last_step = self.global_steps >= self.total_training_steps
+
+ if step_idx == 0 or need_validate:
+ # reset env workers in first step
+ # if validation on last step, the reset was not executed and need to be done here
+ reset_future = self._reset_envs(gen_batch)
+
+ need_validate = (
+ self.val_reward_fn is not None
+ and self.config.trainer.test_freq > 0
+ and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
+ )
+
+ with marked_timer("step", timing_raw):
+ # generate a batch
+ with marked_timer("gen", timing_raw, color="red"):
+ gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch, reset_future)
+
+ # prepare for next batch's env reset
+ if step_idx != dataloader_len - 1 and not need_validate:
+ next_batch: DataProto = DataProto.from_single_dict(next_batch_dict)
+ next_gen_batch = self._get_gen_batch(next_batch)
+ next_gen_batch = next_gen_batch.repeat(
+ repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True
+ )
+ reset_future = self._reset_envs(next_gen_batch)
+
+ # repeat to align with repeated responses in rollout
+ batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)
+ batch = gen_batch_output
+
+ if "response_mask" not in batch.batch.keys():
+ batch.batch["response_mask"] = compute_response_mask(batch)
+
+ with marked_timer("reward", timing_raw, color="yellow"):
+ # compute reward model score
+ reward_tensor, reward_extra_infos_dict = compute_reward(batch, self.reward_fn)
+
+ batch.batch["reward_tensor"] = reward_tensor
+ batch = flatten_trajectories(batch)
+
+ batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
+
+ # recompute old_log_probs
+ with marked_timer("old_log_prob", timing_raw, color="blue"):
+ old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
+ entropys = old_log_prob.batch["entropys"]
+ response_masks = batch.batch["response_mask"]
+ actor_config = self.config.actor_rollout_ref.actor
+ entropy_agg = agg_loss(
+ loss_mat=entropys,
+ loss_mask=response_masks,
+ loss_agg_mode=actor_config.loss_agg_mode,
+ loss_scale_factor=actor_config.loss_scale_factor,
+ )
+ old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()}
+ metrics.update(old_log_prob_metrics)
+ old_log_prob.batch.pop("entropys")
+ batch = batch.union(old_log_prob)
+
+ if "rollout_log_probs" in batch.batch.keys():
+ # TODO: we may want to add diff of probs too.
+ from verl.utils.debug.metrics import calculate_debug_metrics
+
+ metrics.update(calculate_debug_metrics(batch))
+
+ if self.use_reference_policy:
+ # compute reference log_prob
+ with marked_timer("ref", timing_raw, color="olive"):
+ if not self.ref_in_actor:
+ ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
+ else:
+ ref_log_prob = self.actor_rollout_wg.compute_ref_log_prob(batch)
+ batch = batch.union(ref_log_prob)
+
+ # compute values
+ if self.use_critic:
+ with marked_timer("values", timing_raw, color="cyan"):
+ values = self.critic_wg.compute_values(batch)
+ batch = batch.union(values)
+
+ with marked_timer("adv", timing_raw, color="brown"):
+ # we combine with rule-based rm
+ reward_extra_infos_dict: dict[str, list] = None
+
+ token_level_scores = torch.zeros_like(response_masks, dtype=torch.float32)
+ flipped_mask = response_masks.flip(dims=[1])
+ indices_in_flipped = torch.argmax(flipped_mask.long(), dim=1)
+
+ last_true_indices = response_masks.shape[-1] - 1 - indices_in_flipped
+ rows_with_response = response_masks.any(dim=1)
+ effective_rewards = batch.batch["reward_tensor"] * rows_with_response.to(
+ batch.batch["reward_tensor"].dtype
+ )
+ row_indices = torch.arange(response_masks.shape[0], device=token_level_scores.device)
+
+ token_level_scores[row_indices, last_true_indices] = effective_rewards
+ batch.batch["token_level_scores"] = token_level_scores
+ if reward_extra_infos_dict:
+ batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})
+
+ # compute rewards. apply_kl_penalty if available
+ if self.config.algorithm.use_kl_in_reward:
+ batch, kl_metrics = apply_kl_penalty(
+ batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty
+ )
+ metrics.update(kl_metrics)
+ else:
+ batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
+
+ # compute advantages, executed on the driver process
+
+ norm_adv_by_std_in_grpo = self.config.algorithm.get(
+ "norm_adv_by_std_in_grpo", True
+ ) # GRPO adv normalization factor
+
+ batch = compute_advantage(
+ batch,
+ adv_estimator=self.config.algorithm.adv_estimator,
+ gamma=self.config.algorithm.gamma,
+ lam=self.config.algorithm.lam,
+ num_repeat=self.config.actor_rollout_ref.rollout.n,
+ norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,
+ config=self.config.algorithm,
+ )
+
+ # update critic
+ if self.use_critic:
+ with marked_timer("update_critic", timing_raw, color="pink"):
+ critic_output = self.critic_wg.update_critic(batch)
+ critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"])
+ metrics.update(critic_output_metrics)
+
+ # implement critic warmup
+ if self.config.trainer.critic_warmup <= self.global_steps:
+ # update actor
+ with marked_timer("update_actor", timing_raw, color="red"):
+ batch.meta_info["multi_turn"] = self.config.actor_rollout_ref.rollout.multi_turn.enable
+ actor_output = self.actor_rollout_wg.update_actor(batch)
+ actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])
+ metrics.update(actor_output_metrics)
+
+ # Log rollout generations if enabled
+ rollout_data_dir = self.config.trainer.get("rollout_data_dir", None)
+ if rollout_data_dir:
+ with marked_timer("dump_rollout_generations", timing_raw, color="green"):
+ inputs = self.tokenizer.batch_decode(batch.batch["prompts"], skip_special_tokens=True)
+ outputs = self.tokenizer.batch_decode(batch.batch["responses"], skip_special_tokens=True)
+ scores = batch.batch["token_level_scores"].sum(-1).cpu().tolist()
+ sample_gts = [
+ item.non_tensor_batch.get("reward_model", {}).get("ground_truth", None)
+ for item in batch
+ ]
+
+ if "request_id" in batch.non_tensor_batch:
+ reward_extra_infos_dict.setdefault(
+ "request_id",
+ batch.non_tensor_batch["request_id"].tolist(),
+ )
+
+ self._dump_generations(
+ inputs=inputs,
+ outputs=outputs,
+ gts=sample_gts,
+ scores=scores,
+ reward_extra_infos_dict=reward_extra_infos_dict,
+ dump_path=rollout_data_dir,
+ )
+
+ # validate
+ if need_validate:
+ with marked_timer("testing", timing_raw, color="green"):
+ val_metrics: dict = self._validate()
+ if is_last_step:
+ last_val_metrics = val_metrics
+ metrics.update(val_metrics)
+
+ # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.
+ esi_close_to_expiration = should_save_ckpt_esi(
+ max_steps_duration=self.max_steps_duration,
+ redundant_time=self.config.trainer.esi_redundant_time,
+ )
+ # Check if the conditions for saving a checkpoint are met.
+ # The conditions include a mandatory condition (1) and
+ # one of the following optional conditions (2/3/4):
+ # 1. The save frequency is set to a positive value.
+ # 2. It's the last training step.
+ # 3. The current step number is a multiple of the save frequency.
+ # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.
+ if self.config.trainer.save_freq > 0 and (
+ is_last_step or self.global_steps % self.config.trainer.save_freq == 0 or esi_close_to_expiration
+ ):
+ if esi_close_to_expiration:
+ print("Force saving checkpoint: ESI instance expiration approaching.")
+ with marked_timer("save_checkpoint", timing_raw, color="green"):
+ self._save_checkpoint()
+
+ with marked_timer("stop_profile", timing_raw):
+ next_step_profile = (
+ self.global_steps + 1 in self.config.global_profiler.steps
+ if self.config.global_profiler.steps is not None
+ else False
+ )
+ self._stop_profiling(
+ curr_step_profile and not next_step_profile
+ if self.config.global_profiler.profile_continuous_steps
+ else curr_step_profile
+ )
+ prev_step_profile = curr_step_profile
+ curr_step_profile = next_step_profile
+
+ steps_duration = timing_raw["step"]
+ self.max_steps_duration = max(self.max_steps_duration, steps_duration)
+
+ # training metrics
+ metrics.update(
+ {
+ "training/global_step": self.global_steps,
+ "training/epoch": epoch,
+ }
+ )
+ # collect metrics
+ metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
+ metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
+ # TODO: implement actual tflpo and theoretical tflpo
+ n_gpus = self.resource_pool_manager.get_n_gpus()
+ metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))
+
+ # this is experimental and may be changed/removed in the future in favor of a general-purpose one
+ if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):
+ self.train_dataloader.sampler.update(batch=batch)
+
+ # TODO: make a canonical logger that supports various backend
+ logger.log(data=metrics, step=self.global_steps)
+
+ progress_bar.update(1)
+ self.global_steps += 1
+
+ if (
+ hasattr(self.config.actor_rollout_ref.actor, "profiler")
+ and self.config.actor_rollout_ref.actor.profiler.tool == "torch_memory"
+ ):
+ self.actor_rollout_wg.dump_memory_snapshot(
+ tag=f"post_update_step{self.global_steps}", sub_dir=f"step{self.global_steps}"
+ )
+
+ if is_last_step:
+ pprint(f"Final validation metrics: {last_val_metrics}")
+ progress_bar.close()
+ return
+
+ # this is experimental and may be changed/removed in the future
+ # in favor of a general-purpose data buffer pool
+ if hasattr(self.train_dataset, "on_batch_end"):
+ # The dataset may be changed after each training batch
+ self.train_dataset.on_batch_end(batch=batch)
+
+ def _validate(self):
+ data_source_lst = []
+ reward_extra_infos_dict: dict[str, list] = defaultdict(list)
+
+ # Lists to collect samples for the table
+ sample_scores = []
+ sample_turns = []
+ sample_uids = []
+
+ for test_data in self.val_dataloader:
+ test_batch = DataProto.from_single_dict(test_data)
+ if len(test_batch) < self.config.data.val_batch_size:
+ print(f"drop last batch in val_dataloader, len {len(test_batch)}")
+ break
+
+ if "uid" not in test_batch.non_tensor_batch:
+ test_batch.non_tensor_batch["uid"] = np.array(
+ [str(uuid.uuid4()) for _ in range(len(test_batch))], dtype=object
+ )
+
+ test_gen_batch = self._get_gen_batch(test_batch)
+ test_gen_batch.meta_info = {
+ "eos_token_id": self.tokenizer.eos_token_id,
+ "pad_token_id": self.tokenizer.pad_token_id,
+ "prompt_length": self.config.actor_rollout_ref.rollout.prompt_length,
+ "recompute_log_prob": False,
+ "do_sample": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample,
+ "temperature": self.config.actor_rollout_ref.rollout.temperature,
+ "n_samples": self.config.actor_rollout_ref.rollout.n,
+ "validate": True,
+ "global_steps": self.global_steps,
+ }
+
+ test_gen_batch = test_gen_batch.repeat(
+ repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True
+ )
+
+ sample_uids.extend(test_gen_batch.non_tensor_batch["uid"])
+
+ # pad to be divisible by dp_size
+ size_divisor = self.config.env.train.num_envs * self.config.env.rollout.pipeline_stage_num
+ test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor)
+ if not self.async_rollout_mode:
+ test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences(test_gen_batch_padded)
+ else:
+ reset_future = self._reset_envs(test_gen_batch_padded)
+ test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(
+ test_gen_batch_padded, reset_future
+ )
+
+ # unpad
+ test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size)
+
+ print("validation generation end")
+
+ test_batch = test_output_gen_batch
+ test_batch.meta_info["validate"] = True
+
+ # evaluate using reward_function
+ if self.val_reward_fn is None:
+ raise ValueError("val_reward_fn must be provided for validation.")
+ result = self.val_reward_fn(test_batch, return_dict=True)
+ reward_tensor = result["reward_tensor"]
+ scores = reward_tensor.sum(-1).cpu().tolist()
+ sample_scores.extend(scores)
+
+ reward_extra_infos_dict["reward"].extend(scores)
+ print(f"len reward_extra_infos_dict['reward']: {len(reward_extra_infos_dict['reward'])}")
+ if "reward_extra_info" in result:
+ for key, lst in result["reward_extra_info"].items():
+ reward_extra_infos_dict[key].extend(lst)
+ print(f"len reward_extra_infos_dict['{key}']: {len(reward_extra_infos_dict[key])}")
+
+ # collect num_turns of each prompt
+ if "__num_turns__" in test_batch.non_tensor_batch:
+ sample_turns.append(test_batch.non_tensor_batch["__num_turns__"])
+
+ data_source_lst.append(test_batch.non_tensor_batch.get("data_source", ["unknown"] * reward_tensor.shape[0]))
+
+ for key_info, lst in reward_extra_infos_dict.items():
+ assert len(lst) == 0 or len(lst) == len(sample_scores), f"{key_info}: {len(lst)=}, {len(sample_scores)=}"
+
+ data_sources = np.concatenate(data_source_lst, axis=0)
+
+ data_src2var2metric2val = process_validation_metrics(data_sources, sample_uids, reward_extra_infos_dict)
+ metric_dict = {}
+ for data_source, var2metric2val in data_src2var2metric2val.items():
+ core_var = "acc" if "acc" in var2metric2val else "reward"
+ for var_name, metric2val in var2metric2val.items():
+ n_max = max([int(name.split("@")[-1].split("/")[0]) for name in metric2val.keys()])
+ for metric_name, metric_val in metric2val.items():
+ if (
+ (var_name == core_var)
+ and any(metric_name.startswith(pfx) for pfx in ["mean", "maj", "best"])
+ and (f"@{n_max}" in metric_name)
+ ):
+ metric_sec = "val-core"
+ else:
+ metric_sec = "val-aux"
+ pfx = f"{metric_sec}/{data_source}/{var_name}/{metric_name}"
+ metric_dict[pfx] = metric_val
+
+ if len(sample_turns) > 0:
+ sample_turns = np.concatenate(sample_turns)
+ metric_dict["val-aux/num_turns/min"] = sample_turns.min()
+ metric_dict["val-aux/num_turns/max"] = sample_turns.max()
+ metric_dict["val-aux/num_turns/mean"] = sample_turns.mean()
+
+ return metric_dict
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/run_simpleVLA_isaac_disagg.sh b/code/RL_model/verl/verl_train/verl/experimental/vla/run_simpleVLA_isaac_disagg.sh
new file mode 100644
index 0000000000000000000000000000000000000000..c16c6124785a74d246f1ce4f340b768f11aa0c5d
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/run_simpleVLA_isaac_disagg.sh
@@ -0,0 +1,109 @@
+#!/bin/bash
+set -x
+
+echo "remember to set ray param < --resources='{\"sim\"/\"actor_rollout\":1}' > if using disagg sim"
+
+libero_train_path=$HOME/data/libero_rl/train.parquet
+libero_test_path=$HOME/data/libero_rl/test.parquet
+
+train_files=$libero_train_path
+test_files=$libero_test_path
+
+OUTPUT_DIR=${MLP_MODEL_OUTPUT:-"$HOME/models/vla_libero_grpo"}
+VIDEO_OUTPUT=${MLP_MODEL_OUTPUT:-"$HOME"}/video
+SFT_MODEL_PATH=${SFT_MODEL_PATH:-"$HOME/data/Openvla-oft-SFT-libero10-trajall"}
+
+# for rollout and train
+NUM_NODES=1
+# for simulator
+SIM_NODES=1
+NUM_ENV_GPUS=8
+NUM_ROLLOUT_GPUS=8
+STAGE_NUM=2
+BATCH_SIZE=16
+# rollout.n should equal to num_envs for isaac env
+ROLLOUT_N=8
+
+# 512 is required for libero benchmark, but you can reduce it in debugging to run faster
+MAX_EPISODE_STEPS=512
+
+# isaac or libero
+# libero means original libero benchmark with mujoco sim
+# isaac means libero benchmark using isaac sim
+SIM_TYPE=${SIM_TYPE:-"isaac"}
+PROJECT_NAME="vla-disagg-issac"
+EXPERIMENT_NAME="${SIM_TYPE}_rl"
+
+ISSC_PYTHON="/workspace/isaaclab/_isaac_sim/python.sh"
+PYTHON=python
+if [ -f "$ISSC_PYTHON" ]; then
+ PYTHON=$ISSC_PYTHON
+fi
+
+# avoiding warnings
+mkdir /root/LIBERO/libero/libero/../datasets
+
+
+$PYTHON -m verl.experimental.vla.main_ppo \
+ data.train_files="$train_files" \
+ data.val_files="$test_files" \
+ data.train_batch_size=${BATCH_SIZE} \
+ data.val_batch_size=${BATCH_SIZE} \
+ actor_rollout_ref.rollout.n=$ROLLOUT_N \
+ env.train.num_envs=$ROLLOUT_N \
+ data.max_prompt_length=256 \
+ data.max_response_length=128 \
+ env.rollout.pipeline_stage_num=$STAGE_NUM \
+ env.train.simulator_type=$SIM_TYPE \
+ env.actor.model.num_action_chunks=8 \
+ env.actor.model.action_dim=7 \
+ env.train.only_eval=False \
+ env.train.max_episode_steps=$MAX_EPISODE_STEPS \
+ env.train.video_cfg.save_video=True \
+ env.train.video_cfg.video_base_dir=${VIDEO_OUTPUT} \
+ env.train.seed=42 \
+ env.disagg_sim.enable=True \
+ env.disagg_sim.nnodes=$SIM_NODES \
+ actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \
+ actor_rollout_ref.model.path=$SFT_MODEL_PATH \
+ actor_rollout_ref.rollout.mode=async_envloop \
+ actor_rollout_ref.actor.optim.lr=5e-6 \
+ actor_rollout_ref.actor.optim.warmup_style=constant \
+ actor_rollout_ref.actor.ppo_mini_batch_size=128 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \
+ actor_rollout_ref.actor.use_dynamic_bsz=False \
+ actor_rollout_ref.actor.grad_clip=1 \
+ actor_rollout_ref.actor.clip_ratio_high=0.28 \
+ actor_rollout_ref.actor.clip_ratio_low=0.2 \
+ actor_rollout_ref.actor.num_images_in_input=1 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=False \
+ actor_rollout_ref.model.use_remove_padding=False \
+ actor_rollout_ref.model.trust_remote_code=False \
+ actor_rollout_ref.actor.entropy_coeff=0. \
+ actor_rollout_ref.rollout.temperature=1.6 \
+ actor_rollout_ref.rollout.prompt_length=512 \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=hf \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \
+ actor_rollout_ref.rollout.free_cache_engine=False \
+ actor_rollout_ref.ref.log_prob_micro_batch_size=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.kl_ctrl.kl_coef=0.00 \
+ trainer.logger=['console'] \
+ trainer.project_name=$PROJECT_NAME \
+ trainer.experiment_name=$EXPERIMENT_NAME \
+ trainer.default_local_dir=$OUTPUT_DIR \
+ trainer.n_gpus_per_node=$NUM_ROLLOUT_GPUS \
+ +trainer.n_env_gpus_per_node=$NUM_ENV_GPUS \
+ +trainer.n_rollout_gpus_per_node=$NUM_ROLLOUT_GPUS \
+ trainer.nnodes=$NUM_NODES \
+ trainer.save_freq=30 \
+ trainer.test_freq=-1 \
+ trainer.total_epochs=20 \
+ trainer.val_only=False \
+ trainer.total_training_steps=10000 \
+ algorithm.adv_estimator=reinforce_plus_plus \
+ trainer.val_before_train=False $@
+
+
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/run_simpleVLA_libero_grpo.sh b/code/RL_model/verl/verl_train/verl/experimental/vla/run_simpleVLA_libero_grpo.sh
new file mode 100644
index 0000000000000000000000000000000000000000..d8980a02bf3710ef1f6d2ecd46af602d9fa48094
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/run_simpleVLA_libero_grpo.sh
@@ -0,0 +1,102 @@
+set -x
+libero_train_path=$HOME/data/libero_rl/train.parquet
+libero_test_path=$HOME/data/libero_rl/test.parquet
+
+
+train_files=$libero_train_path
+test_files=$libero_test_path
+
+OUTPUT_DIR=${MLP_MODEL_OUTPUT:-"$HOME/models/vla_libero_grpo"}
+VIDEO_OUTPUT=${MLP_MODEL_OUTPUT:-"$HOME"}/video
+SFT_MODEL_PATH=${SFT_MODEL_PATH:-"$HOME/data/Openvla-oft-SFT-libero10-trajall"}
+
+NUM_NODES=1
+NUM_GPUS=8
+NUM_ENV_GPUS=4
+# rollout.n should equal to num_envs for isaac env
+ROLLOUT_N=8
+# isaac or libero
+# libero means original libero benchmark with mujoco sim
+# isaac means libero benchmark using isaac sim
+SIM_TYPE=${SIM_TYPE:-"isaac"}
+PROJECT_NAME="vla_libero_RL"
+EXPERIMENT_NAME="${SIM_TYPE}_reinforce_plus_plus"
+
+ISSC_PYTHON="/workspace/isaaclab/_isaac_sim/python.sh"
+PYTHON=python
+if [ -f "$ISSC_PYTHON" ]; then
+ PYTHON=$ISSC_PYTHON
+fi
+
+# avoiding warnings
+mkdir /root/LIBERO/libero/libero/../datasets
+gpu_name=$(nvidia-smi --query-gpu=name --format=csv,noheader,nounits | head -n 1)
+
+# force osmesa in Hopper
+if echo "$gpu_name" | grep "NVIDIA H"; then
+ echo "enable MUJOCO_GL=osmesa in Hopper"
+ export MUJOCO_GL=osmesa
+fi
+
+
+$PYTHON -m verl.experimental.vla.main_ppo \
+ data.train_files="$train_files" \
+ data.val_files="$test_files" \
+ data.train_batch_size=8 \
+ data.val_batch_size=8 \
+ actor_rollout_ref.rollout.n=$ROLLOUT_N \
+ env.train.num_envs=$ROLLOUT_N \
+ data.max_prompt_length=256 \
+ data.max_response_length=128 \
+ env.rollout.pipeline_stage_num=2 \
+ env.train.simulator_type=$SIM_TYPE \
+ env.actor.model.num_action_chunks=8 \
+ env.actor.model.action_dim=7 \
+ env.train.only_eval=False \
+ env.train.max_episode_steps=512 \
+ env.train.video_cfg.save_video=True \
+ env.train.video_cfg.video_base_dir=${VIDEO_OUTPUT} \
+ env.train.seed=42 \
+ actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \
+ actor_rollout_ref.model.path=$SFT_MODEL_PATH \
+ actor_rollout_ref.rollout.mode=async_envloop \
+ actor_rollout_ref.actor.optim.lr=5e-6 \
+ actor_rollout_ref.actor.optim.warmup_style=constant \
+ actor_rollout_ref.actor.ppo_mini_batch_size=128 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \
+ actor_rollout_ref.actor.use_dynamic_bsz=False \
+ actor_rollout_ref.actor.grad_clip=1 \
+ actor_rollout_ref.actor.clip_ratio_high=0.28 \
+ actor_rollout_ref.actor.clip_ratio_low=0.2 \
+ actor_rollout_ref.actor.num_images_in_input=1 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=False \
+ actor_rollout_ref.model.use_remove_padding=False \
+ actor_rollout_ref.model.trust_remote_code=False \
+ actor_rollout_ref.actor.entropy_coeff=0. \
+ actor_rollout_ref.rollout.temperature=1.6 \
+ actor_rollout_ref.rollout.prompt_length=512 \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=hf \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \
+ actor_rollout_ref.rollout.free_cache_engine=False \
+ actor_rollout_ref.ref.log_prob_micro_batch_size=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.kl_ctrl.kl_coef=0.00 \
+ trainer.logger=['console'] \
+ trainer.project_name=$PROJECT_NAME \
+ trainer.experiment_name=$EXPERIMENT_NAME \
+ trainer.default_local_dir=$OUTPUT_DIR \
+ trainer.n_gpus_per_node=$NUM_GPUS \
+ +trainer.n_env_gpus_per_node=$NUM_ENV_GPUS \
+ +trainer.n_rollout_gpus_per_node=$((NUM_GPUS - NUM_ENV_GPUS)) \
+ trainer.nnodes=$NUM_NODES \
+ trainer.save_freq=30 \
+ trainer.test_freq=30 \
+ trainer.total_epochs=20 \
+ trainer.val_only=False \
+ trainer.total_training_steps=10000 \
+ algorithm.adv_estimator=reinforce_plus_plus \
+ trainer.val_before_train=False $@
+
+
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_loop_wg_test.py b/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_loop_wg_test.py
new file mode 100644
index 0000000000000000000000000000000000000000..71fb6348f6235b3ec26b7ee5ab22c3e74b3beb4b
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_loop_wg_test.py
@@ -0,0 +1,181 @@
+# Copyright 2025 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
+#
+# https://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.
+
+import asyncio
+
+import ray
+from omegaconf import OmegaConf
+
+from verl import DataProto
+from verl.experimental.vla.naive_rollout_rob import NaiveRolloutRob
+
+# from verl.workers.env.env_worker import EnvWorker
+from verl.experimental.vla.workers.env.env_worker import EnvWorker
+from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup
+
+if not ray.is_initialized():
+ ray.init()
+
+ # for debugging
+ # ray.init(
+ # runtime_env={
+ # "env_vars": {"RAY_DEBUG_POST_MORTEM": "1"},
+ # }
+ # )
+
+ENV_WORKERS_NUM = 1
+STAGE_NUM = 1
+# NUM_ENVS_PER_ITER = 32
+
+# NUM_ENVS_PER_STAGE = 8
+# NUM_ENVS_PER_ITER = STAGE_NUM * NUM_ENVS_PER_STAGE
+# NUM_ENVS_PER_ITER = 8
+# NUM_ENVS_PER_ITER = 32
+NUM_ENVS_PER_ITER = 2
+NUM_ENVS_PER_WORKER = NUM_ENVS_PER_ITER // ENV_WORKERS_NUM
+# NUM_ENVS_PER_WORKER_PER_STAGE = NUM_ENVS_PER_STAGE // ENV_WORKERS_NUM
+GROUP_SIZE = 2 # real group size = GROUP_SIZE * STAGE_NUM
+GROUP_NUM_PER_ITER = NUM_ENVS_PER_ITER * STAGE_NUM // GROUP_SIZE
+BATCH_SIZE_PER_GPU = 2
+NUM_ACTS_CHUNKS = 8
+MAX_EPISODE_STEPS = 32
+MAX_INFER_STEPS = MAX_EPISODE_STEPS // NUM_ACTS_CHUNKS
+cfg_dict = {
+ "rollout": {"pipeline_stage_num": STAGE_NUM},
+ "train": {
+ "use_fixed_reset_state_ids": False,
+ "ignore_terminations": False,
+ # "auto_reset": True,
+ "auto_reset": False,
+ "max_episode_steps": MAX_EPISODE_STEPS,
+ "use_rel_reward": False,
+ "reward_coef": 1.0,
+ "only_eval": False,
+ "use_ordered_reset_state_ids": False,
+ # "num_images_in_input": 1,
+ "init_params": {
+ "camera_depths": False,
+ "camera_heights": 256,
+ "camera_widths": 256,
+ "camera_names": ["agentview", "robot0_eye_in_hand"],
+ },
+ "video_cfg": {
+ "save_video": True,
+ "video_base_dir": "/tmp/videos",
+ },
+ "task_suite_name": "libero_10",
+ "num_envs": NUM_ENVS_PER_WORKER,
+ "simulator_type": "isaac",
+ "seed": 0,
+ },
+ "enable_offload": False,
+ "actor": {"model": {"num_action_chunks": NUM_ACTS_CHUNKS, "action_dim": 7}},
+ "runner": {"only_eval": False},
+}
+env_cfg = OmegaConf.create(cfg_dict)
+
+gpu_pool = RayResourcePool([ENV_WORKERS_NUM], use_gpu=True)
+# RayEnvWorker = ray.remote(num_gpus=1)(EnvWorker)
+ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(EnvWorker), config=env_cfg)
+
+env_wg = RayWorkerGroup(gpu_pool, ray_cls_with_init)
+
+
+def restructure_data_proto(data_proto: DataProto) -> list[DataProto]:
+ total_batch_size = len(data_proto)
+ tensors = data_proto.batch
+ non_tensors = data_proto.non_tensor_batch
+
+ full_image_tensor = tensors["full_image"]
+ state_tensor = tensors["state"]
+ task_descriptions_np = non_tensors["task_descriptions"]
+ if total_batch_size != ENV_WORKERS_NUM * STAGE_NUM * NUM_ENVS_PER_WORKER:
+ raise ValueError(
+ f"Total batch size {total_batch_size} does not match the expected size "
+ f"ENV_WORKERS_NUM * STAGE_NUM * NUM_ENVS_PER_WORKER = "
+ f"{ENV_WORKERS_NUM * STAGE_NUM * NUM_ENVS_PER_WORKER}"
+ )
+
+ image_rest_shape = (ENV_WORKERS_NUM, STAGE_NUM, NUM_ENVS_PER_WORKER) + full_image_tensor.shape[1:]
+ state_rest_shape = (ENV_WORKERS_NUM, STAGE_NUM, NUM_ENVS_PER_WORKER) + state_tensor.shape[1:]
+ reshaped_full_image = full_image_tensor.view(image_rest_shape)
+ reshaped_state = state_tensor.view(state_rest_shape)
+
+ reshaped_task_descriptions = task_descriptions_np.reshape(ENV_WORKERS_NUM, STAGE_NUM, NUM_ENVS_PER_WORKER)
+ stages_data_list = []
+ for stage_idx in range(STAGE_NUM):
+ stage_images = reshaped_full_image[:, stage_idx, :]
+ stage_states = reshaped_state[:, stage_idx, :]
+ stage_tasks = reshaped_task_descriptions[:, stage_idx, :]
+ final_images = stage_images.reshape(ENV_WORKERS_NUM * NUM_ENVS_PER_WORKER, *full_image_tensor.shape[1:])
+ final_states = stage_states.reshape(ENV_WORKERS_NUM * NUM_ENVS_PER_WORKER, *state_tensor.shape[1:])
+ final_tasks = stage_tasks.flatten().tolist()
+
+ stage_dp = DataProto.from_dict(
+ tensors={"full_image": final_images, "state": final_states},
+ non_tensors={"task_descriptions": final_tasks},
+ meta_info={"do_sample": True, "temperature": 1.6, "prompt_length": 512},
+ )
+ stages_data_list.append(stage_dp)
+ return stages_data_list
+
+
+async def run():
+ # breakpoint()
+ env_wg.init_worker()
+ env_wg.init_simulator()
+
+ reset_state_ids_tensordict = DataProto.from_dict(
+ non_tensors={"state_ids": [0] * NUM_ENVS_PER_ITER * STAGE_NUM, "task_ids": [0] * NUM_ENVS_PER_ITER * STAGE_NUM}
+ )
+
+ reset_result = env_wg.reset_envs_to_state_ids(reset_state_ids_tensordict)
+ print(f"reset_envs_to_state_ids result: {reset_result}")
+ stages_data_list = restructure_data_proto(reset_result)
+
+ RayNaiveRolloutRob = ray.remote(num_gpus=1)(NaiveRolloutRob)
+
+ model_config = {"path": "Haozhan72/Openvla-oft-SFT-libero10-trajall"}
+ rollout_workers = RayNaiveRolloutRob.remote(model_config)
+
+ env_obs_refs = {}
+ rollout_refs = {}
+ traj = [[], []]
+
+ for _ in range(MAX_INFER_STEPS):
+ for stage_id in range(STAGE_NUM):
+ if _ == 0:
+ rollout_refs[stage_id] = rollout_workers.generate_sequences.remote(stages_data_list[stage_id])
+ else:
+ # env_batch = env_obs_refs[stage_id]
+ env_batch: DataProto = env_obs_refs[stage_id].get()
+ env_batch_traj = env_batch.select(batch_keys=["rews", "terminations", "truncations"])
+ traj[stage_id][-1].update({"env": env_batch_traj})
+ obs = env_batch
+ obs.meta_info.update({"do_sample": True, "temperature": 1.6, "prompt_length": 512})
+ rollout_refs[stage_id] = rollout_workers.generate_sequences.remote(obs)
+ for stage_id in range(STAGE_NUM):
+ batch: DataProto = ray.get(rollout_refs[stage_id])
+ traj[stage_id].append({"model": batch})
+ action = batch.batch["action"]
+ action = action.cpu().numpy()
+ # already in env
+ data = DataProto.from_dict(non_tensors={"actions": action}, meta_info={"stage_id": stage_id})
+ env_obs_refs[stage_id] = env_wg.env_interact_step(data)
+
+ env_wg.finish_rollout()
+
+
+asyncio.run(run())
+ray.timeline(filename="2stage_pipeline_timeline_wg.json")
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_manager.py b/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..7050a517d89a19706a23964e10fdaca50a190578
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_manager.py
@@ -0,0 +1,387 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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.
+
+import gc
+import logging
+import os
+import subprocess
+from typing import Optional
+
+import torch
+import torch.multiprocessing as mp
+
+from verl.utils.device import get_torch_device
+
+logger = logging.getLogger(__name__)
+
+
+def cleanup_device_tensors():
+ gc.collect()
+ get_torch_device().empty_cache()
+
+
+def get_gpu_numa_node(gpu_id: int) -> int:
+ try:
+ try:
+ import pynvml
+
+ pynvml.nvmlInit()
+ handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
+ # Get PCI bus info
+ pci_info = pynvml.nvmlDeviceGetPciInfo(handle)
+ pci_bus_id = pci_info.busId
+ except ImportError:
+ # Fallback to nvidia-smi
+ result = subprocess.run(
+ [
+ "nvidia-smi",
+ "--query-gpu=pci.bus_id",
+ "--format=csv,noheader,nounits",
+ f"--id={gpu_id}",
+ ],
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ pci_bus_id = result.stdout.strip()
+
+ # Extract bus number from PCI bus ID (format: 0000:XX:YY.Z)
+ bus_number = pci_bus_id.split(":")[1]
+
+ # Get NUMA node from sysfs
+ numa_node_path = f"/sys/bus/pci/devices/0000:{bus_number}:00.0/numa_node"
+ if os.path.exists(numa_node_path):
+ with open(numa_node_path) as f:
+ numa_node = int(f.read().strip())
+ if numa_node >= 0:
+ return numa_node
+
+ # Fallback: try to get from lscpu
+ result = subprocess.run(["lscpu"], capture_output=True, text=True, check=True)
+ numa_nodes = 0
+ for line in result.stdout.split("\n"):
+ if "NUMA node(s):" in line:
+ numa_nodes = int(line.split(":")[1].strip())
+ break
+
+ # If we can't determine the exact NUMA node, distribute evenly
+ return gpu_id % numa_nodes if numa_nodes > 0 else 0
+
+ except Exception as e:
+ logger.error(f"Warning: Could not determine NUMA node for GPU {gpu_id}: {e}")
+ return 0
+
+
+def get_numa_cpus(numa_node: int) -> list:
+ try:
+ # Read from sysfs
+ cpulist_path = f"/sys/devices/system/node/node{numa_node}/cpulist"
+ if os.path.exists(cpulist_path):
+ with open(cpulist_path) as f:
+ cpulist = f.read().strip()
+
+ # Parse CPU list (e.g., "0-7,16-23" or "0,1,2,3")
+ cpus = []
+ for part in cpulist.split(","):
+ if "-" in part:
+ start, end = map(int, part.split("-"))
+ cpus.extend(range(start, end + 1))
+ else:
+ cpus.append(int(part))
+ return cpus
+ except Exception as e:
+ logger.error(f"Warning: Could not get CPU list for NUMA node {numa_node}: {e}")
+
+ # Fallback: return all available CPUs
+ return list(range(os.cpu_count() or 1))
+
+
+def set_process_numa_affinity(gpu_id: int) -> None:
+ try:
+ numa_node = get_gpu_numa_node(gpu_id)
+ cpus = get_numa_cpus(numa_node)
+
+ if not cpus:
+ logger.error(f"Warning: No CPUs found for NUMA node {numa_node}")
+ return
+
+ os.sched_setaffinity(0, cpus)
+ try:
+ subprocess.run(
+ ["numactl", "--membind", str(numa_node), "--"],
+ check=False,
+ capture_output=True,
+ )
+ except FileNotFoundError:
+ pass # numactl not available, that's ok
+
+ except Exception as e:
+ logger.error(f"Warning: Could not set NUMA affinity for GPU {gpu_id}: {e}")
+
+
+def recursive_to_own(obj):
+ if isinstance(obj, torch.Tensor):
+ return obj.clone() if obj.is_shared() else obj
+ elif isinstance(obj, list):
+ return [recursive_to_own(elem) for elem in obj]
+ elif isinstance(obj, tuple):
+ return tuple(recursive_to_own(elem) for elem in obj)
+ elif isinstance(obj, dict):
+ return {k: recursive_to_own(v) for k, v in obj.items()}
+ else:
+ return obj
+
+
+class EnvManager:
+ def __init__(self, cfg, rank, world_size, env_cls):
+ self.cfg = cfg
+ self.rank = rank
+ self.world_size = world_size
+ self.process: Optional[mp.Process] = None
+ self.command_queue: Optional[mp.Queue] = None
+ self.result_queue: Optional[mp.Queue] = None
+ self.state_buffer: Optional[bytes] = None
+
+ self.env_cls = env_cls
+
+ def start_simulator(self):
+ """Start simulator process with shared memory queues"""
+ if self.process:
+ logger.info(f"Simulator process already running for rank {self.rank}")
+ return
+
+ self.context = mp.get_context("spawn")
+ # Create shared memory queues
+ self.command_queue = self.context.Queue()
+ self.result_queue = self.context.Queue()
+
+ # Start simulator process
+ self.process = self.context.Process(
+ target=_simulator_worker,
+ args=(
+ self.cfg,
+ self.rank,
+ self.world_size,
+ self.env_cls,
+ self.command_queue,
+ self.result_queue,
+ self.state_buffer,
+ True,
+ ),
+ )
+ self.process.start()
+
+ # Wait for initialization
+ result = self.result_queue.get(timeout=180)
+ if result["status"] != "ready":
+ raise RuntimeError(f"Simulator initialization failed: {result}")
+
+ def stop_simulator(self):
+ if not self.process:
+ return
+
+ # Request state save
+ self.command_queue.put({"method": "get_state", "args": [], "kwargs": {}})
+
+ # Get saved state
+ result = self.result_queue.get(timeout=180)
+ if result["status"] == "success":
+ self.state_buffer = result["data"]
+
+ self.command_queue.put({"method": "shutdown"})
+ self.command_queue.close()
+ self.result_queue.close()
+ self.command_queue = None
+ self.result_queue = None
+ self.process.join(timeout=5)
+
+ self.command_queue = None
+ self.result_queue = None
+ if self.process.is_alive():
+ self.process.terminate()
+ self.process.join()
+
+ self.process = None
+
+ def __getattr__(self, name):
+ if name in [
+ "cfg",
+ "rank",
+ "world_size",
+ "process",
+ "command_queue",
+ "result_queue",
+ "state_buffer",
+ "env_cls",
+ "context",
+ ]:
+ return super().__getattr__(name)
+
+ def method_proxy(*args, **kwargs):
+ if self.process is None or not self.process.is_alive():
+ raise RuntimeError("Simulator not running")
+
+ args = recursive_to_own(args)
+ kwargs = recursive_to_own(kwargs)
+ self.command_queue.put({"method": name, "args": args, "kwargs": kwargs})
+
+ result = self.result_queue.get()
+ result = recursive_to_own(result)
+ if result["status"] == "error":
+ raise Exception(result["error"])
+ return result["data"]
+
+ return method_proxy
+
+ def get_all_state_ids(self):
+ """Get all available state IDs from the environment."""
+ if self.process is None or not self.process.is_alive():
+ raise RuntimeError("Simulator not running")
+
+ self.command_queue.put({"method": "get_all_state_ids", "args": [], "kwargs": {}})
+ result = self.result_queue.get()
+ result = recursive_to_own(result)
+ if result["status"] == "error":
+ raise Exception(result["error"])
+ return result["data"]
+
+ def reset_envs_to_state_ids(self, state_ids_list, task_ids_list):
+ """Reset environments to specified state IDs."""
+ if self.process is None or not self.process.is_alive():
+ raise RuntimeError("Simulator not running")
+
+ state_ids_list = recursive_to_own(state_ids_list)
+ task_ids_list = recursive_to_own(task_ids_list)
+
+ self.command_queue.put(
+ {
+ "method": "reset_envs_to_state_ids",
+ "args": [state_ids_list, task_ids_list],
+ "kwargs": {},
+ }
+ )
+
+ result = self.result_queue.get()
+ result = recursive_to_own(result)
+ if result["status"] == "error":
+ raise Exception(result["error"])
+ return result["data"]
+
+ def __setattr__(self, name, value):
+ # Handle special attributes that should be set on self
+ if name in [
+ "cfg",
+ "rank",
+ "world_size",
+ "process",
+ "command_queue",
+ "result_queue",
+ "state_buffer",
+ "env_cls",
+ "context",
+ ]:
+ super().__setattr__(name, value)
+ return
+
+ if self.process is None or not self.process.is_alive():
+ raise RuntimeError(f"Simulator not running to set attribute {name} to {value}")
+
+ value = recursive_to_own(value)
+ self.command_queue.put(
+ {
+ "method": "__setattr__",
+ "args": [name, value],
+ "kwargs": {},
+ }
+ )
+
+ result = self.result_queue.get()
+ result = recursive_to_own(result)
+ if result["status"] == "error":
+ raise Exception(result["error"])
+
+
+def _simulator_worker(
+ cfg,
+ rank,
+ world_size,
+ env_cls,
+ command_queue,
+ result_queue,
+ state_buffer,
+ bind_numa=True,
+):
+ """Worker process for simulator"""
+ # Set NUMA affinity for the process to match the GPU rank
+ import logging
+ import os
+
+ pid = os.getpid()
+ logger = logging.getLogger(f"simulator_worker_{rank}_{pid}")
+
+ if bind_numa:
+ set_process_numa_affinity(rank)
+ try:
+ env = env_cls(cfg, rank, world_size)
+
+ if state_buffer:
+ env.load_state(state_buffer)
+
+ # Signal ready
+ result_queue.put({"status": "ready"})
+
+ # Main command processing loop
+ while True:
+ try:
+ command = command_queue.get()
+ logger.debug(f"Received command method: {command['method']}")
+
+ if command["method"] == "shutdown":
+ env.close()
+ break
+
+ method_name = command["method"]
+ args = command.get("args", [])
+ kwargs = command.get("kwargs", {})
+ if method_name == "__setattr__":
+ # Handle attribute setting
+ attr_name, attr_value = args
+ setattr(env, attr_name, attr_value)
+ result_queue.put({"status": "success", "data": None})
+ elif hasattr(env, method_name):
+ method = getattr(env, method_name)
+ assert callable(method), f"Method {method_name} is not callable"
+ result = method(*args, **kwargs)
+ result_queue.put({"status": "success", "data": result})
+ else:
+ logger.error(f"Method '{method_name}' not found")
+ result_queue.put(
+ {
+ "status": "error",
+ "error": f"Method '{method_name}' not found",
+ }
+ )
+
+ except Exception as e:
+ logger.exception(e)
+ result_queue.put({"status": "error", "error": str(e)})
+
+ except Exception as e:
+ logger.exception(e)
+ result_queue.put({"status": "error", "error": str(e)})
+
+ finally:
+ command_queue.close()
+ result_queue.close()
diff --git a/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_worker.py b/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..92bb6364a4a89c943cea0f9d791ba5a7908eb09d
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/experimental/vla/workers/env/env_worker.py
@@ -0,0 +1,242 @@
+# Copyright 2025 The RLinf Authors.
+# Copyright 2024 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
+#
+# https://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.
+
+import itertools
+
+import torch
+from omegaconf import DictConfig
+from torch.distributed.device_mesh import init_device_mesh
+
+from verl import DataProto
+from verl.experimental.vla.envs.action_utils import prepare_actions
+from verl.experimental.vla.workers.env.env_manager import EnvManager
+from verl.single_controller.base import Worker
+from verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register
+from verl.utils.config import omega_conf_to_dataclass
+from verl.utils.device import (
+ get_device_name,
+)
+from verl.utils.distributed import initialize_global_process_group_ray
+from verl.utils.profiler import DistProfiler, DistProfilerExtension, ProfilerConfig
+
+
+def put_tensor_cpu(data_dict):
+ for key, value in data_dict.items():
+ if isinstance(value, dict):
+ data_dict[key] = put_tensor_cpu(value)
+ if isinstance(value, torch.Tensor):
+ data_dict[key] = value.cpu().contiguous()
+ return data_dict
+
+
+def create_env_batch(obs, rews, dones, infos, meta=None):
+ ret_dict = {"obs": obs, "rews": rews, "dones": dones, "infos": infos}
+ if meta is not None:
+ ret_dict.update(meta=meta)
+
+ ret_dict = put_tensor_cpu(ret_dict)
+ return ret_dict
+
+
+def create_env_batch_dataproto(obs, rews, terminations, truncations, infos, meta=None):
+ ret_dict = {"obs": obs, "rews": rews, "terminations": terminations, "truncations": truncations, "infos": infos}
+ if meta is not None:
+ ret_dict.update(meta=meta)
+
+ ret_dict = put_tensor_cpu(ret_dict)
+ tensor_batch = {
+ "full_image": ret_dict["obs"]["images_and_states"]["full_image"],
+ "state": ret_dict["obs"]["images_and_states"]["state"],
+ "rews": ret_dict["rews"],
+ "terminations": ret_dict["terminations"],
+ "truncations": ret_dict["truncations"],
+ }
+ non_tensor_batch = {"task_descriptions": obs["task_descriptions"]}
+ output = DataProto.from_dict(tensors=tensor_batch, non_tensors=non_tensor_batch)
+
+ return output
+
+
+class EnvWorker(Worker, DistProfilerExtension):
+ def __init__(self, config: DictConfig):
+ Worker.__init__(self)
+ self.cfg = config
+ self.train_video_cnt = 0
+ self.eval_video_cnt = 0
+
+ self.simulator_list = []
+ self.last_obs_list = []
+ self.last_dones_list = []
+ self.eval_simulator_list = []
+
+ self.stage_num = self.cfg.rollout.pipeline_stage_num
+ initialize_global_process_group_ray(timeout_second=None)
+ device_name = get_device_name()
+ env_device_mesh = init_device_mesh(device_name, mesh_shape=(self.world_size, 1), mesh_dim_names=["dp", "tp"])
+ self._register_dispatch_collect_info("env", dp_rank=env_device_mesh["dp"].get_local_rank(), is_collect=True)
+
+ # Initialize profiler
+ omega_profiler_config = config.train.get("profiler", {})
+ profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig)
+ if omega_profiler_config.get("tool", None) in ["npu", "nsys", "torch", "torch_memory"]:
+ tool_config = omega_conf_to_dataclass(
+ omega_profiler_config.get("tool_config", {}).get(omega_profiler_config.get("tool"))
+ )
+ else:
+ tool_config = None
+ DistProfilerExtension.__init__(
+ self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config)
+ )
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ @DistProfiler.annotate(color="green", role="env_init")
+ def init_worker(self):
+ if self.cfg.train.simulator_type == "libero":
+ from verl.experimental.vla.envs.libero_env.libero_env import LiberoEnv
+
+ for _ in range(self.stage_num):
+ self.simulator_list.append(
+ EnvManager(
+ self.cfg.train,
+ rank=self._rank,
+ world_size=self._world_size,
+ env_cls=LiberoEnv,
+ )
+ )
+
+ elif self.cfg.train.simulator_type == "isaac":
+ from verl.experimental.vla.envs.isaac_env.isaac_env import IsaacEnv
+
+ for _ in range(self.stage_num):
+ self.simulator_list.append(
+ EnvManager(
+ self.cfg.train,
+ rank=self._rank,
+ world_size=self._world_size,
+ env_cls=IsaacEnv,
+ )
+ )
+ else:
+ raise NotImplementedError(f"Simulator type {self.cfg.train.simulator_type} not implemented")
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ @DistProfiler.annotate(color="green", role="env_init_simulator")
+ def init_simulator(self):
+ for i in range(self.stage_num):
+ self.simulator_list[i].start_simulator()
+ return
+
+ @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="env"), blocking=False)
+ @DistProfiler.annotate(color="red", role="env_interact_step")
+ def env_interact_step(self, data: DataProto) -> dict:
+ """
+ This function is used to interact with the environment.
+ """
+ chunk_actions: torch.Tensor = data.non_tensor_batch["actions"]
+ stage_id: int = data.meta_info["stage_id"]
+ chunk_actions = prepare_actions(
+ simulator_type=self.cfg.train.simulator_type,
+ raw_chunk_actions=chunk_actions,
+ num_action_chunks=self.cfg.actor.model.num_action_chunks,
+ action_dim=self.cfg.actor.model.action_dim,
+ )
+ env_info_list = {}
+
+ extracted_obs, chunk_rewards, chunk_terminations, chunk_truncations, infos = self.simulator_list[
+ stage_id
+ ].chunk_step(chunk_actions)
+ chunk_dones = torch.logical_or(chunk_terminations, chunk_truncations)
+
+ if chunk_dones.any():
+ if "final_info" in infos:
+ final_info = infos["final_info"]
+ for key in final_info["episode"]:
+ env_info_list[key] = final_info["episode"][key][chunk_dones[:, -1]].cpu()
+
+ env_batch = create_env_batch_dataproto(
+ obs=extracted_obs,
+ rews=chunk_rewards,
+ terminations=chunk_terminations,
+ truncations=chunk_truncations,
+ infos=infos,
+ meta=env_info_list,
+ )
+ return env_batch
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def get_all_state_ids(self):
+ """Get all available state IDs from the environment."""
+ state_ids = self.simulator_list[0].get_all_state_ids()
+ return state_ids
+
+ @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="env"), blocking=False)
+ @DistProfiler.annotate(color="blue", role="env_reset_envs_to_state_ids")
+ def reset_envs_to_state_ids(self, data: DataProto):
+ """Reset environments to specified state IDs.
+
+ Args:
+ state_ids: State IDs to reset environments to
+ """
+ state_ids_list = list(data.non_tensor_batch["state_ids"])
+ task_ids_list = list(data.non_tensor_batch["task_ids"])
+
+ assert len(state_ids_list) == self.cfg.train.num_envs * self.stage_num, (
+ f"state_ids_list length is {len(state_ids_list)}, but should be {self.cfg.train.num_envs * self.stage_num}"
+ )
+ result_list = []
+ for stage_id in range(self.stage_num):
+ if self.cfg.train.simulator_type == "isaac":
+ assert (
+ len(
+ set(
+ state_ids_list[
+ stage_id * self.cfg.train.num_envs : (stage_id + 1) * self.cfg.train.num_envs
+ ]
+ )
+ )
+ == 1
+ ), "rollout.n should equal to num_envs for isaac"
+
+ result = self.simulator_list[stage_id].reset_envs_to_state_ids(
+ state_ids_list[stage_id * self.cfg.train.num_envs : (stage_id + 1) * self.cfg.train.num_envs],
+ task_ids_list[stage_id * self.cfg.train.num_envs : (stage_id + 1) * self.cfg.train.num_envs],
+ )
+ result_list.append(result)
+ output_tensor_dict = {}
+ output_non_tensor_dict = {}
+
+ # Handle nested 'images_and_states'
+ images_and_states_list = [d[0]["images_and_states"] for d in result_list]
+ if images_and_states_list:
+ # Assuming all dicts in the list have the same keys
+ for k in images_and_states_list[0].keys():
+ if isinstance(images_and_states_list[0][k], torch.Tensor):
+ output_tensor_dict[k] = torch.cat([d[k] for d in images_and_states_list])
+
+ # Handle 'task_descriptions'
+ task_descriptions_list = [d[0]["task_descriptions"] for d in result_list]
+ output_non_tensor_dict["task_descriptions"] = list(itertools.chain.from_iterable(task_descriptions_list))
+
+ output = DataProto.from_dict(tensors=output_tensor_dict, non_tensors=output_non_tensor_dict)
+ return output
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ @DistProfiler.annotate(color="gray", role="env_finish_rollout")
+ def finish_rollout(self, mode="train"):
+ # reset
+ if mode == "train":
+ if self.cfg.train.video_cfg.save_video:
+ for i in range(self.stage_num):
+ self.simulator_list[i].flush_video(video_sub_dir=f"stage_{i}")
diff --git a/code/RL_model/verl/verl_train/verl/single_controller/base/__init__.py b/code/RL_model/verl/verl_train/verl/single_controller/base/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b24bd9942b872b71f4c7b3a2dbfe6db5530cfe25
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/single_controller/base/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2024 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 .worker import Worker
+from .worker_group import ClassWithInitArgs, ResourcePool, WorkerGroup
+
+__all__ = ["Worker", "WorkerGroup", "ClassWithInitArgs", "ResourcePool"]
diff --git a/code/RL_model/verl/verl_train/verl/single_controller/base/decorator.py b/code/RL_model/verl/verl_train/verl/single_controller/base/decorator.py
new file mode 100644
index 0000000000000000000000000000000000000000..540c4e00552ada733d34e020b3f45d6e7fca097d
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/single_controller/base/decorator.py
@@ -0,0 +1,475 @@
+# Copyright 2024 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.
+import inspect
+from functools import partial, wraps
+from types import FunctionType
+
+from tensordict import TensorDict
+
+from verl.protocol import DataProtoFuture, _padding_size_key
+from verl.utils.py_functional import DynamicEnum
+from verl.utils.tensordict_utils import chunk_tensordict, concat_tensordict, contiguous
+from verl.utils.transferqueue_utils import BatchMeta
+
+# here we add a magic number of avoid user-defined function already have this attribute
+MAGIC_ATTR = "attrs_3141562937"
+
+
+class Dispatch(DynamicEnum):
+ """Enum class defining different dispatch modes for distributed computation.
+
+ Each mode represents a specific strategy for distributing data across
+ different ranks in a distributed system. The modes are used to control
+ how data is partitioned and processed across different worker groups.
+ """
+
+ _registry = {}
+ _next_value = 0
+
+
+def init_predefined_dispatch_mode():
+ Dispatch.register("RANK_ZERO")
+ Dispatch.register("ONE_TO_ALL")
+ Dispatch.register("ALL_TO_ALL")
+ Dispatch.register("DP_COMPUTE")
+ Dispatch.register("DP_COMPUTE_PROTO")
+ Dispatch.register("DP_COMPUTE_PROTO_WITH_FUNC")
+ Dispatch.register("DP_COMPUTE_METRIC")
+ # This is a special dispatch mode for vllm ExternalRayDistributedExecutor
+ Dispatch.register("DIRECT_ROLLOUT_METHOD")
+
+
+class Execute(DynamicEnum):
+ """Enum class defining different execution modes for distributed computation.
+
+ These modes control how a function should be executed across different ranks
+ in a distributed system.
+ """
+
+ _registry = {}
+ _next_value = 0
+
+
+def init_predefined_execute_mode():
+ Execute.register("ALL")
+ Execute.register("RANK_ZERO")
+
+
+# Initialize the two Dynamic Enum Classes
+init_predefined_dispatch_mode()
+init_predefined_execute_mode()
+
+
+def _consolidate_tuple_td(chunked_arg):
+ return tuple(contiguous(val).consolidate() for val in chunked_arg)
+
+
+def _split_args_kwargs_data_proto(chunks, *args, **kwargs):
+ from verl.protocol import DataProto, DataProtoFuture
+
+ splitted_args = []
+ for arg in args:
+ assert isinstance(arg, DataProto | DataProtoFuture | BatchMeta | TensorDict)
+ if isinstance(arg, TensorDict):
+ chunked_arg = chunk_tensordict(arg, chunks)
+ chunked_arg = _consolidate_tuple_td(chunked_arg)
+ else:
+ chunked_arg = arg.chunk(chunks=chunks)
+ assert len(chunked_arg) == chunks
+ splitted_args.append(chunked_arg)
+
+ splitted_kwargs = {}
+ for key, val in kwargs.items():
+ assert isinstance(val, DataProto | DataProtoFuture | BatchMeta | TensorDict)
+ if isinstance(val, TensorDict):
+ chunked_kwarg = chunk_tensordict(val, chunks)
+ chunked_kwarg = _consolidate_tuple_td(chunked_kwarg)
+ else:
+ chunked_kwarg = val.chunk(chunks=chunks)
+ assert len(chunked_kwarg) == chunks
+ splitted_kwargs[key] = chunked_kwarg
+
+ return splitted_args, splitted_kwargs
+
+
+def _split_args_kwargs_data_proto_with_auto_padding(chunks, *args, **kwargs):
+ from verl.protocol import DataProto, DataProtoFuture
+
+ data_proto_len = None
+ padding_size = None
+
+ def _padding_and_split_data(obj, chunks):
+ nonlocal data_proto_len, padding_size
+ assert isinstance(obj, DataProto | DataProtoFuture)
+ if isinstance(obj, DataProto) and obj.is_padding_enabled():
+ # for padding, we only support DataProto with same length
+ if data_proto_len is None:
+ data_proto_len = len(obj)
+ padding_size = (chunks - (data_proto_len % chunks)) if (data_proto_len % chunks > 0) else 0
+ else:
+ assert data_proto_len == len(obj), (
+ f"expecting all arg share same length of {data_proto_len}, but got {len(obj)}"
+ )
+ obj.padding(padding_size=padding_size)
+ return obj.chunk(chunks=chunks)
+
+ splitted_args = [_padding_and_split_data(arg, chunks) for arg in args]
+ splitted_kwargs = {key: _padding_and_split_data(val, chunks) for key, val in kwargs.items()}
+ if padding_size is not None:
+ splitted_kwargs[_padding_size_key] = padding_size
+
+ return splitted_args, splitted_kwargs
+
+
+def dispatch_one_to_all(worker_group, *args, **kwargs):
+ args = tuple([arg] * worker_group.world_size for arg in args)
+ kwargs = {k: [v] * worker_group.world_size for k, v in kwargs.items()}
+ return args, kwargs
+
+
+def dummy_direct_rollout_call(worker_group, *args, **kwargs):
+ raise NotImplementedError("Direct rollout call is forbidden.")
+
+
+def dispatch_all_to_all(worker_group, *args, **kwargs):
+ return args, kwargs
+
+
+def collect_all_to_all(worker_group, output):
+ return output
+
+
+def _concat_data_proto_or_future(output: list):
+ import ray
+
+ from verl.protocol import DataProto, DataProtoFuture
+
+ # make sure all the elements in output has the same type
+ for o in output:
+ assert type(o) is type(output[0])
+
+ o = output[0]
+
+ if isinstance(o, DataProto):
+ return DataProto.concat(output)
+ elif isinstance(o, ray.ObjectRef):
+ return DataProtoFuture.concat(output)
+ elif isinstance(o, BatchMeta):
+ return BatchMeta.concat(output)
+ elif isinstance(o, TensorDict):
+ return concat_tensordict(output)
+ else:
+ raise NotImplementedError
+
+
+def dispatch_dp_compute(worker_group, *args, **kwargs):
+ from verl.single_controller.base.worker_group import WorkerGroup
+
+ assert isinstance(worker_group, WorkerGroup)
+ for arg in args:
+ assert isinstance(arg, tuple | list) and len(arg) == worker_group.world_size
+ for k, v in kwargs.items():
+ assert isinstance(v, tuple | list) and len(v) == worker_group.world_size
+ return args, kwargs
+
+
+def collect_dp_compute(worker_group, output):
+ from verl.single_controller.base.worker_group import WorkerGroup
+
+ assert isinstance(worker_group, WorkerGroup)
+ assert len(output) == worker_group.world_size
+ return output
+
+
+def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):
+ from verl.single_controller.base.worker_group import WorkerGroup
+
+ assert isinstance(worker_group, WorkerGroup)
+ # Note: enable auto padding for dp compute DatapProto
+ splitted_args, splitted_kwargs = _split_args_kwargs_data_proto_with_auto_padding(
+ worker_group.world_size,
+ *args,
+ **kwargs,
+ )
+ return splitted_args, splitted_kwargs
+
+
+def dispatch_dp_compute_data_proto_with_func(worker_group, *args, **kwargs):
+ from verl.single_controller.base.worker_group import WorkerGroup
+
+ assert isinstance(worker_group, WorkerGroup)
+ assert isinstance(args[0], FunctionType) # NOTE: The first one args is a function!
+
+ splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args[1:], **kwargs)
+ splitted_args_with_func = [[args[0]] * worker_group.world_size] + splitted_args
+ return splitted_args_with_func, splitted_kwargs
+
+
+def collect_dp_compute_data_proto(worker_group, output):
+ import ray
+
+ from verl.protocol import DataProto
+
+ for o in output:
+ assert isinstance(o, DataProto | ray.ObjectRef), f"expecting {o} to be DataProto, but got {type(o)}"
+
+ output = collect_dp_compute(worker_group, output)
+ return _concat_data_proto_or_future(output)
+
+
+def dispatch_nd_compute(dp_rank_mapping: list[int], dp_size, worker_group, *args, **kwargs):
+ import os
+
+ from verl.single_controller.base.worker_group import WorkerGroup
+ from verl.utils.ray_utils import parallel_put
+
+ assert isinstance(worker_group, WorkerGroup)
+
+ max_workers = max(1, min(len(args[0]), os.cpu_count()))
+
+ args = [parallel_put(arg, max_workers=max_workers) for arg in args]
+ kwargs = {k: parallel_put(v, max_workers=max_workers) for k, v in kwargs.items()}
+
+ all_args = []
+ for arg in args:
+ assert isinstance(arg, tuple | list) and len(arg) == dp_size
+ transformed_args = []
+ for i in range(worker_group.world_size):
+ local_dp_rank = dp_rank_mapping[i]
+ transformed_args.append(arg[local_dp_rank])
+ all_args.append(transformed_args)
+ all_args = tuple(all_args)
+
+ all_kwargs = {}
+ for k, v in kwargs.items():
+ assert isinstance(v, tuple | list) and len(v) == dp_size
+ transformed_v = []
+ for i in range(worker_group.world_size):
+ local_dp_rank = dp_rank_mapping[i]
+ transformed_v.append(v[local_dp_rank])
+ all_kwargs[k] = transformed_v
+ return all_args, all_kwargs
+
+
+def collect_nd_compute(collect_mask: list[bool], worker_group, output):
+ from verl.single_controller.base.worker_group import WorkerGroup
+
+ assert isinstance(worker_group, WorkerGroup)
+ assert len(output) == worker_group.world_size
+
+ output_in_dp = []
+ for global_rank in range(worker_group.world_size):
+ collect_dp_rank = collect_mask[global_rank]
+ if collect_dp_rank:
+ output_in_dp.append(output[global_rank])
+ return output_in_dp
+
+
+def dispatch_nd_compute_dataproto(dp_rank_mapping: list[int], dp_size, worker_group, *args, **kwargs):
+ splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(dp_size, *args, **kwargs)
+ return dispatch_nd_compute(dp_rank_mapping, dp_size, worker_group, *splitted_args, **splitted_kwargs)
+
+
+def collect_nd_compute_dataproto(collect_mask: list[bool], worker_group, output):
+ output = collect_nd_compute(collect_mask, worker_group, output)
+ import ray
+
+ from verl.protocol import DataProto
+
+ for o in output:
+ assert isinstance(o, DataProto | ray.ObjectRef | BatchMeta | TensorDict), (
+ f"expecting {o} to be DataProto | ray.ObjectRef | BatchMeta | TensorDict, but got {type(o)}"
+ )
+ return _concat_data_proto_or_future(output)
+
+
+def dispatch_lazy_compute_data_proto(mesh_name, worker_group, *args, **kwargs):
+ from verl.single_controller.base.worker_group import WorkerGroup
+
+ assert isinstance(worker_group, WorkerGroup)
+
+ # query dispatch info of the worker group
+ if mesh_name not in worker_group._dispatch_info:
+ worker_group._dispatch_info[mesh_name] = worker_group._query_dispatch_info(mesh_name)
+ assert len(worker_group._dispatch_info[mesh_name]) == worker_group.world_size
+
+ dp_rank_mapping = worker_group._dispatch_info[mesh_name]
+ # perform dispatch
+ dp_size = max(dp_rank_mapping) + 1
+ return dispatch_nd_compute_dataproto(dp_rank_mapping, dp_size, worker_group, *args, **kwargs)
+
+
+def collect_lazy_compute_data_proto(mesh_name, worker_group, *args, **kwargs):
+ from verl.single_controller.base.worker_group import WorkerGroup
+
+ assert isinstance(worker_group, WorkerGroup)
+
+ # the dispatch info is stored in the worker group
+ assert mesh_name in worker_group._dispatch_info
+
+ if mesh_name not in worker_group._collect_info:
+ worker_group._collect_info[mesh_name] = worker_group._query_collect_info(mesh_name)
+ assert len(worker_group._collect_info[mesh_name]) == worker_group.world_size
+
+ # a boolean of whether the dp_rank is used for collect
+ collect_mask = worker_group._collect_info[mesh_name]
+ # perform dispatch
+ return collect_nd_compute_dataproto(collect_mask, worker_group, *args, **kwargs)
+
+
+def make_nd_compute_dataproto_dispatch_fn(mesh_name):
+ return {
+ "dispatch_fn": partial(dispatch_lazy_compute_data_proto, mesh_name),
+ "collect_fn": partial(collect_lazy_compute_data_proto, mesh_name),
+ }
+
+
+# Global registry for dispatch mode.
+DISPATCH_MODE_FN_REGISTRY = {
+ Dispatch.ONE_TO_ALL: {
+ "dispatch_fn": dispatch_one_to_all,
+ "collect_fn": collect_all_to_all,
+ },
+ Dispatch.ALL_TO_ALL: {
+ "dispatch_fn": dispatch_all_to_all,
+ "collect_fn": collect_all_to_all,
+ },
+ Dispatch.DP_COMPUTE: {"dispatch_fn": dispatch_dp_compute, "collect_fn": collect_dp_compute},
+ Dispatch.DP_COMPUTE_PROTO: {
+ "dispatch_fn": dispatch_dp_compute_data_proto,
+ "collect_fn": collect_dp_compute_data_proto,
+ },
+ Dispatch.DP_COMPUTE_PROTO_WITH_FUNC: {
+ "dispatch_fn": dispatch_dp_compute_data_proto_with_func,
+ "collect_fn": collect_dp_compute_data_proto,
+ },
+ Dispatch.DP_COMPUTE_METRIC: {"dispatch_fn": dispatch_dp_compute_data_proto, "collect_fn": collect_dp_compute},
+ Dispatch.DIRECT_ROLLOUT_METHOD: {
+ "dispatch_fn": dummy_direct_rollout_call,
+ "collect_fn": dummy_direct_rollout_call,
+ },
+}
+
+
+def get_predefined_dispatch_fn(dispatch_mode):
+ return DISPATCH_MODE_FN_REGISTRY[dispatch_mode]
+
+
+def register_dispatch_mode(dispatch_mode_name, dispatch_fn, collect_fn):
+ """
+ Register a new dispatch mode.
+ """
+ dispatch_mode = Dispatch.register(dispatch_mode_name)
+ _check_dispatch_mode(dispatch_mode)
+ assert dispatch_mode not in DISPATCH_MODE_FN_REGISTRY, f"dispatch_mode_name {dispatch_mode_name} already exists"
+ DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = {"dispatch_fn": dispatch_fn, "collect_fn": collect_fn}
+
+
+def update_dispatch_mode(dispatch_mode, dispatch_fn, collect_fn):
+ """
+ Update the dispatch mode.
+ """
+ _check_dispatch_mode(dispatch_mode)
+ assert dispatch_mode in DISPATCH_MODE_FN_REGISTRY, f"dispatch_mode {dispatch_mode} not found"
+ DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = {"dispatch_fn": dispatch_fn, "collect_fn": collect_fn}
+
+
+def get_predefined_execute_fn(execute_mode):
+ """
+ Note that here we only asks execute_all and execute_rank_zero to be implemented
+ Leave the choice of how these two functions handle argument 'blocking' to users
+ """
+ predefined_execute_mode_fn = {
+ Execute.ALL: {"execute_fn_name": "execute_all"},
+ Execute.RANK_ZERO: {"execute_fn_name": "execute_rank_zero"},
+ }
+ return predefined_execute_mode_fn[execute_mode]
+
+
+def _check_dispatch_mode(dispatch_mode):
+ assert isinstance(dispatch_mode, Dispatch | dict), (
+ f"dispatch_mode must be a Dispatch or a Dict. Got {dispatch_mode}"
+ )
+ if isinstance(dispatch_mode, dict):
+ necessary_keys = ["dispatch_fn", "collect_fn"]
+ for key in necessary_keys:
+ assert key in dispatch_mode, f"key {key} should be in dispatch_mode if it is a dictionary"
+
+
+def _check_execute_mode(execute_mode):
+ assert isinstance(execute_mode, Execute), f"execute_mode must be a Execute. Got {execute_mode}"
+
+
+def _materialize_futures(*args, **kwargs):
+ new_args = []
+ for arg in args:
+ if isinstance(arg, DataProtoFuture):
+ arg = arg.get()
+ # add more type to materialize
+ new_args.append(arg)
+ for k, v in kwargs.items():
+ if isinstance(v, DataProtoFuture):
+ kwargs[k] = v.get()
+
+ new_args = tuple(new_args)
+ return new_args, kwargs
+
+
+def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True):
+ """Register a function with distributed execution configuration.
+
+ This decorator registers a function with specific dispatch and execution modes
+ for distributed computation. It handles both synchronous and asynchronous
+ functions, and optionally materializes futures before execution.
+
+ Args:
+ dispatch_mode:
+ Dispatch mode for computation distribution. Default: Dispatch.ALL_TO_ALL.
+ execute_mode:
+ Execute mode for computation distribution. Default: Execute.ALL.
+ blocking:
+ Whether the execution should be blocking. Defaults to True.
+ materialize_futures:
+ Whether to materialize the data before dispatching. Defaults to True.
+
+ Returns:
+ A decorator that wraps the original function with distributed execution
+ configuration.
+ """
+ from verl.utils.transferqueue_utils import tqbridge
+
+ _check_dispatch_mode(dispatch_mode=dispatch_mode)
+ _check_execute_mode(execute_mode=execute_mode)
+
+ def decorator(func):
+ func = tqbridge(dispatch_mode=dispatch_mode)(func)
+
+ @wraps(func)
+ def inner(*args, **kwargs):
+ if materialize_futures:
+ args, kwargs = _materialize_futures(*args, **kwargs)
+ return func(*args, **kwargs)
+
+ @wraps(func)
+ async def async_inner(*args, **kwargs):
+ if materialize_futures:
+ args, kwargs = _materialize_futures(*args, **kwargs)
+ return await func(*args, **kwargs)
+
+ wrapper = async_inner if inspect.iscoroutinefunction(func) else inner
+ attrs = {"dispatch_mode": dispatch_mode, "execute_mode": execute_mode, "blocking": blocking}
+ setattr(wrapper, MAGIC_ATTR, attrs)
+ return wrapper
+
+ return decorator
diff --git a/code/RL_model/verl/verl_train/verl/single_controller/base/worker.py b/code/RL_model/verl/verl_train/verl/single_controller/base/worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..cffaf5d30a0f966ebc857bfe365a90685c1dd565
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/single_controller/base/worker.py
@@ -0,0 +1,357 @@
+# Copyright 2024 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.
+"""
+the class for Worker
+"""
+
+import os
+import socket
+import warnings
+from dataclasses import dataclass
+
+import ray
+
+from verl.utils.device import (
+ get_torch_device,
+ get_visible_devices_keyword,
+ is_npu_available,
+)
+
+from .decorator import Dispatch, Execute, register
+
+
+@dataclass
+class DistRankInfo:
+ tp_rank: int
+ dp_rank: int
+ pp_rank: int
+ cp_rank: int
+
+
+@dataclass
+class DistGlobalInfo:
+ tp_size: int
+ dp_size: int
+ pp_size: int
+ cp_size: int
+
+
+class WorkerHelper:
+ @staticmethod
+ def _get_node_ip():
+ if os.getenv("WG_BACKEND", None) == "ray":
+ return ray.util.get_node_ip_address()
+ else:
+ raise NotImplementedError("WG_BACKEND now just support ray mode.")
+
+ @staticmethod
+ def _get_free_port():
+ with socket.socket() as sock:
+ sock.bind(("", 0))
+ return sock.getsockname()[1]
+
+ def get_availale_master_addr_port(self):
+ warnings.warn(
+ "This function is deprecated due to typo in name; Please use `get_available_master_addr_port` instead",
+ stacklevel=2,
+ )
+ return self.get_available_master_addr_port()
+
+ def get_available_master_addr_port(self):
+ return self._get_node_ip().strip("[]"), str(self._get_free_port())
+
+
+# we assume that in each WorkerGroup, there is a Master Worker
+class Worker(WorkerHelper):
+ """A distributed worker that handles initialization and configuration for distributed training.
+
+ This class manages worker initialization, configuration, and provides methods for executing
+ distributed operations. It handles communication settings, device configuration, and worker
+ metadata management.
+ """
+
+ fused_worker_attr_name = "fused_worker_dict"
+
+ def _register_dispatch_collect_info(self, mesh_name: str, dp_rank: int, is_collect: bool):
+ """Register the dp_rank for a given mesh name. This function is meant to be called by the worker
+
+ Args:
+ mesh_name (str):
+ Name of the mesh to register dp_rank for.
+ dp_rank (int):
+ dp_rank to register for the given mesh name.
+ is_collect (bool):
+ Whether the dp_rank is used for collect.
+ """
+ if mesh_name in self.__dispatch_dp_rank or mesh_name in self.__collect_dp_rank:
+ raise ValueError(f"mesh_name {mesh_name} has been registered")
+ self.__dispatch_dp_rank[mesh_name] = dp_rank
+ self.__collect_dp_rank[mesh_name] = is_collect
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def _query_dispatch_info(self, mesh_name: str):
+ """Query the dispatch info for a given mesh name.
+
+ Args:
+ mesh_name (str):
+ Name of the mesh to query dispatch info for.
+
+ Returns:
+ int:
+ The dp_rank for the given mesh name.
+ """
+ assert mesh_name in self.__dispatch_dp_rank, f"{mesh_name} is not registered in {self.__class__.__name__}"
+ # note that each rank store its own dp_rank
+ return self.__dispatch_dp_rank[mesh_name]
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL)
+ def _query_collect_info(self, mesh_name: str):
+ return self.query_collect_info(mesh_name)
+
+ def query_collect_info(self, mesh_name: str):
+ """Query the collect info for a given mesh name.
+
+ Args:
+ mesh_name (str):
+ Name of the mesh to query collect info for.
+
+ Returns:
+ bool:
+ Whether the dp_rank is used for collect.
+ """
+ assert mesh_name in self.__collect_dp_rank, f"{mesh_name} is not registered in {self.__class__.__name__}"
+ return self.__collect_dp_rank[mesh_name]
+
+ def get_dispatch_collect(self):
+ """Get all registered dispatch and collect dp_ranks.
+
+ Returns:
+ dict[str, int]:
+ A dictionary mapping mesh names to their dispatch dp_ranks.
+ dict[str, bool]:
+ A dictionary mapping mesh names to whether they are used for collect.
+ """
+ return {"dispatch_dp_rank": self.__dispatch_dp_rank, "collect_dp_rank": self.__collect_dp_rank}
+
+ def set_dispatch_collect(self, mesh_name: str, dispatch_dp_rank: dict[str, int], collect_dp_rank: dict[str, bool]):
+ """Set the dispatch and collect dp_ranks for all registered meshes.
+
+ Args:
+ mesh_name (str): Mesh name to set dispatch and collect dp_ranks for.
+ dispatch_dp_rank (dict[str, int]):
+ A dictionary mapping mesh names to their dispatch dp_ranks.
+ collect_dp_rank (dict[str, bool]):
+ A dictionary mapping mesh names to whether they are used for collect.
+ """
+ assert mesh_name not in self.__dispatch_dp_rank, (
+ f"{mesh_name} is already registered, {self.__dispatch_dp_rank.keys()}"
+ )
+ assert mesh_name not in self.__collect_dp_rank, (
+ f"{mesh_name} is already registered, {self.__collect_dp_rank.keys()}"
+ )
+ for dp_rank in dispatch_dp_rank.values():
+ self.__dispatch_dp_rank[mesh_name] = dp_rank
+ for is_collect in collect_dp_rank.values():
+ self.__collect_dp_rank[mesh_name] = is_collect
+
+ @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=True)
+ def create_transferqueue_client(self, config):
+ from verl.utils.transferqueue_utils import create_transferqueue_client
+
+ create_transferqueue_client(
+ client_id=f"worker_{self.rank}",
+ config=config.transfer_queue,
+ )
+
+ @classmethod
+ def env_keys(cls):
+ """The keys of the environment variables that are used to configure the Worker."""
+ return [
+ "WORLD_SIZE",
+ "RANK",
+ "LOCAL_WORLD_SIZE",
+ "LOCAL_RANK",
+ "MASTER_ADDR",
+ "MASTER_PORT",
+ get_visible_devices_keyword().upper(),
+ ]
+
+ def __init__(self, cuda_visible_devices=None) -> None:
+ """Initialize the worker with environment settings and device configuration.
+
+ Args:
+ cuda_visible_devices (str, optional):
+ CUDA visible devices configuration. Defaults to None.
+ """
+ # construct a meta from environment variable. Note that the import must be inside the class because
+ # it is executed remotely
+ import os
+
+ self._setup_env_cuda_visible_devices()
+
+ world_size = int(os.environ["WORLD_SIZE"])
+ rank = int(os.environ["RANK"])
+ self._rank = rank
+ self._world_size = world_size
+
+ master_addr = os.environ["MASTER_ADDR"]
+ master_port = os.environ["MASTER_PORT"]
+
+ local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", "1"))
+ local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+ store = {
+ "_world_size": world_size,
+ "_rank": rank,
+ "_local_world_size": local_world_size,
+ "_local_rank": local_rank,
+ "_master_addr": master_addr,
+ "_master_port": master_port,
+ }
+ if cuda_visible_devices is not None:
+ store[f"_{get_visible_devices_keyword()}".lower()] = cuda_visible_devices
+
+ self._configure_with_store(store=store)
+
+ self.fused_worker_dict = {}
+ self.__dispatch_dp_rank = {}
+ self.__collect_dp_rank = {}
+
+ def get_fused_worker_by_name(self, worker_name: str):
+ """Get a fused worker by its name.
+
+ Args:
+ worker_name (str):
+ Name of the worker to retrieve
+ """
+ return self.fused_worker_dict.get(worker_name, None)
+
+ def _setup_env_cuda_visible_devices(self):
+ from verl.utils.ray_utils import ray_noset_visible_devices
+
+ is_ray_noset_visible_devices = ray_noset_visible_devices()
+
+ # Prevent use of clashing `{CUDA/HIP/ROCR}_VISIBLE_DEVICES``
+ rocr_val = os.environ.get("ROCR_VISIBLE_DEVICES", None)
+ hip_val = os.environ.get("HIP_VISIBLE_DEVICES", None)
+ cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+ if hip_val:
+ # Switch the use of HIP_VISIBLE_DEVICES to CUDA_VISIBLE_DEVICES for consistency.
+ # Make sure that the HIP_VISIBLE_DEVICES is set to the same value as CUDA_VISIBLE_DEVICES
+ # at this point.
+ val = os.environ.pop("HIP_VISIBLE_DEVICES")
+ hip_val = None
+ if cuda_val:
+ assert val == cuda_val, (
+ f"Please use the same HIP_VISIBLE_DEVICES or CUDA_VISIBLE_DEVICES, inconsistant values "
+ f"found: {val} and {cuda_val}."
+ )
+ else:
+ cuda_val = val
+ os.environ["CUDA_VISIBLE_DEVICES"] = val
+ # os.environ["HIP_VISIBLE_DEVICES"] = val
+
+ if rocr_val:
+ # You must take care if both HIP/CUDA and ROCR env vars are set as they have
+ # different meanings. Both env vars accept either a list of ints or a
+ # list of UUIDs. The ROCR env var is processed first which then reduces
+ # the number of GPUs that HIP can select from.
+ # https://github.com/pytorch/pytorch/pull/144026
+ # To avoid the complexity of this, we simply gives out error if both are set
+ # (Also to keep consistency with ray's practice with 2.45.0).
+ # Otherwise, we will set ROCR_VISIBLE_DEVICES to CUDA_VISIBLE_DEVICES
+ # and remove ROCR_VISIBLE_DEVICES.
+ if cuda_val:
+ raise ValueError("Please don't set ROCR_VISIBLE_DEVICES when HIP/CUDA_VISIBLE_DEVICES is set.")
+
+ cuda_val = os.environ.pop("ROCR_VISIBLE_DEVICES")
+ os.environ["CUDA_VISIBLE_DEVICES"] = cuda_val
+ rocr_val = None
+
+ if is_ray_noset_visible_devices:
+ # NOTE: Ray will automatically set the *_VISIBLE_DEVICES
+ # environment variable for each actor, unless
+ # RAY_EXPERIMENTAL_NOSET_*_VISIBLE_DEVICES is set,
+ # so we need to set local rank when the flag is set.
+ device_name = "NPU" if is_npu_available else "GPU"
+ local_rank = ray.get_runtime_context().get_accelerator_ids()[device_name][0]
+ os.environ["LOCAL_RANK"] = local_rank
+ get_torch_device().set_device(int(local_rank))
+
+ def _configure_with_store(self, store: dict):
+ """
+ This function should only be called inside by WorkerGroup
+ """
+ store_env_dict = {f"_{key.lower()}": store.get(f"_{key.lower()}", None) for key in type(self).env_keys()}
+ self.__dict__.update(store_env_dict) # this is hacky
+ # print(f"__dict__: {self.__dict__}")
+ for key in type(self).env_keys():
+ val = self.__dict__.get(f"_{key.lower()}", None)
+ if val is not None:
+ # print(f"set {key} to {val}")
+ os.environ[key] = str(val)
+ os.environ["REDIS_STORE_SERVER_HOST"] = (
+ str(self._master_addr).replace("[", "").replace("]", "") if self._master_addr else ""
+ )
+
+ def get_master_addr_port(self):
+ """Get the master address and port for distributed communication."""
+ return self._master_addr, self._master_port
+
+ def get_cuda_visible_devices(self):
+ """Get the CUDA visible devices configuration."""
+ import os
+
+ visible_devices = os.environ.get(get_visible_devices_keyword().upper(), "not set")
+ return visible_devices
+
+ @property
+ def world_size(self):
+ """Get the total number of workers in the distributed setup."""
+ return self._world_size
+
+ @property
+ def rank(self):
+ """Get the rank of this worker in the distributed setup."""
+ return self._rank
+
+ @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO_WITH_FUNC)
+ def execute_with_func_generator(self, func, *args, **kwargs):
+ """Execute a function with function generator dispatch mode.
+
+ Args:
+ func:
+ Function to execute
+ *args:
+ Positional arguments for the function
+ **kwargs:
+ Keyword arguments for the function
+ """
+ ret_proto = func(self, *args, **kwargs)
+ return ret_proto
+
+ @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO)
+ def execute_func_rank_zero(self, func, *args, **kwargs):
+ """Execute a function in rank zero execution mode.
+
+ Args:
+ func:
+ Function to execute
+ *args:
+ Positional arguments for the function
+ **kwargs:
+ Keyword arguments for the function
+ """
+ result = func(*args, **kwargs)
+ return result
diff --git a/code/RL_model/verl/verl_train/verl/single_controller/base/worker_group.py b/code/RL_model/verl/verl_train/verl/single_controller/base/worker_group.py
new file mode 100644
index 0000000000000000000000000000000000000000..f5df3d6b31b32ce216dfcc6595a1e835171fb097
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/single_controller/base/worker_group.py
@@ -0,0 +1,255 @@
+# Copyright 2024 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.
+"""
+the class of WorkerGroup
+"""
+
+import logging
+import signal
+import threading
+import time
+from typing import Any, Callable
+
+from .decorator import MAGIC_ATTR, Dispatch, get_predefined_dispatch_fn, get_predefined_execute_fn
+
+
+class ResourcePool:
+ """
+ Manages a pool of resources across multiple nodes, tracking process counts and GPU allocations.
+ The class provides methods to calculate world size, local world sizes, and local ranks
+ across all nodes in the pool.
+ """
+
+ def __init__(self, process_on_nodes=None, max_colocate_count: int = 10, n_gpus_per_node=8) -> None:
+ """Initialize the ResourcePool with node processes and GPU configuration.
+
+ Args:
+ process_on_nodes (List[int], optional): List of process counts per node. Defaults to empty list.
+ max_colocate_count (int, optional): Maximum number of processes that can be colocated. Defaults to 10.
+ n_gpus_per_node (int, optional): Number of GPUs available per node. Defaults to 8.
+ """
+ if process_on_nodes is None:
+ process_on_nodes = []
+ self._store = process_on_nodes
+ self.max_colocate_count = max_colocate_count
+ self.n_gpus_per_node = n_gpus_per_node # this is left for future huawei GPU that contains 16 GPUs per node
+
+ def add_node(self, process_count):
+ self._store.append(process_count)
+
+ @property
+ def world_size(self):
+ """Total number of processes across all nodes in the pool."""
+ return sum(self._store)
+
+ def __call__(self) -> Any:
+ return self._store
+
+ @property
+ def store(self):
+ return self._store
+
+ def local_world_size_list(self) -> list[int]:
+ """Returns a flat list where each process has its local world size."""
+ nested_local_world_size_list = [
+ [local_world_size for _ in range(local_world_size)] for local_world_size in self._store
+ ]
+ return [item for row in nested_local_world_size_list for item in row]
+
+ def local_rank_list(self) -> list[int]:
+ """Returns a flat list of local ranks for all processes across all nodes."""
+ nested_local_rank_list = [[i for i in range(local_world_size)] for local_world_size in self._store]
+ return [item for row in nested_local_rank_list for item in row]
+
+
+class ClassWithInitArgs:
+ """
+ Wrapper class that stores constructor arguments for deferred instantiation.
+ This class is particularly useful for remote class instantiation where
+ the actual construction needs to happen at a different time or location.
+ """
+
+ def __init__(self, cls, *args, **kwargs) -> None:
+ """Initialize the ClassWithInitArgs instance.
+
+ Args:
+ cls: The class to be instantiated later
+ *args: Positional arguments for the class constructor
+ **kwargs: Keyword arguments for the class constructor
+ """
+ self.cls = cls
+ self.args = args
+ self.kwargs = kwargs
+
+ self.fused_worker_used = False
+
+ def __call__(self) -> Any:
+ """Instantiate the stored class with the stored arguments."""
+ return self.cls(*self.args, **self.kwargs)
+
+
+def check_workers_alive(workers: list, is_alive: Callable, gap_time: float = 1) -> None:
+ """Continuously monitors worker processes and raises SIGABRT if any worker dies.
+
+ Args:
+ workers (List):
+ List of worker objects to monitor
+ is_alive (Callable):
+ Function to check if a worker is alive
+ gap_time (float):
+ Time interval between checks
+ """
+ import time
+
+ while True:
+ for worker in workers:
+ if not is_alive(worker):
+ logging.warning(f"worker {worker} is not alive sending signal to main thread")
+ signal.raise_signal(signal.SIGABRT)
+ time.sleep(gap_time)
+
+
+class WorkerGroup:
+ """
+ Base class for managing a group of workers in a distributed system.
+ The class provides methods for worker management, aliveness checking, and method binding.
+ """
+
+ fused_worker_execute_fn_name = "_fuw_execute"
+
+ def __init__(self, resource_pool: ResourcePool, **kwargs) -> None:
+ self._is_init_with_detached_workers = resource_pool is None
+
+ self.fused_worker_used = False
+
+ if resource_pool is not None:
+ # handle the case when WorkGroup is attached to an existing one
+ self._procecss_dispatch_config = resource_pool()
+ else:
+ self._procecss_dispatch_config = None
+
+ self._workers = []
+ self._worker_names = []
+
+ self._dispatch_info = {}
+ self._collect_info = {}
+
+ self._master_addr = None
+ self._master_port = None
+
+ self._checker_thread: threading.Thread = None
+
+ def _is_worker_alive(self, worker):
+ """Check if a worker is alive. Must be implemented by derived classes."""
+ raise NotImplementedError("WorkerGroup._is_worker_alive called, should be implemented in derived class.")
+
+ def _block_until_all_workers_alive(self) -> None:
+ """Blocks until all workers in the group are alive."""
+ while True:
+ all_state = [self._is_worker_alive(worker) for worker in self._workers]
+ if False in all_state:
+ time.sleep(1)
+ else:
+ break
+
+ def start_worker_aliveness_check(self, every_n_seconds=1) -> None:
+ """Starts a background thread to monitor worker aliveness.
+
+ Args:
+ every_n_seconds (int): Interval between aliveness checks
+ """
+ # before starting checking worker aliveness, make sure all workers are already alive
+ self._block_until_all_workers_alive()
+
+ self._checker_thread = threading.Thread(
+ target=check_workers_alive, args=(self._workers, self._is_worker_alive, every_n_seconds)
+ )
+ self._checker_thread.start()
+
+ @property
+ def world_size(self):
+ """Number of workers in the group."""
+ return len(self._workers)
+
+ def _bind_worker_method(self, user_defined_cls, func_generator):
+ """Binds worker methods to the WorkerGroup based on registered attributes.
+
+ Args:
+ user_defined_cls (type): The class containing methods to bind
+ func_generator (Callable): Function that generates the bound method
+
+ Returns:
+ List[str]: List of method names that were successfully bound
+ """
+ method_names = []
+ for method_name in dir(user_defined_cls):
+ try:
+ method = getattr(user_defined_cls, method_name)
+ assert callable(method), f"{method_name} in {user_defined_cls} is not callable"
+ except Exception:
+ # if it is a property, it will fail because Class doesn't have instance property
+ continue
+
+ if hasattr(method, MAGIC_ATTR):
+ # this method is decorated by register
+ attribute = getattr(method, MAGIC_ATTR)
+ assert isinstance(attribute, dict), f"attribute must be a dictionary. Got {type(attribute)}"
+ assert "dispatch_mode" in attribute, "attribute must contain dispatch_mode in its key"
+
+ dispatch_mode = attribute["dispatch_mode"]
+ execute_mode = attribute["execute_mode"]
+ blocking = attribute["blocking"]
+
+ # get dispatch fn
+ if isinstance(dispatch_mode, Dispatch):
+ # get default dispatch fn
+ fn = get_predefined_dispatch_fn(dispatch_mode=dispatch_mode)
+ dispatch_fn = fn["dispatch_fn"]
+ collect_fn = fn["collect_fn"]
+ else:
+ assert isinstance(dispatch_mode, dict)
+ assert "dispatch_fn" in dispatch_mode
+ assert "collect_fn" in dispatch_mode
+ dispatch_fn = dispatch_mode["dispatch_fn"]
+ collect_fn = dispatch_mode["collect_fn"]
+
+ # get execute_fn_name
+ execute_mode = get_predefined_execute_fn(execute_mode=execute_mode)
+ wg_execute_fn_name = execute_mode["execute_fn_name"]
+
+ # get execute_fn from string
+ try:
+ execute_fn = getattr(self, wg_execute_fn_name)
+ assert callable(execute_fn), "execute_fn must be callable"
+ except Exception:
+ print(f"execute_fn {wg_execute_fn_name} is invalid")
+ raise
+
+ # bind a new method to the RayWorkerGroup
+ func = func_generator(
+ self,
+ method_name,
+ dispatch_fn=dispatch_fn,
+ collect_fn=collect_fn,
+ execute_fn=execute_fn,
+ blocking=blocking,
+ )
+
+ try:
+ setattr(self, method_name, func)
+ method_names.append(method_name)
+ except Exception as e:
+ raise ValueError(f"Fail to set method_name {method_name}") from e
+
+ return method_names
diff --git a/code/RL_model/verl/verl_train/verl/single_controller/ray/__init__.py b/code/RL_model/verl/verl_train/verl/single_controller/ray/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b60291d23acf0cde6480b3caa988dc1c872fddc5
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/single_controller/ray/__init__.py
@@ -0,0 +1,33 @@
+# Copyright 2024 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 .base import (
+ RayClassWithInitArgs,
+ RayResourcePool,
+ RayWorkerGroup,
+ ResourcePoolManager,
+ SubRayResourcePool,
+ create_colocated_worker_cls,
+ create_colocated_worker_cls_fused,
+)
+
+__all__ = [
+ "RayClassWithInitArgs",
+ "RayResourcePool",
+ "SubRayResourcePool",
+ "RayWorkerGroup",
+ "ResourcePoolManager",
+ "create_colocated_worker_cls",
+ "create_colocated_worker_cls_fused",
+]
diff --git a/code/RL_model/verl/verl_train/verl/single_controller/ray/base.py b/code/RL_model/verl/verl_train/verl/single_controller/ray/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..d632be4f6fb6e9cadace8c7f59ef936389a80432
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/single_controller/ray/base.py
@@ -0,0 +1,1098 @@
+# Copyright 2024 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.
+import inspect
+import logging
+import os
+import socket
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+import numpy as np
+import ray
+from ray.experimental.state.api import get_actor
+from ray.util.placement_group import PlacementGroup, placement_group
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, PlacementGroupSchedulingStrategy
+
+from verl.protocol import DataProto, _padding_size_key
+from verl.single_controller.base import ClassWithInitArgs, ResourcePool, Worker, WorkerGroup
+from verl.single_controller.base.decorator import MAGIC_ATTR, Dispatch
+from verl.utils.device import get_device_name
+from verl.utils.py_functional import temp_env_var
+
+__all__ = ["Worker"]
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+def get_random_string(length: int) -> str:
+ import random
+ import string
+
+ letters_digits = string.ascii_letters + string.digits
+ return "".join(random.choice(letters_digits) for _ in range(length))
+
+
+def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking):
+ class Functor:
+ def __call__(this, *args, **kwargs):
+ args, kwargs = dispatch_fn(self, *args, **kwargs)
+ padding_count = kwargs.pop(_padding_size_key, 0)
+ output = execute_fn(method_name, *args, **kwargs)
+ if blocking:
+ output = ray.get(output)
+ output = collect_fn(self, output)
+ if padding_count > 0:
+ if isinstance(output, DataProto):
+ indices = [i for i in range(len(output))][:-padding_count]
+ output = output.select_idxs(indices)
+ elif isinstance(output, list):
+ output = output[:-padding_count]
+ return output
+
+ # use class type to pass the method_name to get a better observability
+ return type(method_name, (Functor,), {})()
+
+
+def sort_placement_group_by_node_ip(pgs: list[PlacementGroup]) -> list[PlacementGroup]:
+ """
+ Sort the placement groups by node ip, all bundles in a single placement group should be on the same node.
+
+ FSDPCheckpointManager saves sharded model states and optimizer states in local storage, which requires RANK
+ to be consistent across nodes when resume from checkpoint.
+
+ With this function, if there's only one resource pool and there's no node change, RANK should be consistent
+ across nodes in multiple ray jobs, even if the whole ray cluster is restarted.
+ """
+ node_ip = {node["NodeID"]: node["NodeManagerAddress"] for node in ray.nodes()}
+ pg_ip = {}
+ for pg in pgs:
+ specs = ray._private.state.state.placement_group_table(pg.id)
+ # all bunles should be on the same node
+ node_id = specs["bundles_to_node_id"][0]
+ pg_ip[pg.id] = node_ip[node_id]
+ return sorted(pgs, key=lambda pg: pg_ip[pg.id])
+
+
+@ray.remote
+def get_master_addr_port() -> tuple[str, str]:
+ addr = ray.util.get_node_ip_address().strip("[]")
+ with socket.socket() as sock:
+ sock.bind(("", 0))
+ port = sock.getsockname()[1]
+ return addr, str(port)
+
+
+class RayResourcePool(ResourcePool):
+ def __init__(
+ self,
+ process_on_nodes: Optional[list[int]] = None,
+ use_gpu: bool = True,
+ name_prefix: str = None,
+ max_colocate_count: int = 10,
+ detached=False,
+ accelerator_type: Optional[str] = None,
+ ) -> None:
+ super().__init__(process_on_nodes, max_colocate_count)
+ self.use_gpu = use_gpu
+ # print(f"in RayProcessDispatchConfiguration: name_prefix = {name_prefix}")
+ self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix
+ self.pgs = None
+ self.detached = detached
+ self.accelerator_type = accelerator_type
+
+ def get_placement_groups(self, strategy="STRICT_PACK", name=None, device_name="cuda"):
+ if self.pgs is not None:
+ return self.pgs
+
+ pg_name_prefix = (
+ name if name else f"{self.name_prefix}verl_group_{'_'.join([str(count) for count in self._store])}:"
+ )
+ # print(f"pg_name_prefix = {pg_name_prefix}")
+ if device_name == "npu":
+ device_name = "NPU"
+ elif device_name == "cuda":
+ device_name = "GPU"
+
+ bundle = {"CPU": self.max_colocate_count}
+ if self.use_gpu:
+ bundle[device_name] = 1
+ if self.accelerator_type is not None:
+ bundle[self.accelerator_type] = 1e-4
+ pg_scheme = [[bundle.copy() for _ in range(process_count)] for process_count in self._store]
+
+ lifetime = "detached" if self.detached else None
+
+ pgs = [
+ placement_group(bundles=bundles, strategy=strategy, name=pg_name_prefix + str(idx), lifetime=lifetime)
+ for idx, bundles in enumerate(pg_scheme)
+ ]
+
+ ray.get([pg.ready() for pg in pgs])
+
+ self.pgs = sort_placement_group_by_node_ip(pgs)
+ return pgs
+
+
+class SubRayResourcePool(RayResourcePool):
+ def __init__(
+ self,
+ placement_groups: list[PlacementGroup],
+ start_bundle_index: int,
+ subgroup_world_size: int,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.pgs = placement_groups
+ self.start_bundle_index = start_bundle_index
+ self.subgroup_world_size = subgroup_world_size
+
+ @property
+ def world_size(self):
+ return self.subgroup_world_size
+
+
+@dataclass
+class ResourcePoolManager:
+ """
+ Define a resource pool specification. Resource pool will be initialized first.
+ """
+
+ resource_pool_spec: dict[str, list[int]]
+ mapping: dict[int, str]
+ resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict)
+
+ def create_resource_pool(self):
+ """Create Ray resource pools for distributed training.
+
+ Initializes resource pools based on the resource pool specification,
+ with each pool managing GPU resources across multiple nodes.
+ For FSDP backend, uses max_colocate_count=1 to merge WorkerGroups.
+ For Megatron backend, uses max_colocate_count>1 for different models.
+ """
+ for resource_pool_name, process_on_nodes in self.resource_pool_spec.items():
+ # max_colocate_count means the number of WorkerGroups (i.e. processes) in each RayResourcePool
+ # For FSDP backend, using max_colocate_count=3: actor_critic_ref, rollout, reward model (optional)
+ # For Megatron backend, we recommend using max_colocate_count>1
+ # that can utilize different WorkerGroup for differnt models
+ resource_pool = RayResourcePool(
+ process_on_nodes=process_on_nodes, use_gpu=True, max_colocate_count=3, name_prefix=resource_pool_name
+ )
+ self.resource_pool_dict[resource_pool_name] = resource_pool
+
+ self._check_resource_available()
+
+ def get_resource_pool(self, role) -> RayResourcePool:
+ """Get the resource pool of the worker_cls"""
+ return self.resource_pool_dict[self.mapping[role]]
+
+ def get_n_gpus(self) -> int:
+ """Get the number of gpus in this cluster."""
+ return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes])
+
+ def _check_resource_available(self):
+ """Check if the resource pool can be satisfied in this ray cluster."""
+ node_available_resources = ray._private.state.available_resources_per_node()
+ node_available_gpus = {
+ node: node_info.get("GPU", 0) if "GPU" in node_info else node_info.get("NPU", 0)
+ for node, node_info in node_available_resources.items()
+ }
+
+ # check total required gpus can be satisfied
+ total_available_gpus = sum(node_available_gpus.values())
+ total_required_gpus = sum(
+ [n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes]
+ )
+ if total_available_gpus < total_required_gpus:
+ raise ValueError(
+ f"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}"
+ )
+
+
+def extract_pg_from_exist(
+ resource_pools: dict[str, RayResourcePool], src_role_names: list[str], resource_pool: RayResourcePool
+) -> list:
+ src_pgs = [
+ pg
+ for role_name, resource_pool in resource_pools.items()
+ for pg in resource_pool.get_placement_groups()
+ if role_name in src_role_names
+ ]
+
+ sorted_src_pgs = sorted(src_pgs, key=lambda pg: pg.bundle_count, reverse=True)
+ sorted_process_on_nodes = sorted([(val, idx) for idx, val in enumerate(resource_pool.store)], reverse=True)
+
+ unsorted_pgs: list[tuple[int, PlacementGroup]] = []
+ searching_idx = 0
+ for request_process, original_idx in sorted_process_on_nodes:
+ assert searching_idx < len(sorted_src_pgs), f"no enough nodes for request: searching {searching_idx} th node"
+ assert request_process <= sorted_src_pgs[searching_idx].bundle_count, (
+ f"requesting {request_process} processes, bundle count cannot satisfy"
+ )
+ unsorted_pgs.append((original_idx, sorted_src_pgs[searching_idx]))
+ searching_idx += 1
+
+ return [pg for _, pg in sorted(unsorted_pgs)]
+
+
+# split a RayResourcePool or SubRayResourcePool into multiple SubRayResourcePool
+def split_resource_pool(
+ resource_pool: RayResourcePool | SubRayResourcePool, split_size: int | list[int]
+) -> list[SubRayResourcePool]:
+ """
+ Split a RayResourcePool into multiple SubRayResourcePool.
+ resouce_pool can also be a SubRayResourcePool (have been splited) for multiple-time spliting.
+
+ Args:
+ resource_pool (RayResourcePool | SubRayResourcePool): The resource pool to split.
+ split_size (int | list[int]): The size of each split. If int, all splits will have the same size.
+ If list[int], each element in the list represents the size of a split.
+
+ Returns:
+ list[SubRayResourcePool]: A list of SubRayResourcePool after splitting.
+ """
+ # convert split_size to list[int]
+ if isinstance(split_size, int):
+ assert resource_pool.world_size % split_size == 0, "split_size must be a divisor of world_size"
+ num_replica = resource_pool.world_size // split_size
+ split_size_list = [split_size] * num_replica
+ else:
+ split_size_list = split_size
+
+ assert sum(split_size_list) == resource_pool.world_size, "split_size must sum up to world_size"
+
+ # judge if this resource pool has been splited
+ if isinstance(resource_pool, SubRayResourcePool):
+ start_bundle_idx_list = np.cumsum([resource_pool.start_bundle_index] + split_size_list[:-1])
+ else:
+ start_bundle_idx_list = np.cumsum([0] + split_size_list[:-1])
+
+ # ensure resource_pool.pgs has been initialized
+ placement_groups = resource_pool.get_placement_groups()
+ split_resource_pools = [
+ SubRayResourcePool(
+ process_on_nodes=resource_pool.store,
+ use_gpu=resource_pool.use_gpu,
+ name_prefix=f"{resource_pool.name_prefix}_split_{split_idx}",
+ max_colocate_count=resource_pool.max_colocate_count,
+ placement_groups=placement_groups,
+ start_bundle_index=start_bundle_idx_list[split_idx],
+ subgroup_world_size=split_size_list[split_idx],
+ )
+ for split_idx in range(len(split_size_list))
+ ]
+ return split_resource_pools
+
+
+def merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> RayResourcePool:
+ assert rp1.use_gpu == rp2.use_gpu, "Both RayResourcePool must either use_gpu or not"
+ assert rp1.max_colocate_count == rp2.max_colocate_count, "Both RayResourcePool must has the same max_colocate_count"
+ assert rp1.n_gpus_per_node == rp2.n_gpus_per_node, "Both RayResourcePool must has the same n_gpus_per_node"
+ assert rp1.detached == rp2.detached, "Detached ResourcePool cannot be merged with non-detached ResourcePool"
+
+ new_store = rp1.store + rp2.store
+
+ merged = type(rp1)(
+ new_store, rp1.use_gpu, f"{rp1.name_prefix}_{rp2.name_prefix}", rp1.max_colocate_count, rp1.detached
+ )
+ merged.pgs = rp1.get_placement_groups(device_name=get_device_name()) + rp2.get_placement_groups(
+ device_name=get_device_name()
+ )
+
+ return merged
+
+
+class RayClassWithInitArgs(ClassWithInitArgs):
+ """A wrapper class for Ray actors with initialization arguments.
+
+ This class extends ClassWithInitArgs to provide additional functionality for
+ configuring and creating Ray actors with specific resource requirements and
+ scheduling strategies.
+ """
+
+ def __init__(self, cls, *args, **kwargs) -> None:
+ # self._options = kwargs.pop('options', dict())
+ super().__init__(cls, *args, **kwargs)
+ self._options = {}
+ self._additional_resource = {}
+
+ def set_additional_resource(self, additional_resource):
+ """Set additional resource requirements for the actor.
+
+ Args:
+ additional_resource: Dictionary specifying additional resource requirements
+ """
+ self._additional_resource = additional_resource
+
+ def update_options(self, options: dict):
+ """Update the Ray actor creation options.
+
+ Args:
+ options: Dictionary of options to update
+ """
+ self._options.update(options)
+
+ def __call__(
+ self,
+ placement_group,
+ placement_group_bundle_idx,
+ use_gpu: bool = True,
+ num_gpus=1,
+ sharing_with=None,
+ device_name="cuda",
+ ) -> Any:
+ """Create and return a Ray actor with the configured options.
+
+ Args:
+ placement_group: Ray placement group for scheduling
+ placement_group_bundle_idx: Index of the bundle in the placement group
+ use_gpu: Whether to use GPU resources
+ num_gpus: Number of GPUs to allocate
+ sharing_with: Actor to share resources with
+ device_name: Device for training
+
+ Returns:
+ A Ray actor handle with the configured options
+ """
+ if sharing_with is not None:
+ target_node_id = ray.get(sharing_with.get_node_id.remote())
+ visible_devices = ray.get(sharing_with.get_cuda_visible_devices.remote())
+ options = {"scheduling_strategy": NodeAffinitySchedulingStrategy(node_id=target_node_id, soft=False)}
+ return self.cls.options(**options).remote(*self.args, cuda_visible_devices=visible_devices, **self.kwargs)
+
+ options = {
+ "scheduling_strategy": PlacementGroupSchedulingStrategy(
+ placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_idx
+ )
+ }
+ options.update(self._options)
+
+ if use_gpu and device_name == "cuda":
+ options["num_gpus"] = num_gpus
+ if use_gpu and device_name == "npu":
+ options["resources"] = {"NPU": num_gpus}
+
+ if len(self._additional_resource) > 1:
+ for k, v in self._additional_resource.items():
+ options[k] = v
+
+ # print("cls:", self.cls)
+ # print("args: ", self.args)
+ # print("kwargs: ", self.kwargs)
+ return self.cls.options(**options).remote(*self.args, **self.kwargs)
+
+
+class RayWorkerGroup(WorkerGroup):
+ """A group of Ray workers that can be managed collectively.
+
+ This class extends WorkerGroup to provide Ray-specific functionality for
+ creating and managing groups of Ray actors with specific resource requirements
+ and scheduling strategies.
+ """
+
+ def __init__(
+ self,
+ resource_pool: RayResourcePool = None,
+ ray_cls_with_init: RayClassWithInitArgs = None,
+ bin_pack: bool = True,
+ name_prefix: str = None,
+ detached=False,
+ worker_names=None,
+ worker_handles: list[ray.actor.ActorHandle] = None,
+ ray_wait_register_center_timeout: int = 300,
+ **kwargs,
+ ) -> None:
+ """Initialize a RayWorkerGroup.
+
+ Args:
+ resource_pool: Resource pool for worker allocation
+ ray_cls_with_init: Class with initialization arguments for workers
+ bin_pack: Whether to use strict bin packing for resource allocation
+ name_prefix: Prefix for worker names
+ detached: Whether workers should be detached
+ worker_names: Names of existing workers to attach to
+ ray_wait_register_center_timeout: Timeout for waiting on register center
+ **kwargs: Additional keyword arguments
+ """
+ self._master_addr = kwargs.pop("master_addr", None)
+ self._master_port = kwargs.pop("master_port", None)
+ self.use_gpu = kwargs.pop("use_gpu", resource_pool.use_gpu if resource_pool is not None else True)
+ super().__init__(resource_pool=resource_pool, **kwargs)
+ self.ray_cls_with_init = ray_cls_with_init
+ self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix
+ self._ray_wait_register_center_timeout = ray_wait_register_center_timeout
+ # Whether the WorkerGroup is a Colocate WorkerGroup created by FusedWorker.
+ self.fused_worker_used = False if ray_cls_with_init is None else ray_cls_with_init.fused_worker_used
+ # if a WorkerGroup is spawned from Colocate WorkerGroup, this indicates which sub-class is binded to
+ # this WorkerGroup.
+ self.sub_cls_name = ""
+ self.device_name = kwargs.get("device_name", "cuda")
+ self.profile_steps = kwargs.get("profile_steps", None)
+ self.worker_nsight_options = kwargs.get("worker_nsight_options", None)
+ self.customized_worker_env = kwargs.get("worker_env", {})
+ if self.worker_nsight_options is not None and self.worker_nsight_options["capture-range-end"] is None:
+ self.worker_nsight_options["capture-range-end"] = f"repeat-shutdown:{6 * len(self.profile_steps)}"
+
+ if worker_names is not None and (not self.fused_worker_used):
+ assert self._is_init_with_detached_workers
+ self._worker_names = worker_names
+
+ if self._is_init_with_detached_workers:
+ self._init_with_detached_workers(worker_names=worker_names, worker_handles=worker_handles)
+ elif isinstance(resource_pool, SubRayResourcePool):
+ self._init_with_subresource_pool(
+ resource_pool=resource_pool,
+ ray_cls_with_init=ray_cls_with_init,
+ bin_pack=bin_pack,
+ detached=detached,
+ worker_env=self.customized_worker_env,
+ )
+ else:
+ self._init_with_resource_pool(
+ resource_pool=resource_pool,
+ ray_cls_with_init=ray_cls_with_init,
+ bin_pack=bin_pack,
+ detached=detached,
+ worker_env=self.customized_worker_env,
+ )
+
+ if ray_cls_with_init is not None:
+ self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)
+
+ self.wg_dict = None
+ self.method_names = []
+
+ def _is_worker_alive(self, worker: ray.actor.ActorHandle):
+ """Check if a worker actor is still alive.
+
+ Args:
+ worker: Ray actor handle to check
+
+ Returns:
+ bool: True if the worker is alive, False otherwise
+ """
+ worker_state_dict = get_actor(worker._actor_id.hex())
+ return worker_state_dict.get("state", "undefined") == "ALIVE" if worker_state_dict is not None else False
+
+ def _init_with_detached_workers(self, worker_names, worker_handles):
+ # ray.get_actor holds a weak reference to the actor, which causes actors garbage collected unexpectedly
+ # if we only hold spawn RayWorkerGroup. By passing actor handle explicitly, spawn RayWorkerGroup have
+ # strong reference to these actors.
+ # https://github.com/ray-project/ray/pull/45699
+ workers = worker_handles if worker_handles else [ray.get_actor(name=name) for name in worker_names]
+ self._workers = workers
+ self._world_size = len(workers)
+
+ def _get_master_addr_port(self, pg, bundle_index=0):
+ """Get master addr and port for this worker group"""
+ if self._master_addr is None and self._master_port is None:
+ self._master_addr, self._master_port = ray.get(
+ get_master_addr_port.options(
+ scheduling_strategy=PlacementGroupSchedulingStrategy(
+ placement_group=pg, placement_group_bundle_index=bundle_index
+ ),
+ ).remote()
+ )
+ elif self._master_addr is not None and self._master_port is not None:
+ logger.debug(f"{self._master_addr=} {self._master_port=}")
+ else:
+ raise ValueError(
+ "Both 'master_addr' and 'master_port' must be provided if you intend to manually specify them, "
+ "or neither should be provided to use Ray's default assignment."
+ )
+
+ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached, worker_env=None):
+ """Initialize the worker group by creating new workers from a resource pool.
+
+ Args:
+ resource_pool: Resource pool for worker allocation
+ ray_cls_with_init: Class with initialization arguments for workers
+ bin_pack: Whether to use strict bin packing for resource allocation
+ detached: Whether workers should be detached
+ """
+ self.resource_pool = resource_pool
+
+ strategy = "PACK"
+ if bin_pack:
+ strategy = "STRICT_PACK"
+ pgs = resource_pool.get_placement_groups(strategy=strategy, device_name=self.device_name)
+ world_size = resource_pool.world_size
+ self._world_size = world_size
+ # cia.add_kwarg("_world_size", world_size)
+
+ rank = -1
+ local_world_size = resource_pool.store[0]
+ for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)):
+ assert local_world_size <= pg.bundle_count, f"when generating for {self.name_prefix}, for the "
+ if pg_idx == 0:
+ self._get_master_addr_port(pg)
+
+ for local_rank in range(local_world_size):
+ rank += 1
+ self._create_worker(
+ rank=rank,
+ pg_idx=pg_idx,
+ pg=pg,
+ local_rank=local_rank,
+ resource_pool=resource_pool,
+ ray_cls_with_init=ray_cls_with_init,
+ worker_env=worker_env,
+ detached=detached,
+ )
+
+ def _init_with_subresource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached, worker_env=None):
+ """Initialize the worker group by creating new workers from a resource pool or sub resource pool.
+ Args:
+ resource_pool: Resource pool for worker allocation
+ ray_cls_with_init: Class with initialization arguments for workers
+ bin_pack: Whether to use strict bin packing for resource allocation
+ detached: Whether workers should be detached
+ """
+ strategy = "PACK"
+ if bin_pack:
+ strategy = "STRICT_PACK"
+ pgs = resource_pool.get_placement_groups(strategy=strategy, device_name=self.device_name)
+ world_size = resource_pool.world_size
+ self._world_size = world_size
+
+ rank = -1
+ local_world_size = resource_pool.store[0]
+ self._get_master_addr_port(
+ pgs[resource_pool.start_bundle_index // local_world_size],
+ resource_pool.start_bundle_index % local_world_size,
+ )
+ for curr_rank in range(resource_pool.start_bundle_index, resource_pool.start_bundle_index + world_size):
+ pg_idx = curr_rank // local_world_size
+ pg = pgs[pg_idx]
+ local_rank = curr_rank % local_world_size
+ assert local_world_size <= pg.bundle_count, f"when generating for {self.name_prefix}, for the "
+
+ rank += 1
+ self._create_worker(
+ rank=rank,
+ pg_idx=pg_idx,
+ pg=pg,
+ local_rank=local_rank,
+ resource_pool=resource_pool,
+ ray_cls_with_init=ray_cls_with_init,
+ worker_env=worker_env,
+ detached=detached,
+ )
+
+ def _create_worker(self, rank, pg_idx, pg, local_rank, resource_pool, ray_cls_with_init, worker_env, detached):
+ world_size = resource_pool.world_size
+ use_gpu = resource_pool.use_gpu
+ if self.use_gpu and not use_gpu:
+ raise ValueError("use_gpu is True but resource_pool.use_gpu is False")
+ local_world_size = resource_pool.store[0]
+ num_gpus = 1 / resource_pool.max_colocate_count
+
+ # we pass in environment variable at option so that Worker can use environment variable to set
+ env_vars = {
+ "WORLD_SIZE": str(world_size),
+ "RANK": str(rank),
+ "WG_PREFIX": self.name_prefix,
+ "WG_BACKEND": "ray",
+ "RAY_LOCAL_WORLD_SIZE": str(local_world_size),
+ "MASTER_ADDR": self._master_addr,
+ "MASTER_PORT": self._master_port,
+ }
+ if worker_env is not None:
+ logging.debug(f"Appending ray class env, origin: {env_vars}, customized env: {worker_env}")
+ conflict_env_vars = set(env_vars.keys()) & set(worker_env.keys())
+ if len(conflict_env_vars) > 0:
+ logging.error(
+ f"User customized env vars conflict with system env: {conflict_env_vars} "
+ f"Overriding may cause unexpected behavior."
+ )
+ raise ValueError(f"Cannot override protected system env: {conflict_env_vars}")
+ env_vars.update(worker_env)
+ import re
+
+ cia_name = type(ray_cls_with_init.cls).__name__
+ match = re.search(r"ActorClass\(([^)]+)\)", cia_name) # ray.remote(Obj) -> "ActorClass(Obj)"
+ cia_name = match.group(1) if match else cia_name # "ActorClass(Obj)" -> "Obj"
+ name = f"{self.name_prefix}{cia_name}_{pg_idx}:{local_rank}" # e.g. Worker_2:5
+
+ if self.profile_steps and self.device_name == "cuda":
+ ray_cls_with_init.update_options(
+ {
+ "runtime_env": {
+ "env_vars": env_vars,
+ "nsight": self.worker_nsight_options,
+ },
+ "name": name,
+ }
+ )
+ else:
+ ray_cls_with_init.update_options({"runtime_env": {"env_vars": env_vars}, "name": name})
+
+ if detached:
+ ray_cls_with_init.update_options({"lifetime": "detached"})
+
+ # create a worker
+ worker = ray_cls_with_init(
+ placement_group=pg,
+ placement_group_bundle_idx=local_rank,
+ use_gpu=self.use_gpu,
+ num_gpus=num_gpus,
+ device_name=self.device_name,
+ )
+ self._workers.append(worker)
+ self._worker_names.append(name)
+
+ @property
+ def worker_names(self):
+ return self._worker_names
+
+ @classmethod
+ def from_detached(
+ cls,
+ name_prefix=None,
+ worker_names=None,
+ worker_handles=None,
+ ray_cls_with_init=None,
+ **kwargs,
+ ):
+ """Create a worker group from existing detached workers.
+
+ Args:
+ name_prefix: Prefix for worker names
+ worker_names: Names of existing workers to attach to
+ ray_cls_with_init: Class with initialization arguments for workers
+
+ Returns:
+ A new RayWorkerGroup instance
+ """
+ worker_group = cls(
+ resource_pool=None,
+ ray_cls_with_init=ray_cls_with_init,
+ name_prefix=name_prefix,
+ worker_names=worker_names,
+ worker_handles=worker_handles,
+ **kwargs,
+ )
+ return worker_group
+
+ def spawn(self, prefix_set):
+ """Spawn to a dictionary of worker groups, each with a subset of method with prefix.
+
+ Args:
+ prefix_set: Set of prefixes to create worker groups for
+
+ Returns:
+ Dictionary of worker groups keyed by prefix
+ """
+ if self.fused_worker_used:
+ return self.spawn_fused(prefix_set)
+
+ def _rebind_actor_methods(worker_group, actor_name):
+ prefix: str = actor_name + "_"
+ for method_name in dir(worker_group):
+ if method_name.startswith(prefix):
+ original_method_name = method_name.removeprefix(prefix)
+ method = getattr(worker_group, method_name)
+ setattr(worker_group, original_method_name, method)
+
+ new_worker_group_dict = {}
+ for prefix in prefix_set:
+ new_worker_group = self.from_detached(
+ name_prefix=self.name_prefix,
+ worker_names=self._worker_names,
+ worker_handles=self._workers,
+ ray_cls_with_init=self.ray_cls_with_init,
+ profile_steps=self.profile_steps,
+ worker_nsight_options=self.worker_nsight_options,
+ )
+
+ _rebind_actor_methods(new_worker_group, prefix)
+ new_worker_group_dict[prefix] = new_worker_group
+ return new_worker_group_dict
+
+ def spawn_fused(self, prefix_set):
+ """Create a dictionary of worker groups for fused workers.
+
+ Args:
+ prefix_set: Set of prefixes to create worker groups for
+
+ Returns:
+ Dictionary of worker groups keyed by prefix
+ """
+ wg_dict = dict()
+ for key in prefix_set:
+ new_wg = deepcopy(self)
+ new_wg._bind_worker_method(self.ray_cls_with_init.cls.raw_cls_dict[key], func_generator)
+ new_wg.sub_cls_name = key
+ wg_dict[key] = new_wg
+ return wg_dict
+
+ def fuse(self, prefix_set):
+ """Fuse multiple worker groups into the current worker group.
+
+ Args:
+ prefix_set: Set of prefixes to fuse into the worker group
+ """
+ if self.wg_dict is None:
+ self.wg_dict = self.spawn(prefix_set)
+ for role_name, role_wg in self.wg_dict.items():
+ setattr(self, role_name, role_wg)
+ self.method_names = self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)
+
+ def _execute_remote_single_worker(self, worker, method_name: str, *args, **kwargs):
+ """Execute a method on a single worker remotely.
+
+ Args:
+ worker: The worker actor handle
+ method_name: Name of the method to execute
+ *args: Positional arguments for the method
+ **kwargs: Keyword arguments for the method
+
+ Returns:
+ Remote object reference to the method execution
+ """
+ if self.fused_worker_used and method_name not in self.method_names:
+ remote_call = getattr(worker, self.fused_worker_execute_fn_name)
+ return remote_call.remote(f"{self.sub_cls_name}_fwmn_{method_name}", *args, **kwargs)
+ # fused worker not used
+ remote_call = getattr(worker, method_name)
+ return remote_call.remote(*args, **kwargs)
+
+ def execute_rank_zero_sync(self, method_name: str, *args, **kwargs):
+ """Execute a method on rank zero worker synchronously.
+
+ Args:
+ method_name: Name of the method to execute
+ *args: Positional arguments for the method
+ **kwargs: Keyword arguments for the method
+
+ Returns:
+ Result of the method execution
+ """
+ return ray.get(self.execute_rank_zero_async(method_name, *args, **kwargs))
+
+ def execute_rank_zero_async(self, method_name: str, *args, **kwargs):
+ """Execute a method on rank zero worker asynchronously.
+
+ Args:
+ method_name: Name of the method to execute
+ *args: Positional arguments for the method
+ **kwargs: Keyword arguments for the method
+
+ Returns:
+ Remote object reference to the method execution
+ """
+ return self._execute_remote_single_worker(self._workers[0], method_name, *args, **kwargs)
+
+ def execute_rank_zero(self, method_name: str, *args, **kwargs):
+ """Alias for execute_rank_zero_async.
+
+ Args:
+ method_name: Name of the method to execute
+ *args: Positional arguments for the method
+ **kwargs: Keyword arguments for the method
+
+ Returns:
+ Remote object reference to the method execution
+ """
+ return self.execute_rank_zero_async(method_name, *args, **kwargs)
+
+ def execute_all(self, method_name: str, *args, **kwargs):
+ """Alias for execute_all_async.
+
+ Args:
+ method_name: Name of the method to execute
+ *args: Positional arguments for the method
+ **kwargs: Keyword arguments for the method
+
+ Returns:
+ List of remote object references to the method executions
+ """
+ return self.execute_all_async(method_name, *args, **kwargs)
+
+ def execute_all_sync(self, method_name: str, *args, **kwargs):
+ """Execute a method on all workers synchronously.
+
+ Args:
+ method_name: Name of the method to execute
+ *args: Positional arguments for the method
+ **kwargs: Keyword arguments for the method
+
+ Returns:
+ List of results from all workers
+ """
+ return ray.get(self.execute_all_async(method_name, *args, **kwargs))
+
+ def execute_all_async(self, method_name: str, *args, **kwargs):
+ """Execute a method on all workers asynchronously.
+
+ Args:
+ method_name: Name of the method to execute
+ *args: Positional arguments for the method
+ **kwargs: Keyword arguments for the method
+
+ Returns:
+ List of remote object references to the method executions
+ """
+ # Here, we assume that if all arguments in args and kwargs are lists,
+ # and their lengths match len(self._workers), we'll distribute each
+ # element in these lists to the corresponding worker
+ # print(f"execute_all_async: method {method_name}({args}, {kwargs})")
+ length = len(self._workers)
+ if all(isinstance(arg, list) for arg in args) and all(isinstance(kwarg, list) for kwarg in kwargs.values()):
+ if all(len(arg) == length for arg in args) and all(len(kwarg) == length for kwarg in kwargs.values()):
+ # print(f"splitting args and kwargs into {length} shards")
+ result = []
+ for i in range(length):
+ sliced_args = tuple(arg[i] for arg in args)
+ sliced_kwargs = {k: v[i] for k, v in kwargs.items()}
+ result.append(
+ self._execute_remote_single_worker(self._workers[i], method_name, *sliced_args, **sliced_kwargs)
+ )
+ return result
+
+ return [self._execute_remote_single_worker(worker, method_name, *args, **kwargs) for worker in self._workers]
+
+ @property
+ def master_address(self):
+ return self._master_addr
+
+ @property
+ def master_port(self):
+ return self._master_port
+
+ @property
+ def workers(self):
+ return self._workers
+
+ @property
+ def world_size(self):
+ return self._world_size
+
+
+"""
+Utilities that enables creating workers inside the same ray.Actor,
+with code written in separate ray.Actors.
+"""
+
+
+# deprecated, switching to FusedWorker
+def _bind_workers_method_to_parent(cls, key, user_defined_cls):
+ """
+ Binds the methods of each worker to the WorkerDict.
+ Note that we only bind public methods that are decorated by register
+ """
+
+ for method_name in dir(user_defined_cls):
+ try:
+ method = getattr(user_defined_cls, method_name)
+ assert callable(method), f"{method_name} in {user_defined_cls} is not callable"
+ except Exception:
+ # if it is a property, it will fail because Class doesn't have instance property
+ continue
+
+ if hasattr(method, MAGIC_ATTR):
+
+ def generate_function(name, key=key):
+ def func(self, *args, **kwargs):
+ # dispatch to the actual worker
+ return getattr(self.worker_dict[key], name)(*args, **kwargs)
+
+ async def async_func(self, *args, **kwargs):
+ # dispatch to the actual worker
+ return await getattr(self.worker_dict[key], name)(*args, **kwargs)
+
+ wrapper = async_func if inspect.iscoroutinefunction(method) else func # noqa: B023
+
+ return wrapper
+
+ func = generate_function(method_name)
+ # pass MAGIC_ATTR for outer worker group
+ attrs = getattr(method, MAGIC_ATTR)
+ setattr(func, MAGIC_ATTR, attrs)
+ try:
+ # bind direct rollout method to class without prefix
+ if attrs["dispatch_mode"] == Dispatch.DIRECT_ROLLOUT_METHOD and "rollout" in key:
+ assert not hasattr(cls, method_name), (
+ f"conflict direct rollout method {method_name} with role {key}"
+ )
+ setattr(cls, method_name, func)
+ print(f"bind role {key} method {method_name} to class {cls}")
+ else:
+ method_name_with_prefix = key + "_" + method_name
+ setattr(cls, method_name_with_prefix, func)
+ except Exception as e:
+ raise ValueError(f"Fail to set method_name {method_name}") from e
+
+
+def _unwrap_ray_remote(cls):
+ if hasattr(cls, "__ray_actor_class__"):
+ cls = cls.__ray_actor_class__
+ return cls
+
+
+def _determine_fsdp_megatron_base_class(mros: list):
+ """
+ - megatron: base class should be MegatronWorker
+ - fsdp: base class should be Worker
+ """
+ for cls in mros[0]:
+ if cls.__name__ == "MegatronWorker":
+ return cls
+ if cls.__name__ == "Worker":
+ return cls
+ raise ValueError(f"Cannot determine base class for {mros}")
+
+
+# deprecated, switching to FusedWorker
+def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]):
+ """
+ This function should return a class instance that delegates the calls to every
+ cls in cls_dict
+ """
+ cls_dict = {}
+ init_args_dict = {}
+ worker_cls = _determine_fsdp_megatron_base_class(
+ [cls.cls.__ray_actor_class__.__mro__ for cls in class_dict.values()]
+ )
+ assert issubclass(worker_cls, Worker), f"worker_cls {worker_cls} should be a subclass of Worker"
+ print(f"colocated worker base class {worker_cls}")
+
+ for key, cls in class_dict.items():
+ cls_dict[key] = cls.cls
+ init_args_dict[key] = {"args": cls.args, "kwargs": cls.kwargs}
+
+ assert cls_dict.keys() == init_args_dict.keys()
+
+ # TODO: create a class with customizable name
+ class WorkerDict(worker_cls):
+ def __init__(self):
+ super().__init__()
+ self.worker_dict = {}
+ for key, user_defined_cls in cls_dict.items():
+ user_defined_cls = _unwrap_ray_remote(user_defined_cls)
+ # directly instantiate the class without remote
+ # in worker class, e.g.
+ # when DISABLE_WORKER_INIT == 1 it will return immediately
+ with temp_env_var("DISABLE_WORKER_INIT", "1"):
+ self.worker_dict[key] = user_defined_cls(
+ *init_args_dict[key].get("args", ()), **init_args_dict[key].get("kwargs", {})
+ )
+
+ # now monkey-patch the methods from inner class to WorkerDict
+ for key, user_defined_cls in cls_dict.items():
+ user_defined_cls = _unwrap_ray_remote(user_defined_cls)
+ _bind_workers_method_to_parent(WorkerDict, key, user_defined_cls)
+
+ remote_cls = ray.remote(WorkerDict)
+ remote_cls = RayClassWithInitArgs(cls=remote_cls)
+ return remote_cls
+
+
+FusedWorkerCLSName = "FusedWorker"
+
+
+def create_colocated_worker_raw_cls(class_dict: dict[str, RayClassWithInitArgs]):
+ """
+ This function returns a FusedWorker class.
+
+ `FusedWorker.{class_name}` -> FusedClass
+ Use `class_name` as a param to directly access the underlying class.
+
+ `FusedWorker._fuw_execute("{class_name}_fwmn_{method_name}", *args, **kwargs)`
+ First param must be "{class_name}_fwmn_{method_name}" in order to access `method_name`
+ of underlying class `{class_name}`.
+
+ `FusedWorker.fused_worker_dict` -> {"class_name": FusedClass}
+ Stores all underlying classes.
+
+ `FusedClass.fused_worker_dict` -> {"class_name": FusedClass}
+ The same as `FusedWorker.fused_worker_dict`, enables underlying class to access other
+ underlying classes.
+ """
+ raw_cls_dict = {cls_name: _unwrap_ray_remote(cia.cls) for cls_name, cia in class_dict.items()}
+ init_args_dict = {cls_name: cia.args for cls_name, cia in class_dict.items()}
+ init_kwargs_dict = {cls_name: cia.kwargs for cls_name, cia in class_dict.items()}
+ cls_names = list(class_dict.keys())
+
+ # FusedWorker_Actor_Critic
+ class_name_renamed = "_".join([FusedWorkerCLSName] + cls_names)
+
+ class FusedWorker(Worker):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.cls_names = cls_names
+ self.raw_cls_dict = raw_cls_dict
+ self.init_args_dict = init_args_dict
+ self.init_kwargs_dict = init_kwargs_dict
+
+ for cls_name, udc, ud_args, ud_kwargs in zip(
+ self.cls_names,
+ self.raw_cls_dict.values(),
+ self.init_args_dict.values(),
+ self.init_kwargs_dict.values(),
+ strict=True,
+ ):
+ with temp_env_var("DISABLE_WORKER_INIT", "1"):
+ udc._get_ray_actor_cls_name = lambda x, name_renamed=class_name_renamed: name_renamed
+ udc._get_ray_method_prefix = lambda x, name_prefixed=cls_name: f"{name_prefixed}_"
+ # cls_name = "actor", "critic", udc = ActorWorker, CriticWorker
+ self.fused_worker_dict[cls_name] = udc(*ud_args, **ud_kwargs)
+ setattr(self, cls_name, self.fused_worker_dict[cls_name])
+
+ # injecting fused_worker to each sub worker so they can be aware of existence of each other
+ for _, worker in self.fused_worker_dict.items():
+ setattr(worker, Worker.fused_worker_attr_name, self.fused_worker_dict)
+
+ def _fuw_execute(self, method_name: str, *args, **kwargs):
+ # for fused_worker, method_name is in a form of "{cls_name}_fwmn_{method_name}"
+ # where fwmn stands "fused worker method name"
+ names = method_name.split("_fwmn_")
+ cls_name = names[0]
+ method_name = names[1]
+
+ assert cls_name in self.fused_worker_dict, (
+ f"calling {cls_name}'s {method_name}, but {cls_name} not in fused_worker_dict"
+ )
+ udc_method = getattr(self.fused_worker_dict[cls_name], method_name)
+ return udc_method(*args, **kwargs)
+
+ renamed_fused_worker_cls = type(class_name_renamed, (FusedWorker,), {})
+ renamed_fused_worker_cls.is_fused_worker = True
+ renamed_fused_worker_cls.raw_cls_dict = raw_cls_dict
+
+ return renamed_fused_worker_cls
+
+
+def create_colocated_worker_cls_fused(class_dict: dict[str, RayClassWithInitArgs]):
+ """
+ This function returns a RayClassWithInitArgs instance of FusedWorker, which is an replacement
+ of `create_colocated_worker_cls`. WorkerGroup constructed using this class will be a colocated
+ WorkerGroup, which will be referenced as `ColocateWorkerGroup` below.
+
+ `ColocateWorkerGroup.spawn(prefix_set)`
+ returns a dict of WorkerGroup {"class_name": WorkerGroup}, WorkerGroup in this dict will
+ have methods of underlying class `class_name` attached.
+
+ `ColocateWorkerGroup.fuse(prefix_set)`
+ After executing this function, `ColocateWorkerGroup.{class_name}` will return WorkerGroup
+ with methods of underlying class `class_name` attached.
+ """
+ raw_colocated_worker_cls = create_colocated_worker_raw_cls(class_dict)
+
+ remote_cls = ray.remote(raw_colocated_worker_cls)
+ cia = RayClassWithInitArgs(cls=remote_cls)
+ cia.fused_worker_used = True
+
+ return cia
diff --git a/code/RL_model/verl/verl_train/verl/tools/utils/__init__.py b/code/RL_model/verl/verl_train/verl/tools/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4b932b1ae7eeeb4c53c98c684cf0ba9b670a86b
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/tools/utils/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2023-2024 SGLang Team
+# Copyright 2025 ModelBest Inc. 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.
diff --git a/code/RL_model/verl/verl_train/verl/tools/utils/mcp_clients/McpClientManager.py b/code/RL_model/verl/verl_train/verl/tools/utils/mcp_clients/McpClientManager.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee5fe31191321f653230f6dc0cfb9e71a42e1722
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/tools/utils/mcp_clients/McpClientManager.py
@@ -0,0 +1,97 @@
+# Copyright 2025 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.
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+from fastmcp import Client
+from fastmcp.client.transports import SSETransport
+
+from verl.tools.utils.mcp_clients.utils import TokenBucket, mcp2openai
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClientManager:
+ rootServerName = "mcpServers"
+ initialized = False
+ clients = []
+ tool_client_mapping = {}
+ rate_limiter = None
+
+ async def initialize(self, config_path, rate_limit: float = 10.0):
+ if self.initialized:
+ return
+ """Initialize the MCP Client Manager and start all clients"""
+ result = self._load_config(config_path)
+ servers = result[self.rootServerName]
+ exclude_sse_servers = {self.rootServerName: {}}
+ for server_name in servers.keys():
+ server = servers[server_name]
+ if "auth_token" in server:
+ transport = SSETransport(url=server["url"], headers={"Authorization": f"Bearer {server['auth_token']}"})
+ client = Client(transport)
+ self.clients.append(client)
+ else:
+ exclude_sse_servers[self.rootServerName][server_name] = server
+
+ if exclude_sse_servers[self.rootServerName]:
+ self.clients.append(Client(exclude_sse_servers))
+
+ # Initialize rate limiter
+ self.rate_limiter = TokenBucket(rate_limit)
+ self.initialized = True
+
+ async def call_tool(self, tool_name, parameters, timeout):
+ # Apply rate limiting
+ while not self.rate_limiter.acquire():
+ await asyncio.sleep(0.1)
+
+ client = self.get_client_with_tool_name(tool_name)
+ async with client:
+ return await client.call_tool_mcp(tool_name, parameters)
+
+ async def fetch_tool_schemas(self, tool_selected_list: list[str]) -> list[dict]:
+ tool_schemas = []
+ for client in self.clients:
+ async with client:
+ tools = await client.list_tools_mcp()
+ for tool in tools.tools:
+ if not tool_selected_list:
+ self.tool_client_mapping[tool.name] = client
+ tool_schemas.append(mcp2openai(tool))
+ elif tool.name in tool_selected_list:
+ self.tool_client_mapping[tool.name] = client
+ tool_schemas.append(mcp2openai(tool))
+
+ return tool_schemas
+
+ def get_client_with_tool_name(self, tool_name: str):
+ return self.tool_client_mapping[tool_name]
+
+ def _load_config(self, file: str) -> dict[str, Any]:
+ try:
+ with open(file) as f:
+ return json.load(f)
+ except FileNotFoundError:
+ logger.warning(f'the "{file}" file was not found')
+ except Exception:
+ logger.error(f'there was an error reading the "{file}" file')
+
+ return {}
+
+
+ClientManager = MCPClientManager()
diff --git a/code/RL_model/verl/verl_train/verl/tools/utils/mcp_clients/utils.py b/code/RL_model/verl/verl_train/verl/tools/utils/mcp_clients/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..22a5f63532713dcb895b0a940bf9bc9dfe42cfdf
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/tools/utils/mcp_clients/utils.py
@@ -0,0 +1,58 @@
+# Copyright 2025 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.
+
+import logging
+import threading
+import time
+
+from mcp import Tool
+
+logger = logging.getLogger(__file__)
+
+
+class TokenBucket:
+ def __init__(self, rate_limit: float):
+ self.rate_limit = rate_limit # tokens per second
+ self.tokens = rate_limit
+ self.last_update = time.time()
+ self.lock = threading.Lock()
+
+ def acquire(self) -> bool:
+ with self.lock:
+ now = time.time()
+ # Add new tokens based on time elapsed
+ new_tokens = (now - self.last_update) * self.rate_limit
+ self.tokens = min(self.rate_limit, self.tokens + new_tokens)
+ self.last_update = now
+
+ if self.tokens >= 1:
+ self.tokens -= 1
+ return True
+ return False
+
+
+def mcp2openai(mcp_tool: Tool) -> dict:
+ """Convert a MCP Tool to an OpenAI ChatCompletionTool."""
+ openai_format = {
+ "type": "function",
+ "function": {
+ "name": mcp_tool.name,
+ "description": mcp_tool.description,
+ "parameters": mcp_tool.inputSchema,
+ "strict": False,
+ },
+ }
+ if not openai_format["function"]["parameters"].get("required", None):
+ openai_format["function"]["parameters"]["required"] = []
+ return openai_format
diff --git a/code/RL_model/verl/verl_train/verl/tools/utils/search_r1_like_utils.py b/code/RL_model/verl/verl_train/verl/tools/utils/search_r1_like_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..610698e3b602d44b1bc19919e397a2d4cfb08bc9
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/tools/utils/search_r1_like_utils.py
@@ -0,0 +1,245 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2023-2024 SGLang Team
+#
+# 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.
+
+import json
+import logging
+import threading
+import time
+import traceback
+import uuid
+from typing import Any, Optional
+
+import requests
+
+DEFAULT_TIMEOUT = 30 # Default search request timeout
+MAX_RETRIES = 10
+INITIAL_RETRY_DELAY = 1
+API_TIMEOUT = 10
+
+logger = logging.getLogger(__name__)
+
+
+def call_search_api(
+ retrieval_service_url: str,
+ query_list: list[str],
+ topk: int = 3,
+ return_scores: bool = True,
+ timeout: int = DEFAULT_TIMEOUT,
+) -> tuple[Optional[dict[str, Any]], Optional[str]]:
+ """
+ Calls the remote search API to perform retrieval with retry logic for various errors,
+ using increasing delay between retries. Logs internal calls with a unique ID.
+
+ Args:
+ retrieval_service_url: The URL of the retrieval service API.
+ query_list: List of search queries.
+ topk: Number of top results to return.
+ return_scores: Whether to return scores.
+ timeout: Request timeout in seconds.
+
+ Returns:
+ A tuple (response_json, error_message).
+ If successful, response_json is the API's returned JSON object, error_message is None.
+ If failed after retries, response_json is None, error_message contains the error information.
+ """
+ request_id = str(uuid.uuid4())
+ log_prefix = f"[Search Request ID: {request_id}] "
+
+ payload = {"queries": query_list, "topk": topk, "return_scores": return_scores}
+
+ headers = {"Content-Type": "application/json", "Accept": "application/json"}
+
+ last_error = None
+
+ for attempt in range(MAX_RETRIES):
+ try:
+ logger.info(
+ f"{log_prefix}Attempt {attempt + 1}/{MAX_RETRIES}: Calling search API at {retrieval_service_url}"
+ )
+ response = requests.post(
+ retrieval_service_url,
+ headers=headers,
+ json=payload,
+ timeout=timeout,
+ )
+
+ # Check for Gateway Timeout (504) and other server errors for retrying
+ if response.status_code in [500, 502, 503, 504]:
+ last_error = (
+ f"{log_prefix}API Request Error: Server Error ({response.status_code}) on attempt "
+ f"{attempt + 1}/{MAX_RETRIES}"
+ )
+ logger.warning(last_error)
+ if attempt < MAX_RETRIES - 1:
+ delay = INITIAL_RETRY_DELAY * (attempt + 1)
+ logger.info(f"{log_prefix}Retrying after {delay} seconds...")
+ time.sleep(delay)
+ continue
+
+ # Check for other HTTP errors (e.g., 4xx)
+ response.raise_for_status()
+
+ # If successful (status code 2xx)
+ logger.info(f"{log_prefix}Search API call successful on attempt {attempt + 1}")
+ return response.json(), None
+
+ except requests.exceptions.ConnectionError as e:
+ last_error = f"{log_prefix}Connection Error: {e}"
+ logger.warning(last_error)
+ if attempt < MAX_RETRIES - 1:
+ delay = INITIAL_RETRY_DELAY * (attempt + 1)
+ logger.info(f"{log_prefix}Retrying after {delay} seconds...")
+ time.sleep(delay)
+ continue
+ except requests.exceptions.Timeout as e:
+ last_error = f"{log_prefix}Timeout Error: {e}"
+ logger.warning(last_error)
+ if attempt < MAX_RETRIES - 1:
+ delay = INITIAL_RETRY_DELAY * (attempt + 1)
+ logger.info(f"{log_prefix}Retrying after {delay} seconds...")
+ time.sleep(delay)
+ continue
+ except requests.exceptions.RequestException as e:
+ last_error = f"{log_prefix}API Request Error: {e}"
+ break # Exit retry loop on other request errors
+ except json.JSONDecodeError as e:
+ raw_response_text = response.text if "response" in locals() else "N/A"
+ last_error = f"{log_prefix}API Response JSON Decode Error: {e}, Response: {raw_response_text[:200]}"
+ break # Exit retry loop on JSON decode errors
+ except Exception as e:
+ last_error = f"{log_prefix}Unexpected Error: {e}"
+ break # Exit retry loop on other unexpected errors
+
+ # If loop finishes without returning success, return the last recorded error
+ logger.error(f"{log_prefix}Search API call failed. Last error: {last_error}")
+ return None, last_error.replace(log_prefix, "API Call Failed: ") if last_error else "API Call Failed after retries"
+
+
+def _passages2string(retrieval_result):
+ """Convert retrieval results to formatted string."""
+ format_reference = ""
+ for idx, doc_item in enumerate(retrieval_result):
+ content = doc_item["document"]["contents"]
+ title = content.split("\n")[0]
+ text = "\n".join(content.split("\n")[1:])
+ format_reference += f"Doc {idx + 1} (Title: {title})\n{text}\n\n"
+ return format_reference.strip()
+
+
+def perform_single_search_batch(
+ retrieval_service_url: str,
+ query_list: list[str],
+ topk: int = 3,
+ concurrent_semaphore: Optional[threading.Semaphore] = None,
+ timeout: int = DEFAULT_TIMEOUT,
+) -> tuple[str, dict[str, Any]]:
+ """
+ Performs a single batch search for multiple queries (original search tool behavior).
+
+ Args:
+ retrieval_service_url: The URL of the retrieval service API.
+ query_list: List of search queries.
+ topk: Number of top results to return.
+ concurrent_semaphore: Optional semaphore for concurrency control.
+ timeout: Request timeout in seconds.
+
+ Returns:
+ A tuple (result_text, metadata).
+ result_text: The search result JSON string.
+ metadata: Metadata dictionary for the batch search.
+ """
+ logger.info(f"Starting batch search for {len(query_list)} queries.")
+
+ api_response = None
+ error_msg = None
+
+ try:
+ if concurrent_semaphore:
+ with concurrent_semaphore:
+ api_response, error_msg = call_search_api(
+ retrieval_service_url=retrieval_service_url,
+ query_list=query_list,
+ topk=topk,
+ return_scores=True,
+ timeout=timeout,
+ )
+ else:
+ api_response, error_msg = call_search_api(
+ retrieval_service_url=retrieval_service_url,
+ query_list=query_list,
+ topk=topk,
+ return_scores=True,
+ timeout=timeout,
+ )
+ except Exception as e:
+ error_msg = f"API Request Exception during batch search: {e}"
+ logger.error(f"Batch search: {error_msg}")
+ traceback.print_exc()
+
+ metadata = {
+ "query_count": len(query_list),
+ "queries": query_list,
+ "api_request_error": error_msg,
+ "api_response": None,
+ "status": "unknown",
+ "total_results": 0,
+ "formatted_result": None,
+ }
+
+ result_text = json.dumps({"result": "Search request failed or timed out after retries."}, ensure_ascii=False)
+
+ if error_msg:
+ metadata["status"] = "api_error"
+ result_text = json.dumps({"result": f"Search error: {error_msg}"}, ensure_ascii=False)
+ logger.error(f"Batch search: API error occurred: {error_msg}")
+ elif api_response:
+ logger.debug(f"Batch search: API Response: {api_response}")
+ metadata["api_response"] = api_response
+
+ try:
+ raw_results = api_response.get("result", [])
+ if raw_results:
+ pretty_results = []
+ total_results = 0
+
+ for retrieval in raw_results:
+ formatted = _passages2string(retrieval)
+ pretty_results.append(formatted)
+ total_results += len(retrieval) if isinstance(retrieval, list) else 1
+
+ final_result = "\n---\n".join(pretty_results)
+ result_text = json.dumps({"result": final_result}, ensure_ascii=False)
+ metadata["status"] = "success"
+ metadata["total_results"] = total_results
+ metadata["formatted_result"] = final_result
+ logger.info(f"Batch search: Successful, got {total_results} total results")
+ else:
+ result_text = json.dumps({"result": "No search results found."}, ensure_ascii=False)
+ metadata["status"] = "no_results"
+ metadata["total_results"] = 0
+ logger.info("Batch search: No results found")
+ except Exception as e:
+ error_msg = f"Error processing search results: {e}"
+ result_text = json.dumps({"result": error_msg}, ensure_ascii=False)
+ metadata["status"] = "processing_error"
+ logger.error(f"Batch search: {error_msg}")
+ else:
+ metadata["status"] = "unknown_api_state"
+ result_text = json.dumps(
+ {"result": "Unknown API state (no response and no error message)."}, ensure_ascii=False
+ )
+ logger.error("Batch search: Unknown API state.")
+
+ return result_text, metadata
diff --git a/code/RL_model/verl/verl_train/verl/tools/utils/tool_registry.py b/code/RL_model/verl/verl_train/verl/tools/utils/tool_registry.py
new file mode 100644
index 0000000000000000000000000000000000000000..2b20fa48b96fab68da86a02932960d7b88d81928
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/tools/utils/tool_registry.py
@@ -0,0 +1,142 @@
+# Copyright 2025 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.
+
+import asyncio
+import importlib
+import logging
+import os
+import sys
+import threading
+from enum import Enum
+
+from omegaconf import OmegaConf
+
+from verl.tools.schemas import OpenAIFunctionToolSchema
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+class ToolType(Enum):
+ NATIVE = "native"
+ MCP = "mcp"
+
+
+async def initialize_mcp_tool(tool_cls, tool_config) -> list:
+ from verl.tools.utils.mcp_clients.McpClientManager import ClientManager
+
+ tool_list = []
+ mcp_servers_config_path = tool_config.mcp.mcp_servers_config_path
+ tool_selected_list = tool_config.mcp.tool_selected_list if "tool_selected_list" in tool_config.mcp else None
+ await ClientManager.initialize(mcp_servers_config_path, tool_config.config.rate_limit)
+ # Wait for MCP client to be ready
+ max_retries = 10
+ retry_interval = 2 # seconds
+ for i in range(max_retries):
+ tool_schemas = await ClientManager.fetch_tool_schemas(tool_selected_list)
+ if tool_schemas:
+ break
+ if i < max_retries - 1:
+ logger.debug(f"Waiting for MCP client to be ready, attempt {i + 1}/{max_retries}")
+ await asyncio.sleep(retry_interval)
+ else:
+ raise RuntimeError("Failed to initialize MCP tools after maximum retries")
+ # mcp registry
+ assert len(tool_schemas), "mcp tool is empty"
+ for tool_schema_dict in tool_schemas:
+ logger.debug(f"tool_schema_dict: {tool_schema_dict}")
+ tool_schema = OpenAIFunctionToolSchema.model_validate(tool_schema_dict)
+ tool = tool_cls(
+ config=OmegaConf.to_container(tool_config.config, resolve=True),
+ tool_schema=tool_schema,
+ )
+ tool_list.append(tool)
+ return tool_list
+
+
+def get_tool_class(cls_name):
+ module_name, class_name = cls_name.rsplit(".", 1)
+ if module_name not in sys.modules:
+ spec = importlib.util.find_spec(module_name)
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[module_name] = module
+ spec.loader.exec_module(module)
+ else:
+ module = sys.modules[module_name]
+
+ tool_cls = getattr(module, class_name)
+ return tool_cls
+
+
+def initialize_tools_from_config(tools_config_file):
+ """Initialize tools from config file.
+
+ Supports both NATIVE and MCP tool types. For MCP tools, a temporary event loop
+ is created only when needed and properly closed after use to prevent memory leaks.
+ """
+ tools_config = OmegaConf.load(tools_config_file)
+ tool_list = []
+
+ # Lazy initialization for MCP support - only create event loop when needed
+ tmp_event_loop = None
+ thread = None
+
+ def get_mcp_event_loop():
+ """Lazily create event loop and thread for MCP tools."""
+ nonlocal tmp_event_loop, thread
+ if tmp_event_loop is None:
+ tmp_event_loop = asyncio.new_event_loop()
+ thread = threading.Thread(target=tmp_event_loop.run_forever, name="mcp tool list fetcher", daemon=True)
+ thread.start()
+ return tmp_event_loop
+
+ def run_coroutine(coroutine):
+ """Run coroutine in the MCP event loop."""
+ loop = get_mcp_event_loop()
+ future = asyncio.run_coroutine_threadsafe(coroutine, loop)
+ return future.result()
+
+ try:
+ for tool_config in tools_config.tools:
+ cls_name = tool_config.class_name
+ tool_type = ToolType(tool_config.config.type)
+ tool_cls = get_tool_class(cls_name)
+
+ match tool_type:
+ case ToolType.NATIVE:
+ if tool_config.get("tool_schema", None) is None:
+ tool_schema = None
+ else:
+ tool_schema_dict = OmegaConf.to_container(tool_config.tool_schema, resolve=True)
+ tool_schema = OpenAIFunctionToolSchema.model_validate(tool_schema_dict)
+ tool = tool_cls(
+ config=OmegaConf.to_container(tool_config.config, resolve=True),
+ tool_schema=tool_schema,
+ )
+ tool_list.append(tool)
+ case ToolType.MCP:
+ mcp_tools = run_coroutine(initialize_mcp_tool(tool_cls, tool_config))
+ tool_list.extend(mcp_tools)
+ case _:
+ raise NotImplementedError
+ finally:
+ # Properly cleanup event loop if it was created
+ if tmp_event_loop is not None:
+ # stop first and then close
+ tmp_event_loop.call_soon_threadsafe(tmp_event_loop.stop)
+ if thread is not None and thread.is_alive():
+ thread.join(timeout=5.0)
+ tmp_event_loop.close()
+
+ return tool_list
diff --git a/code/RL_model/verl/verl_train/verl/utils/checkpoint/__init__.py b/code/RL_model/verl/verl_train/verl/utils/checkpoint/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..df9275830f0654a585435ea6ac74659e03a1cbb4
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/checkpoint/__init__.py
@@ -0,0 +1,17 @@
+# Copyright 2024 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 .checkpoint_handler import CheckpointHandler, OrchestrationMode
+
+__all__ = ["CheckpointHandler", "OrchestrationMode"]
diff --git a/code/RL_model/verl/verl_train/verl/utils/checkpoint/checkpoint_handler.py b/code/RL_model/verl/verl_train/verl/utils/checkpoint/checkpoint_handler.py
new file mode 100644
index 0000000000000000000000000000000000000000..aee7a5c6c0825d6b9cf83209befbd2211bd1fb71
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/checkpoint/checkpoint_handler.py
@@ -0,0 +1,224 @@
+# Copyright 2025 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.
+
+
+# TODO: add unit tests
+
+import logging
+import os
+import re
+from enum import Enum
+
+import torch
+
+import verl.utils.hdfs_io as hdfs_io
+from verl.single_controller import WorkerGroup
+from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, get_checkpoint_tracker_filename
+from verl.utils.logger import log_with_rank
+from verl.workers.engine import BaseEngine
+
+
+def extract_step(path):
+ match = re.search(r"global_step_(\d+)", path)
+ if match:
+ return int(match.group(1))
+ return None
+
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_SFT_LOGGING_LEVEL", "WARN"))
+
+
+class OrchestrationMode(Enum):
+ SPMD = 0
+ RAY = 1
+
+
+class CheckpointHandler:
+ """
+ Checkpoint handler handles the path, global_step of a checkpoint folder.
+ Currently, it only works with a single model.
+ We can expand it to support multiple models. It is expected to be used with SPMD style (e.g., torchrun)
+ """
+
+ def __init__(
+ self,
+ engine: BaseEngine | WorkerGroup,
+ train_dataloader,
+ *,
+ default_local_dir,
+ max_ckpt_to_keep=None,
+ default_hdfs_dir=None,
+ resume_mode="auto",
+ resume_from_path=None,
+ mode=OrchestrationMode.SPMD,
+ ):
+ self.default_local_dir = default_local_dir
+ self.max_ckpt_to_keep = max_ckpt_to_keep
+ self.default_hdfs_dir = default_hdfs_dir
+ self.resume_mode = resume_mode
+ self.resume_from_path = resume_from_path
+ self.engine = engine
+ self.train_dataloader = train_dataloader
+ self.mode = mode
+
+ if self.mode == OrchestrationMode.SPMD:
+ self.rank = torch.distributed.get_rank()
+ self.is_mp_src_rank_with_outputs = self.engine.is_mp_src_rank_with_outputs()
+ self.dp_rank = self.engine.get_data_parallel_rank()
+ elif self.mode == OrchestrationMode.RAY:
+ self.rank = 0
+ self.is_mp_src_rank_with_outputs = True
+ self.dp_rank = 0
+ else:
+ raise ValueError(f"Unknown {self.mode=}")
+
+ def save_checkpoint(self, step):
+ """Save checkpoint using FSDPCheckpointManager with improved tracking"""
+ from verl.utils.fs import local_mkdir_safe
+
+ # Determine checkpoint path
+ local_global_step_folder = os.path.join(self.default_local_dir, f"global_step_{step}")
+ if self.rank == 0:
+ print(f"Saving checkpoint to: {local_global_step_folder}")
+
+ # Get max checkpoints to keep
+ max_ckpt_to_keep = self.max_ckpt_to_keep
+
+ # Use checkpoint manager to save
+ self.engine.save_checkpoint(
+ local_path=local_global_step_folder, global_step=step, max_ckpt_to_keep=max_ckpt_to_keep
+ )
+
+ # Save dataloader state. Note that we only save the iterator in the train_dataloader.
+ # So it's identical in each dp rank.
+ if self.is_mp_src_rank_with_outputs:
+ dp_rank = self.dp_rank
+ local_mkdir_safe(local_global_step_folder)
+ dataloader_local_path = os.path.join(local_global_step_folder, f"data_{dp_rank}.pt")
+
+ # Use StatefulDataLoader's built-in state dict functionality
+ dataloader_state_dict = self.train_dataloader.state_dict()
+ torch.save(dataloader_state_dict, dataloader_local_path)
+ print(f"Saved dataloader state to: {dataloader_local_path}")
+
+ if self.rank == 0:
+ # Update latest checkpoint tracker (atomic write)
+ tracker_file = get_checkpoint_tracker_filename(self.default_local_dir)
+ temp_tracker_file = tracker_file + ".tmp"
+ with open(temp_tracker_file, "w") as f:
+ f.write(str(step))
+ os.rename(temp_tracker_file, tracker_file)
+ print(f"Updated checkpoint tracker: {tracker_file}")
+
+ # Copy to HDFS if configured
+ if self.rank == 0 and self.default_hdfs_dir:
+ hdfs_io.makedirs(self.default_hdfs_dir, exist_ok=True)
+ hdfs_io.copy(src=local_global_step_folder, dst=self.default_hdfs_dir, dirs_exist_ok=True)
+
+ if self.mode == OrchestrationMode.SPMD:
+ torch.distributed.barrier()
+
+ def load_checkpoint(self):
+ # Determine resume path based on configuration
+ checkpoint_path = self._determine_resume_path()
+
+ if checkpoint_path is None:
+ return 0
+
+ # extract resume step from checkpoint path
+ resume_step = extract_step(checkpoint_path)
+ if resume_step is None:
+ log_with_rank(
+ f"Warning: Could not extract step number from {checkpoint_path}, starting from step 0",
+ logger=logger,
+ rank=self.rank,
+ level=logging.WARNING,
+ log_only_rank_0=True,
+ )
+ return 0
+ self.resume_global_step = resume_step
+
+ # Use checkpoint manager to load model state
+ self.engine.load_checkpoint(checkpoint_path)
+ # Always load dataloader state for StatefulDataLoader
+ self._load_dataloader_state(checkpoint_path)
+
+ return resume_step
+
+ def _load_dataloader_state(self, checkpoint_path: str):
+ """Load dataloader state from checkpoint"""
+ dp_rank = self.dp_rank
+ dataloader_path = os.path.join(checkpoint_path, f"data_{dp_rank}.pt")
+
+ if os.path.exists(dataloader_path):
+ # Use StatefulDataLoader's built-in state dict functionality
+ dataloader_state_dict = torch.load(dataloader_path, map_location="cpu", weights_only=False)
+ self.train_dataloader.load_state_dict(dataloader_state_dict)
+
+ log_with_rank(
+ f"Successfully loaded dataloader state from {dataloader_path}",
+ logger=logger,
+ rank=self.rank,
+ log_only_rank_0=True,
+ )
+
+ else:
+ log_with_rank(
+ f"Warning: No dataloader state found at {dataloader_path}, will start from scratch",
+ logger=logger,
+ rank=self.rank,
+ level=logging.WARNING,
+ log_only_rank_0=True,
+ )
+
+ def _determine_resume_path(self):
+ """Determine the path to resume from based on resume_mode configuration"""
+ resume_mode = self.resume_mode
+ resume_from_path = self.resume_from_path
+
+ if resume_mode == "disable":
+ return None
+ elif resume_mode == "auto":
+ if resume_from_path is not None:
+ assert os.path.exists(resume_from_path), (
+ "resume_from_path must be null or an existing path when resume_mode is 'auto'"
+ )
+ assert "global_step_" in resume_from_path, "resume_from_path must specify the global_steps"
+ return resume_from_path
+ # Try to find the latest checkpoint in the default directory
+ return self._find_latest_checkpoint()
+ elif resume_mode == "resume_path":
+ assert os.path.exists(resume_from_path), (
+ "resume_from_path must be an existing path when resume_mode is 'resume_path'"
+ )
+ assert "global_step_" in resume_from_path, "resume_from_path must specify the global_steps"
+ return resume_from_path
+ else:
+ raise ValueError(f"Invalid resume_mode: {resume_mode}. Must be 'auto', 'disable', or 'resume_path'")
+
+ def _find_latest_checkpoint(self):
+ """Find the latest checkpoint in the default local directory"""
+ checkpoint_dir = self.default_local_dir
+
+ if not os.path.exists(checkpoint_dir):
+ return None
+
+ latest_checkpoint = find_latest_ckpt_path(checkpoint_dir)
+
+ if latest_checkpoint and self.rank == 0:
+ step_num = extract_step(latest_checkpoint)
+ print(f"Found latest checkpoint: {latest_checkpoint} (step {step_num})")
+
+ return latest_checkpoint
diff --git a/code/RL_model/verl/verl_train/verl/utils/checkpoint/checkpoint_manager.py b/code/RL_model/verl/verl_train/verl/utils/checkpoint/checkpoint_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f48147b8f538af45921798be1790b79a805dbda
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/checkpoint/checkpoint_manager.py
@@ -0,0 +1,268 @@
+# Copyright 2024 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.
+
+import os
+import random
+import shutil
+
+import numpy as np
+import torch
+import torch.distributed
+from omegaconf import DictConfig
+from transformers import PreTrainedTokenizer, ProcessorMixin
+
+from verl.trainer.config import CheckpointConfig
+from verl.utils.device import get_device_name, get_torch_device
+
+
+class BaseCheckpointManager:
+ """
+ A checkpoint manager that saves and loads the following states in a SPMD way:
+ - model
+ - optimizer
+ - lr_scheduler
+ - extra_states
+
+ We save
+ - sharded model states and optimizer states
+ - full lr_scheduler states
+ - huggingface tokenizer and config for ckpt merge
+ """
+
+ def __init__(
+ self,
+ model,
+ optimizer: torch.optim.Optimizer,
+ lr_scheduler: torch.optim.lr_scheduler.LRScheduler = None,
+ processing_class: PreTrainedTokenizer | ProcessorMixin = None,
+ checkpoint_config: DictConfig | CheckpointConfig = None,
+ ):
+ self.checkpoint_config = checkpoint_config
+ checkpoint_load_contents = checkpoint_config.get("load_contents", None) if checkpoint_config else None
+ checkpoint_save_contents = checkpoint_config.get("save_contents", None) if checkpoint_config else None
+ if checkpoint_load_contents is None:
+ checkpoint_load_contents = ["model", "optimizer", "extra"]
+ if checkpoint_save_contents is None:
+ checkpoint_save_contents = ["model", "optimizer", "extra"]
+ self.previous_global_step = None
+ self.previous_saved_paths = []
+
+ self.model = model
+ self.optimizer = optimizer
+ self.lr_scheduler = lr_scheduler
+ self.processing_class = processing_class
+ self.checkpoint_load_contents = checkpoint_load_contents
+ self.checkpoint_save_contents = checkpoint_save_contents
+
+ self.rank = torch.distributed.get_rank()
+ self.world_size = torch.distributed.get_world_size()
+
+ @property
+ def should_save_model(self) -> bool:
+ """
+ Returns True if 'model' is in checkpoint_save_contents, indicating the model state should be saved.
+ """
+ return "model" in self.checkpoint_save_contents
+
+ @property
+ def should_save_optimizer(self) -> bool:
+ """
+ Returns True if 'optimizer' is in checkpoint_save_contents, indicating the optimizer state should be saved.
+ """
+ return "optimizer" in self.checkpoint_save_contents
+
+ @property
+ def should_save_extra(self) -> bool:
+ """
+ Returns True if 'extra' is in checkpoint_save_contents, indicating the extra state should be saved.
+ """
+ return "extra" in self.checkpoint_save_contents
+
+ @property
+ def should_save_hf_model(self) -> bool:
+ """
+ Returns True if 'hf_model' is in checkpoint_save_contents, indicating the model should be converted to hf
+ model and saved.
+ """
+ return "hf_model" in self.checkpoint_save_contents
+
+ @property
+ def should_load_model(self) -> bool:
+ """
+ Returns True if 'model' is in checkpoint_load_contents, indicating the model state should be loaded.
+ """
+ return "model" in self.checkpoint_load_contents
+
+ @property
+ def should_load_optimizer(self) -> bool:
+ """
+ Returns True if 'optimizer' is in checkpoint_load_contents, indicating the optimizer state should be loaded.
+ """
+ return "optimizer" in self.checkpoint_load_contents
+
+ @property
+ def should_load_extra(self) -> bool:
+ """
+ Returns True if 'extra' is in checkpoint_load_contents, indicating the extra state should be loaded.
+ """
+ return "extra" in self.checkpoint_load_contents
+
+ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load: bool = False):
+ raise NotImplementedError
+
+ def save_checkpoint(
+ self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep: int = None
+ ):
+ raise NotImplementedError
+
+ @staticmethod
+ def checkpath(local_path: str, hdfs_path: str):
+ assert local_path is not None or hdfs_path is not None, "local_path and hdfs_path cannot be both None"
+ return local_path is not None, local_path if local_path is not None else hdfs_path
+
+ def remove_previous_save_local_path(self, path):
+ if isinstance(path, str):
+ path = [path]
+ for p in path:
+ abs_path = os.path.abspath(p)
+ print(f"Checkpoint manager remove previous save local path: {abs_path}")
+ if not os.path.exists(abs_path):
+ continue
+ shutil.rmtree(abs_path, ignore_errors=True)
+
+ def ensure_checkpoint_capacity(self, max_ckpt_to_keep: int):
+ """
+ Remove old checkpoints to make room for a new one, keeping a safety buffer.
+
+ With max_ckpt_to_keep=1, this does nothing - we keep the existing checkpoint
+ until the new save completes successfully (handled by register_checkpoint).
+ For max_ckpt_to_keep >= 2, we keep (max_ckpt_to_keep - 1) checkpoints before save.
+ """
+ if not (max_ckpt_to_keep and isinstance(max_ckpt_to_keep, int) and max_ckpt_to_keep > 1):
+ return
+ if len(self.previous_saved_paths) >= max_ckpt_to_keep:
+ keep_start = len(self.previous_saved_paths) - max_ckpt_to_keep + 1
+ self.remove_previous_save_local_path(self.previous_saved_paths[:keep_start])
+ self.previous_saved_paths = self.previous_saved_paths[keep_start:]
+
+ def register_checkpoint(self, new_path: str, max_ckpt_to_keep: int):
+ """
+ Register a successfully saved checkpoint and enforce retention limit.
+
+ Adds the new checkpoint path to tracking and removes excess old
+ checkpoints beyond max_ckpt_to_keep.
+ """
+ self.previous_saved_paths.append(new_path)
+ if not (max_ckpt_to_keep and isinstance(max_ckpt_to_keep, int) and max_ckpt_to_keep > 0):
+ return
+ if len(self.previous_saved_paths) > max_ckpt_to_keep:
+ keep_start = len(self.previous_saved_paths) - max_ckpt_to_keep
+ self.remove_previous_save_local_path(self.previous_saved_paths[:keep_start])
+ self.previous_saved_paths = self.previous_saved_paths[keep_start:]
+
+ @staticmethod
+ def get_rng_state():
+ rng_state = {
+ "cpu": torch.get_rng_state(),
+ "numpy": np.random.get_state(),
+ "random": random.getstate(),
+ }
+
+ if get_device_name() != "cpu":
+ rng_state[get_device_name()] = get_torch_device().get_rng_state()
+
+ return rng_state
+
+ @staticmethod
+ def load_rng_state(rng_state):
+ torch.set_rng_state(rng_state["cpu"])
+ np.random.set_state(rng_state["numpy"])
+ random.setstate(rng_state["random"])
+
+ if get_device_name() != "cpu":
+ get_torch_device().set_rng_state(rng_state[get_device_name()])
+
+
+def find_latest_ckpt_path(path, directory_format="global_step_{}"):
+ """
+ Return the most recent checkpoint directory based on a tracker file.
+
+ Args:
+ path (str): Base directory containing the checkpoint tracker.
+ directory_format (str): Template for checkpoint subfolders with one
+ placeholder for the iteration number (default "global_step_{}").
+
+ Returns:
+ str or None: Full path to the latest checkpoint directory, or
+ None if the tracker or checkpoint folder is missing.
+ """
+ if path is None:
+ return None
+
+ tracker_file = get_checkpoint_tracker_filename(path)
+ if not os.path.exists(tracker_file):
+ if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+ print(f"Checkpoint tracker file does not exist: {tracker_file}")
+ return None
+
+ with open(tracker_file, "rb") as f:
+ iteration = int(f.read().decode())
+ ckpt_path = os.path.join(path, directory_format.format(iteration))
+ if not os.path.exists(ckpt_path):
+ print("Checkpoint does not exist: %s", ckpt_path)
+ return None
+
+ print("Found checkpoint: %s", ckpt_path)
+ return ckpt_path
+
+
+def get_checkpoint_tracker_filename(root_path: str):
+ """
+ Tracker file rescords the latest chckpoint during training to restart from.
+ """
+ return os.path.join(root_path, "latest_checkpointed_iteration.txt")
+
+
+def should_save_ckpt_esi(max_steps_duration: float, save_ckpt_duration: float = 60, redundant_time: float = 0) -> bool:
+ """
+ Determine if checkpoint should be saved based on capacity esi expiration.
+
+ Args:
+ max_steps_duration: Max estimated time (seconds) required to complete one training step
+ save_ckpt_duration: Estimated time (seconds) required to save checkpoint (default: 60)
+ redundant_time: Additional buffer time (seconds) for unexpected delays (default: 0)
+ """
+ exp_ts_mlp = os.getenv("MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP") # vemlp
+ exp_ts_aws = os.getenv("SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP") # aws
+ if exp_ts_mlp:
+ try:
+ import time
+
+ remaining = float(exp_ts_mlp) - time.time()
+ except ValueError:
+ return False
+ return (
+ remaining > 0
+ and max_steps_duration > 0
+ and remaining <= save_ckpt_duration + max_steps_duration + redundant_time
+ )
+ elif exp_ts_aws:
+ from datetime import datetime, timedelta
+
+ expiration_time = datetime.fromtimestamp(int(exp_ts_aws))
+ time_difference = expiration_time - datetime.now()
+ threshold_minutes = (save_ckpt_duration + max_steps_duration + redundant_time) / 60
+ return time_difference < timedelta(minutes=threshold_minutes)
+ else:
+ return False
diff --git a/code/RL_model/verl/verl_train/verl/utils/checkpoint/fsdp_checkpoint_manager.py b/code/RL_model/verl/verl_train/verl/utils/checkpoint/fsdp_checkpoint_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bd57d907aee38f44e08a4d34c2e8c50037f7363
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/checkpoint/fsdp_checkpoint_manager.py
@@ -0,0 +1,362 @@
+# Copyright 2024 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.
+
+import json
+import logging
+import os
+import warnings
+from dataclasses import asdict, dataclass
+from typing import Optional
+
+import torch
+import torch.distributed
+from accelerate import init_empty_weights
+from omegaconf import DictConfig
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import ShardedOptimStateDictConfig, ShardedStateDictConfig, StateDictType
+from transformers import GenerationConfig, PreTrainedTokenizer, ProcessorMixin
+from transformers.dynamic_module_utils import custom_object_save
+
+from verl.utils.device import is_cuda_available
+from verl.utils.fs import copy_to_local, is_non_local, local_mkdir_safe
+from verl.utils.fsdp_utils import fsdp_version, get_fsdp_full_state_dict, get_fsdp_state_ctx
+from verl.utils.logger import log_with_rank
+
+from .checkpoint_manager import BaseCheckpointManager
+
+# Setup logging
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO"))
+
+
+@dataclass
+class FSDPConfig:
+ """Configuration for FSDP checkpointing.
+
+ Args:
+ FSDP_version (int): Version of FSDP being used.
+ world_size (int): Number of processes in the distributed training setup.
+ """
+
+ FSDP_version: int
+ world_size: int
+
+
+class FSDPCheckpointManager(BaseCheckpointManager):
+ """
+ Manage FSDP checkpointing in SPMD training.
+
+ - Saves/loads per-rank sharded model & optimizer states
+ - Persists full lr_scheduler and RNG state
+ - Stores HF tokenizer/processor and model/config for unified restore
+
+ Args:
+ model (FSDP): Wrapped model instance.
+ optimizer (Optimizer): Training optimizer.
+ lr_scheduler (LRScheduler): Learning-rate scheduler.
+ processing_class (PreTrainedTokenizer or ProcessorMixin, optional):
+ Pre-/post-processing artifact handler.
+ checkpoint_contents DictConfig: Configuration for checkpoint contents.
+ - 'load': Components to load; must contain 'model'. Defaults to ['model', 'optimizer', 'extra'].
+ - 'save': Components to save; must contain 'model'. Defaults to ['model', 'optimizer', 'extra'].
+ """
+
+ def __init__(
+ self,
+ model: FSDP,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ lr_scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
+ processing_class: PreTrainedTokenizer | ProcessorMixin = None,
+ checkpoint_config: DictConfig = None,
+ **kwargs,
+ ):
+ if processing_class is None and "tokenizer" in kwargs:
+ warnings.warn(
+ "`tokenizer` is deprecated. use `processing_class` instead.", DeprecationWarning, stacklevel=2
+ )
+ processing_class = kwargs.pop("tokenizer")
+
+ super().__init__(
+ model,
+ optimizer,
+ lr_scheduler=lr_scheduler,
+ processing_class=processing_class,
+ checkpoint_config=checkpoint_config,
+ )
+
+ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load=False):
+ """
+ Load an FSDP checkpoint for this rank.
+
+ Downloads and loads:
+ - model and optimizer shards
+ - extra state dict (scheduler + RNG)
+
+ Args:
+ local_path: Directory with per-rank checkpoint files.
+ hdfs_path: Unused (for API compatibility).
+ del_local_after_load: Remove local files after loading.
+ """
+ if local_path is None:
+ return
+
+ # check if the checkpoint_load_contents is valid
+ if self.should_load_model:
+ assert self.model is not None, "model must be provided when checkpoint_contents.load includes ['model']"
+ if self.should_load_optimizer:
+ assert self.optimizer is not None, (
+ "optimizer must be provided when checkpoint_contents.load includes ['optimizer']"
+ )
+
+ # every rank download its own checkpoint
+ state_dict_cfg = (
+ ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
+ if self.should_load_model
+ else None
+ )
+ optim_cfg = (
+ ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
+ if self.should_load_optimizer
+ else None
+ )
+ with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
+ if self.should_load_model:
+ remote_model_path = os.path.join(local_path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt")
+ local_model_path = copy_to_local(remote_model_path)
+ model_state_dict = torch.load(local_model_path, weights_only=False)
+ self.model.load_state_dict(model_state_dict)
+ log_with_rank(f"Loaded model from {remote_model_path}", rank=self.rank, logger=logger)
+
+ if self.should_load_optimizer:
+ remote_optim_path = os.path.join(local_path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt")
+ local_optim_path = copy_to_local(remote_optim_path)
+ optimizer_state_dict = torch.load(local_optim_path, weights_only=False)
+ self.optimizer.load_state_dict(optimizer_state_dict)
+ log_with_rank(f"Loaded optimizer from {remote_optim_path}", rank=self.rank, logger=logger)
+
+ if self.should_load_extra:
+ remote_extra_state_path = os.path.join(
+ local_path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt"
+ )
+ local_extra_state_path = copy_to_local(remote_extra_state_path)
+ extra_state_dict = torch.load(local_extra_state_path, weights_only=False)
+ # recover random state
+ if "rng" in extra_state_dict:
+ # 'rng' may not exist for backward compatibility
+ self.load_rng_state(extra_state_dict["rng"])
+ log_with_rank(f"Loaded rng from {remote_extra_state_path}", rank=self.rank, logger=logger)
+
+ lr_scheduler_state_dict = extra_state_dict["lr_scheduler"]
+ if lr_scheduler_state_dict is not None and self.lr_scheduler is not None:
+ self.lr_scheduler.load_state_dict(lr_scheduler_state_dict)
+ log_with_rank(f"Loaded lr_scheduler from {remote_extra_state_path}", rank=self.rank, logger=logger)
+
+ if self.rank == 0 and del_local_after_load:
+ try:
+ os.remove(local_model_path) if is_non_local(local_model_path) else None
+ os.remove(local_optim_path) if is_non_local(local_optim_path) else None
+ os.remove(local_extra_state_path) if is_non_local(local_extra_state_path) else None
+ except Exception as e:
+ log_with_rank(
+ f"remove local resume ckpt file after loading failed, exception {e} will be ignored",
+ rank=self.rank,
+ logger=logger,
+ )
+
+ # wait for everyone to load checkpoints
+ torch.distributed.barrier()
+
+ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep=None):
+ """
+ Save an FSDP checkpoint for this rank.
+
+ Writes:
+ - model & optimizer shard files
+ - extra state dict (scheduler + RNG)
+ - HF tokenizer/processor and model/config on rank 0
+ - optional full HF model under 'huggingface/' if requested
+
+ Rotates old checkpoints, keeping at most `max_ckpt_to_keep`.
+
+ Args:
+ local_path: Target directory for checkpoint files.
+ hdfs_path: Unused (for API compatibility).
+ global_step: Current training step (used for bookkeeping).
+ max_ckpt_to_keep: Number of recent checkpoints to retain.
+ """
+ if local_path is None:
+ return
+
+ # record the previous global step
+ self.previous_global_step = global_step
+
+ if self.rank == 0:
+ self.ensure_checkpoint_capacity(max_ckpt_to_keep)
+
+ local_path = local_mkdir_safe(local_path)
+ torch.distributed.barrier()
+
+ # check if the checkpoint_save_contents is valid
+ if self.should_save_model:
+ assert self.model is not None, "model must be provided when checkpoint_contents.save includes ['model']"
+ if self.should_save_optimizer:
+ assert self.optimizer is not None, (
+ "optimizer must be provided when checkpoint_contents.save includes ['optimizer']"
+ )
+
+ # every rank will save its own model and optim shard
+ state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
+ optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False)
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):
+ model_path = os.path.join(local_path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt")
+ optim_path = os.path.join(local_path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt")
+ extra_path = os.path.join(local_path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt")
+
+ if self.should_save_model:
+ model_state_dict = self.model.state_dict()
+ torch.save(model_state_dict, model_path)
+ log_with_rank(f"Saved model to {os.path.abspath(model_path)}", rank=self.rank, logger=logger)
+
+ if self.should_save_optimizer:
+ optimizer_state_dict = self.optimizer.state_dict()
+ torch.save(optimizer_state_dict, optim_path)
+ log_with_rank(f"Saved optim to {os.path.abspath(optim_path)}", rank=self.rank, logger=logger)
+
+ if self.should_save_extra:
+ lr_scheduler_state_dict = self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None
+ extra_state_dict = {
+ "lr_scheduler": lr_scheduler_state_dict,
+ "rng": self.get_rng_state(),
+ }
+ torch.save(extra_state_dict, extra_path)
+ log_with_rank(f"Saved extra_state to {os.path.abspath(extra_path)}", rank=self.rank, logger=logger)
+
+ if self.rank == 0:
+ # Save HF tokenizer/processor and model config on rank 0 to huggingface/ directory, no matter whether
+ # huggingface model is requested to be saved or not.
+
+ if fsdp_version(self.model) == 1:
+ unwrap_model = self.model._fsdp_wrapped_module
+ else:
+ unwrap_model = self.model
+
+ hf_config_tokenizer_path = os.path.join(local_path, "huggingface")
+ local_mkdir_safe(hf_config_tokenizer_path)
+ model_config = unwrap_model.config
+ generation_config = None
+ if unwrap_model.can_generate() and hasattr(model_config, "name_or_path") and model_config.name_or_path:
+ try:
+ # Some model's name_or_path is empty if not initialized from pretrained,
+ # in this cases, we don't save generation config.
+ generation_config = GenerationConfig.from_pretrained(model_config.name_or_path)
+ generation_config.save_pretrained(hf_config_tokenizer_path)
+ except Exception:
+ # if the generation config isn't available, we don't save it
+ pass
+
+ if hasattr(model_config, "auto_map") and None in model_config.auto_map:
+ model_config.auto_map = {k: v for k, v in model_config.auto_map.items() if k is not None}
+
+ model_config.save_pretrained(hf_config_tokenizer_path)
+ if self.processing_class is not None:
+ self.processing_class.save_pretrained(hf_config_tokenizer_path)
+ log_with_rank(
+ f"Saved model config and tokenizer class to {os.path.abspath(hf_config_tokenizer_path)}",
+ rank=self.rank,
+ logger=logger,
+ log_only_rank_0=True,
+ )
+
+ # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
+ # loaded from the Hub.
+ if hasattr(model_config, "auto_map"):
+ custom_object_save(unwrap_model, hf_config_tokenizer_path, config=model_config)
+
+ # Also save runtime FSDP config
+ fsdp_config_path = os.path.join(local_path, "fsdp_config.json")
+ fsdp_config = FSDPConfig(
+ FSDP_version=fsdp_version(self.model),
+ world_size=self.world_size,
+ )
+ with open(fsdp_config_path, "w") as f:
+ json.dump(asdict(fsdp_config), f, indent=4)
+
+ # wait for everyone to dump to local
+ torch.distributed.barrier()
+
+ if self.should_save_hf_model:
+ # Only rank 0 will save hf model and,
+ # offload to cpu to save LLMs which may be too large to fit in one GPU
+ state_dict = get_fsdp_full_state_dict(self.model, offload_to_cpu=True, rank0_only=True)
+
+ if self.rank == 0:
+ hf_local_path = os.path.join(local_path, "huggingface")
+ os.makedirs(hf_local_path, exist_ok=True)
+
+ if "ForTokenClassification" in model_config.architectures[0]:
+ from transformers import AutoModelForTokenClassification
+
+ auto_model_cls = AutoModelForTokenClassification
+ elif "ForCausalLM" in model_config.architectures[0]:
+ from transformers import AutoModelForCausalLM
+
+ auto_model_cls = AutoModelForCausalLM
+ elif "ForConditionalGeneration" in model_config.architectures[0]:
+ # Handle different transformers versions for Vision2Seq models
+ import transformers
+ from packaging import version
+
+ if version.parse(transformers.__version__) >= version.parse("4.54.0"):
+ # transformers >= 4.54.0 uses AutoModelForImageTextToText
+ from transformers import AutoModelForImageTextToText
+
+ auto_model_cls = AutoModelForImageTextToText
+ else:
+ # transformers < 4.54.0 uses AutoModelForVision2Seq
+ from transformers import AutoModelForVision2Seq
+
+ auto_model_cls = AutoModelForVision2Seq
+ else:
+ raise NotImplementedError(f"Unknown architecture {model_config['architectures']}")
+
+ with init_empty_weights():
+ save_model = auto_model_cls.from_config(model_config, torch_dtype=torch.bfloat16)
+ save_model.to_empty(device="cpu")
+
+ if save_model.can_generate():
+ if generation_config is not None:
+ save_model.generation_config = generation_config
+ else:
+ print(
+ f"Warning: {self.__class__.__name__}.save_checkpoint: Generation config file not found "
+ f"in, using a generation config created from the model config when saving hf_model."
+ )
+
+ save_model.save_pretrained(hf_local_path, state_dict=state_dict)
+ log_with_rank(
+ f"Saved hf_model to {os.path.abspath(hf_local_path)}",
+ rank=self.rank,
+ logger=logger,
+ log_only_rank_0=True,
+ )
+ del state_dict
+ del save_model
+
+ # wait for rank0 to dump hf_model to local
+ torch.distributed.barrier()
+
+ if self.rank == 0:
+ self.register_checkpoint(local_path, max_ckpt_to_keep)
diff --git a/code/RL_model/verl/verl_train/verl/utils/checkpoint/megatron_checkpoint_manager.py b/code/RL_model/verl/verl_train/verl/utils/checkpoint/megatron_checkpoint_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..b763e64432d6859de79ba29d37bb7d6005be664f
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/checkpoint/megatron_checkpoint_manager.py
@@ -0,0 +1,666 @@
+# Copyright 2024 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.
+
+import json
+import logging
+import os
+import random
+from collections.abc import Callable
+from dataclasses import asdict
+
+import numpy as np
+import torch
+import torch.distributed
+from megatron.core import mpu, tensor_parallel
+from megatron.core.dist_checkpointing.mapping import ShardedObject
+from megatron.core.transformer.enums import AttnBackend
+from transformers import GenerationConfig
+
+from verl.models.weight_loader_registry import get_weight_saver
+from verl.utils.device import get_device_name, get_torch_device
+from verl.utils.fs import is_non_local, local_mkdir_safe
+from verl.utils.logger import log_with_rank
+from verl.utils.megatron.dist_checkpointing import load_dist_checkpointing, save_dist_checkpointing
+from verl.utils.megatron_utils import (
+ get_dist_checkpoint_path,
+ get_hf_model_checkpoint_path,
+ get_transformer_config_checkpoint_path,
+)
+
+from .checkpoint_manager import BaseCheckpointManager
+
+# Setup logging
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO"))
+
+
+class MegatronCheckpointManager(BaseCheckpointManager):
+ """
+ Checkpoint manager for Megatron-LM distributed training.
+
+ This class manages the saving and loading of model checkpoints in a Megatron-LM
+ distributed training environment. It handles various aspects of checkpointing
+ including model states, optimizer states, learning rate schedulers, and random
+ number generator states, ensuring compatibility with HuggingFace formats.
+
+ Key features:
+ - Distributed checkpoint saving and loading using Megatron's dist_checkpointing
+ - Support for tensor parallel, pipeline parallel, and data parallel configurations
+ - Automatic handling of model state dictionaries across multiple pipeline stages
+ - Integration with HuggingFace model configurations and tokenizers
+ - Random number generator state management for reproducibility
+ - Support for both synchronous and asynchronous checkpoint operations
+
+ The manager automatically handles:
+ - Directory structure creation based on global steps and process ranks
+ - Model configuration and tokenizer saving in HuggingFace format
+ - Optimizer and scheduler state persistence
+ - CUDA RNG state management for deterministic training
+ - Checkpoint cleanup and retention policies
+
+ Args:
+ model: The Megatron model instance to checkpoint
+ optimizer: The optimizer instance (optional)
+ lr_scheduler: The learning rate scheduler instance (optional)
+
+ Attributes:
+ model: Reference to the Megatron model being checkpointed
+ optimizer: Reference to the optimizer (if provided)
+ lr_scheduler: Reference to the learning rate scheduler (if provided)
+ rank: Current process rank in the distributed setup
+
+ Example:
+ ```python
+ checkpoint_manager = MegatronCheckpointManager(
+ model=megatron_model,
+ optimizer=optimizer,
+ lr_scheduler=scheduler
+ )
+
+ checkpoint_manager.save_checkpoint(
+ local_path="checkpoints/step_1000",
+ global_step=1000
+ )
+
+ checkpoint_manager.load_checkpoint(
+ local_path="checkpoints/step_1000"
+ )
+ ```
+ """
+
+ def __init__(
+ self,
+ config,
+ checkpoint_config,
+ model_config,
+ transformer_config,
+ role,
+ model: torch.nn.ModuleList,
+ arch: str,
+ hf_config,
+ param_dtype: torch.dtype,
+ share_embeddings_and_output_weights: bool,
+ processing_class,
+ optimizer,
+ optimizer_scheduler,
+ use_distributed_optimizer: bool,
+ use_checkpoint_opt_param_scheduler: bool = False,
+ use_dist_checkpointing: bool = True,
+ bridge=None,
+ provider=None,
+ peft_cls=None,
+ **kwargs,
+ ):
+ super().__init__(
+ model,
+ optimizer=optimizer,
+ lr_scheduler=optimizer_scheduler,
+ processing_class=processing_class,
+ checkpoint_config=checkpoint_config,
+ )
+ self.arch = arch
+ self.config = config
+ self.transformer_config = transformer_config
+ self.role = role
+ self.is_value_model = False
+ if self.role in ["reward", "critic"]:
+ self.is_value_model = True
+ self.model_config = model_config
+ self.hf_config = hf_config
+ self.param_dtype = param_dtype
+ self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+ self.model_path = self.config.model.path
+ self.use_distributed_optimizer = use_distributed_optimizer
+ self.use_checkpoint_opt_param_scheduler = use_checkpoint_opt_param_scheduler
+ self.bridge = bridge
+ self.provider = provider
+ self.vanilla_bridge = self.provider is None
+ self.peft_cls = peft_cls
+ self.rank = torch.distributed.get_rank()
+ # Megatron-Bridge is Okay to load/save HF checkpoint for value model as well
+ self.use_dist_checkpointing = (
+ use_dist_checkpointing or not self.bridge or (self.vanilla_bridge and self.is_value_model)
+ )
+ self.use_hf_checkpoint = not self.use_dist_checkpointing
+
+ self.weight_saver = None
+ if self.bridge is None:
+ self.weight_saver = get_weight_saver(self.arch)
+
+ def get_rng_state(self, use_dist_ckpt: bool = True, data_parallel_random_init: bool = False):
+ """collect rng state across data parallel ranks"""
+ rng_state = {
+ "random_rng_state": random.getstate(),
+ "np_rng_state": np.random.get_state(),
+ "torch_rng_state": torch.get_rng_state(),
+ "rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states(),
+ }
+
+ if get_device_name() != "cpu":
+ rng_state[f"{get_device_name()}_rng_state"] = get_torch_device().get_rng_state()
+
+ rng_state_list = None
+ if torch.distributed.is_initialized() and mpu.get_data_parallel_world_size() > 1 and data_parallel_random_init:
+ rng_state_list = [None for i in range(mpu.get_data_parallel_world_size())]
+ torch.distributed.all_gather_object(rng_state_list, rng_state, group=mpu.get_data_parallel_group())
+ else:
+ rng_state_list = [rng_state]
+
+ if use_dist_ckpt:
+ pp_rank = mpu.get_pipeline_model_parallel_rank()
+ pp_size = mpu.get_pipeline_model_parallel_world_size()
+ tp_rank = mpu.get_tensor_model_parallel_rank()
+ tp_size = mpu.get_tensor_model_parallel_world_size()
+ rng_state_list = ShardedObject(
+ "rng_state",
+ rng_state_list,
+ (pp_size, tp_size),
+ (pp_rank, tp_rank),
+ replica_id=mpu.get_data_parallel_rank(with_context_parallel=True),
+ )
+
+ return rng_state_list
+
+ def get_checkpoint_name(
+ self,
+ checkpoints_path,
+ pipeline_parallel=None,
+ tensor_rank=None,
+ pipeline_rank=None,
+ cp_rank=None,
+ expert_parallel=None,
+ expert_rank=None,
+ return_base_dir=True,
+ basename="model.pt",
+ ):
+ """Determine the directory name for this rank's checkpoint."""
+ # Use both the tensor and pipeline MP rank.
+ if pipeline_parallel is None:
+ pipeline_parallel = mpu.get_pipeline_model_parallel_world_size() > 1
+ if tensor_rank is None:
+ tensor_rank = mpu.get_tensor_model_parallel_rank()
+ if pipeline_rank is None:
+ pipeline_rank = mpu.get_pipeline_model_parallel_rank()
+ if cp_rank is None:
+ cp_rank = mpu.get_context_parallel_rank()
+ if expert_parallel is None:
+ expert_parallel = mpu.get_expert_model_parallel_world_size() > 1
+ if expert_rank is None:
+ expert_rank = mpu.get_expert_model_parallel_rank()
+
+ # Use both the tensor and pipeline MP rank. If using the distributed
+ # optimizer, then the optimizer's path must additionally include the
+ # data parallel rank.
+
+ # due to the fact that models are identical across cp ranks, cp rank is not used in the checkpoint path
+ if not pipeline_parallel:
+ common_path = os.path.join(checkpoints_path, f"mp_rank_{tensor_rank:02d}")
+ else:
+ common_path = os.path.join(checkpoints_path, f"mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}")
+
+ if expert_parallel:
+ common_path = common_path + f"_{expert_rank:03d}"
+
+ os.makedirs(common_path, exist_ok=True)
+
+ if return_base_dir:
+ return common_path
+ return os.path.join(common_path, basename)
+
+ def generate_state_dict(
+ self,
+ generate_model: bool = True,
+ generate_optimizer: bool = True,
+ generate_extra: bool = True,
+ is_loading: bool = False,
+ ):
+ # For save dist checkpointing
+ state_dict = {}
+
+ # Should always generate model state dict
+ # All ranks Save Model to reduce memory pressure
+ # Get sharded state dict, notice that state_dict will collect among dp groups, causing memory pressure
+ for vpp_rank, model in enumerate(self.model):
+ if len(self.model) > 1:
+ mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank)
+ key = f"model{vpp_rank}" if len(self.model) > 1 else "model"
+ else:
+ key = "model"
+ if hasattr(model, "module"):
+ model = model.module
+
+ # GPTModel's sharded_state_dict function when having mtp requires metadata['dp_cp_group']
+ kwargs = {"metadata": {"dp_cp_group": mpu.get_data_parallel_group(with_context_parallel=True)}}
+ state_dict[key] = model.sharded_state_dict(**kwargs)
+
+ # Optimizer State Dict
+ if generate_optimizer:
+ torch.distributed.barrier()
+ optimizer_sharded_states = self.optimizer.sharded_state_dict(state_dict, is_loading=is_loading)
+ state_dict["optimizer"] = optimizer_sharded_states
+
+ if self.lr_scheduler is not None:
+ lr_state_dict = self.lr_scheduler.state_dict()
+ state_dict["lr_scheduler"] = lr_state_dict
+
+ if not generate_model:
+ state_dict.pop("model", None)
+
+ # RNG States State Dict
+ if generate_extra:
+ torch.distributed.barrier()
+ rng_state = self.get_rng_state()
+ state_dict["rng_state"] = rng_state
+
+ return state_dict
+
+ def load_rng_states(self, rng_states, data_parallel_random_init=False, use_dist_ckpt=True):
+ # access rng_state for data parallel rank
+ if data_parallel_random_init:
+ rng_states = rng_states[mpu.get_data_parallel_rank()]
+ else:
+ rng_states = rng_states[0]
+ random.setstate(rng_states["random_rng_state"])
+ np.random.set_state(rng_states["np_rng_state"])
+ torch.set_rng_state(rng_states["torch_rng_state"])
+
+ if get_device_name() != "cpu":
+ get_torch_device().set_rng_state(rng_states[f"{get_device_name()}_rng_state"])
+
+ # Check for empty states array
+ if not rng_states["rng_tracker_states"]:
+ raise KeyError
+ tensor_parallel.get_cuda_rng_tracker().set_states(rng_states["rng_tracker_states"])
+
+ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load=False):
+ if local_path is not None:
+ assert os.path.exists(local_path), f"Checkpoint path {local_path} does not exist."
+
+ # For load optimizer dist_ckpt
+ try:
+ import transformer_engine
+
+ torch.serialization.add_safe_globals([torch.optim.AdamW])
+ torch.serialization.add_safe_globals([transformer_engine.pytorch.optimizers.fused_adam.FusedAdam])
+ except Exception:
+ pass
+
+ dist_checkpoint_path = get_dist_checkpoint_path(local_path)
+
+ # Get State Dict for loading
+ sharded_state_dict = self.generate_state_dict(
+ self.should_load_model and self.use_dist_checkpointing,
+ self.should_load_optimizer,
+ self.should_load_extra,
+ is_loading=True,
+ )
+ log_with_rank(f"Generated state dict for loading: {sharded_state_dict.keys()}", rank=self.rank, logger=logger)
+
+ # Load Dist Checkpointing
+ state_dict = load_dist_checkpointing(
+ sharded_state_dict=sharded_state_dict,
+ ckpt_dir=dist_checkpoint_path,
+ )
+
+ if self.should_load_model and self.use_dist_checkpointing:
+ assert "model" in state_dict or any(
+ f"model{vpp_rank}" in state_dict for vpp_rank in range(len(self.model))
+ ), f"Model state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}."
+ for vpp_rank, model in enumerate(self.model):
+ if len(self.model) == 1:
+ model_state_dict = state_dict["model"]
+ else:
+ assert f"model{vpp_rank}" in state_dict, f"model{vpp_rank} not found in state_dict"
+ model_state_dict = state_dict[f"model{vpp_rank}"]
+ mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank)
+ self.model[vpp_rank].load_state_dict(model_state_dict)
+ log_with_rank(f"Loaded sharded model checkpoint from {local_path}", rank=self.rank, logger=logger)
+
+ # Skip HF checkpoint loading if PEFT is used
+ elif self.should_load_model and self.use_hf_checkpoint and self.peft_cls is None:
+ hf_model_path = get_hf_model_checkpoint_path(local_path)
+ if self.vanilla_bridge:
+ self.bridge.load_weights(self.model, hf_model_path)
+ else:
+ self.bridge.load_hf_weights(self.model, hf_model_path)
+ log_with_rank(f"Loaded HF model checkpoint from {hf_model_path} with bridge", rank=self.rank, logger=logger)
+ # Load PEFT adapter checkpoint if available
+ if self.should_load_model and self.peft_cls is not None:
+ adapter_ckpt_path = os.path.join(local_path, "adapter_checkpoint")
+ if os.path.exists(adapter_ckpt_path):
+ from verl.utils.megatron_peft_utils import load_adapter_checkpoint
+
+ # TODO: a better format for adapter checkpoint, waiting megatron-bridge support
+
+ load_adapter_checkpoint(
+ self.model,
+ adapter_ckpt_path,
+ )
+ log_with_rank(
+ f"Loaded adapter checkpoint from {adapter_ckpt_path}",
+ rank=self.rank,
+ logger=logger,
+ )
+ else:
+ log_with_rank(
+ f"PEFT config is set but no adapter checkpoint found at {adapter_ckpt_path}",
+ rank=self.rank,
+ logger=logger,
+ )
+
+ if self.should_load_optimizer:
+ assert "optimizer" in state_dict, (
+ f"Optimizer state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}."
+ )
+ optimizer_state_dict = state_dict["optimizer"]
+ self.optimizer.load_state_dict(optimizer_state_dict)
+ log_with_rank(f"Loaded optimizer checkpoint from {local_path}", rank=self.rank, logger=logger)
+ if self.use_checkpoint_opt_param_scheduler:
+ assert "lr_scheduler" in state_dict, (
+ f"LR scheduler state dict not found in {state_dict.keys()}. Please check the checkpoint file "
+ f"{local_path}."
+ )
+ lr_scheduler_state_dict = state_dict["lr_scheduler"]
+ if self.lr_scheduler is not None:
+ self.lr_scheduler.load_state_dict(lr_scheduler_state_dict)
+ log_with_rank(f"Loaded LR scheduler checkpoint from {local_path}", rank=self.rank, logger=logger)
+
+ if self.should_load_extra:
+ assert "rng_state" in state_dict, (
+ f"RNG state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}."
+ )
+ rng_state = state_dict["rng_state"]
+ self.load_rng_states(rng_state)
+ log_with_rank(f"Loaded RNG states from {local_path}", rank=self.rank, logger=logger)
+
+ if del_local_after_load:
+ try:
+ os.remove(local_path) if is_non_local(local_path) else None
+ except Exception as e:
+ log_with_rank(
+ f"remove local resume ckpt file after loading failed, exception {e} will be ignored",
+ rank=self.rank,
+ logger=logger,
+ )
+
+ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep=None):
+ # record the previous global step
+ self.previous_global_step = global_step
+
+ if not self.checkpoint_config.async_save:
+ self.ensure_checkpoint_capacity(max_ckpt_to_keep)
+
+ local_path = local_mkdir_safe(local_path)
+ dist_checkpoint_path = get_dist_checkpoint_path(local_path)
+
+ # Note that model weights, optimizer states, and extra states are generated
+ # together in a state dict, we save them in one time
+ if self.use_dist_checkpointing:
+ # Generate state dict for saving
+ state_dict = self.generate_state_dict(
+ self.should_save_model, self.should_save_optimizer, self.should_save_extra
+ )
+ log_with_rank(f"Generated state dict for saving: {state_dict.keys()}", rank=self.rank, logger=logger)
+ for vpp_rank, model in enumerate(self.model):
+ if len(self.model) > 1:
+ model_i_keys = state_dict[f"model{vpp_rank}"].keys()
+ log_with_rank(f"Generated state dict for saving: {model_i_keys}", rank=self.rank, logger=logger)
+ else:
+ log_with_rank(
+ f"Generated state dict for saving: {state_dict['model'].keys()}", rank=self.rank, logger=logger
+ )
+ # Start Async save if enabled
+ async_save_request = save_dist_checkpointing(
+ sharded_state_dict=state_dict,
+ ckpt_path=dist_checkpoint_path,
+ async_save=self.checkpoint_config.async_save,
+ )
+
+ # Synchronize all async save requests
+ if not self.checkpoint_config.async_save:
+ assert async_save_request is None, "Async save request should be None when not using async save."
+ torch.distributed.barrier()
+ else:
+ assert self.use_hf_checkpoint, "When not using distributed checkpointing, use_hf_checkpoint should be True."
+ # Generate optimizer and exra state dicts
+ state_dict = self.generate_state_dict(
+ generate_model=False,
+ generate_optimizer=self.should_save_optimizer,
+ generate_extra=self.should_save_extra,
+ )
+ # Save optimizer and extra states to local path
+ # Start Async save if enabled
+ async_save_request = save_dist_checkpointing(
+ sharded_state_dict=state_dict,
+ ckpt_path=dist_checkpoint_path,
+ async_save=self.checkpoint_config.async_save,
+ )
+
+ # Synchronize all async save requests
+ if not self.checkpoint_config.async_save:
+ assert async_save_request is None, "Async save request should be None when not using async save."
+ torch.distributed.barrier()
+
+ if self.should_save_model:
+ # Save adapter-only checkpoint if PEFT is enabled
+ if self.peft_cls is not None:
+ from verl.utils.megatron_peft_utils import save_adapter_checkpoint
+
+ adapter_ckpt_path = os.path.join(local_path, "adapter_checkpoint")
+
+ # Save adapter weights only (much smaller than full model)
+ save_adapter_checkpoint(
+ self.model,
+ adapter_ckpt_path,
+ self.rank,
+ )
+
+ log_with_rank(
+ f"Saved adapter-only checkpoint to {adapter_ckpt_path}",
+ rank=self.rank,
+ logger=logger,
+ log_only_rank_0=True,
+ )
+ elif self.use_hf_checkpoint:
+ # Use mbridge to save HF model checkpoint
+ log_with_rank(f"Saving HF model checkpoint to {local_path} with bridge", rank=self.rank, logger=logger)
+ hf_ckpt_path = get_hf_model_checkpoint_path(local_path)
+ if self.vanilla_bridge:
+ self.bridge.save_weights(
+ self.model, hf_ckpt_path, distributed_filesystem=True, memory_efficient=True
+ )
+ else:
+ self.bridge.save_hf_weights(self.model, hf_ckpt_path)
+
+ log_with_rank(f"Saved bridge checkpoint to {hf_ckpt_path}", rank=self.rank, logger=logger)
+
+ # Only rank 0 saves the hf config and tokenizer to huggingface path
+ # No matter whether we save hf model or not
+ if self.rank == 0:
+ # Save tokenizer
+ hf_config_tokenizer_path = get_hf_model_checkpoint_path(local_path)
+ if self.processing_class is not None:
+ self.processing_class.save_pretrained(hf_config_tokenizer_path)
+ # Save huggingface config
+ self.hf_config.save_pretrained(hf_config_tokenizer_path)
+ if hasattr(self.hf_config, "name_or_path") and self.hf_config.name_or_path:
+ try:
+ generation_config = GenerationConfig.from_pretrained(self.hf_config.name_or_path)
+ generation_config.save_pretrained(hf_config_tokenizer_path)
+ except Exception:
+ # if the generation config isn't available, we don't save it
+ pass
+ log_with_rank(
+ f"Saved Huggingface config and tokenizer to {hf_config_tokenizer_path}",
+ rank=self.rank,
+ logger=logger,
+ log_only_rank_0=True,
+ )
+
+ if self.should_save_extra:
+ if self.rank == 0:
+ # Save transformer config
+ print(self.transformer_config)
+ bypass_keys = [
+ "finalize_model_grads_func",
+ "grad_scale_func",
+ "no_sync_func",
+ "grad_sync_func",
+ "param_sync_func",
+ "generation_config",
+ "_pg_collection",
+ ]
+ backup = {}
+ for k in bypass_keys:
+ if hasattr(self.transformer_config, k):
+ backup[k] = getattr(self.transformer_config, k, None)
+ delattr(self.transformer_config, k)
+ transformer_config_dict = asdict(self.transformer_config)
+ for k in backup:
+ setattr(self.transformer_config, k, backup[k])
+ to_convert_types = {torch.dtype: str, AttnBackend: str}
+ ignore_types = [Callable]
+ pop_keys = []
+ for key, value in transformer_config_dict.items():
+ if type(value) in to_convert_types:
+ transformer_config_dict[key] = to_convert_types[type(value)](value)
+ if type(value) in ignore_types:
+ pop_keys.append(key)
+ if callable(value):
+ pop_keys.append(key)
+ for key in pop_keys:
+ transformer_config_dict.pop(key)
+ transformer_config_path = get_transformer_config_checkpoint_path(local_path)
+ with open(transformer_config_path, "w") as f:
+ json.dump(transformer_config_dict, f, indent=2)
+
+ if self.should_save_hf_model and not self.use_hf_checkpoint:
+ # wait for everyone to dump to local
+ if self.bridge is not None:
+ hf_model_ckpt_path = get_hf_model_checkpoint_path(local_path)
+ if self.vanilla_bridge:
+ self.bridge.save_weights(
+ self.model, hf_model_ckpt_path, distributed_filesystem=True, memory_efficient=True
+ )
+ else:
+ self.bridge.save_hf_weights(self.model, hf_model_ckpt_path)
+ else:
+ state_dict = self.weight_saver(
+ self.model,
+ self.hf_config,
+ dtype=self.param_dtype,
+ is_value_model=self.is_value_model,
+ tie_word_embeddings=self.share_embeddings_and_output_weights,
+ )
+
+ torch.distributed.barrier()
+ if self.rank == 0:
+ hf_model_ckpt_path = get_hf_model_checkpoint_path(local_path)
+ import warnings
+
+ from accelerate import init_empty_weights
+
+ with init_empty_weights(), warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ if "mistral7b-rm" in self.config.model.path:
+ from transformers import MistralForSequenceClassification
+
+ model = MistralForSequenceClassification.from_pretrained(
+ self.config.model.path
+ ) # use score head instead of lm_head
+ state_dict["score.weight"] = state_dict["score.weight"]
+ else:
+ from transformers import AutoModelForCausalLM
+
+ model = AutoModelForCausalLM.from_pretrained(self.config.model.path, torch_dtype="auto")
+ model.save_pretrained(hf_model_ckpt_path, state_dict=state_dict)
+ log_with_rank(
+ f"Saved Huggingface config and tokenizer to {hf_model_ckpt_path}",
+ rank=self.rank,
+ logger=logger,
+ log_only_rank_0=True,
+ )
+
+ if hdfs_path is not None:
+ log_with_rank(
+ f"Uploading checkpoint to {hdfs_path}", rank=self.rank, logger=logger, log_only_rank_0=True
+ )
+ from verl.utils import hdfs_io
+
+ hdfs_io.makedirs(hdfs_path, exist_ok=True)
+ hdfs_io.copy(src=hf_model_ckpt_path, dst=hdfs_path, dirs_exist_ok=True)
+ log_with_rank(
+ f"HDFS checkpoint uploaded to {hdfs_path}",
+ rank=self.rank,
+ logger=logger,
+ log_only_rank_0=True,
+ )
+
+ def finalize_save_fn():
+ # Rank 0 uploads checkpoint to HDFS if hdfs_path is provided
+ log_with_rank(
+ f"Dist checkpointing save completed for {dist_checkpoint_path}", rank=self.rank, logger=logger
+ )
+ if self.rank == 0:
+ if hdfs_path is not None:
+ log_with_rank(f"Uploading checkpoint to {hdfs_path}", rank=self.rank, logger=logger)
+ from verl.utils import hdfs_io
+
+ hdfs_io.makedirs(hdfs_path, exist_ok=True)
+ hdfs_io.copy(src=dist_checkpoint_path, dst=hdfs_path, dirs_exist_ok=True)
+ hdfs_io.copy(src=hf_config_tokenizer_path, dst=hdfs_path, dirs_exist_ok=True)
+
+ # update latest_checkpointed_iteration.txt when async_save is True
+ if self.checkpoint_config.async_save and self.rank == 0:
+ log_with_rank(
+ f"Update latest_checkpointed_iteration.txt to step {global_step}",
+ rank=self.rank,
+ logger=logger,
+ )
+ local_latest_checkpointed_iteration = os.path.join(
+ os.path.dirname(os.path.dirname(local_path)), "latest_checkpointed_iteration.txt"
+ )
+ with open(local_latest_checkpointed_iteration, "w") as f:
+ f.write(str(global_step))
+
+ self.register_checkpoint(local_path, max_ckpt_to_keep)
+
+ if self.checkpoint_config.async_save:
+ assert async_save_request is not None, "Async save request should not be None when using async save."
+ async_save_request.add_finalize_fn(finalize_save_fn)
+ from megatron.core.dist_checkpointing.strategies.base import async_calls
+
+ async_calls.schedule_async_request(async_save_request)
+ else:
+ finalize_save_fn()
diff --git a/code/RL_model/verl/verl_train/verl/utils/dataset/README.md b/code/RL_model/verl/verl_train/verl/utils/dataset/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..f886a70aabf443fb167453d667529b62f3311765
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/dataset/README.md
@@ -0,0 +1,16 @@
+# Dataset Format
+## RLHF dataset
+We combine all the data sources into a single parquet files. We directly organize the prompt into the chat format so that multi-turn chats can be easily incorporated. In the prompt, we may add instruction following texts to guide the model output the answers in a particular format so that we can extract the answers.
+
+Math problems
+```json
+{
+ "data_source": "openai/gsm8k",
+ "prompt": [{"role": "user", "content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? Let's think step by step and output the final answer after \"####\""}],
+ "ability": "math",
+ "reward_model": {
+ "style": "rule",
+ "ground_truth": ["72"]
+ },
+}
+```
diff --git a/code/RL_model/verl/verl_train/verl/utils/dataset/__init__.py b/code/RL_model/verl/verl_train/verl/utils/dataset/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6032d68c86423f0e6c57afba684dff5e1b8362c0
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/dataset/__init__.py
@@ -0,0 +1,19 @@
+# Copyright 2024 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 .rl_dataset import RLHFDataset
+from .rm_dataset import RMDataset
+from .sft_dataset import SFTDataset
+
+__all__ = ["RLHFDataset", "RMDataset", "SFTDataset"]
diff --git a/code/RL_model/verl/verl_train/verl/utils/dataset/dataset_utils.py b/code/RL_model/verl/verl_train/verl/utils/dataset/dataset_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..03bde7b01d2e5ec37af8adf535aaaa199ce5f90e
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/dataset/dataset_utils.py
@@ -0,0 +1,75 @@
+# Copyright 2025 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 enum import Enum
+
+import torch
+from tensordict.tensorclass import NonTensorData
+
+
+class DatasetPadMode(str, Enum):
+ """Padding mode for dataset"""
+
+ RIGHT = "right"
+ LEFT_RIGHT = "left_right"
+ NO_PADDING = "no_padding"
+
+
+class SFTTensorCollator:
+ """
+ A custom collate_fn that handles batching of sequences.
+ 1. for variable-length sequences, convert them into NestedTensors.
+ 2. for fixed-length sequences, use default_collate.
+ """
+
+ def __init__(self, pad_mode: DatasetPadMode = DatasetPadMode.LEFT_RIGHT):
+ self.pad_mode = pad_mode
+
+ def __call__(self, batch: list[dict[str, any]]) -> dict[str, any]:
+ if self.pad_mode == DatasetPadMode.NO_PADDING:
+ return self.collate_variable_batch(batch)
+ elif self.pad_mode in [DatasetPadMode.RIGHT, DatasetPadMode.LEFT_RIGHT]:
+ from torch.utils.data import default_collate
+
+ return default_collate(batch)
+ else:
+ raise NotImplementedError(f"pad_mode {self.pad_mode} not implemented")
+
+ def collate_variable_batch(self, batch: list[dict[str, any]]) -> dict[str, any]:
+ """
+ Collates a list of samples into a single batch.
+
+ Args:
+ batch: A list of dictionary samples from the dataset.
+
+ Returns:
+ A dictionary representing the batched data, with variable-length
+ sequences converted to NestedTensors.
+ """
+
+ final_batch = {}
+
+ tensor_keys = set().union(*(d.keys() for d in batch))
+
+ # Handle tensor values by creating a NestedTensor.
+ for key in tensor_keys:
+ if isinstance(batch[0][key], torch.Tensor):
+ tensors = [item[key] for item in batch]
+ final_batch[key] = torch.nested.as_nested_tensor(tensors, layout=torch.jagged)
+ else:
+ tensors = [NonTensorData(item.get(key)) for item in batch]
+ final_batch[key] = torch.stack(tensors, dim=0)
+
+ return final_batch
diff --git a/code/RL_model/verl/verl_train/verl/utils/dataset/multiturn_sft_dataset.py b/code/RL_model/verl/verl_train/verl/utils/dataset/multiturn_sft_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..9da33228e216ba35e4c6b534289e65a029150799
--- /dev/null
+++ b/code/RL_model/verl/verl_train/verl/utils/dataset/multiturn_sft_dataset.py
@@ -0,0 +1,455 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2025 ModelBest Inc. 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.
+"""
+Multi-turn SFT dataset that supports training on conversation data with multiple turns
+"""
+
+import logging
+import os
+import re
+from functools import wraps
+from typing import Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from omegaconf import DictConfig, ListConfig
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer, ProcessorMixin
+
+from verl.models.transformers.qwen2_vl import get_rope_index
+from verl.utils import hf_tokenizer
+from verl.utils.chat_template import extract_system_prompt_and_generation
+from verl.utils.dataset.dataset_utils import DatasetPadMode
+from verl.utils.dataset.vision_utils import process_image, process_video
+from verl.utils.fs import copy_local_path_from_hdfs
+
+logger = logging.getLogger(__file__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
+
+
+def once(func):
+ """Decorator to ensure a function runs only once. Subsequent calls do nothing."""
+
+ @wraps(func)
+ def wrapper(*args, **kwargs):
+ if not hasattr(wrapper, "called"):
+ wrapper.called = True
+ return func(*args, **kwargs)
+
+ return wrapper
+
+
+@once
+def print_assembled_message(tokenizer, message_list, input_ids, loss_mask, attn_mask, tools):
+ """
+ Print the message after applying the chat template
+ """
+
+ tokenized = tokenizer.apply_chat_template(message_list, add_generation_prompt=False, tokenize=False, tools=tools)
+ sep = "\n\n"
+ str = f"tokenized entire message:\n{tokenized}"
+ str += sep
+ str += f"tokenized seperately :\n{tokenizer.decode(input_ids)}"
+
+ logger.debug(str)
+
+
+def convert_nested_value_to_list_recursive(data_item):
+ if isinstance(data_item, dict):
+ return {k: convert_nested_value_to_list_recursive(v) for k, v in data_item.items()}
+ elif isinstance(data_item, list):
+ return [convert_nested_value_to_list_recursive(elem) for elem in data_item]
+ elif isinstance(data_item, np.ndarray):
+ # Convert to list, then recursively process the elements of the new list
+ return convert_nested_value_to_list_recursive(data_item.tolist())
+ else:
+ # Base case: item is already a primitive type (int, str, float, bool, etc.)
+ return data_item
+
+
+class MultiTurnSFTDataset(Dataset):
+ """
+ Dataset for multi-turn conversations where each assistant response should be trained
+
+ Args:
+ data_files (str or list): Path(s) to Parquet file(s).
+ tokenizer (PreTrainedTokenizer): For the tokenization of text to token IDs.
+ config (DictConfig): Options like cache_dir, prompt_key, max_prompt_length, truncation, etc.
+ processor (ProcessorMixin, optional): Multimodal preprocessor for images/videos.
+ max_samples (int, optional): Limit the number of samples. Defaults to -1 (use all).
+ """
+
+ def __init__(
+ self,
+ parquet_files: str | list[str],
+ tokenizer: PreTrainedTokenizer,
+ config: DictConfig,
+ processor: Optional[ProcessorMixin] = None,
+ max_samples: int = -1,
+ ):
+ # Set defaults and extract parameters from config if provided
+ config = config or {}
+ self.pad_mode = config.get("pad_mode", "right")
+ assert self.pad_mode in ["right", "no_padding"], (
+ f"Expect pad_mode to be 'right' or 'no_padding'. Got {self.pad_mode}"
+ )
+ self.truncation = config.get("truncation", "error")
+ # for right padding
+ self.max_length = config.get("max_length", 1024)
+ # Get messages_key from the new multiturn config structure
+ self.messages_key = config.get("messages_key", "messages")
+ self.image_key = config.get("image_key", "images")
+ self.video_key = config.get("video_key", "videos")
+ self.image_patch_size = config.get(
+ "image_patch_size", processor.image_processor.patch_size if processor else None
+ )
+ self.tools_key = config.get("tools_key", "tools")
+ self.enable_thinking_key = config.get("enable_thinking_key", "enable_thinking")
+ self.enable_thinking_default = config.get("enable_thinking_default", None)
+ self.apply_chat_template_kwargs = config.get("apply_chat_template_kwargs", {})
+ self.shuffle = config.get("shuffle", False)
+ self.seed = config.get("seed")
+ self.max_samples = max_samples
+ self.ignore_input_ids_mismatch = config.get("ignore_input_ids_mismatch", False)
+ assert self.truncation in ["error", "left", "right"]
+
+ if not isinstance(parquet_files, list | ListConfig):
+ parquet_files = [parquet_files]
+
+ self.parquet_files = parquet_files
+ if isinstance(tokenizer, str):
+ tokenizer = hf_tokenizer(tokenizer)
+ self.tokenizer: PreTrainedTokenizer = tokenizer
+ self.processor = processor
+
+ self._download()
+ self._read_files_and_process()
+
+ def _download(self):
+ for i, parquet_file in enumerate(self.parquet_files):
+ self.parquet_files[i] = copy_local_path_from_hdfs(parquet_file, verbose=True)
+
+ def _read_files_and_process(self):
+ def series_to_item(ls):
+ import numpy
+ import pandas
+
+ while isinstance(ls, pandas.core.series.Series | numpy.ndarray) and len(ls) == 1:
+ ls = ls[0]
+ return ls
+
+ dataframes = []
+ for parquet_file in self.parquet_files:
+ # default loader loads some list as np.ndarray, which fails the tokenizer
+ dataframe = pd.read_parquet(parquet_file, dtype_backend="pyarrow")
+ dataframes.append(dataframe)
+ self.dataframe = pd.concat(dataframes)
+
+ total = len(self.dataframe)
+ print(f"dataset len: {len(self.dataframe)}")
+
+ if self.max_samples > 0 and self.max_samples < total:
+ if self.shuffle:
+ rngs_args = (self.seed,) if self.seed is not None else ()
+ rng = np.random.default_rng(*rngs_args)
+ indices = rng.choice(total, size=self.max_samples, replace=False)
+ else:
+ indices = np.arange(self.max_samples)
+ self.dataframe = self.dataframe.iloc[indices.tolist()]
+ print(f"selected {self.max_samples} random samples out of {total}")
+
+ # Extract messages list from dataframe
+ self.messages = self.dataframe[self.messages_key].apply(convert_nested_value_to_list_recursive).tolist()
+
+ # Extract tools list from dataframe
+ if self.tools_key in self.dataframe.columns:
+ self.tools = self.dataframe[self.tools_key].apply(convert_nested_value_to_list_recursive).tolist()
+ else:
+ self.tools = None
+ # Extract enable_thinking list from dataframe
+ if self.enable_thinking_key in self.dataframe.columns:
+ self.enable_thinking = self.dataframe[self.enable_thinking_key].tolist()
+ else:
+ self.enable_thinking = None
+
+ # system prompt: <|im_start|>system\nYou are a helpful assistant.<|im_end|>\n
+ # generation prompt: <|im_start|>assistant\n
+ self.system_prompt, self.generation_prompt = extract_system_prompt_and_generation(self.tokenizer)
+
+ def __len__(self):
+ return len(self.messages)
+
+ def _process_single_message(
+ self,
+ index: int,
+ message: dict[str, Any],
+ full_message: list,
+ tools: Optional[list[dict[str, Any]]] = None,
+ enable_thinking: Optional[bool] = None,
+ ) -> tuple[list[int], list[int], list[int]]:
+ """
+ Process a single message and return its tokenized representation.
+
+ Args:
+ index: turn index in the conversation
+ message: A single message dictionary
+ images: List of images to be used
+ videos: List of videos to be used
+ tools: List of tools to be used
+ enable_thinking: Whether to enable thinking mode
+
+ Returns:
+ Tuple of (input_ids, loss_mask, attention_mask, dict[str, torch.Tensor])
+ """
+ processor = self.processor if self.processor is not None else self.tokenizer
+ apply_chat_template_kwargs = {**self.apply_chat_template_kwargs}
+ if enable_thinking is not None:
+ apply_chat_template_kwargs["enable_thinking"] = enable_thinking
+
+ inputs = processor.apply_chat_template(
+ [message],
+ tools=tools,
+ add_generation_prompt=False,
+ tokenize=True,
+ return_dict=True,
+ return_tensors="pt",
+ **apply_chat_template_kwargs,
+ )
+
+ inputs = dict(inputs)
+ input_ids = inputs.pop("input_ids")[0]
+ attention_mask = inputs.pop("attention_mask")[0]
+
+ # remove system prompt if exists
+ if index != 0 and message["role"] != "system":
+ input_ids = input_ids[len(self.system_prompt) :]
+ attention_mask = attention_mask[len(self.system_prompt) :]
+
+ if message["role"] == "assistant":
+ loss_mask = torch.ones_like(attention_mask)
+ # mask out generation prompt if assistant message
+ loss_mask[: len(self.generation_prompt)] = 0
+ else:
+ loss_mask = torch.zeros_like(attention_mask)
+
+ return input_ids, loss_mask, attention_mask, inputs
+
+ def _build_messages(self, example: dict):
+ """Replace and