diff --git a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py index bd3f2e099..dc701af7a 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py +++ b/tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py @@ -21,10 +21,12 @@ Usage: from __future__ import annotations +import asyncio import uuid from contextlib import ExitStack import pytest +import torch from tests.helpers.mark import hardware_test from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -150,3 +152,101 @@ async def test_generate_after_list_loras_inline_mode(): assert last_output is not None assert isinstance(last_output, OmniRequestOutput) assert last_output.images, "Expected at least one generated image" + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.asyncio +async def test_sleep_memory_reclaimed_custom_pipeline(): + """sleep(level=1) must physically reclaim CuMemAllocator-tracked memory for + custom_pipeline. + + Regression test for: custom pipelines constructed under ``with target_device:`` + (CUDA default-device context) caused safetensors >=0.20.0 to use a + direct-to-GPU fast path (cudaMalloc via the driver API) that bypasses + CuMemAllocator, leaving weights invisible to sleep() and pinned in GPU + memory after the call. + + The fix moves custom_pipeline init outside the CUDA context so all weights + go through the caching allocator and are therefore fully reclaimed by + sleep(level=1). A non-zero ``CuMemAllocator.get_current_usage()`` after + sleep is the direct signal that the bypass is still occurring. + """ + with ExitStack() as after: + engine = AsyncOmni( + model=MODEL, + custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS}, + worker_extension_cls=WORKER_EXTENSION_CLASS, + enforce_eager=True, + enable_sleep_mode=True, + ) + after.callback(engine.shutdown) + + assert not await engine.is_sleeping(), "Engine should be awake after creation" + + # Measure global VRAM before sleep (driver view; includes inline worker + # thread since inline mode runs in the same process). + torch.accelerator.synchronize() + free_before, total = torch.cuda.mem_get_info() + used_before_gib = (total - free_before) / 1024**3 + + # Measure CuMemAllocator-tracked usage before sleep. In inline mode + # the worker runs in a thread pool inside this process, so the allocator + # singleton is shared and can be read directly. + allocator = None + tracked_before = 0 + try: + from vllm.device_allocator.cumem import CuMemAllocator + + allocator = CuMemAllocator.get_instance() + tracked_before = allocator.get_current_usage() + except Exception: + pass + + # Put the engine to sleep; all weights should be offloaded via the pool. + acks = await engine.sleep(level=1) + await asyncio.sleep(0.5) # allow the CUDA driver to settle + torch.accelerator.synchronize() + + # Measure after sleep. + free_after, _ = torch.cuda.mem_get_info() + used_after_gib = (total - free_after) / 1024**3 + drop_gib = used_before_gib - used_after_gib + + # --- Primary assertion: allocator reports zero tracked memory. --- + # If this fails it means weights were allocated outside the CuMem pool + # (safetensors direct-to-GPU bypass) — the exact regression this test + # is designed to catch. + if allocator is not None: + tracked_after = allocator.get_current_usage() + assert tracked_after == 0, ( + f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB " + f"after sleep(level=1) on custom_pipeline path " + f"(was {tracked_before / 1024**3:.3f} GiB before sleep). " + "Weights were allocated outside the CuMem pool via the " + "safetensors direct-to-GPU fast path — loader-context fix " + "may not be applied." + ) + + # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms + # reclamation at the driver level. + total_freed_bytes = sum( + (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0)) + for ack in acks + if ack is not None + ) + freed_gib = total_freed_bytes / 1024**3 + assert freed_gib > 0 or drop_gib > 0, ( + f"Expected GPU memory to be reclaimed after sleep(level=1) on " + f"custom_pipeline + enable_sleep_mode=True. " + f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, " + f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB." + ) + + # Engine must report it is sleeping. + assert await engine.is_sleeping() + + # Wake up and confirm the engine is functional again. + await engine.wake_up() + assert not await engine.is_sleeping() diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0675c226a..433ea5577 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1007,6 +1007,28 @@ class DiffusionOutput: # memory usage info peak_memory_mb: float = 0.0 + # When True, move all tensor fields (including tensors inside + # ``custom_output``) to CPU at construction time. Useful when the output + # is shipped across process boundaries (e.g. step-execution mode) and the + # receiving side must not initialise a stray CUDA context. + to_cpu: bool = False + + def __post_init__(self) -> None: + if not self.to_cpu: + return + + def _maybe_to_cpu(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu() + return value + + self.output = _maybe_to_cpu(self.output) + self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps) + self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents) + self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs) + if self.custom_output: + self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()} + class DiffusionRequestAbortedError(RuntimeError): """Raised when a diffusion request ends via user-visible abort.""" diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 9fee2b67e..d59c7df0a 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -328,19 +328,34 @@ class DiffusersPipelineLoader: od_config, target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name ) else: - with target_device: - if load_format == "default": - model = initialize_model(od_config) - elif load_format == "diffusers": - model = DiffusersAdapterPipeline(od_config=od_config, device=target_device) - elif load_format == "custom_pipeline": - from vllm_omni.diffusion.config import set_current_diffusion_config + if load_format == "custom_pipeline": + # NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)` + # internally. If we construct them under `with target_device:` (CUDA), + # safetensors takes a direct-to-GPU fast path that calls `cudaMalloc` + # via the driver API and BYPASSES PyTorch's caching allocator. + # That makes those bytes invisible to CuMemAllocator, so `sleep()` + # cannot offload/unmap them and GPU memory stays pinned. + # + # Fix: build the custom pipeline on CPU first (no default device + # context), then explicitly move it to the target device. The + # subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)` + # + `copy_`, which goes through the caching allocator and is fully + # tracked by CuMemAllocator. + from vllm_omni.diffusion.config import set_current_diffusion_config - model_cls = resolve_obj_by_qualname(custom_pipeline_name) - with set_current_diffusion_config(od_config): - model = model_cls(od_config=od_config) - else: - raise ValueError(f"Unknown load_format: {load_format}") + model_cls = resolve_obj_by_qualname(custom_pipeline_name) + with set_current_diffusion_config(od_config): + model = model_cls(od_config=od_config) + if target_device.type != "cpu": + model.to(target_device) + else: + with target_device: + if load_format == "default": + model = initialize_model(od_config) + elif load_format == "diffusers": + model = DiffusersAdapterPipeline(od_config=od_config, device=target_device) + else: + raise ValueError(f"Unknown load_format: {load_format}") logger.debug("Loading weights on %s ...", load_device) if load_format == "diffusers": # DiffusersAdapterPipeline.load_weights() calls diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 8f0577b49..93bf71730 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -522,11 +522,11 @@ class DiffusionWorker: logger.info(f"[Worker {self.rank}] Handshake Received: Task {task.task_id}") current_omni_platform.synchronize() - usage_before = current_omni_platform.get_current_memory_usage(self.device) + free_before = current_omni_platform.get_free_memory(self.device) self.sleep(level=task.level) current_omni_platform.synchronize() - usage_after = current_omni_platform.get_current_memory_usage(self.device) - real_freed = max(0, usage_before - usage_after) + free_after = current_omni_platform.get_free_memory(self.device) + real_freed = max(0, free_after - free_before) logger.info(f"[Worker {self.rank}] Preparing ACK: freed_bytes={real_freed / GiB_bytes:.2f} GiB.") # Ensure all ranks have completed sleep before measuring memory and sending ACK @@ -538,6 +538,11 @@ class DiffusionWorker: if self.rank != 0: return None + try: + total_mem = current_omni_platform.get_device_total_memory() + except (NotImplementedError, AttributeError): + total_mem = torch.cuda.get_device_properties(self.device).total_memory + residual_gib = (total_mem - free_after) / GiB_bytes ack = OmniACK( task_id=task.task_id, status="SUCCESS", @@ -548,7 +553,7 @@ class DiffusionWorker: metadata={ "source": f"Platform_{current_omni_platform.get_device_name()}", "total_freed_gib": f"{real_freed / GiB_bytes:.2f}", - "rank_residual_gib": f"{usage_after / GiB_bytes:.2f}", + "rank_residual_gib": f"{residual_gib:.2f}", }, ) logger.info(f"[Worker {self.rank}] ACK emitted. Freed {real_freed / GiB_bytes:.2f} GiB.") @@ -576,8 +581,12 @@ class DiffusionWorker: torch.distributed.barrier() current_omni_platform.synchronize() - usage_now = current_omni_platform.get_current_memory_usage(self.device) - current_used_gib = usage_now / (1024**3) + free_now = current_omni_platform.get_free_memory(self.device) + try: + total_mem = current_omni_platform.get_device_total_memory() + except (NotImplementedError, AttributeError): + total_mem = torch.cuda.get_device_properties(self.device).total_memory + current_used_gib = (total_mem - free_now) / (1024**3) if self.rank != 0: return None diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 67b6f77b3..dd431d16b 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -232,10 +232,10 @@ class OmniGPUWorkerBase(GPUWorker): if hasattr(self.model_runner, "graph_runners"): self.model_runner.graph_runners.clear() logger.info(f"[LLM Worker {self.rank}] LLM CUDA Graphs cleared.") - mem_before = current_omni_platform.get_current_memory_usage(self.device) + mem_before = current_omni_platform.get_free_memory(self.device) self.sleep(level=task.level) - mem_after = current_omni_platform.get_current_memory_usage(self.device) - rank_freed = max(0, mem_before - mem_after) + mem_after = current_omni_platform.get_free_memory(self.device) + rank_freed = max(0, mem_after - mem_before) if torch.distributed.is_initialized(): t_freed = torch.tensor([float(rank_freed)], device=self.device) torch.distributed.all_reduce(t_freed) @@ -246,6 +246,11 @@ class OmniGPUWorkerBase(GPUWorker): if self.rank != 0: return None current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0) + try: + total_mem = current_omni_platform.get_device_total_memory() + except (NotImplementedError, AttributeError): + total_mem = torch.cuda.get_device_properties(self.device).total_memory + residual_gib = (total_mem - mem_after) / 1024**3 ack = OmniACK( task_id=task.task_id, status="SUCCESS", @@ -255,7 +260,7 @@ class OmniGPUWorkerBase(GPUWorker): metadata={ "source": "omni_platform_audit", "total_freed_gib": f"{total_freed / 1024**3:.2f}", - "rank_residual_gib": f"{mem_after / 1024**3:.2f}", + "rank_residual_gib": f"{residual_gib:.2f}", }, ) if hasattr(self, "result_mq") and self.result_mq: @@ -281,7 +286,12 @@ class OmniGPUWorkerBase(GPUWorker): torch.distributed.barrier() gc.collect() current_omni_platform.synchronize() - usage_now = current_omni_platform.get_current_memory_usage(self.device) + free_now = current_omni_platform.get_free_memory(self.device) + try: + total_mem = current_omni_platform.get_device_total_memory() + except (NotImplementedError, AttributeError): + total_mem = torch.cuda.get_device_properties(self.device).total_memory + current_used_gib = (total_mem - free_now) / 1024**3 if self.rank != 0: return None current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0) @@ -290,7 +300,7 @@ class OmniGPUWorkerBase(GPUWorker): status="SUCCESS", stage_id=current_stage_id, rank=self.rank, - metadata={"state": "WARM", "current_vram_gib": f"{usage_now / 1024**3:.2f}"}, + metadata={"state": "WARM", "current_vram_gib": f"{current_used_gib:.2f}"}, ) if hasattr(self, "result_mq") and self.result_mq: self.result_mq.put(ack)