[MRV2] Release weight offloader on shutdown (#57834)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
This commit is contained in:
Taneem Ibrahim
2026-09-22 15:01:20 -04:00
committed by GitHub
parent f944491db9
commit bc162b3f92
2 changed files with 6 additions and 5 deletions
+5 -5
View File
@@ -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:
+1
View File
@@ -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