diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 2780ecdd08..3206df34e8 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -129,14 +129,14 @@ def get_offloader() -> BaseOffloader: return _instance -def set_offloader(instance: BaseOffloader) -> None: - """Set the global offloader instance.""" +def set_offloader(instance: BaseOffloader | None) -> None: + """Set or reset the global offloader instance.""" global _instance - _instance = instance - if isinstance(instance, NoopOffloader): + _instance = NoopOffloader() if instance is None else instance + if isinstance(_instance, NoopOffloader): logger.debug_once("Offloader set to NoopOffloader (no offloading).") else: - logger.info_once("Offloader set to %s", type(instance).__name__) + logger.info_once("Offloader set to %s", type(_instance).__name__) def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a5e5ac2497..57f0180edc 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2248,6 +2248,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): torch.accelerator.synchronize() if self.aux_output_connector is not None: self.aux_output_connector.close() + set_offloader(None) self.cudagraph_manager = None self.fast_prefill = None self.pooling_runner = None