apply review suggestions

Author: ji-huazhong

recipe/transfer_queue/fsdp_workers.py

@@ -113,10 +113,7 @@ 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 .fsdp_workers import (
-                ActorRolloutRefWorker,
-                AsyncActorRolloutRefWorker,
-            )
+            from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker
 
             actor_rollout_cls = (
                 AsyncActorRolloutRefWorker
@@ -126,10 +123,7 @@ def add_actor_rollout_worker(self, config):
             ray_worker_group_cls = RayWorkerGroup
 
         elif config.actor_rollout_ref.actor.strategy == "megatron":
-            from .megatron_workers import (
-                ActorRolloutRefWorker,
-                AsyncActorRolloutRefWorker,
-            )
+            from verl.workers.megatron_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker
 
             actor_rollout_cls = (
                 AsyncActorRolloutRefWorker
@@ -152,7 +146,7 @@ def add_critic_worker(self, config):
         if config.critic.strategy in {"fsdp", "fsdp2"}:
             use_legacy_worker_impl = config.trainer.get("use_legacy_worker_impl", "auto")
             if use_legacy_worker_impl in ["auto", "enable"]:
-                from .fsdp_workers import CriticWorker
+                from verl.workers.fsdp_workers import CriticWorker
             elif use_legacy_worker_impl == "disable":
                 from verl.workers.roles import CriticWorker
 
@@ -161,7 +155,7 @@ def add_critic_worker(self, config):
                 raise ValueError(f"Invalid use_legacy_worker_impl: {use_legacy_worker_impl}")
 
         elif config.critic.strategy == "megatron":
-            from .megatron_workers import CriticWorker
+            from verl.workers.megatron_workers import CriticWorker
 
         else:
             raise NotImplementedError
@@ -203,9 +197,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 .fsdp_workers import RewardModelWorker
+                    from verl.workers.fsdp_workers import RewardModelWorker
                 elif config.reward_model.strategy == "megatron":
-                    from .megatron_workers import RewardModelWorker
+                    from verl.workers.megatron_workers import RewardModelWorker
                 else:
                     raise NotImplementedError
             elif use_legacy_worker_impl == "disable":

recipe/transfer_queue/main_ppo.py

@@ -23,5 +23,4 @@ uvicorn
 fastapi
 latex2sympy2_extended
 math_verify
-tensorboard
-git+https://github.com/TransferQueue/TransferQueue.git@a2ddb30
+tensorboard

recipe/transfer_queue/megatron_workers.py

@@ -0,0 +1,2 @@
+# requirements.txt records the full set of dependencies for development
+git+https://github.com/TransferQueue/TransferQueue.git@a2ddb30

requirements.txt

@@ -42,7 +42,6 @@
     "wandb",
     "packaging>=20.0",
     "tensorboard",
-    "TransferQueue @ git+https://github.com/TransferQueue/TransferQueue.git@a2ddb30",
 ]
 
 TEST_REQUIRES = ["pytest", "pre-commit", "py-spy", "pytest-asyncio"]
@@ -58,6 +57,7 @@
 ]
 TRL_REQUIRES = ["trl<=0.9.6"]
 MCORE_REQUIRES = ["mbridge"]
+TRANSFERQUEUE_REQUIRES = ["TransferQueue @ git+https://github.com/TransferQueue/TransferQueue.git@a2ddb30"]
 
 extras_require = {
     "test": TEST_REQUIRES,
@@ -69,6 +69,7 @@
     "sglang": SGLANG_REQUIRES,
     "trl": TRL_REQUIRES,
     "mcore": MCORE_REQUIRES,
+    "transferqueue": TRANSFERQUEUE_REQUIRES,
 }
 
 

requirements_transferqueue.txt

@@ -129,6 +129,16 @@ 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=True)
+    def create_transferqueue_client(self, controller_infos, storage_infos, role="train"):
+        from verl.utils.transferqueue_utils import create_transferqueue_client
+
+        create_transferqueue_client(
+            client_id=f"{role}_worker_{self.rank}",
+            controller_infos=controller_infos,
+            storage_infos=storage_infos,
+        )
 
     @classmethod
     def env_keys(cls):