[Refactor] Extend CFG Parallel to support 3 or 4 branch dispatch across M GPUs (#2423)

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})")