| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test Dual Batch Overlap (DBO) with Data Parallelism + Expert Parallelism. |
| 4 | |
| 5 | DBO is specifically designed for DP+EP scenarios to hide communication latency |
| 6 | by overlapping computation of two batches. This test validates that DBO works |
| 7 | correctly with the DeepSeek-V2-Lite model using GSM8K evaluation. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from tests.utils import RemoteOpenAIServer |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.import_utils import has_deep_ep |
| 16 | |
| 17 | # Detect Blackwell / B200 (compute capability 10.x) |
| 18 | try: |
| 19 | if torch.cuda.is_available(): |
| 20 | cap = torch.cuda.get_device_capability(0) |
| 21 | IS_BLACKWELL = cap[0] >= 10 |
| 22 | else: |
| 23 | IS_BLACKWELL = False |
| 24 | except Exception: |
| 25 | # Be conservative: if we can't detect, don't xfail by default |
| 26 | IS_BLACKWELL = False |
| 27 | |
| 28 | MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat" |
| 29 | DP_SIZE = 2 |
| 30 | |
| 31 | # GSM8K eval configuration |
| 32 | NUM_QUESTIONS = 256 # Fast eval for CI; but must be large enough to hit dbo thresholds |
| 33 | NUM_SHOTS = 5 # Few-shot examples |
| 34 | |
| 35 | # Expected 0.64 with 2% buffer (based on vLLM test data) |
| 36 | # On ROCm, widened to 0.61 after observing 0.617 amd-ci nightly flake twice. |
| 37 | MIN_ACCURACY = 0.61 if current_platform.is_rocm() else 0.62 |
| 38 | |
| 39 | # Increase max_num_seqs to trigger DBO for decode batches |
| 40 | # With 64 seqs, decode batches should exceed the 32 token threshold |
| 41 | MAX_NUM_SEQS = 64 # Increased from 16 to trigger decode DBO |
| 42 | |
| 43 | # DeepEP backends to test |
| 44 | DEEPEP_BACKENDS = [ |
| 45 | "deepep_low_latency", |
| 46 | "deepep_high_throughput", |
| 47 | ] |
| 48 | |
| 49 | |
| 50 | @pytest.mark.skipif(not has_deep_ep(), reason="These tests require deep_ep to run") |
| 51 | @pytest.mark.parametrize("all2all_backend", DEEPEP_BACKENDS) |
| 52 | @pytest.mark.xfail( |
| 53 | IS_BLACKWELL, |
| 54 | reason=( |
| 55 | "Temporary: DBO accuracy unstable on Blackwell " |
| 56 | "(doesn't meet expectation of MIN_ACCURACY = 0.62)" |
| 57 | ), |
| 58 | ) |
| 59 | def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available): |
| 60 | """Test DBO with DP+EP using GSM8K evaluation.""" |
| 61 | lm_eval = pytest.importorskip("lm_eval") |
| 62 | required_gpus = DP_SIZE |
| 63 | |
| 64 | if num_gpus_available < required_gpus: |
| 65 | pytest.skip(f"Need at least {required_gpus} GPUs (DP={DP_SIZE})") |
| 66 | |
| 67 | # Server arguments for DBO + DP + EP |
| 68 | server_args = [ |
| 69 | "--max-model-len", |
| 70 | "4096", |
| 71 | "--max-num-seqs", |
| 72 | str(MAX_NUM_SEQS), # Use larger batch to trigger decode DBO |
| 73 | "--trust-remote-code", |
| 74 | # Note: Not using --enforce-eager to test DBO's alternate CUDA graph dispatching |
| 75 | "--data-parallel-size", |
| 76 | str(DP_SIZE), |
| 77 | "--enable-expert-parallel", |
| 78 | "--enable-dbo", |
| 79 | # Fix threshold so we know we trigger DBO |
| 80 | "--dbo-decode-token-threshold", |
| 81 | "16", |
| 82 | "--dbo-prefill-token-threshold", |
| 83 | "32", |
| 84 | "--all2all-backend", |
| 85 | all2all_backend, |
| 86 | ] |
| 87 | |
| 88 | with RemoteOpenAIServer( |
| 89 | MODEL_NAME, |
| 90 | server_args, |
| 91 | max_wait_seconds=600, # Allow time for model loading with DP+EP |
| 92 | ) as remote_server: |
| 93 | # Use host and port directly from RemoteOpenAIServer |
| 94 | |
| 95 | base_url = f"http://{remote_server.host}:{remote_server.port}/v1/completions" |
| 96 | |
| 97 | # Run GSM8K evaluation |
| 98 | results = lm_eval.simple_evaluate( |
| 99 | model="local-completions", |
| 100 | model_args=( |
| 101 | f"pretrained={MODEL_NAME}," |
| 102 | f"base_url={base_url}," |
| 103 | "num_concurrent=512,max_retries=3" |
| 104 | ), |
| 105 | tasks=["gsm8k"], |
| 106 | num_fewshot=NUM_SHOTS, |
| 107 | limit=NUM_QUESTIONS, |
| 108 | ) |
| 109 | # Validate accuracy is reasonable |
| 110 | gsm8k = results["results"]["gsm8k"] |
| 111 | accuracy = gsm8k.get( |
| 112 | "exact_match,strict-match", |
| 113 | gsm8k.get("exact_match,flexible-extract"), |
| 114 | ) |
| 115 | assert accuracy is not None, f"gsm8k exact_match missing: {gsm8k}" |
| 116 | assert accuracy >= MIN_ACCURACY, ( |
| 117 | f"DBO+DP+EP accuracy too low ({all2all_backend}): " |
| 118 | f"{accuracy:.3f} < {MIN_ACCURACY:.3f} " |
| 119 | ) |