[grpo] update vllm weight sync & wake up (#4770)

hjh0119 · web-flow · commit 348b11f90771 · 2025-07-01T10:58:07.000+08:00
* less memory during load and rollout

* deprecate gc_collect_after_offload

* offload context

* fix without optimizer

* rm gc_collect_after_offload in docs

* rm gc_collect_after_offload in scripts
diff --git a/docs/source/BestPractices/Qwen3最佳实践.md b/docs/source/BestPractices/Qwen3最佳实践.md
@@ -312,7 +312,6 @@ swift rlhf \
     --sleep_level 1 \
     --offload_model true \
     --offload_optimizer true \
-    --gc_collect_after_offload true \
     --deepspeed zero3 \
     --num_infer_workers 8 \
     --tensor_parallel_size 1 \
diff --git a/docs/source/Instruction/GRPO/DeveloperGuide/奖励模型.md b/docs/source/Instruction/GRPO/DeveloperGuide/奖励模型.md
@@ -80,7 +80,6 @@ swift rlhf \
     --sleep_level 1 \
     --offload_model true \
     --offload_optimizer true \
-    --gc_collect_after_offload true \
     --log_completions true \
     --deepspeed zero2
 ```
diff --git a/docs/source/Instruction/GRPO/GetStarted/GRPO.md b/docs/source/Instruction/GRPO/GetStarted/GRPO.md
@@ -154,7 +154,6 @@ GRPO 训练框架支持集成高性能推理引擎（如 vLLM）来加速采样
 ```bash
 --offload_optimizer true \
 --offload_model true \
---gc_collect_after_offload true \
 ```
 
 4. 在vLLM中使用 Tensor Parallel 技术：
diff --git a/docs/source/Instruction/命令行参数.md b/docs/source/Instruction/命令行参数.md
@@ -476,7 +476,6 @@ reward模型参数将在PPO、GRPO中使用。
   - sleep_level: 训练时释放 vLLM 显存，可选项为[0, 1], 默认为0，不释放
   - offload_optimizer: 是否在vLLM推理时offload optimizer参数，默认为False。
   - offload_model: 是否在vLLM推理时 offload 模型，默认为False。
-  - gc_collect_after_offload: 是否在offload结束时进行gc（python gc和GPU gc），默认为False。
   - completion_length_limit_scope: 在多轮对话中，`max_completion_length` 的限制范围。
   `total`限制所有对话轮次的总输出长度不超过`max_completion_length`, `per_round`限制每一轮的输出长度。
 - num_iterations: 每个批次代更新次数，默认为1。
diff --git a/docs/source_en/BestPractices/Qwen3-Best-Practice.md b/docs/source_en/BestPractices/Qwen3-Best-Practice.md
@@ -316,7 +316,6 @@ swift rlhf \
     --sleep_level 1 \
     --offload_model true \
     --offload_optimizer true \
-    --gc_collect_after_offload true \
     --deepspeed zero3 \
     --num_infer_workers 8 \
     --tensor_parallel_size 1 \
diff --git a/docs/source_en/Instruction/Command-line-parameters.md b/docs/source_en/Instruction/Command-line-parameters.md
@@ -488,7 +488,6 @@ The meanings of the following parameters can be referenced [here](https://huggin
   - sleep_level: make vllm sleep when model is training. Options are 0 or 1, default is 0, no sleep
   - offload_optimizer: Whether to offload optimizer parameters during inference with vLLM. The default is `False`.
   - offload_model: Whether to offload the model during inference with vLLM. The default is `False`.
-  - gc_collect_after_offload: Whether to perform garbage collection (both Python GC and GPU GC) after offloading. The default is `False`.
   - completion_length_limit_scope: Specifies the scope of the `max_completion_length` limit in multi-turn conversations.
   When set to `total`, the total output length across all turns must not exceed `max_completion_length`.
   When set to `per_round`, each individual turn's output length is limited separately.
diff --git a/docs/source_en/Instruction/GRPO/DeveloperGuide/reward_model.md b/docs/source_en/Instruction/GRPO/DeveloperGuide/reward_model.md
@@ -79,7 +79,6 @@ swift rlhf \
     --sleep_level 1 \
     --offload_model true \
     --offload_optimizer true \
-    --gc_collect_after_offload true \
     --log_completions true \
     --deepspeed zero2
 ```
diff --git a/docs/source_en/Instruction/GRPO/GetStarted/GRPO.md b/docs/source_en/Instruction/GRPO/GetStarted/GRPO.md
@@ -153,7 +153,6 @@ When running in Colocate mode, out-of-memory (OOM) issues may frequently occur.
 ```bash
 --offload_optimizer true \
 --offload_model true \
---gc_collect_after_offload true \
 ```
 
 4. Use Tensor Parallelism in vLLM:
diff --git a/examples/train/grpo/internal/vllm_72b_4gpu.sh b/examples/train/grpo/internal/vllm_72b_4gpu.sh
@@ -8,7 +8,7 @@ swift rlhf \
     --train_type lora \
     --use_vllm true \
     --vllm_mode colocate \
-    --vllm_gpu_memory_utilization 0.5 \
+    --vllm_gpu_memory_utilization 0.55 \
     --vllm_max_model_len 2048 \
     --vllm_tensor_parallel_size 4 \
     --dataset AI-MO/NuminaMath-TIR#10000 \
@@ -39,5 +39,4 @@ swift rlhf \
     --move_model_batches 16 \
     --offload_optimizer true \
     --offload_model true \
-    --gc_collect_after_offload true \
     --sleep_level 1
diff --git a/examples/train/grpo/internal/vllm_lora_qwenvl72b.sh b/examples/train/grpo/internal/vllm_lora_qwenvl72b.sh
@@ -42,7 +42,6 @@ swift rlhf \
   --async_generate false \
   --offload_optimizer true \
   --offload_model true \
-  --gc_collect_after_offload true \
   --move_model_batches 40 \
   --sleep_level 1 \
   --report_to wandb \
diff --git a/examples/train/grpo/internal/vllm_multi_turn.sh b/examples/train/grpo/internal/vllm_multi_turn.sh
@@ -35,7 +35,6 @@ swift rlhf \
     --async_generate false \
     --offload_optimizer true \
     --offload_model true \
-    --gc_collect_after_offload true \
     --sleep_level 1 \
     --multi_turn_scheduler math_tip_trick_multi_turn \
     --max_turns 3
diff --git a/examples/train/grpo/plugin/run_external_reward_model.sh b/examples/train/grpo/plugin/run_external_reward_model.sh
@@ -18,6 +18,5 @@ swift rlhf \
     --sleep_level 1 \
     --offload_model true \
     --offload_optimizer true \
-    --gc_collect_after_offload true \
     --log_completions true \
     --deepspeed zero2
diff --git a/examples/train/long_text/sequence_parallel_grpo.sh b/examples/train/long_text/sequence_parallel_grpo.sh
@@ -39,6 +39,5 @@ swift rlhf \
     --offload_model true \
     --padding_free true \
     --sequence_parallel_size 4 \
-    --gc_collect_after_offload true \
     --dataloader_drop_last true \
     --sleep_level 1
diff --git a/swift/llm/argument/rlhf_args.py b/swift/llm/argument/rlhf_args.py
@@ -329,3 +329,7 @@ def _deprecated_warning(self):
                            "Please use 'multi_turn_scheduler' instead")
 
             self.multi_turn_scheduler = self.multi_turn_func
+
+        if self.gc_collect_after_offload:
+            logger.warning(
+                "The parameter 'gc_collect_after_offload' has been deprecated and will be removed in version 3.7. ")
diff --git a/swift/trainers/arguments.py b/swift/trainers/arguments.py
@@ -201,7 +201,7 @@ class GRPOArgumentsMixin:
     move_model_batches: Optional[int] = None
     offload_optimizer: bool = False
     offload_model: bool = False
-    gc_collect_after_offload: bool = False
+    gc_collect_after_offload: bool = False  # deprecated
 
     # multi turn
     multi_turn_func: Optional[str] = None  # deprecated
diff --git a/swift/trainers/rlhf_trainer/grpo_trainer.py b/swift/trainers/rlhf_trainer/grpo_trainer.py
@@ -39,7 +39,7 @@
 from swift.llm.template.template_inputs import StdTemplateInputs
 from swift.plugin import loss_scale_map, multi_turns, orms, rm_plugins
 from swift.plugin.multi_turn import MultiTurnScheduler
-from swift.utils import (JsonlWriter, gc_collect, get_current_device, get_device, get_logger, is_vllm_available,
+from swift.utils import (JsonlWriter, empty_cache, get_current_device, get_device, get_logger, is_vllm_available,
                          is_wandb_available, seed_worker, unwrap_model_for_generation)
 from ..mixin import SwiftMixin
 from .rlhf_mixin import RLHFTrainerMixin
@@ -230,6 +230,7 @@ def __init__(self,
             self.parameter_groups, self.parameter_groups_no_lora = self.split_batches()
         self.use_fast_infer = self.use_vllm  # whether to use the PT backend
         self.vllm_use_async_engine = False
+        self.enable_offload = False
         if self.use_vllm:
             if not is_vllm_available():
                 raise ImportError('vLLM is not available and `use_vllm` is set to True. '
@@ -255,8 +256,14 @@ def __init__(self,
                         list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size))
                         for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size)
                     ])
+                self.enable_offload = self.args.offload_model or self.args.offload_optimizer
+                context = self.offload_context if self.enable_offload else nullcontext
+
+                with context():
+                    self.engine = self.prepare_vllm(model)
+                    if self.args.sleep_level > 0:
+                        self.engine.engine.sleep(self.args.sleep_level)
 
-                self.engine = self.prepare_vllm(model)
         else:
             from swift.llm import PtEngine
             self.engine = PtEngine.from_model_template(self.model, self.template, max_batch_size=0)  # 0: no limit
@@ -513,6 +520,7 @@ def _move_model_to_vllm(self, skip_async_check=False):
                         llm_model.load_weights(state_dict.items())
                     with patch_lora_unmerge(self.model):
                         self.model.unmerge_adapter()
+                    del state_dict
         else:
             for name, param in self.model.named_parameters():
                 with gather_if_zero3([param]):
@@ -774,53 +782,54 @@ def _prefetch(self, dataloader: DataLoader):
         self._queue.put(DataCache(all_inputs, outputs))
 
     def _fast_infer(self, inputs: InputsType) -> Tuple[InputsType, OutputsType]:
+        # Skip the first wake_up to avoid the warning "Executor is not sleeping"
+
         if self.vllm_mode == 'colocate' and self.args.sleep_level > 0:
-            if self.args.offload_model:
-                self.offload_model(self.accelerator.unwrap_model(self.model))
-                if self.ref_model:
-                    self.offload_model(self.ref_model)
-            if self.args.offload_optimizer:
-                self.offload_optimizer()
-            if self.args.gc_collect_after_offload:
-                gc_collect()
-            # Skip the first wake_up to avoid the warning "Executor is not sleeping"
             if self.engine.inner_model_executor.is_sleeping:
-                self.engine.engine.wake_up()
+                # First, load weights only, https://github.com/vllm-project/vllm/pull/15500
+                if 'tags' in inspect.signature(self.engine.engine.wake_up).parameters:
+                    self.engine.engine.wake_up(tags=['weights'])
+                else:
+                    logger.info('We recommend installing vLLM >= 0.8.3, (ideally 0.8.5.post1)'
+                                'to help reduce memory peaks during engine wake-up.')
+                    self.engine.engine.wake_up()
+
         # First, have main process load weights if needed
         if self.state.global_step != self._last_loaded_step:
             self._move_model_to_vllm()
             self._last_loaded_step = self.state.global_step
 
-        if self.async_generate:
-            # send this step data to server
-            # we gather inputs outside the thread for prevent potential gather deadlock
-            all_inputs = gather_object(inputs)
-            self.async_infer(all_inputs)
-            # cached data from last step
-            data_cache = self._queue.get()
-            all_inputs = data_cache.inputs
-            all_outputs = gather_object(data_cache.outputs)
-            process_slice = slice(
-                self.accelerator.process_index * len(inputs),
-                (self.accelerator.process_index + 1) * len(inputs),
-            )
-            inputs = all_inputs[process_slice]
-            outputs = all_outputs[process_slice]
+        context = self.offload_context if self.enable_offload else nullcontext
+        with context():
+            if self.vllm_mode == 'colocate' and self.engine.inner_model_executor.is_sleeping and \
+                    'tags' in inspect.signature(self.engine.engine.wake_up).parameters:
+                # Load the kv_cache only after updating and offload the weights.
+                self.engine.engine.wake_up(tags=['kv_cache'])
 
-        else:
-            with self.multi_turn_completion_length_context():
-                outputs = self._infer_single_or_multi_turn(inputs, self.request_config)
+            if self.async_generate:
+                # send this step data to server
+                # we gather inputs outside the thread for prevent potential gather deadlock
+                all_inputs = gather_object(inputs)
+                self.async_infer(all_inputs)
+                # cached data from last step
+                data_cache = self._queue.get()
+                all_inputs = data_cache.inputs
+                all_outputs = gather_object(data_cache.outputs)
+                process_slice = slice(
+                    self.accelerator.process_index * len(inputs),
+                    (self.accelerator.process_index + 1) * len(inputs),
+                )
+                inputs = all_inputs[process_slice]
+                outputs = all_outputs[process_slice]
+
+            else:
+                with self.multi_turn_completion_length_context():
+                    outputs = self._infer_single_or_multi_turn(inputs, self.request_config)
+
+            if self.vllm_mode == 'colocate' and self.args.sleep_level > 0:
+                self.engine.engine.sleep(level=self.args.sleep_level)
+                empty_cache()
 
-        if self.vllm_mode == 'colocate' and self.args.sleep_level > 0:
-            self.engine.engine.sleep(level=self.args.sleep_level)
-            if self.args.gc_collect_after_offload:
-                gc_collect()
-            if self.args.offload_model:
-                self.load_model(self.accelerator.unwrap_model(self.model))
-                if self.ref_model:
-                    self.load_model(self.ref_model)
-            if self.args.offload_optimizer:
-                self.load_optimizer()
         return inputs, outputs
 
     def _generate_completions(self, inputs: InputsType) -> InputsType:
@@ -1577,3 +1586,25 @@ def inputs_to_rolloutrequest(self, inputs: InputsType) -> RolloutInferRequest:
         ]
 
         return infer_requests
+
+    @contextmanager
+    def offload_context(self):
+        if self.args.offload_model:
+            self.offload_model(self.accelerator.unwrap_model(self.model))
+            if self.ref_model:
+                self.offload_model(self.ref_model)
+        if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
+            self.offload_optimizer()
+        empty_cache()
+
+        try:
+            yield
+        finally:
+            # reload (load back) model when exiting context
+            if self.args.offload_model:
+                self.load_model(self.accelerator.unwrap_model(self.model))
+                if self.ref_model:
+                    self.load_model(self.ref_model)
+            if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
+                self.load_optimizer()
+            empty_cache()
