From e24f86ae2e269ddab24d90f9026b40b440b18cff Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Fri, 24 Apr 2026 16:30:24 +0800 Subject: [PATCH] [Enhancement]modify profiling.md (#3051) Signed-off-by: bjf-frz --- .../model/adding_diffusion_model.md | 18 +- docs/contributing/profiling.md | 222 ++++++++++++------ .../image_to_image/image_edit.py | 21 +- .../image_to_image/image_edit.py | 22 +- .../image_to_video/image_to_video.py | 23 +- .../text_to_image/text_to_image.py | 22 +- .../text_to_video/text_to_video.py | 23 +- vllm_omni/profiler/omni_torch_profiler.py | 7 +- 8 files changed, 261 insertions(+), 97 deletions(-) diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 8d85eb4f6..6d5782a6e 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -813,16 +813,16 @@ class WanTransformer3DModel(nn.Module): --- -### Diffusion Timing (Performance Profiling) +### Diffusion Pipeline Profiler (Performance Profiling) When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding. vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks. !!! info - `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps. + `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps. This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs. -**Enabling Diffusion Timing** +**Enabling Diffusion Pipeline Profiler** Enable timing by setting: @@ -843,7 +843,7 @@ If not specified, the default targets are used: **Adding DiffusionPipelineProfilerMixin to a Pipeline** To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin. ```python -from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.profiler import DiffusionPipelineProfilerMixin class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin): # Optional: Specify custom timing targets @@ -862,7 +862,9 @@ class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin): ... # initialize timing profiler - self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler) + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) ``` The mixin dynamically wraps selected methods and records their execution time during inference. @@ -906,9 +908,9 @@ tokenizer.forward When enabled, timing logs appear like this: ``` -[DiffusionTiming] text_encoder.forward took 0.018s -[DiffusionTiming] diffuse took 2.412s -[DiffusionTiming] vae.decode took 0.063s +[DiffusionPipelineProfiler] text_encoder.forward took 0.018s +[DiffusionPipelineProfiler] diffuse took 2.412s +[DiffusionPipelineProfiler] vae.decode took 0.063s ``` These measurements help identify bottlenecks during model adaptation and optimization diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index a7b508302..e1dbc8234 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -1,15 +1,17 @@ -# Profiling vLLM-Omni +# Profiling Diffusion Models > **Warning:** Profiling is for development and debugging only. It adds significant overhead and should not be enabled in production. -vLLM-Omni supports two profiler backends through `profiler_config`: +Diffusion profiling supports two backends through `profiler_config`: -- `torch`: detailed CPU/CUDA traces written to `torch_profiler_dir` +- `torch`: detailed CPU/CUDA traces, operator tables, and optional memory snapshots - `cuda`: low-overhead CUDA range control for NVIDIA Nsight Systems (`nsys`) -## 1. Configure Profiling +## 1. Configure `profiler_config` -Use the same `profiler_config` shape everywhere: +Use `profiler_config` to enable profiling for a diffusion model. For diffusion usage, pass it directly to `Omni(...)` or `vllm serve`. + +Minimal torch-profiler config: ```yaml profiler_config: @@ -21,54 +23,114 @@ Supported fields: | Field | Description | |---|---| -| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. | -| `torch_profiler_dir` | Output directory for torch traces. Required when `profiler: torch`. | +| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. | +| `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. | +| `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. | +| `torch_profiler_record_shapes` | Record input shapes and add a `by_shape` sheet to `ops_rank*.xlsx`. | +| `torch_profiler_with_stack` | Record call stacks, add a `by_stack` sheet to `ops_rank*.xlsx`, and export `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt`. | +| `torch_profiler_with_memory` | Enable memory profiling and attempt to dump `memory_snapshot_rank*.pickle`. The pickle is only generated when the current backend supports memory history and snapshot APIs. | +| `torch_profiler_with_flops` | Enable FLOPs collection in `torch.profiler`. This does not add a separate output file. | +| `torch_profiler_dump_cuda_time_total` | Export an additional text summary `profiler_out_.txt` sorted by `self_cuda_time_total`. | | `delay_iterations` | Number of worker iterations to skip before profiling starts. | | `max_iterations` | Maximum number of worker iterations to capture before auto-stop. | +| `wait_iterations` | Torch-profiler wait iterations before warmup. | | `warmup_iterations` | Torch-profiler warmup iterations. | | `active_iterations` | Torch-profiler active iterations. | -| `wait_iterations` | Torch-profiler wait iterations before warmup. | -For multi-stage omni pipelines, put `profiler_config` under the target stage's `engine_args`. +### Minimal configurations by output -```yaml -stage_args: - - stage_id: 0 - stage_type: llm - engine_args: - profiler_config: - profiler: torch - torch_profiler_dir: ./perf -``` - -For single-stage diffusion usage, pass `profiler_config` directly to `Omni(...)` or `vllm serve`. - -## 2. Profiling Omni Pipelines - -It is usually best to profile only the stages you need. +Only collect trace output: ```python -# Profile all stages. -omni.start_profile() - -# Profile selected stages only. -omni.start_profile(stages=[0, 2]) -... -omni.stop_profile(stages=[0, 2]) +profiler_config = { + "profiler": "torch", + "torch_profiler_dir": "./perf", +} ``` -Always stop the same stage set that you started. If only some stages have `profiler_config`, pass an explicit `stages=[...]` list instead of relying on the default "all stages" behavior. +Outputs: -Examples: +- `trace_rank*.json` +- `ops_rank*.xlsx` with a `summary` sheet -1. [Qwen2.5-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py) -2. [Qwen3-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py) +Collect compressed trace output: -## 3. Profiling Single-Stage Diffusion +```python +profiler_config = { + "profiler": "torch", + "torch_profiler_dir": "./perf", + "torch_profiler_use_gzip": True, +} +``` -Single-stage diffusion models use the same `start_profile()` / `stop_profile()` controls, but you must provide `profiler_config` explicitly. +Outputs: -### PyTorch profiler +- `trace_rank*.json.gz` +- `ops_rank*.xlsx` with a `summary` sheet + +Collect trace and full operator tables: + +```python +profiler_config = { + "profiler": "torch", + "torch_profiler_dir": "./perf", + "torch_profiler_record_shapes": True, + "torch_profiler_with_stack": True, +} +``` + +Outputs: + +- `trace_rank*.json` +- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack` +- `stacks_cpu_rank*.txt` +- `stacks_cuda_rank*.txt` + +Collect trace, operator tables, and memory snapshots: + +```python +profiler_config = { + "profiler": "torch", + "torch_profiler_dir": "./perf", + "torch_profiler_record_shapes": True, + "torch_profiler_with_stack": True, + "torch_profiler_with_memory": True, +} +``` + +Outputs: + +- `trace_rank*.json` +- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack` +- `stacks_cpu_rank*.txt` +- `stacks_cuda_rank*.txt` +- `memory_snapshot_rank*.pickle` when supported by the current backend + +### Full torch-profiler configuration + +If you want to enable the commonly used torch-profiler options together: + +```python +profiler_config = { + "profiler": "torch", + "torch_profiler_dir": "./perf", + "torch_profiler_use_gzip": False, + "torch_profiler_record_shapes": True, + "torch_profiler_with_stack": True, + "torch_profiler_with_memory": True, + "torch_profiler_with_flops": False, + "torch_profiler_dump_cuda_time_total": False, + "delay_iterations": 0, + "max_iterations": 0, + "wait_iterations": 0, + "warmup_iterations": 0, + "active_iterations": 0, +} +``` + +## 2. Profiling Diffusion with PyTorch Profiler + +Single-stage diffusion models use `start_profile()` / `stop_profile()` controls. The profiler only writes artifacts after profiling has been started and then stopped. ```python from vllm_omni import Omni @@ -86,7 +148,29 @@ omni.start_profile() omni.stop_profile() ``` -### Nsight Systems (`nsys`) +For diffusion offline example scripts under `examples/offline_inference/`, pass `--profiler-config` as a JSON object. The script enables profiling when this argument is set and wraps generation with `start_profile()` / `stop_profile()`. + +Example: + +```bash +python examples/offline_inference/image_to_video/image_to_video.py \ + --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \ + --image input.jpg \ + --prompt "A cat playing with yarn" \ + --profiler-config '{ + "profiler": "torch", + "torch_profiler_dir": "./perf", + "torch_profiler_record_shapes": true, + "torch_profiler_with_stack": true + }' +``` + +Examples: + +1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) +2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video) + +## 3. Profiling Diffusion with Nsight Systems (`nsys`) For Nsight Systems, use `profiler: cuda` and wrap the process with `nsys profile`. @@ -103,35 +187,20 @@ nsys profile \ The Python process being profiled must create the diffusion engine with: ```python -profiler_config={"profiler": "cuda"} +profiler_config = {"profiler": "cuda"} ``` Then call `start_profile()` before the requests you want to capture and `stop_profile()` after them. The diffusion worker processes open and close the CUDA capture range themselves, so `nsys` sees the actual GPU work instead of only the parent process. -Examples: - -1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) -2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video) - ## 4. Profiling Online Serving -When any stage has `profiler_config.profiler` set, the server exposes: +When `profiler_config.profiler` is set for a diffusion model, the server exposes: - `POST /start_profile` - `POST /stop_profile` ### Start the server -Multi-stage omni serving: - -```bash -vllm serve Qwen/Qwen2.5-Omni-7B \ - --omni \ - --port 8091 -``` - -(The default deploy config at `vllm_omni/deploy/qwen2_5_omni.yaml` is loaded automatically. Pass `--deploy-config /path/to/custom.yaml` to override.) - Single-stage diffusion serving with torch profiler: ```bash @@ -144,21 +213,17 @@ vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \ "torch_profiler_with_stack": true, "torch_profiler_with_flops": false, "torch_profiler_use_gzip": true, - "torch_profiler_dump_cuda_time_total": true, + "torch_profiler_dump_cuda_time_total": false, "torch_profiler_record_shapes": true, "torch_profiler_with_memory": true, "delay_iterations": 0, - "max_iterations": 0 + "max_iterations": 0, + "wait_iterations": 0, + "warmup_iterations": 0, + "active_iterations": 0 }' ``` -Useful optional fields: - -- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files. -- `torch_profiler_record_shapes`: export `by_shape` operator views. -- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history. -- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`. - Single-stage diffusion serving with Nsight Systems: ```bash @@ -176,30 +241,37 @@ nsys profile \ ### Control capture +Example profiling flow for an online Qwen-Image request: + ```bash -# Start profiling on all profiled stages. +# Start profiling. curl -X POST http://localhost:8091/start_profile -# Start profiling on selected stages. -curl -X POST http://localhost:8091/start_profile \ +# Send a Qwen-Image generation request while profiling is active. +curl http://localhost:8091/v1/images/generations \ -H "Content-Type: application/json" \ - -d '{"stages": [0]}' + -d '{ + "model": "Qwen/Qwen-Image", + "prompt": "A red vintage bicycle parked beside a quiet canal at sunset" + }' -# Stop profiling. +# Stop profiling and flush profiler artifacts. curl -X POST http://localhost:8091/stop_profile ``` -For mixed-stage pipelines, use explicit `stages` and pass the same stage list to both endpoints. +## 5. Diffusion Pipeline Profiler -## 5. Analyze Results +For lightweight per-stage pipeline timing such as `vae.decode` or `diffuse`, see [Diffusion Pipeline Profiler](model/adding_diffusion_model.md#diffusion-pipeline-profiler-performance-profiling). That utility logs stage durations only and does not generate torch-profiler artifacts such as `trace.json`, Excel tables, or memory snapshots. -Torch profiler output: +## 6. Analyze Results + +Torch-profiler output: - Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz` - Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets - Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled - Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend -- Optional aggregated CUDA-time tables under the same session directory +- Optional CUDA-time text summary: `profiler_out_.txt` when `torch_profiler_dump_cuda_time_total` is enabled CUDA profiler / Nsight Systems output: @@ -211,4 +283,4 @@ Recommended viewers: - `nsys stats .nsys-rep` for CLI summaries - Nsight Systems GUI for CUDA kernel timelines -vLLM-Omni reuses the vLLM profiling infrastructure where possible. For the upstream reference, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/). +For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/). diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py index 8ab5e0d9a..7c91c0323 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -44,9 +44,11 @@ Layered RGBA output: import argparse import asyncio +import json import os import time from pathlib import Path +from typing import Any import torch from PIL import Image @@ -58,6 +60,16 @@ from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform +def parse_profiler_config(value: str) -> dict[str, Any]: + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + return config + + # =========================== # Argument Parser # =========================== @@ -99,6 +111,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--vae-use-slicing", action="store_true") parser.add_argument("--vae-use-tiling", action="store_true") parser.add_argument("--enable-cpu-offload", action="store_true") + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) return parser.parse_args() @@ -158,12 +176,13 @@ async def main(): enable_cpu_offload=args.enable_cpu_offload, diffusion_load_format="dummy", custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"}, + profiler_config=args.profiler_config, ) print(">>> Pipeline loaded successfully") # ---- Profiling + Info ---- - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + profiler_enabled = args.profiler_config is not None print(f"\n{'=' * 60}") print("Generation Configuration") print(f"Model: {args.model}") diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 2a73d875d..7837fe16d 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -87,9 +87,11 @@ For more options, run: """ import argparse +import json import os import time from pathlib import Path +from typing import Any import torch from PIL import Image @@ -101,6 +103,16 @@ from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform +def parse_profiler_config(value: str) -> dict[str, Any]: + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + return config + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.") parser.add_argument( @@ -360,6 +372,12 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) return parser.parse_args() @@ -425,11 +443,11 @@ def main(): enforce_eager=args.enforce_eager, enable_cpu_offload=args.enable_cpu_offload, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + profiler_config=args.profiler_config, ) print("Pipeline loaded") - # Check if profiling is requested via environment variable - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + profiler_enabled = args.profiler_config is not None # Time profiling for generation print(f"\n{'=' * 60}") diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 53319c822..e46db4de4 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -33,9 +33,10 @@ Usage: """ import argparse -import os +import json import time from pathlib import Path +from typing import Any import numpy as np import PIL.Image @@ -48,6 +49,16 @@ from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform +def parse_profiler_config(value: str) -> dict[str, Any]: + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + return config + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).") parser.add_argument( @@ -194,6 +205,12 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) return parser.parse_args() @@ -274,8 +291,7 @@ def main(): "rel_l1_thresh": 0.2, } - # Check if profiling is requested via environment variable - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + profiler_enabled = args.profiler_config is not None parallel_config = DiffusionParallelConfig( ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, @@ -300,6 +316,7 @@ def main(): cache_backend=args.cache_backend, cache_config=cache_config, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + profiler_config=args.profiler_config, ) if profiler_enabled: diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index ff54d9fa1..2c74ddfa8 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse -import os +import json import time from pathlib import Path from typing import Any @@ -30,6 +30,16 @@ def is_nextstep_model(model_name: str) -> bool: return False +def parse_profiler_config(value: str) -> dict[str, Any]: + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + return config + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.") parser.add_argument( @@ -254,6 +264,12 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) parser.add_argument( "--log-stats", action="store_true", @@ -333,8 +349,7 @@ def main(): enable_expert_parallel=args.enable_expert_parallel, ) - # Check if profiling is requested via environment variable - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + profiler_enabled = args.profiler_config is not None # Prepare LoRA kwargs for Omni initialization lora_args: dict[str, Any] = {} @@ -375,6 +390,7 @@ def main(): "mode": "text-to-image", "log_stats": args.log_stats, "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler, + "profiler_config": args.profiler_config, "init_timeout": args.init_timeout, "stage_init_timeout": args.stage_init_timeout, **lora_args, diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index ac9c7b091..1dba97c20 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse -import os +import json import time from pathlib import Path +from typing import Any import numpy as np import torch @@ -44,6 +45,16 @@ def _detect_preset(model: str) -> dict: return _MODEL_PRESETS["wan"] +def parse_profiler_config(value: str) -> dict[str, Any]: + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + return config + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate a video from a text prompt. " @@ -185,6 +196,12 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) parser.add_argument( "--quantization", type=str, @@ -248,8 +265,7 @@ def main(): enable_expert_parallel=args.enable_expert_parallel, ) - # Check if profiling is requested via environment variable - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + profiler_enabled = args.profiler_config is not None omni_kwargs = dict( model=args.model, @@ -263,6 +279,7 @@ def main(): cache_backend=args.cache_backend, cache_config=cache_config, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + profiler_config=args.profiler_config, ) if args.boundary_ratio is not None: omni_kwargs["boundary_ratio"] = args.boundary_ratio diff --git a/vllm_omni/profiler/omni_torch_profiler.py b/vllm_omni/profiler/omni_torch_profiler.py index 023ad5009..69f893b05 100644 --- a/vllm_omni/profiler/omni_torch_profiler.py +++ b/vllm_omni/profiler/omni_torch_profiler.py @@ -474,10 +474,13 @@ class OmniTorchProfilerWrapper(WorkerProfiler): table = self.profiler.key_averages().table(sort_by=sort_key) if not _is_uri_path(profiler_dir): - table_file = os.path.join(profiler_dir, f"profiler_out_{rank}.txt") + table_file = os.path.join( + self._ensure_session_dir(), + f"profiler_out_{rank}.txt", + ) with open(table_file, "w") as f: print(table, file=f) - self._table_path = table_file + self._artifact_paths["profiler_out"] = table_file if rank == 0: print(table)