From a52867abde4fae10203c5d3d649eda77824376d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Thu, 12 Mar 2026 19:19:57 +0800 Subject: [PATCH] [BugFix]: Revert #1582 (#1842) Signed-off-by: princepride --- .../diffusion/worker/diffusion_worker.py | 66 ------------------- vllm_omni/entrypoints/omni.py | 60 ----------------- 2 files changed, 126 deletions(-) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 79a1a9ef9..1b9dc5667 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -9,12 +9,10 @@ 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 @@ -47,16 +45,6 @@ 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: """ @@ -71,60 +59,6 @@ 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 f96b4e9ea..c414b972b 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -313,63 +313,6 @@ 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 @@ -392,9 +335,6 @@ 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