From bc162b3f9280c2dc05dfd5aa282d29551f8f6d33 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 22 Sep 2026 15:01:20 -0400 Subject: [PATCH] [MRV2] Release weight offloader on shutdown (#57834) Signed-off-by: Taneem Ibrahim --- vllm/model_executor/offloader/base.py | 10 +++++----- vllm/v1/worker/gpu/model_runner.py | 1 + 2 files changed, 6 insertions(+), 5 deletions(-) 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