[bugfix] fix grpo move_model_batches (#8091)

Author: hjh0119

swift/trainers/rlhf_trainer/rollout_mixin.py

@@ -647,31 +647,34 @@ def _collect_state_dict_for_vllm(self, parameter_group=None, parameter_group_no_
     def _move_full_model_to_vllm(self):
         """Transfer full model weights to vLLM engine.
 
-        Manages the lifecycle of gather and merge/unmerge:
-        - gather_if_zero3: once for the entire sync (DeepSpeed Zero3)
+        Manages the lifecycle of gather and merge/unmerge per parameter_group:
+        - gather_if_zero3: per parameter_group batch (DeepSpeed Zero3)
         - merge/unmerge: per parameter_group (must be within gather context)
         - No clone needed: unmerge happens after load completes
         """
         is_peft = is_peft_model(self.model)
-        # For DeepSpeed, merge within gather context; FSDP2 uses tensor-level merge
         should_merge = is_peft and not self._is_fsdp2
 
         gather_if_zero3 = get_gather_if_zero3_context(self)
-        parameters = [] if self._is_fsdp2 else list(self.model.parameters())
 
-        with gather_if_zero3(parameters):
-            for i, parameter_group in enumerate(self.parameter_groups):
-                parameter_group_no_lora = self.parameter_groups_no_lora[i]
+        for i, parameter_group in enumerate(self.parameter_groups):
+            parameter_group_no_lora = self.parameter_groups_no_lora[i]
+
+            if not self._is_fsdp2:
+                parameters = [
+                    parameter for name, parameter in self.model.named_parameters()
+                    if not parameter_group or name in parameter_group
+                ]
+            else:
+                parameters = []
 
-                # Merge must be within gather context (needs full parameters)
+            with gather_if_zero3(parameters):
                 if should_merge:
                     with patch_lora_merge(self.model, parameter_group):
                         self.model.merge_adapter()
 
                 try:
-                    # Collect without clone - unmerge happens after load
                     state_dict = self._collect_state_dict_for_vllm(parameter_group, parameter_group_no_lora)
-                    # Data is copied here (FlattenedTensorBucket.copy_ or vLLM load_weights)
                     self._load_state_dict_to_vllm(state_dict)
                 finally:
                     if should_merge: