diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index 1374a4752..7152b6986 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/tests/e2e/online_serving/test_bagel_expansion.py @@ -6,7 +6,6 @@ Coverage: - TeaCache - Cache-DiT - CFG-Parallel -- Tensor-Parallel - Ulysses-SP - Ring-Attention - Layerwise Offloading @@ -15,62 +14,23 @@ assert_diffusion_response validates successful generation and the expected 512x512 resolution. """ -import json - import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data -from tests.helpers.stage_config import get_deploy_config_path pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] -_BAGEL_DEFAULT_YAML = get_deploy_config_path("ci/bagel.yaml") - PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution." NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) -TP_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=3) - - -def _make_tp_cases(model: str, tp_size: int): - """Build Bagel TP test cases with devices auto-derived from tp_size. - Devices can not be set through CLI args, so we set them in the YAML. - """ - # Stage 0 uses GPU 0, so place DiT TP ranks on GPU 1..N to - # avoid AR/DiT memory contention on one device. - devices = ",".join(str(i + 1) for i in range(tp_size)) - stage_overrides = json.dumps( - { - "0": {"tensor_parallel_size": 1}, - "1": {"devices": devices}, - } - ) - return [ - pytest.param( - OmniServerParams( - model=model, - stage_config_path=_BAGEL_DEFAULT_YAML, - server_args=[ - "--stage-overrides", - stage_overrides, - "--cache-backend", - "cache_dit", - "--tensor-parallel-size", - str(tp_size), - ], - ), - id=f"parallel_tp_{tp_size}", - marks=TP_FEATURE_MARKS, - ), - ] def _get_diffusion_feature_cases(model: str): """Return L4 diffusion feature cases for Bagel. - TeaCache, Cache-DiT, CFG-Parallel, Tensor-Parallel, + TeaCache, Cache-DiT, CFG-Parallel, Ulysses-SP, Ring-Attention, Layerwise Offloading. """ @@ -113,10 +73,6 @@ def _get_diffusion_feature_cases(model: str): id="parallel_cfg_2", marks=PARALLEL_FEATURE_MARKS, ), - # Tensor-Parallel size 2 (2 GPUs, Cache-DiT backend) - # Stage 1 (DiT) needs visible GPUs matching TP size; the default YAML - # only exposes device "0", so we patch it here. - *_make_tp_cases(model, tp_size=2), # Ulysses-SP degree=2 (2 GPUs) pytest.param( OmniServerParams( @@ -181,7 +137,6 @@ def test_bagel( - TeaCache - Cache-DiT - CFG-Parallel (size=2) - - Tensor-Parallel (size=2) - Ulysses-SP (degree=2) - Ring-Attention (degree=2) - Layerwise Offloading