update

Author: ji-huazhong

recipe/transfer_queue/main_ppo.py

@@ -113,7 +113,10 @@ def add_actor_rollout_worker(self, config):
         from verl.single_controller.ray import RayWorkerGroup
 
         if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}:
-            from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker
+            from .fsdp_workers import (
+                ActorRolloutRefWorker,
+                AsyncActorRolloutRefWorker,
+            )
 
             actor_rollout_cls = (
                 AsyncActorRolloutRefWorker
@@ -123,7 +126,10 @@ def add_actor_rollout_worker(self, config):
             ray_worker_group_cls = RayWorkerGroup
 
         elif config.actor_rollout_ref.actor.strategy == "megatron":
-            from verl.workers.megatron_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker
+            from .megatron_workers import (
+                ActorRolloutRefWorker,
+                AsyncActorRolloutRefWorker,
+            )
 
             actor_rollout_cls = (
                 AsyncActorRolloutRefWorker
@@ -197,9 +203,9 @@ def add_reward_model_worker(self, config):
             use_legacy_worker_impl = config.trainer.get("use_legacy_worker_impl", "auto")
             if use_legacy_worker_impl in ["auto", "enable"]:
                 if config.reward_model.strategy in {"fsdp", "fsdp2"}:
-                    from verl.workers.fsdp_workers import RewardModelWorker
+                    from .fsdp_workers import RewardModelWorker
                 elif config.reward_model.strategy == "megatron":
-                    from verl.workers.megatron_workers import RewardModelWorker
+                    from .megatron_workers import RewardModelWorker
                 else:
                     raise NotImplementedError
             elif use_legacy_worker_impl == "disable":

recipe/transfer_queue/ray_trainer.py

@@ -847,7 +847,7 @@ def init_workers(self):
         self.actor_rollout_wg = all_wg["actor_rollout"]
         self.actor_rollout_wg.init_model()
 
-        # set transferqueue server info for each worker group
+        # set transferqueue server info for each worker
         for _, wg in all_wg.items():
             wg.set_transferqueue_server_info(self.data_system_controller_infos, self.data_system_storage_unit_infos)
 

verl/single_controller/base/worker.py

@@ -129,20 +129,6 @@ def _query_collect_info(self, mesh_name: str):
         """
         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]
-    
-    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
-    def set_transferqueue_server_info(self, controller_infos, storage_infos):
-        """Set the transfer queue server information for the worker.
-
-        Args:
-            controller_infos (list):
-                List of controller server information.
-            storage_infos (list):
-                List of storage unit server information.
-        """
-        from verl.utils.transferqueue_utils import set_transferqueue_server_info
-
-        set_transferqueue_server_info(controller_infos, storage_infos)
 
     @classmethod
     def env_keys(cls):

verl/utils/transferqueue_utils.py

@@ -12,13 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import inspect
 from functools import wraps
 from typing import Any
 
+import numpy as np
+import torch
+from tensordict import NonTensorData, NonTensorStack, TensorDict
 from transfer_queue import BatchMeta
 
 from verl.experimental.transfer_queue import ZMQServerInfo
+from verl.protocol import DataProto
 
 _TRANSFER_QUEUE_CONTROLLER_INFOS = None
 _TRANSFER_QUEUE_STORAGE_INFOS = None
@@ -50,15 +55,55 @@ def _find_batchmeta(*args, **kwargs):
 
 
 def _batchmeta_to_dataproto(batchmeta: BatchMeta):
-    ...
+    tensordict = asyncio.run(client.async_get_data(batchmeta))
+
+    batch = {}
+    non_tensor_batch = {}
+    batch_size = None
+    for k, v in tensordict.items():
+        if isinstance(v, torch.Tensor):
+            batch[k] = v
+            if batch_size is None:
+                batch_size = v.shape[:1]
+        elif isinstance(v, NonTensorStack):
+            non_tensor_batch[k] = np.array([elem.data for elem in v], dtype=object)
+        else:
+            non_tensor_batch[k] = v
+    return DataProto(
+        batch=TensorDict(batch, batch_size=batch_size),
+        non_tensor_batch=non_tensor_batch,
+        meta_info=batchmeta.extra_info.copy(),
+    )
+
+
+def _dataproto_to_tensordict(data: DataProto):
+    result_dict = {}
 
+    if data.batch is not None:
+        result_dict.update(data.batch)
 
-def _update_batchmeta_with_output(output, batchmeta: BatchMeta):
-    ...
+    batch_size = data.batch.batch_size if data.batch is not None else (len(list(data.non_tensor_batch.values())[0]),)    
+    if data.non_tensor_batch is not None:
+        for k, v in data.non_tensor_batch.items():
+            result_dict[k] = NonTensorData(data=v, batch_size=batch_size)
+    
+    if data.meta_info == {} or data.meta_info is None:
+        result_dict["meta_info"] = NonTensorData(data=[None] * batch_size[0], batch_size=batch_size)
+    else:
+        result_dict["meta_info"] = NonTensorData(data=[data.meta_info] * batch_size[0], batch_size=batch_size)
+    return TensorDict(result_dict, batch_size=batch_size)
+
+
+def _update_batchmeta_with_output(output: DataProto, batchmeta: BatchMeta):
+    tensordict = _dataproto_to_tensordict(output)
+    batchmeta.add_fields(tensordict)
+    asyncio.run(client.async_put(data=tensordict, metadata=batchmeta))
 
 
 async def _async_update_batchmeta_with_output(output, batchmeta: BatchMeta):
-    ...
+    tensordict = _dataproto_to_tensordict(output)
+    batchmeta.add_fields(tensordict)
+    await client.async_put(data=tensordict, metadata=batchmeta)
 
 
 def batchmeta_dataproto_pipe():
@@ -90,4 +135,3 @@ async def async_inner(*args, **kwargs):
         wrapper = async_inner if inspect.iscoroutinefunction(func) else inner
         return wrapper
     return decorator
-