diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index b4639a35f..f3830f02e 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -70,6 +70,11 @@ def test_cpu_offload_diffusion_model(model_name: str): pytest.fail("Inference failed") print(f"Offload peak memory: {offload_peak_memory} MB") print(f"No offload peak memory: {no_offload_peak_memory} MB") - assert offload_peak_memory + 2500 < no_offload_peak_memory, ( - f"Offload peak memory {offload_peak_memory} MB should be less than no offload peak memory {no_offload_peak_memory} MB" + # Set platform-specific VRAM saving thresholds to account + # for varying runtime memory overhead and fragmentation between CUDA and ROCm. + is_rocm = torch.version.hip is not None + threshold = 2500 if not is_rocm else 2100 + assert offload_peak_memory + threshold < no_offload_peak_memory, ( + f"Offload peak memory {offload_peak_memory} MB should be less than " + f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB" ) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 1b9dc5667..79a1a9ef9 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -9,10 +9,12 @@ to DiffusionModelRunner. """ import gc +import json import multiprocessing as mp import os from collections.abc import Iterable from contextlib import AbstractContextManager, nullcontext +from pathlib import Path from typing import Any import torch @@ -45,6 +47,16 @@ from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory logger = init_logger(__name__) +MODEL_PARAM_COUNTS = { + "bagel": 13.5e9, # BAGEL-7B-MoT + "flux": 12.0e9, # FLUX.1-dev + "default": 10.0e9, +} + +# 2.5 GiB is the empirical peak activation footprint (intermediate tensors + +# attention workspace) for models like Bagel-7B/Flux at 1024x1024 resolution. +ACTIVATION_MEMORY_MULTIPLIER = 2.5 + class DiffusionWorker: """ @@ -59,6 +71,60 @@ class DiffusionWorker: delegated to DiffusionModelRunner. """ + @staticmethod + def predict_resource_usage(od_config: OmniDiffusionConfig) -> dict[str, float]: + from vllm.utils.mem_utils import GiB_bytes + + total_params = 0 + estimate_source = "fallback" + try: + model_path = Path(getattr(od_config, "model", "")) + index_file = model_path / "model.safetensors.index.json" + if index_file.exists(): + with open(index_file) as f: + data = json.load(f) + # metadata.total_size + total_params = int(data.get("metadata", {}).get("total_size", 0)) + if total_params > 0: + estimate_source = "safetensors index" + except Exception as e: + logger.debug(f"Failed to parse safetensors metadata: {e}") + if total_params == 0: + m_name = str(getattr(od_config, "model", "")).lower() + if "bagel" in m_name: + total_params = MODEL_PARAM_COUNTS["bagel"] + estimate_source = "bagel" + elif "flux" in m_name: + total_params = MODEL_PARAM_COUNTS["flux"] + estimate_source = "flux" + else: + total_params = MODEL_PARAM_COUNTS["default"] + estimate_source = "default" + logger.info(f"VRAM Quota: Estimated {total_params / 1e9:.2f}B params using {estimate_source} logic") + + dtype = getattr(od_config, "dtype", torch.bfloat16) + dtype_str = str(dtype).lower() + # Calculate the number of bytes per parameter + # based on different levels of precision. + if "int4" in dtype_str: + bytes_per_param = 0.5 + elif any(kw in dtype_str for kw in ["float8", "int8"]): + bytes_per_param = 1 + elif any(kw in dtype_str for kw in ["float16", "bfloat16", "half"]): + bytes_per_param = 2 + elif any(kw in dtype_str for kw in ["float32", "int32"]): + bytes_per_param = 4 + else: + bytes_per_param = 4 + static_gb = (total_params * bytes_per_param) / GiB_bytes + h, w = getattr(od_config, "height", 1024), getattr(od_config, "width", 1024) + dynamic_gb = ACTIVATION_MEMORY_MULTIPLIER * (h * w / (1024 * 1024)) + return { + "static_gb": round(static_gb, 2), + "dynamic_gb": round(dynamic_gb, 2), + "total_gb": round(static_gb + dynamic_gb, 2), + } + def __init__( self, local_rank: int, diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index c414b972b..f96b4e9ea 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -313,6 +313,63 @@ class OmniBase: return config_path, stage_configs + def _coordinate_vram_resources(self) -> None: + """ + Coordinate VRAM resource reservations across stages based on their type and configs. + Note: Current implementation supports per-device isolation. Advanced multi-GPU + scenarios like TP or PP are planned as sleep mode ack Test Scenario. + """ + import torch + from vllm.utils.mem_utils import GiB_bytes + + from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker + + # Device ID + reserved_gb_per_device: dict[int, float] = {} + active_configs = self.stage_configs + if hasattr(self, "_single_stage_id") and self._single_stage_id is not None: + active_configs = [cfg for cfg in self.stage_configs if cfg.stage_id == self._single_stage_id] + for cfg in active_configs: + s_type = getattr(cfg, "stage_type", None) + if s_type == "diffusion": + prediction = DiffusionWorker.predict_resource_usage(cfg.engine_args) + device_str = getattr(cfg.runtime, "devices", "0") + try: + devices = [int(d.strip()) for d in device_str.split(",")] + except (ValueError, AttributeError): + devices = [0] + for d_id in devices: + reserved_gb_per_device[d_id] = reserved_gb_per_device.get(d_id, 0.0) + prediction["total_gb"] + logger.info( + f"[Coordinator] Stage-{cfg.stage_id} ({s_type.capitalize()}) " + f"on devices {devices} predicted budget: {prediction['total_gb']:.2f} GiB" + ) + if not torch.cuda.is_available(): + # TODO: Add support for other accelerators (NPU/XPU) + # as their specific memory management APIs are integrated. + return + for cfg in active_configs: + if getattr(cfg, "stage_type", None) == "llm": + # Get the master device ID where the LLM is located + llm_device_str = getattr(cfg.runtime, "devices", "0") + try: + target_device = int(llm_device_str.split(",")[0].strip()) + except (ValueError, AttributeError): + target_device = 0 + # The physical total of the card containing LLM + physical_vram_gb = torch.cuda.get_device_properties(target_device).total_memory / GiB_bytes + total_reserved_on_this_device = reserved_gb_per_device.get(target_device, 0.0) + if total_reserved_on_this_device > 0: + original_util = cfg.engine_args.get("gpu_memory_utilization", 0.9) + reserved_util_ratio = total_reserved_on_this_device / physical_vram_gb + adjusted_util = min(0.95, original_util + reserved_util_ratio) + cfg.engine_args["gpu_memory_utilization"] = round(adjusted_util, 3) + logger.info( + f"[Coordinator] LLM Stage-{cfg.stage_id} on Device {target_device} dynamic boost: " + f"{original_util} -> {cfg.engine_args['gpu_memory_utilization']} " + f"(Compensating {reserved_util_ratio:.2f} ratio for resource domain isolation)" + ) + def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None: """Initialize stage list management.""" self._inline_diffusion = False @@ -335,6 +392,9 @@ class OmniBase: # Resolve stage configs shared by orchestrator/headless paths. self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs) + # This ensures all engines receive the corrected gpu_memory_utilization through engine_args. + self._coordinate_vram_resources() + # Initialize connectors self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors( self.config_path, worker_backend=worker_backend, shm_threshold_bytes=shm_threshold_bytes