[Bugfix] fix GPU VRAM calculation problem (#328)

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>

docs/configuration/README.md

@@ -10,6 +10,10 @@ For specific example, please refer to [Qwen2.5-omni stage config](stage_configs/
 
 For introduction, please check [Introduction for stage config](./stage_configs.md)
 
+## Memory Configuration
+
+- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
+
 ## Optimization Features
 
 - **[TeaCache Configuration](../user_guide/acceleration/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss

docs/configuration/gpu_memory_utilization.md

@@ -0,0 +1,207 @@
+# GPU Memory Calculation and Configuration
+
+This guide explains how to calculate GPU memory requirements and properly configure `gpu_memory_utilization` for vLLM-Omni stages.
+
+## Overview
+
+`gpu_memory_utilization` is a critical parameter that controls how much GPU memory each stage can use. It's specified as a fraction between 0.0 and 1.0, where:
+- `0.8` means 80% of the GPU's total memory
+- `1.0` means 100% of the GPU's total memory (not recommended, leaves no buffer)
+
+## How Memory is Calculated
+
+### Memory Allocation Formula
+
+For each stage, vLLM-Omni calculates the requested memory as:
+
+```
+requested_memory = total_gpu_memory × gpu_memory_utilization
+```
+
+The system checks that:
+```
+free_memory ≥ requested_memory
+```
+
+If this condition is not met, the stage will fail to initialize with an error message showing the memory requirements.
+
+### Memory Components
+
+The total memory used by a stage includes:
+
+1. **Model Weights**: The size of the model parameters loaded on the GPU
+2. **KV Cache**: Memory for storing key-value cache during generation
+3. **Activation Memory**: Temporary memory for intermediate computations
+4. **System Overhead**: Memory used by CUDA, PyTorch, and other system components
+5. **Non-Torch Memory**: Memory allocated outside of PyTorch (e.g., CUDA graphs)
+
+### Example Calculation
+
+For a GPU with 80GB total memory:
+- `gpu_memory_utilization: 0.8` → 64GB available for the stage
+- `gpu_memory_utilization: 0.6` → 48GB available for the stage
+- `gpu_memory_utilization: 0.15` → 12GB available for the stage
+
+## Setting Up `gpu_memory_utilization`
+
+### Step 1: Determine GPU Memory
+
+First, check your GPU's total memory:
+
+```bash
+# Using nvidia-smi
+nvidia-smi --query-gpu=memory.total --format=csv
+
+# Or using Python
+python -c "import torch; print(f'{torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')"
+```
+
+### Step 2: Estimate Model Memory Requirements
+
+#### For Autoregressive (AR) Stages
+
+AR stages typically need more memory due to:
+- Large model weights
+- KV cache for attention
+- Activation buffers
+
+#### For Diffusion/Generation Stages
+
+Diffusion stages (like code2wav) typically need less memory:
+- Smaller model components
+- Different memory access patterns
+
+**Typical values:**
+- `0.1 - 0.3` for most diffusion stages
+
+### Step 3: Consider Multi-Stage Scenarios
+
+When multiple stages share the same GPU, you must ensure the sum of their `gpu_memory_utilization` values doesn't exceed 1.0.
+
+**Example: Two stages on GPU 0**
+```yaml
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0"
+    engine_args:
+      gpu_memory_utilization: 0.6  # Uses 60% of GPU 0
+
+  - stage_id: 1
+    runtime:
+      devices: "0"
+    engine_args:
+      gpu_memory_utilization: 0.3  # Uses 30% of GPU 0
+      # Total: 90% of GPU 0 (safe, leaves 10% buffer)
+```
+
+**Important:** If stages run on different GPUs, each can use up to 1.0 independently.
+
+### Step 4: Account for Tensor Parallelism
+
+When using `tensor_parallel_size > 1`, the model is split across multiple GPUs, so each GPU needs less memory.
+
+**Example: 2-way tensor parallelism**
+```yaml
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"  # Uses both GPUs
+    engine_args:
+      tensor_parallel_size: 2
+      gpu_memory_utilization: 0.6  # 60% per GPU
+      # Model is split, so each GPU uses ~30% of model memory
+```
+
+## Examples
+
+### Qwen3-Omni-MoE on 2x H100-80GB
+
+```yaml
+stage_args:
+  - stage_id: 0  # Thinker stage with TP=2
+    runtime:
+      devices: "0,1"
+    engine_args:
+      tensor_parallel_size: 2
+      gpu_memory_utilization: 0.6  # 48GB per GPU
+
+  - stage_id: 1  # Talker stage
+    runtime:
+      devices: "1"
+    engine_args:
+      gpu_memory_utilization: 0.3  # 24GB on GPU 1
+
+  - stage_id: 2  # Code2Wav stage
+    runtime:
+      devices: "0"
+    engine_args:
+      gpu_memory_utilization: 0.1  # 8GB on GPU 0
+```
+**Note:** In this configuration, stages 0 and 2 share GPU 0, but they run at different times in the pipeline, so their memory usage doesn't overlap.
+
+## Troubleshooting
+
+### Error: "Free memory is less than desired GPU memory utilization"
+
+This means the GPU doesn't have enough free memory when the stage starts.
+
+**Solutions:**
+1. Free up memory by closing other processes
+2. Reduce `gpu_memory_utilization` for this stage
+3. Use a GPU with more memory
+4. Move the stage to a different GPU
+
+### Error: OOM during inference
+
+The stage initialized but ran out of memory during processing.
+
+**Solutions:**
+1. Reduce `max_num_batched_tokens`
+2. Reduce `max_batch_size` in runtime config
+3. Lower `gpu_memory_utilization` slightly
+4. Enable quantization if supported
+
+### Memory Not Fully Utilized
+
+If you see low memory usage, you can:
+1. Increase `gpu_memory_utilization` to allow larger KV cache
+2. Increase `max_num_batched_tokens` for better batching
+3. Check if other stages are limiting throughput
+
+## Useful formula for Memory Calculation
+
+### KV Cache Memory
+
+The KV cache size depends on:
+- Number of sequences in batch
+- Sequence length (prompt + generation)
+- Model hidden size
+- Number of attention heads
+- Number of layers
+
+approximate Formula:
+```
+kv_cache_memory ≈ batch_size × seq_len × hidden_size × num_layers × 2 × dtype_size
+```
+2 for k & v
+
+### Model Weight Memory
+
+```
+model_memory ≈ num_parameters × dtype_size
+```
+
+For example:
+- 7B parameters in FP16: ~14GB
+- 7B parameters in FP32: ~28GB
+- 7B parameters in INT8: ~7GB
+
+### Activation Memory
+
+Activation memory is typically smaller but varies with:
+- Batch size
+- Sequence length
+- Model architecture
+
+It's usually 10-30% of model weight memory during inference.

docs/configuration/stage_configs.md

@@ -183,6 +183,9 @@ The fraction of GPU memory to allocate for this stage within a single GPU chip.
 
 Default: `0.8`
 
+!!! tip "Memory Configuration Guide"
+    For detailed information on how to calculate memory requirements and properly configure `gpu_memory_utilization`, see the [GPU Memory Calculation and Configuration Guide](./gpu_memory_utilization.md).
+
 #### `engine_args.enforce_eager`
 
 Whether to enforce eager execution mode. When set to `true`, the engine will run in eager mode without using CUDA graphs or other compilation optimizations. Currently, vLLM-Omni only supports eager mode.

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -44,7 +44,7 @@ stage_args:
        model_arch: Qwen3OmniMoeForConditionalGeneration
        worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.32
+       gpu_memory_utilization: 0.3
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav

vllm_omni/entrypoints/omni_stage.py

@@ -11,6 +11,7 @@ the original input processing utilities for cross-stage data wiring.
 """
 
 import asyncio
+import fcntl
 import importlib
 import logging
 import multiprocessing as mp
@@ -475,6 +476,7 @@ def _stage_worker(
     _batch_seq = 0
 
     # Device mapping
+    device_type = None
     try:
         from vllm_omni.utils import detect_device_type
 
@@ -483,11 +485,146 @@ def _stage_worker(
     except Exception as e:
         _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
 
+    # Sequential initialization on the same device to avoid memory calculation errors
+    # when multiple instances start simultaneously
+    # For TP/PP/DP/SP, we need to lock ALL devices that will be used by this stage
+    lock_files = []
+    if device_type == "cuda":
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                # Get all parallel sizes from engine_args (defaults to 1)
+                tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
+                pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
+                data_parallel_size = engine_args.get("data_parallel_size", 1)
+                prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+
+                # Calculate total number of devices needed for this stage
+                # For a single stage worker:
+                # - TP: splits model across GPUs (always needed)
+                # - PP: splits layers across pipelinestages, but each stage uses TP devices
+                # - DP: replicates model, but each replica uses TP devices
+                # - PCP: context parallelism, typically uses TP devices
+                # The number of devices per stage is determined by TP * PP * DP * PCP size
+                # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
+                num_devices_per_stage = (
+                    tensor_parallel_size * pipeline_parallel_size * data_parallel_size * prefill_context_parallel_size
+                )
+
+                # Get physical device IDs from CUDA_VISIBLE_DEVICES
+                # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s)
+                cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES")
+                physical_devices = []
+
+                if cuda_visible_devices:
+                    try:
+                        physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()]
+                    except (ValueError, IndexError):
+                        pass
+
+                if not physical_devices:
+                    # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set
+                    num_devices = torch.cuda.device_count()
+                    physical_devices = list(range(num_devices))
+
+                # Determine which devices will be used (min of devices per stage and available devices)
+                num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+                devices_to_lock = physical_devices[:num_devices_to_lock]
+
+                # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
+                devices_to_lock = sorted(devices_to_lock)
+
+                _logging.getLogger(__name__).debug(
+                    "[Stage-%s] Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
+                    stage_id,
+                    tensor_parallel_size,
+                    pipeline_parallel_size,
+                    data_parallel_size,
+                    prefill_context_parallel_size,
+                    num_devices_to_lock,
+                    devices_to_lock,
+                )
+
+                # Acquire exclusive locks for all devices using fcntl.flock
+                # Locks are automatically released when process dies
+                max_wait_time = 300  # 5 minutes max wait
+                wait_start = _time.time()
+                acquired_lock_fds = []  # Store file descriptors to keep locks alive
+
+                for device_id in devices_to_lock:
+                    lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+                    lock_acquired = False
+
+                    while not lock_acquired:
+                        try:
+                            # Open or create the lock file
+                            lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
+
+                            # Try to acquire exclusive lock (non-blocking first)
+                            try:
+                                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                                # Successfully acquired lock - write PID
+                                _os.ftruncate(lock_fd, 0)  # Clear file
+                                _os.write(lock_fd, f"{_os.getpid()}\n".encode())
+                                _os.fsync(lock_fd)  # Ensure written to disk
+                                lock_acquired = True
+                                acquired_lock_fds.append(lock_fd)
+                                _logging.getLogger(__name__).debug(
+                                    "[Stage-%s] Acquired exclusive lock for device %s", stage_id, device_id
+                                )
+                            except BlockingIOError:
+                                # Lock is held by another process
+                                _os.close(lock_fd)
+
+                                # Check if we've been waiting too long
+                                if _time.time() - wait_start > max_wait_time:
+                                    _logging.getLogger(__name__).warning(
+                                        "[Stage-%s] Timeout waiting for device %s "
+                                        "initialization lock, proceeding anyway",
+                                        stage_id,
+                                        device_id,
+                                    )
+                                    break
+
+                                # Wait a bit before retrying
+                                _time.sleep(0.1)
+                        except OSError as e:
+                            # Other error - log and continue without lock
+                            _logging.getLogger(__name__).debug(
+                                "[Stage-%s] Failed to acquire lock for device %s: %s, continuing anyway",
+                                stage_id,
+                                device_id,
+                                e,
+                            )
+                            try:
+                                _os.close(lock_fd)
+                            except (OSError, NameError):
+                                pass
+                            break
+
+                lock_files = acquired_lock_fds
+        except Exception as e:
+            _logging.getLogger(__name__).debug(
+                "[Stage-%s] Failed to set up sequential initialization lock: %s", stage_id, e
+            )
+
     # Init LLM
     _logging.getLogger(__name__).debug(
         "[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys())
     )
-    stage_engine = OmniStageLLM(model=model, **engine_args)
+    try:
+        stage_engine = OmniStageLLM(model=model, **engine_args)
+    finally:
+        # Release all locks by closing file descriptors
+        # Locks are automatically released when file descriptors are closed
+        # or when process dies
+        for lock_fd in lock_files:
+            try:
+                _os.close(lock_fd)
+                _logging.getLogger(__name__).debug("[Stage-%s] Released initialization lock (fd=%s)", stage_id, lock_fd)
+            except (OSError, ValueError):
+                pass
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
 
     # Initialize OmniConnectors if configured
@@ -747,6 +884,7 @@ async def _stage_worker_async(
 ) -> None:
     """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
     import logging as _logging
+    import os as _os
     import time as _time
 
     from vllm_omni.distributed.omni_connectors import build_stage_connectors
@@ -802,6 +940,7 @@ async def _stage_worker_async(
     _batch_seq = 0
 
     # Device mapping
+    device_type = None
     try:
         from vllm_omni.utils import detect_device_type
 
@@ -821,20 +960,155 @@ async def _stage_worker_async(
             return
         connectors = built_connectors
 
+    # Sequential initialization on the same device to avoid memory calculation errors
+    # when multiple instances start simultaneously
+    # For TP, we need to lock ALL devices that will be used by this stage
+    lock_files = []
+    if device_type == "cuda":
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                # Get all parallel sizes from engine_args (defaults to 1)
+                tensor_parallel_size = engine_args.get("tensor_parallel_size", 1)
+                pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1)
+                data_parallel_size = engine_args.get("data_parallel_size", 1)
+                prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1)
+
+                # Calculate total number of devices needed for this stage
+                # For a single stage worker in omni:
+                # - TP: splits model across GPUs (always needed)
+                # - PP: splits layers across stages, but each stage uses TP devices
+                # - DP: replicates model, but each replica uses TP devices
+                # - PCP: context parallelism, typically uses TP devices
+                # The number of devices per stage is determined by TP * PP * DP * PCP size
+                # (PP/DP/PCP are higher-level parallelism that don't add devices per stage)
+                num_devices_per_stage = (
+                    tensor_parallel_size * pipeline_parallel_size * data_parallel_size * prefill_context_parallel_size
+                )
+
+                # Get physical device IDs from CUDA_VISIBLE_DEVICES
+                # After set_stage_devices, CUDA_VISIBLE_DEVICES is set to physical device(s)
+                cuda_visible_devices = _os.environ.get("CUDA_VISIBLE_DEVICES")
+                physical_devices = []
+
+                if cuda_visible_devices:
+                    try:
+                        physical_devices = [int(x.strip()) for x in cuda_visible_devices.split(",") if x.strip()]
+                    except (ValueError, IndexError):
+                        pass
+
+                if not physical_devices:
+                    # Fallback: use logical device count if CUDA_VISIBLE_DEVICES not set
+                    num_devices = torch.cuda.device_count()
+                    physical_devices = list(range(num_devices))
+
+                # Determine which devices will be used (min of devices per stage and available devices)
+                num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+                devices_to_lock = physical_devices[:num_devices_to_lock]
+
+                # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
+                devices_to_lock = sorted(devices_to_lock)
+
+                _logging.getLogger(__name__).debug(
+                    "[Stage-%s] Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
+                    stage_id,
+                    tensor_parallel_size,
+                    pipeline_parallel_size,
+                    data_parallel_size,
+                    prefill_context_parallel_size,
+                    num_devices_to_lock,
+                    devices_to_lock,
+                )
+
+                # Acquire exclusive locks for all devices using fcntl.flock
+                # Locks are automatically released when process dies
+                max_wait_time = 300  # 5 minutes max wait
+                wait_start = _time.time()
+                acquired_lock_fds = []  # Store file descriptors to keep locks alive
+
+                for device_id in devices_to_lock:
+                    lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+                    lock_acquired = False
+
+                    while not lock_acquired:
+                        try:
+                            # Open or create the lock file
+                            lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644)
+
+                            # Try to acquire exclusive lock (non-blocking first)
+                            try:
+                                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                                # Successfully acquired lock - write PID
+                                _os.ftruncate(lock_fd, 0)  # Clear file
+                                _os.write(lock_fd, f"{_os.getpid()}\n".encode())
+                                _os.fsync(lock_fd)  # Ensure written to disk
+                                lock_acquired = True
+                                acquired_lock_fds.append(lock_fd)
+                                _logging.getLogger(__name__).debug(
+                                    "[Stage-%s] Acquired exclusive lock for device %s", stage_id, device_id
+                                )
+                            except BlockingIOError:
+                                # Lock is held by another process
+                                _os.close(lock_fd)
+
+                                # Check if we've been waiting too long
+                                if _time.time() - wait_start > max_wait_time:
+                                    _logging.getLogger(__name__).warning(
+                                        "[Stage-%s] Timeout waiting for device %s "
+                                        "initialization lock, proceeding anyway",
+                                        stage_id,
+                                        device_id,
+                                    )
+                                    break
+
+                                # Wait a bit before retrying
+                                _time.sleep(0.1)
+                        except OSError as e:
+                            # Other error - log and continue without lock
+                            _logging.getLogger(__name__).debug(
+                                "[Stage-%s] Failed to acquire lock for device %s: %s, continuing anyway",
+                                stage_id,
+                                device_id,
+                                e,
+                            )
+                            try:
+                                _os.close(lock_fd)
+                            except (OSError, NameError):
+                                pass
+                            break
+
+                lock_files = acquired_lock_fds
+        except Exception as e:
+            _logging.getLogger(__name__).debug(
+                "[Stage-%s] Failed to set up sequential initialization lock: %s", stage_id, e
+            )
+
     # Init LLM
     _logging.getLogger(__name__).debug(
         "[Stage-%s] Initializing engine with args keys=%s",
         stage_id,
         list(engine_args.keys()),
     )
-    omni_engine_args = AsyncOmniEngineArgs(**engine_args)
-    usage_context = UsageContext.OPENAI_API_SERVER
-    vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context)
-    stage_engine = AsyncOmniStageLLM.from_vllm_config(
-        vllm_config=vllm_config,
-        usage_context=usage_context,
-        engine_args=omni_engine_args,
-    )
+    try:
+        omni_engine_args = AsyncOmniEngineArgs(**engine_args)
+        usage_context = UsageContext.OPENAI_API_SERVER
+        vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context)
+        stage_engine = AsyncOmniStageLLM.from_vllm_config(
+            vllm_config=vllm_config,
+            usage_context=usage_context,
+            engine_args=omni_engine_args,
+        )
+    finally:
+        # Release all locks by closing file descriptors
+        # Locks are automatically released when file descriptors are closed
+        # or when process dies
+        for lock_fd in lock_files:
+            try:
+                _os.close(lock_fd)
+                _logging.getLogger(__name__).debug("[Stage-%s] Released initialization lock (fd=%s)", stage_id, lock_fd)
+            except (OSError, ValueError):
+                pass
     omni_stage.set_async_engine(stage_engine)
     # Don't keep the dummy data in memory
     await stage_engine.reset_mm_cache()