From 16041ab550608b429ca96ea3f9fff100f128ca37 Mon Sep 17 00:00:00 2001 From: zhou zhuoxin Date: Sun, 12 Apr 2026 22:06:49 +0800 Subject: [PATCH] [Refactor] Extend CFG Parallel to support 3 or 4 branch dispatch across M GPUs (#2423) --- docs/design/feature/cfg_parallel.md | 70 +++- docs/user_guide/diffusion_features.md | 2 +- .../image_to_image/image_edit.py | 4 +- .../x_to_video_audio/x_to_video_audio.py | 4 +- .../distributed/test_cfg_parallel.py | 342 +++++++++++++++++- .../diffusion/distributed/cfg_parallel.py | 180 +++++++++ .../dreamid_omni/pipeline_dreamid_omni.py | 107 ++---- .../models/omnigen2/pipeline_omnigen2.py | 83 +++-- 8 files changed, 670 insertions(+), 122 deletions(-) diff --git a/docs/design/feature/cfg_parallel.md b/docs/design/feature/cfg_parallel.md index 64decbe95..c73a87749 100644 --- a/docs/design/feature/cfg_parallel.md +++ b/docs/design/feature/cfg_parallel.md @@ -25,7 +25,9 @@ In standard Classifier-Free Guidance, each diffusion step requires two forward p 1. **Positive/Conditional**: Guided by the text prompt 2. **Negative/Unconditional**: Typically using empty or negative prompt -CFG-Parallel eliminates this bottleneck by distributing the two forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially. +Some models require 3 or more CFG branches (see [N-Branch CFG](#n-branch-cfg-3-branches)). + +CFG-Parallel eliminates this bottleneck by distributing the forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially. ### Architecture @@ -33,9 +35,11 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. | Method | Purpose | Automatic Behavior | |--------|---------|-------------------| -| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results | +| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with 2-branch CFG | Detects parallel mode, distributes computation, gathers results | +| [`predict_noise_with_multi_branch_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with N-branch CFG | Round-robin dispatches N branches across M GPUs | | [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler | All ranks step locally (no broadcast needed) | -| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization | +| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine 2-branch predictions | Applies CFG formula with optional normalization | +| [`combine_multi_branch_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine N-branch predictions | Override for custom multi-branch combine logic | | [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls | | [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization | @@ -57,6 +61,22 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. - All ranks compute the scheduler step locally — no broadcast needed because `predict_noise_maybe_with_cfg` already ensures all ranks have identical noise predictions after `all_gather` + local combine. +### N-Branch CFG (3+ branches) + +Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches. + +`predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`): + +| Branches (N) | GPUs (M) | Dispatch | +|:---:|:---:|:---| +| 3 | 2 | `[[0, 2], [1]]` | +| 3 | 3 | `[[0], [1], [2]]` | +| 4 | 2 | `[[0, 2], [1, 3]]` | +| 4 | 3 | `[[0, 3], [1], [2]]` | +| 4 | 4 | `[[0], [1], [2], [3]]` | + +When a rank handles multiple branches, it runs them sequentially. After `all_gather`, all ranks execute `combine_multi_branch_cfg_noise()` locally, producing identical results. + --- ## Step-by-Step Implementation @@ -98,6 +118,7 @@ class YourModelPipeline(nn.Module, CFGParallelMixin): - `positive_kwargs`: transformer arguments for conditional (text-guided) prediction - `negative_kwargs`: transformer arguments for unconditional prediction (set to `None` if CFG disabled) - For image editing pipelines, add `output_slice=image_seq_len` to extract the generative image portion +- For models with 3+ CFG branches, see [Multi-Branch CFG](#multi-branch-cfg-3-branches) in the Customization section ### Step 2: Call `diffuse` @@ -171,20 +192,42 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin): ``` -### Override `combine_cfg_noise()` for Multi-Output Models +### Multi-Branch CFG (3+ branches) -When `predict_noise()` returns a tuple (e.g., video + audio), the default `combine_cfg_noise()` applies CFG to every element. Override it to apply different logic per element — for example, CFG on video but positive-only on audio: +For models with 3 or more CFG branches, use `predict_noise_with_multi_branch_cfg()` instead of `predict_noise_maybe_with_cfg()`, and override `combine_multi_branch_cfg_noise()` for custom combine logic. This interface also works for standard 2-branch CFG — just pass 2 branches in `branches_kwargs`. + +**Example (3-branch with dual guidance scale):** ```python -class MyVideoAudioPipeline(nn.Module, CFGParallelMixin): - def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize): - (video_pos, audio_pos) = positive_noise_pred - (video_neg, audio_neg) = negative_noise_pred - video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize) - return (video_combined, audio_pos) # audio: positive only, no CFG +class YourMultiBranchPipeline(nn.Module, CFGParallelMixin): + def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False): + text_scale = true_cfg_scale["text"] + image_scale = true_cfg_scale["image"] + pos, ref, uncond = predictions + return uncond + image_scale * (ref - uncond) + text_scale * (pos - ref) + + def diffuse(self, ...): + for i, t in enumerate(timesteps): + positive_kwargs = {...} # conditional prompt + ref_neg_kwargs = {...} # negative prompt + reference + uncond_kwargs = {...} # unconditional + + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale={"text": text_guidance_scale, "image": image_guidance_scale}, + branches_kwargs=[positive_kwargs, ref_neg_kwargs, uncond_kwargs], + ) + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + return latents ``` -This also requires `predict_noise()` to return a tuple (see [Override predict_noise](#override-predict_noise-for-custom-transformer-calls) above). +### Override Combine Functions + +There are two combine functions for different scenarios: + +- **`combine_cfg_noise()`** — Used by `predict_noise_maybe_with_cfg()`. Override when `predict_noise()` returns a tuple (e.g., video + audio) and you need per-element CFG logic. +- **`combine_multi_branch_cfg_noise()`** — Used by `predict_noise_with_multi_branch_cfg()`. Override to implement custom multi-branch combine formulas (see [Multi-Branch CFG](#multi-branch-cfg-3-branches) above). ### Implement a Composite Scheduler for Multi-Output Models @@ -303,4 +346,5 @@ Adding CFG-Parallel support: 1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method 2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior -3. ✅ **Test** - Verify with `--cfg-parallel-size 2` and compare performance +3. ✅ **(Optional) Multi-branch** - For 3+ branch models, use `predict_noise_with_multi_branch_cfg()` and override `combine_multi_branch_cfg_noise()` +4. ✅ **Test** - Verify with `--cfg-parallel-size 2` (or 3/4 for multi-branch) and compare performance diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 7e0885181..2f28131ee 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -118,7 +118,7 @@ The following tables show which models support each feature: | **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | | **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index a8035a3fd..1a7e86f13 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -297,8 +297,8 @@ def parse_args() -> argparse.Namespace: "--cfg-parallel-size", type=int, default=1, - choices=[1, 2], - help="Number of GPUs used for classifier free guidance parallel size.", + choices=[1, 2, 3], + help="Number of GPUs used for classifier free guidance parallel size (max 3 branches).", ) parser.add_argument( "--enforce-eager", diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index fb77b2148..49a0f496f 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py @@ -36,8 +36,8 @@ def parse_args() -> argparse.Namespace: "--cfg-parallel-size", type=int, default=1, - choices=[1, 2], - help="Number of GPUs used for classifier free guidance parallel size.", + choices=[1, 2, 3, 4], + help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).", ) parser.add_argument( "--video-negative-prompt", diff --git a/tests/diffusion/distributed/test_cfg_parallel.py b/tests/diffusion/distributed/test_cfg_parallel.py index 79dbe9e6d..bf709618d 100644 --- a/tests/diffusion/distributed/test_cfg_parallel.py +++ b/tests/diffusion/distributed/test_cfg_parallel.py @@ -2,8 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for CFG (Classifier-Free Guidance) parallel functionality. -This test verifies that predict_noise_maybe_with_cfg produces numerically -equivalent results with and without CFG parallel using fixed random inputs. +This test verifies that predict_noise_maybe_with_cfg and +predict_noise_with_multi_branch_cfg produce numerically equivalent results +with and without CFG parallel using fixed random inputs. """ import os @@ -429,3 +430,340 @@ def test_predict_noise_without_cfg(dtype: torch.dtype): assert noise_pred.shape == (1, 4, 16, 16) print(f"✓ Test passed: predict_noise without CFG (dtype={dtype})") + + +class MultiBranchTestPipeline(CFGParallelMixin): + """Test pipeline with custom 3-branch combine logic (like OmniGen2).""" + + def __init__(self, in_channels: int = 4, hidden_dim: int = 128, seed: int = 42): + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + self.transformer = SimpleTransformer(in_channels, hidden_dim) + + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + for param in self.transformer.parameters(): + torch.nn.init.normal_(param, mean=0.0, std=0.02) + + def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False): + """N-branch combine with weighted sum for testing. + + - 2-branch: standard CFG formula (true_cfg_scale is float) + - 3-branch: OmniGen2-style dual guidance scale (true_cfg_scale is dict) + - 4-branch: DreamID-style weighted sum (true_cfg_scale is dict) + """ + if len(predictions) == 4: + text_scale = true_cfg_scale["text"] + image_scale = true_cfg_scale["image"] + vid_ref_scale = true_cfg_scale["vid_ref"] + pos, neg, vid_neg, audio_neg = predictions + combined = ( + audio_neg + + vid_ref_scale * (vid_neg - audio_neg) + + image_scale * (neg - vid_neg) + + text_scale * (pos - neg) + ) + elif len(predictions) == 3: + text_scale = true_cfg_scale["text"] + image_scale = true_cfg_scale["image"] + pos, ref, uncond = predictions + combined = uncond + image_scale * (ref - uncond) + text_scale * (pos - ref) + else: + pos, neg = predictions[0], predictions[1] + combined = neg + true_cfg_scale * (pos - neg) + + if cfg_normalize: + combined = self.cfg_normalize_function(pos, combined) + return combined + + +def _test_multi_branch_parallel_worker( + local_rank: int, + world_size: int, + cfg_parallel_size: int, + dtype: torch.dtype, + test_config: dict, + result_queue: torch.multiprocessing.Queue, +): + """Worker function for multi-branch CFG parallel test.""" + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + + update_environment_variables( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": "29504", + } + ) + + init_distributed_environment() + initialize_model_parallel(cfg_parallel_size=cfg_parallel_size) + + cfg_rank = get_classifier_free_guidance_rank() + cfg_world_size = get_classifier_free_guidance_world_size() + assert cfg_world_size == cfg_parallel_size + + pipeline = MultiBranchTestPipeline( + in_channels=test_config["channels"], + hidden_dim=test_config["hidden_dim"], + seed=test_config["model_seed"], + ) + pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) + pipeline.transformer.eval() + + n_branches = test_config["n_branches"] + batch_size = test_config["batch_size"] + channels = test_config["channels"] + height = test_config["height"] + width = test_config["width"] + + # Create N branch inputs with distinct seeds + branches_kwargs = [] + for b in range(n_branches): + torch.manual_seed(test_config["input_seed"] + b) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(test_config["input_seed"] + b) + x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) + branches_kwargs.append({"x": x}) + + with torch.no_grad(): + noise_pred = pipeline.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale=test_config["cfg_scale"], + branches_kwargs=branches_kwargs, + cfg_normalize=test_config["cfg_normalize"], + ) + + assert noise_pred is not None + result_queue.put((cfg_rank, noise_pred.cpu())) + + destroy_distributed_env() + + +def _test_multi_branch_sequential_worker( + local_rank: int, + world_size: int, + dtype: torch.dtype, + test_config: dict, + result_queue: torch.multiprocessing.Queue, +): + """Worker function for sequential multi-branch CFG test (baseline).""" + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + + update_environment_variables( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": "29505", + } + ) + + init_distributed_environment() + initialize_model_parallel(cfg_parallel_size=1) + + cfg_world_size = get_classifier_free_guidance_world_size() + assert cfg_world_size == 1 + + pipeline = MultiBranchTestPipeline( + in_channels=test_config["channels"], + hidden_dim=test_config["hidden_dim"], + seed=test_config["model_seed"], + ) + pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) + pipeline.transformer.eval() + + n_branches = test_config["n_branches"] + batch_size = test_config["batch_size"] + channels = test_config["channels"] + height = test_config["height"] + width = test_config["width"] + + branches_kwargs = [] + for b in range(n_branches): + torch.manual_seed(test_config["input_seed"] + b) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(test_config["input_seed"] + b) + x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device) + branches_kwargs.append({"x": x}) + + with torch.no_grad(): + noise_pred = pipeline.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale=test_config["cfg_scale"], + branches_kwargs=branches_kwargs, + cfg_normalize=test_config["cfg_normalize"], + ) + + assert noise_pred is not None + result_queue.put(noise_pred.cpu()) + + destroy_distributed_env() + + +@pytest.mark.parametrize( + "cfg_parallel_size,n_branches", + [ + (2, 2), # 2 branches on 2 GPUs: [[0],[1]] + (2, 3), # 3 branches on 2 GPUs: [[0,2],[1]] + (3, 3), # 3 branches on 3 GPUs: [[0],[1],[2]] + (2, 4), # 4 branches on 2 GPUs: [[0,2],[1,3]] + ], +) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("batch_size", [2]) +@pytest.mark.parametrize("cfg_normalize", [False, True]) +def test_predict_noise_with_multi_branch_cfg( + cfg_parallel_size: int, + n_branches: int, + dtype: torch.dtype, + batch_size: int, + cfg_normalize: bool, +): + """ + Test that predict_noise_with_multi_branch_cfg produces identical results + with and without CFG parallel for N-branch models. + + Args: + cfg_parallel_size: Number of GPUs for CFG parallel + n_branches: Number of CFG branches + dtype: Data type for computation + batch_size: Batch size for testing + cfg_normalize: Whether to normalize CFG output + """ + available_gpus = current_omni_platform.get_device_count() + if available_gpus < cfg_parallel_size: + pytest.skip(f"Test requires {cfg_parallel_size} GPUs but only {available_gpus} available") + + if n_branches == 2: + cfg_scale = 5.0 + elif n_branches == 3: + cfg_scale = {"text": 5.0, "image": 2.0} + else: + cfg_scale = {"text": 5.0, "image": 2.0, "vid_ref": 1.5} + + test_config = { + "batch_size": batch_size, + "channels": 4, + "height": 16, + "width": 16, + "hidden_dim": 128, + "cfg_scale": cfg_scale, + "cfg_normalize": cfg_normalize, + "model_seed": 42, + "input_seed": 123, + "n_branches": n_branches, + } + + mp_context = torch.multiprocessing.get_context("spawn") + manager = mp_context.Manager() + baseline_queue = manager.Queue() + cfg_parallel_queue = manager.Queue() + + # Run baseline (sequential, cfgp=1) + torch.multiprocessing.spawn( + _test_multi_branch_sequential_worker, + args=(1, dtype, test_config, baseline_queue), + nprocs=1, + ) + + # Run CFG parallel + torch.multiprocessing.spawn( + _test_multi_branch_parallel_worker, + args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue), + nprocs=cfg_parallel_size, + ) + + baseline_output = baseline_queue.get() + cfg_parallel_outputs = [cfg_parallel_queue.get() for _ in range(cfg_parallel_size)] + cfg_parallel_outputs.sort(key=lambda item: item[0]) + cfg_parallel_output = cfg_parallel_outputs[0][1] + + # All ranks should produce identical output + for cfg_rank, rank_output in cfg_parallel_outputs[1:]: + torch.testing.assert_close( + rank_output, + cfg_parallel_output, + rtol=0, + atol=0, + msg=f"Multi-branch CFG parallel ranks differ (rank 0 vs rank {cfg_rank})", + ) + + assert baseline_output.shape == cfg_parallel_output.shape, ( + f"Shape mismatch: baseline {baseline_output.shape} vs CFG parallel {cfg_parallel_output.shape}" + ) + + if dtype == torch.float32: + rtol, atol = 1e-5, 1e-5 + elif dtype == torch.bfloat16: + rtol, atol = 1e-2, 1e-2 + else: + rtol, atol = 1e-3, 1e-3 + + torch.testing.assert_close( + cfg_parallel_output, + baseline_output, + rtol=rtol, + atol=atol, + msg=( + f"Multi-branch CFG parallel output differs from sequential\n" + f" n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size}\n" + f" dtype={dtype}, cfg_normalize={cfg_normalize}\n" + f" Max diff: {(cfg_parallel_output - baseline_output).abs().max().item():.6e}" + ), + ) + + print( + f"✓ Test passed: multi_branch n_branches={n_branches}, " + f"cfg_size={cfg_parallel_size}, dtype={dtype}, cfg_normalize={cfg_normalize}" + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_multi_branch_without_cfg(dtype: torch.dtype): + """ + Test predict_noise_with_multi_branch_cfg when do_true_cfg=False. + + When CFG is disabled, only the first branch (positive) should be computed. + This test runs on a single GPU without distributed environment. + """ + available_gpus = current_omni_platform.get_device_count() + if available_gpus < 1: + pytest.skip("Test requires at least 1 GPU") + + device = torch.device(f"{current_omni_platform.device_type}:0") + current_omni_platform.set_device(device) + + pipeline = MultiBranchTestPipeline(in_channels=4, hidden_dim=128, seed=42) + pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype) + pipeline.transformer.eval() + + # Create 3 branch inputs (only first should be used) + branches_kwargs = [] + for b in range(3): + torch.manual_seed(123 + b) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(123 + b) + x = torch.randn(1, 4, 16, 16, dtype=dtype, device=device) + branches_kwargs.append({"x": x}) + + with torch.no_grad(): + noise_pred = pipeline.predict_noise_with_multi_branch_cfg( + do_true_cfg=False, # No CFG + true_cfg_scale=5.0, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + + assert noise_pred is not None + assert noise_pred.shape == (1, 4, 16, 16) + + print(f"✓ Test passed: multi_branch predict_noise without CFG (dtype={dtype})") diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index a8b0012f6..98757006b 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/vllm_omni/diffusion/distributed/cfg_parallel.py @@ -9,6 +9,7 @@ from abc import ABCMeta from typing import Any import torch +from vllm.logger import init_logger from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, @@ -16,6 +17,8 @@ from vllm_omni.diffusion.distributed.parallel_state import ( get_classifier_free_guidance_world_size, ) +logger = init_logger(__name__) + def _wrap(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: """Normalize prediction to tuple form.""" @@ -32,6 +35,24 @@ def _slice_pred(pred: tuple[torch.Tensor, ...], output_slice: int) -> tuple[torc return tuple(p[:, :output_slice] for p in pred) +def _dispatch_branches(n_branches: int, n_ranks: int) -> list[list[int]]: + """ + Round-robin dispatch N branches to M ranks. + + Rule: branch i → rank (i % n_ranks). + + Examples: + _dispatch_branches(3, 2) -> [[0, 2], [1]] + _dispatch_branches(3, 3) -> [[0], [1], [2]] + _dispatch_branches(4, 2) -> [[0, 2], [1, 3]] + _dispatch_branches(4, 4) -> [[0], [1], [2], [3]] + """ + assignments: list[list[int]] = [[] for _ in range(n_ranks)] + for i in range(n_branches): + assignments[i % n_ranks].append(i) + return assignments + + class CFGParallelMixin(metaclass=ABCMeta): """ Base Mixin class for Diffusion pipelines providing shared CFG methods. @@ -189,6 +210,165 @@ class CFGParallelMixin(metaclass=ABCMeta): results.append(comb) return _unwrap(tuple(results)) + # ── N-branch CFG interface (for 3+ branch models) ── + + def predict_noise_with_multi_branch_cfg( + self, + do_true_cfg: bool, + true_cfg_scale: float | dict[str, float], + branches_kwargs: list[dict[str, Any]], + cfg_normalize: bool = False, + output_slice: int | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """ + Predict noise with N-branch CFG dispatch across M GPUs. + + This is the multi-branch counterpart of predict_noise_maybe_with_cfg(). + Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel, + DreamID). Existing 2-branch models should continue using + predict_noise_maybe_with_cfg(). + + Args: + do_true_cfg: Whether to apply CFG. + true_cfg_scale: CFG scale factor (passed to combine_multi_branch_cfg_noise). + branches_kwargs: List of N dicts, each containing kwargs for one + predict_noise() call. branches_kwargs[0] is always the + positive/conditional branch. + cfg_normalize: Whether to normalize (passed to combine_multi_branch_cfg_noise). + output_slice: If set, slice each output to [:, :output_slice]. + + Returns: + Combined noise prediction, identical on all ranks in CFG parallel. + """ + if do_true_cfg: + n_branches = len(branches_kwargs) + cfg_world_size = get_classifier_free_guidance_world_size() + cfg_parallel_ready = cfg_world_size > 1 + + if cfg_parallel_ready: + return self._predict_multi_branch_parallel( + branches_kwargs, + n_branches, + cfg_world_size, + true_cfg_scale, + cfg_normalize, + output_slice, + ) + else: + # Sequential: run all N branches on single device + preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = [] + for kw in branches_kwargs: + pred = _wrap(self.predict_noise(**kw)) + if output_slice is not None: + pred = _slice_pred(pred, output_slice) + preds.append(_unwrap(pred)) + return self.combine_multi_branch_cfg_noise(preds, true_cfg_scale, cfg_normalize) + else: + # No CFG: only compute positive/conditional prediction + pred = self.predict_noise(**branches_kwargs[0]) + if output_slice is not None: + pred = _unwrap(_slice_pred(_wrap(pred), output_slice)) + return pred + + def _predict_multi_branch_parallel( + self, + branches_kwargs: list[dict[str, Any]], + n_branches: int, + cfg_world_size: int, + true_cfg_scale: float, + cfg_normalize: bool, + output_slice: int | None, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Dispatch N branches across M ranks, all_gather, then combine.""" + cfg_group = get_cfg_group() + cfg_rank = get_classifier_free_guidance_rank() + + if cfg_world_size > n_branches: + logger.warning_once( + "cfg_parallel_size=%d > n_branches=%d, %d GPU(s) will be idle for CFG", + cfg_world_size, + n_branches, + cfg_world_size - n_branches, + ) + + # Assign branches to ranks via round-robin + assignments = _dispatch_branches(n_branches, cfg_world_size) + my_branch_ids = assignments[cfg_rank] + max_per_rank = max(len(a) for a in assignments) + + # Run assigned branches + my_preds: list[tuple[torch.Tensor, ...]] = [] + for bid in my_branch_ids: + pred = _wrap(self.predict_noise(**branches_kwargs[bid])) + if output_slice is not None: + pred = _slice_pred(pred, output_slice) + my_preds.append(pred) + + # Idle ranks (cfg_world_size > n_branches) run a forward pass to get the output shape for all_gather. + # Output shape cannot be inferred from kwargs — may be tuple, sliced, etc. + if not my_preds: + pred = _wrap(self.predict_noise(**branches_kwargs[0])) + if output_slice is not None: + pred = _slice_pred(pred, output_slice) + my_preds.append(pred) + + # Pad to max_per_rank with zeros so all ranks have same size + ref_pred = my_preds[0] + while len(my_preds) < max_per_rank: + my_preds.append(tuple(torch.zeros_like(t) for t in ref_pred)) + + # All-gather each output element separately (like predict_noise_maybe_with_cfg) + # For each slot, gather across ranks; then pick valid results by owner_rank + # all_slots[slot][elem_idx] = [rank0_tensor, rank1_tensor, ...] + all_slots: list[list[list[torch.Tensor]]] = [] + for slot in range(max_per_rank): + slot_results: list[list[torch.Tensor]] = [] + for p in my_preds[slot]: + gathered = cfg_group.all_gather(p, separate_tensors=True) + slot_results.append(gathered) + all_slots.append(slot_results) + + # Reconstruct final_preds in branch order + final_preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = [] + for bid in range(n_branches): + owner_rank = bid % cfg_world_size + slot_idx = bid // cfg_world_size + elements = tuple(all_slots[slot_idx][elem_idx][owner_rank] for elem_idx in range(len(ref_pred))) + final_preds.append(_unwrap(elements)) + + return self.combine_multi_branch_cfg_noise(final_preds, true_cfg_scale, cfg_normalize) + + def combine_multi_branch_cfg_noise( + self, + predictions: list[torch.Tensor | tuple[torch.Tensor, ...]], + true_cfg_scale: float | dict[str, float], + cfg_normalize: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """ + Combine N branch predictions. Default: standard 2-branch CFG formula. + + Override this method for custom multi-branch combine logic. + + Args: + predictions: List of N predictions, where predictions[0] is always + the positive/conditional branch. + true_cfg_scale: CFG scale factor (float for 2-branch, dict for multi-branch). + cfg_normalize: Whether to normalize the combined prediction. + + Returns: + Combined noise prediction. + """ + positive = _wrap(predictions[0]) + negative = _wrap(predictions[1]) + + results = [] + for p, n in zip(positive, negative): + comb = n + true_cfg_scale * (p - n) + if cfg_normalize: + comb = self.cfg_normalize_function(p, comb) + results.append(comb) + return _unwrap(tuple(results)) + def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Forward pass through transformer to predict noise. diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py index e22765f80..974cc582f 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py @@ -15,11 +15,6 @@ from tqdm import tqdm from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.parallel_state import ( - get_cfg_group, - get_classifier_free_guidance_rank, - get_classifier_free_guidance_world_size, -) from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -249,6 +244,28 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor return sample_scheduler, timesteps + def predict_noise(self, **kwargs): + pred_vid, pred_audio = self.model(**kwargs) + return (pred_vid[0], pred_audio[0]) + + def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False): + vid_pos, audio_pos = predictions[0] + vid_neg, audio_neg = predictions[1] + vid_ip_neg, _ = predictions[2] + _, refaudio_neg = predictions[3] + + pred_video = ( + vid_neg + + true_cfg_scale["video_cfg_scale"] * (vid_pos - vid_neg) + + true_cfg_scale["video_ref_cfg_scale"] * (vid_pos - vid_ip_neg) + ) + pred_audio = ( + audio_neg + + true_cfg_scale["audio_cfg_scale"] * (audio_pos - audio_neg) + + true_cfg_scale["audio_ref_cfg_scale"] * (audio_pos - refaudio_neg) + ) + return (pred_video, pred_audio) + def diffuse( self, video_noise: torch.Tensor, @@ -306,72 +323,22 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor "vid_context": [text_embeddings_video_neg], } - if get_classifier_free_guidance_world_size() > 1: - # Enable CFG-parallel: rank0 computes positive, rank1 computes negative. - cfg_group = get_cfg_group() - cfg_rank = get_classifier_free_guidance_rank() + branches_kwargs = [ + {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **pos_args}, + {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **neg_args}, + {"vid": [model_input_video_neg], "audio": [model_input_audio], "t": timestep_input, **pos_args}, + {"vid": [model_input_video], "audio": [model_input_audio_neg], "t": timestep_input, **pos_args}, + ] - if cfg_rank == 0: - pred_vid, pred_audio = self.model( - vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args - ) - pre_vid_ip_neg, _ = self.model( - vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args - ) - pred_vid_0 = pred_vid[0] - pred_audio_0 = pred_audio[0] - pre_vid_ip_0 = pre_vid_ip_neg[0] - pred_refaudio_0 = torch.zeros_like(pred_audio_0) # dummy tensor - else: - pred_vid, pred_audio = self.model( - vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args - ) - _, pred_refaudio_neg = self.model( - vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args - ) - pred_vid_0 = pred_vid[0] - pred_audio_0 = pred_audio[0] - pre_vid_ip_0 = torch.zeros_like(pred_vid_0) # dummy tensor - pred_refaudio_0 = pred_refaudio_neg[0] - - pred_vid_gathered = cfg_group.all_gather(pred_vid_0, separate_tensors=True) - pred_audio_gathered = cfg_group.all_gather(pred_audio_0, separate_tensors=True) - pre_vid_ip_gathered = cfg_group.all_gather(pre_vid_ip_0, separate_tensors=True) - pred_refaudio_gathered = cfg_group.all_gather(pred_refaudio_0, separate_tensors=True) - - pred_vid_pos = [pred_vid_gathered[0]] - pred_vid_neg = [pred_vid_gathered[1]] - pred_audio_pos = [pred_audio_gathered[0]] - pred_audio_neg = [pred_audio_gathered[1]] - pre_vid_ip_neg = [pre_vid_ip_gathered[0]] - pred_refaudio_neg = [pred_refaudio_gathered[1]] - else: - pred_vid_pos, pred_audio_pos = self.model( - vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args - ) - - pred_vid_neg, pred_audio_neg = self.model( - vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args - ) - - pre_vid_ip_neg, _ = self.model( - vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args - ) - - _, pred_refaudio_neg = self.model( - vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args - ) - - pred_video_guided = ( - pred_vid_neg[0] - + self.video_cfg_scale * (pred_vid_pos[0] - pred_vid_neg[0]) - + self.video_ref_cfg_scale * (pred_vid_pos[0] - pre_vid_ip_neg[0]) - ) - - pred_audio_guided = ( - pred_audio_neg[0] - + self.audio_cfg_scale * (pred_audio_pos[0] - pred_audio_neg[0]) - + self.audio_ref_cfg_scale * (pred_audio_pos[0] - pred_refaudio_neg[0]) + pred_video_guided, pred_audio_guided = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "video_cfg_scale": self.video_cfg_scale, + "video_ref_cfg_scale": self.video_ref_cfg_scale, + "audio_cfg_scale": self.audio_cfg_scale, + "audio_ref_cfg_scale": self.audio_ref_cfg_scale, + }, + branches_kwargs=branches_kwargs, ) video_noise = scheduler_video.step( pred_video_guided.unsqueeze(0), t_v, video_noise.unsqueeze(0), return_dict=False diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py index 2d370aea1..e8e307b87 100644 --- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py +++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py @@ -29,6 +29,7 @@ from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import ( @@ -619,7 +620,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OmniGen2Pipeline(nn.Module): +class OmniGen2Pipeline(CFGParallelMixin, nn.Module): """ Pipeline for text-to-image generation using OmniGen2. @@ -1171,14 +1172,6 @@ class OmniGen2Pipeline(nn.Module): self._num_timesteps = len(timesteps) for i, t in enumerate(timesteps): - model_pred = self.predict( - t=t, - latents=latents, - prompt_embeds=prompt_embeds, - freqs_cis=freqs_cis, - prompt_attention_mask=prompt_attention_mask, - ref_image_hidden_states=ref_latents, - ) text_guidance_scale = ( self.text_guidance_scale if self.cfg_range[0] <= i / len(timesteps) <= self.cfg_range[1] else 1.0 ) @@ -1186,8 +1179,26 @@ class OmniGen2Pipeline(nn.Module): self.image_guidance_scale if self.cfg_range[0] <= i / len(timesteps) <= self.cfg_range[1] else 1.0 ) + positive_kwargs = dict( + t=t, + latents=latents, + prompt_embeds=prompt_embeds, + freqs_cis=freqs_cis, + prompt_attention_mask=prompt_attention_mask, + ref_image_hidden_states=ref_latents, + ) + uncond_kwargs = dict( + t=t, + latents=latents, + prompt_embeds=negative_prompt_embeds, + freqs_cis=freqs_cis, + prompt_attention_mask=negative_prompt_attention_mask, + ref_image_hidden_states=None, + ) + if text_guidance_scale > 1.0 and image_guidance_scale > 1.0: - model_pred_ref = self.predict( + # 3-branch CFG: pos + ref_neg + uncond + ref_neg_kwargs = dict( t=t, latents=latents, prompt_embeds=negative_prompt_embeds, @@ -1195,31 +1206,24 @@ class OmniGen2Pipeline(nn.Module): prompt_attention_mask=negative_prompt_attention_mask, ref_image_hidden_states=ref_latents, ) - - model_pred_uncond = self.predict( - t=t, - latents=latents, - prompt_embeds=negative_prompt_embeds, - freqs_cis=freqs_cis, - prompt_attention_mask=negative_prompt_attention_mask, - ref_image_hidden_states=None, - ) - - model_pred = ( - model_pred_uncond - + image_guidance_scale * (model_pred_ref - model_pred_uncond) - + text_guidance_scale * (model_pred - model_pred_ref) + model_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "text": text_guidance_scale, + "image": image_guidance_scale, + }, + branches_kwargs=[positive_kwargs, ref_neg_kwargs, uncond_kwargs], ) elif text_guidance_scale > 1.0: - model_pred_uncond = self.predict( - t=t, - latents=latents, - prompt_embeds=negative_prompt_embeds, - freqs_cis=freqs_cis, - prompt_attention_mask=negative_prompt_attention_mask, - ref_image_hidden_states=None, + # 2-branch CFG: pos + uncond + model_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale=text_guidance_scale, + branches_kwargs=[positive_kwargs, uncond_kwargs], ) - model_pred = model_pred_uncond + text_guidance_scale * (model_pred - model_pred_uncond) + else: + # No CFG + model_pred = self.predict_noise(**positive_kwargs) latents = self.scheduler.step(model_pred, t, latents, return_dict=False)[0] @@ -1265,6 +1269,21 @@ class OmniGen2Pipeline(nn.Module): ) return model_pred + def predict_noise(self, **kwargs): + """Override CFGParallelMixin.predict_noise to use self.predict.""" + return self.predict(**kwargs) + + def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False): + """Override: 3-branch dual scale or 2-branch standard CFG.""" + if len(predictions) == 3: + text_scale = true_cfg_scale["text"] + image_scale = true_cfg_scale["image"] + pos, ref, uncond = predictions[0], predictions[1], predictions[2] + return uncond + image_scale * (ref - uncond) + text_scale * (pos - ref) + # 2-branch: standard CFG + pos, neg = predictions[0], predictions[1] + return neg + true_cfg_scale * (pos - neg) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return loader.load_weights(weights)