diff --git a/docs/configuration/README.md b/docs/configuration/README.md index 1ceb9f827..173381011 100644 --- a/docs/configuration/README.md +++ b/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 diff --git a/docs/configuration/gpu_memory_utilization.md b/docs/configuration/gpu_memory_utilization.md new file mode 100644 index 000000000..19fc042aa --- /dev/null +++ b/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. diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 9b5ea8928..18972eb79 100644 --- a/docs/configuration/stage_configs.md +++ b/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. diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml index bcaea6703..fec6169d8 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml +++ b/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 diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 8b2050be6..95d9a492b 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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()