fix blocking flag in fsdp workers

volcengine · Dec 11, 2024 · 5406550 · 5406550
1 parent 3919ab1
commit 5406550
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/verl/trainer/ppo/workers/fsdp_workers.py b/verl/trainer/ppo/workers/fsdp_workers.py
@@ -304,7 +304,7 @@ def init_model(self):
 
         torch.cuda.empty_cache()
 
-    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO, blocking=False)
+    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
     def update_actor(self, data: DataProto):
         data = data.to('cuda')
 
@@ -592,7 +592,7 @@ def compute_values(self, data: DataProto):
         torch.cuda.empty_cache()
         return output
 
-    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO, blocking=False)
+    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
     def update_critic(self, data: DataProto):
         data = data.to('cuda')
         if self._is_offload_param: