[Diffusion] Document torch.compile graph-break checks in diffusion benchmark skills (#20681)
This commit is contained in:
Xiaoyu Zhang
@@ -47,7 +47,7 @@ check "plotly" python3 -c "import plotly"
|
||||
**Minimum for benchmarking**: `sglang`, `torch` with CUDA.
|
||||
**Level 1 profiling**: `torch.profiler` (bundled with torch).
|
||||
**Level 2 profiling**: `nsys`, `pandas`, `plotly` + `gputrc2graph.py` from the sglang repo.
|
||||
All commands below assume you are inside the configured diffusion container shell, already `cd`'d to the repo root derived from `sglang.__file__`, with `FLASHINFER_DISABLE_VERSION_CHECK=1` exported. Re-run `print-idle-gpus` before each perf command if GPU availability may have changed. For 8-GPU commands, request eight idle GPUs and export the comma-separated list to `CUDA_VISIBLE_DEVICES` first.
|
||||
All commands below assume you are inside the configured diffusion container shell, already `cd`'d to the repo root derived from `sglang.__file__`, with `FLASHINFER_DISABLE_VERSION_CHECK=1` exported. Re-run `print-idle-gpus` before each perf command if GPU availability may have changed. Keep benchmark commands within 4 GPUs or fewer.
|
||||
|
||||
Download input images required by some models:
|
||||
```bash
|
||||
@@ -137,18 +137,17 @@ sglang generate \
|
||||
--dit-cpu-offload false --text-encoder-cpu-offload false
|
||||
```
|
||||
|
||||
### Wan2.2-T2V-A14B 720P (8 GPUs, 81 frames, 40 steps)
|
||||
### Wan2.2-T2V-A14B 720P (4 GPUs, 81 frames, 2 steps)
|
||||
```bash
|
||||
# Select eight idle GPUs first:
|
||||
# export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 8)
|
||||
# Select four idle GPUs first:
|
||||
# export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4)
|
||||
sglang generate \
|
||||
--model-path=Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--prompt="A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon." \
|
||||
--negative-prompt=" " --720p --num-inference-steps=40 --num-frames=81 \
|
||||
--negative-prompt=" " --720p --num-inference-steps=2 --num-frames=81 \
|
||||
--guidance-scale=5.0 --seed=42 --save-output \
|
||||
--num-gpus=8 --enable-cfg-parallel --ulysses-degree=4 \
|
||||
--dit-layerwise-offload true --dit-cpu-offload false \
|
||||
--vae-cpu-offload false --text-encoder-cpu-offload true \
|
||||
--num-gpus=4 --ulysses-degree=4 \
|
||||
--text-encoder-cpu-offload --pin-cpu-memory \
|
||||
--warmup --enable-torch-compile
|
||||
```
|
||||
|
||||
@@ -177,7 +176,7 @@ sglang generate \
|
||||
--warmup --enable-torch-compile
|
||||
```
|
||||
|
||||
### MOVA-720p (4 GPUs, 193 frames, 24 steps)
|
||||
### MOVA-720p (4 GPUs, 193 frames, 2 steps)
|
||||
```bash
|
||||
# Select four idle GPUs first:
|
||||
# export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4)
|
||||
@@ -188,7 +187,7 @@ sglang generate \
|
||||
--adjust-frames=false \
|
||||
--num-gpus=4 --ring-degree=1 --ulysses-degree=4 \
|
||||
--num-frames=193 --fps=24 \
|
||||
--num-inference-steps=24 \
|
||||
--num-inference-steps=2 \
|
||||
--enable-torch-compile --save-output --warmup
|
||||
```
|
||||
|
||||
@@ -206,6 +205,13 @@ Add `--log-level=info` and observe:
|
||||
|
||||
### Step 2: Profile with torch.profiler (Level 1)
|
||||
|
||||
**Compile-safety rule for fused or rewritten kernels**
|
||||
- Any new kernel must be checked for `torch.compile` graph breaks before trusting its benchmark result.
|
||||
- If a direct Python/library call triggers tracing issues, wrap it as a custom op first.
|
||||
- For external libraries, use `register_custom_op_from_extern(...)`.
|
||||
- For SGLang JIT kernels, use `@register_custom_op(...)` and keep the JIT/module loading inside the custom op body.
|
||||
- Re-run `torch._dynamo.explain` on representative shapes and verify the optimized path still gets `graph_count=1` and `graph_break_count=0`.
|
||||
|
||||
```bash
|
||||
SGLANG_TORCH_PROFILER_DIR="${PROFILE_DIR}/torch" \
|
||||
sglang generate \
|
||||
@@ -464,11 +470,11 @@ TORCH_COMPILE_DEBUG=1 sglang generate ...
|
||||
- Dynamic shape changes trigger recompilation → fix resolution and frame count when benchmarking
|
||||
- `tensor.item()` in conditional branches causes graph breaks → rewrite as tensor ops
|
||||
|
||||
### Step 6: Multi-GPU Efficiency (Wan2.2-T2V-A14B)
|
||||
### Step 6: Multi-GPU Efficiency (Wan2.2-T2V-A14B / MOVA)
|
||||
|
||||
- Verify `--ulysses-degree` evenly divides `--num-gpus`
|
||||
- Confirm `--enable-cfg-parallel` is active (requires `guidance_scale > 1`)
|
||||
- `--dit-layerwise-offload true` introduces CPU↔GPU transfer overhead; disable when memory permits
|
||||
- Keep the command shape fixed when comparing kernels; for quick checks, reduce only `--num-inference-steps`
|
||||
- If a run OOMs or jitters because of host contention, first confirm there are no leaked scheduler processes on the chosen GPU set
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ Usage:
|
||||
# Side-by-side comparison
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --compare
|
||||
|
||||
# All 9 models, comparison
|
||||
# All 10 models, comparison
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --all --compare
|
||||
|
||||
Input images required for image-guided models:
|
||||
@@ -136,27 +136,20 @@ MODELS = {
|
||||
"false",
|
||||
],
|
||||
},
|
||||
# 6. Wan-AI/Wan2.2-T2V-A14B-Diffusers — Text-to-Video, 720P, 8 GPUs, 81 frames, 40 steps
|
||||
# 6. Wan-AI/Wan2.2-T2V-A14B-Diffusers — Text-to-Video, 720P, 4 GPUs, 81 frames, 2 steps
|
||||
"wan-t2v": {
|
||||
"path": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"prompt": "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon.",
|
||||
"negative_prompt": " ",
|
||||
"extra_args": [
|
||||
"--720p",
|
||||
"--num-inference-steps=40",
|
||||
"--num-inference-steps=2",
|
||||
"--num-frames=81",
|
||||
"--guidance-scale=5.0",
|
||||
"--num-gpus=8",
|
||||
"--enable-cfg-parallel",
|
||||
"--num-gpus=4",
|
||||
"--ulysses-degree=4",
|
||||
"--dit-layerwise-offload",
|
||||
"true",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--vae-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"true",
|
||||
"--pin-cpu-memory",
|
||||
],
|
||||
},
|
||||
# 7. Wan-AI/Wan2.2-TI2V-5B-Diffusers — Text-Image-to-Video, 720P, 1 GPU, 81 frames, 50 steps
|
||||
@@ -197,7 +190,7 @@ MODELS = {
|
||||
"--num-inference-steps=30",
|
||||
],
|
||||
},
|
||||
# 9. OpenMOSS-Team/MOVA-720p — Image-to-Video, 4 GPUs, 193 frames, 24 steps
|
||||
# 9. OpenMOSS-Team/MOVA-720p — Image-to-Video, 4 GPUs, 193 frames, 2 steps
|
||||
# Requires: <repo>/inputs/diffusion_benchmark/figs/mova_single_person.jpg
|
||||
"mova-720p": {
|
||||
"path": "OpenMOSS-Team/MOVA-720p",
|
||||
@@ -210,7 +203,25 @@ MODELS = {
|
||||
"--ulysses-degree=4",
|
||||
"--num-frames=193",
|
||||
"--fps=24",
|
||||
"--num-inference-steps=24",
|
||||
"--num-inference-steps=2",
|
||||
],
|
||||
},
|
||||
# 10. BestWishYsh/Helios-Base — Text-to-Video, 640×384, 33 frames
|
||||
"helios": {
|
||||
"path": "BestWishYsh/Helios-Base",
|
||||
"prompt": "A curious raccoon",
|
||||
"extra_args": [
|
||||
"--width=640",
|
||||
"--height=384",
|
||||
"--num-frames=33",
|
||||
"--dit-layerwise-offload",
|
||||
"false",
|
||||
"--dit-cpu-offload",
|
||||
"false",
|
||||
"--text-encoder-cpu-offload",
|
||||
"false",
|
||||
"--vae-cpu-offload",
|
||||
"false",
|
||||
],
|
||||
},
|
||||
}
|
||||
@@ -218,7 +229,7 @@ MODELS = {
|
||||
|
||||
def required_gpus_for_model(model_key: str) -> int:
|
||||
if model_key == "wan-t2v":
|
||||
return 8
|
||||
return 4
|
||||
if model_key == "mova-720p":
|
||||
return 4
|
||||
return 1
|
||||
|
||||
Reference in New Issue
Block a user