[grpo] fix incorrect placement of data in eval_queue during async_generate (#3573)

* fix

* move to cur_queue to property

* fix

---------

Co-authored-by: hjh <[email protected]>

Author: hjh0119

swift/trainers/rlhf_trainer/grpo_trainer.py

@@ -102,7 +102,6 @@ def __init__(self,
         from swift.trainers.rlhf_arguments import GRPOConfig
         args: GRPOConfig = kwargs['args']
         self.args = args
-        self.queue = None
         self.train_queue = Queue()
         self.eval_queue = Queue()
         self.processing_class = kwargs.get('template').tokenizer
@@ -598,7 +597,7 @@ def patch_merge(model):
                     unwrapped_model.unmerge_adapter()
 
     def _wait_queue(self):
-        while self.queue.empty():
+        while self._queue.empty():
             time.sleep(0.01)
 
     @staticmethod
@@ -621,9 +620,11 @@ def infer_task():
                 return result
 
         future: Future = self.executor.submit(infer_task)
+        # pre-fetch the queue to avoid switching back to eval_queue at the end of training sample sampling
+        current_queue = self._queue
 
         def done(_self):
-            self.queue.put(DataCache(inputs, _self.result(), distributed_idx))
+            current_queue.put(DataCache(inputs, _self.result(), distributed_idx))
 
         future.add_done_callback(done)
 
@@ -634,9 +635,9 @@ def _prefetch(self, dataloader):
         if self.infer_rank >= 0:
             _input_slice = np.array(all_inputs)[distributed_idx[self.infer_rank]]
             outputs = self.engine.infer(_input_slice, self.request_config, use_tqdm=False)
-            self.queue.put(DataCache(inputs, outputs, distributed_idx))
+            self._queue.put(DataCache(inputs, outputs, distributed_idx))
         else:
-            self.queue.put(DataCache(inputs, [], distributed_idx))
+            self._queue.put(DataCache(inputs, [], distributed_idx))
         if self.accelerator.num_processes > 1:
             self.accelerator.wait_for_everyone()
 
@@ -666,7 +667,7 @@ def _fast_infer(self, inputs):
             _input_slice = np.array(all_inputs)[distributed_idx[self.infer_rank]]
             if self.args.async_generate:
                 self.async_infer(inputs, _input_slice, distributed_idx)
-                data_cache = self.queue.get()
+                data_cache = self._queue.get()
                 inputs = data_cache.inputs
                 outputs = data_cache.outputs
                 distributed_idx = data_cache.distributed_idx
@@ -690,8 +691,8 @@ def _fast_infer(self, inputs):
         else:
             if self.args.async_generate:
                 # using old model to generate, which will ignore the `clip` of advantages.
-                self.queue.put(DataCache(inputs, [], distributed_idx))
-                data_cache = self.queue.get()
+                self._queue.put(DataCache(inputs, [], distributed_idx))
+                data_cache = self._queue.get()
                 inputs = data_cache.inputs
                 distributed_idx = data_cache.distributed_idx
             outputs = []
@@ -907,12 +908,17 @@ def _get_per_token_logps(self, model, inputs):
         return selective_log_softmax(logits, input_ids)  # compute logprobs for the input tokens
 
     def evaluation_loop(self, dataloader, *args, **kwargs):
-        self.queue = self.eval_queue
-        if self.queue.empty() and self.args.async_generate:
+        if self._queue.empty() and self.args.async_generate:
             self._prefetch(dataloader)
         metric_key_prefix = kwargs['metric_key_prefix']
         output = super().evaluation_loop(dataloader, *args, **kwargs)
         metrics = {f'{metric_key_prefix}_{key}': sum(val) / len(val) for key, val in self._metrics['eval'].items()}
         output.metrics.update(metrics)
-        self.queue = self.train_queue
         return output
+
+    @property
+    def _queue(self):
+        if self.control.should_evaluate:
+            return self.eval_queue
+        else:
+            return self.train_queue