"docs/source/deployment/frameworks/triton.md" did not exist on "32aa2059addd97be1afce7a199d228191710c294"
test_pipeline_parallel.py 3.89 KB
Newer Older
1
2
3
4
5
6
7
"""
WARNING: This test runs in both single-node (4 GPUs) and multi-node
 (2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
 important to set the distributed backend to "mp" to avoid Ray scheduling
 all workers in a node other than the head node, which can cause the test
 to fail.
"""
8
9
import os

10
11
import pytest

12
13
from vllm.logger import init_logger

14
from ..utils import compare_two_settings, fork_new_process_for_each_test
15

16
17
logger = init_logger("test_pipeline_parallel")

18
19
VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"

20

21
22
23
24
25
26
27
28
29
30
31
32
33
34
@pytest.mark.parametrize(
    ("TP_SIZE, PP_SIZE, EAGER_MODE, CHUNKED_PREFILL, TRUST_REMOTE_CODE, "
     "MODEL_NAME, DIST_BACKEND"),
    [
        (2, 2, 0, 1, 0, "meta-llama/Meta-Llama-3-8B", "mp"),
        (2, 2, 1, 0, 0, "meta-llama/Meta-Llama-3-8B", "mp"),
        (1, 3, 0, 0, 0, "meta-llama/Meta-Llama-3-8B", "mp"),
        (1, 4, 0, 1, 0, "meta-llama/Meta-Llama-3-8B", "mp"),
        (1, 4, 1, 0, 0, "meta-llama/Meta-Llama-3-8B", "mp"),
        (1, 3, 0, 0, 0, "meta-llama/Meta-Llama-3-8B", "ray"),
        (1, 4, 0, 1, 0, "meta-llama/Meta-Llama-3-8B", "ray"),
        (1, 4, 1, 0, 0, "meta-llama/Meta-Llama-3-8B", "ray"),
        (2, 2, 1, 0, 0, "meta-llama/Meta-Llama-3-8B", "ray"),
        (2, 2, 0, 1, 0, "meta-llama/Meta-Llama-3-8B", "ray"),
35
36
37
        (1, 2, 1, 1, 1, "OpenGVLab/InternVL2-1B", "ray"),
        (1, 2, 1, 1, 1, "OpenGVLab/InternVL2-2B", "ray"),
        (1, 2, 1, 0, 1, "OpenGVLab/InternVL2-4B", "ray"),
38
39
    ],
)
40
@fork_new_process_for_each_test
41
42
def test_compare_tp(TP_SIZE, PP_SIZE, EAGER_MODE, CHUNKED_PREFILL,
                    TRUST_REMOTE_CODE, MODEL_NAME, DIST_BACKEND):
43
44
45
    if VLLM_MULTI_NODE and DIST_BACKEND == "mp":
        pytest.skip("Skipping multi-node pipeline parallel test for "
                    "multiprocessing distributed backend")
46

47
    pp_args = [
48
49
        # use half precision for speed and memory savings in CI environment
        "--dtype",
50
        "float16",
51
52
        "--max-model-len",
        "8192",
53
54
55
56
57
        "--pipeline-parallel-size",
        str(PP_SIZE),
        "--tensor-parallel-size",
        str(TP_SIZE),
        "--distributed-executor-backend",
58
        DIST_BACKEND,
59
    ]
60
61
62
63
64
65
66
67
68

    # compare without pipeline parallelism
    # NOTE: use mp backend for TP
    # PP tests might involve multiple nodes, and ray might
    #  schedule all workers in a node other than the head node,
    #  which can cause the test to fail.
    tp_args = [
        # use half precision for speed and memory savings in CI environment
        "--dtype",
69
70
71
        "float16",
        "--max-model-len",
        "8192",
72
        "--tensor-parallel-size",
73
        str(max(TP_SIZE, 2)),  # We only use 2 GPUs in the CI.
74
75
76
        "--distributed-executor-backend",
        "mp",
    ]
77
    if CHUNKED_PREFILL:
78
79
        pp_args.append("--enable-chunked-prefill")
        tp_args.append("--enable-chunked-prefill")
80
    if EAGER_MODE:
81
82
        pp_args.append("--enforce-eager")
        tp_args.append("--enforce-eager")
83
84
85
    if TRUST_REMOTE_CODE:
        pp_args.append("--trust-remote-code")
        tp_args.append("--trust-remote-code")
86
    pp_env = None
87
88
89
    if (DIST_BACKEND == "ray" and TP_SIZE == 2 and PP_SIZE == 2
            and CHUNKED_PREFILL):
        # Test Ray ADAG for a subset of the tests
90
91
92
        pp_env = {
            "VLLM_USE_RAY_COMPILED_DAG": "1",
            "VLLM_USE_RAY_SPMD_WORKER": "1",
93
            "VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL": "1",
94
        }
95
96
97
98
        # Temporary. Currently when zeromq + SPMD is used, it does not properly
        # terminate because of aDAG issue.
        pp_args.append("--disable-frontend-multiprocessing")
        tp_args.append("--disable-frontend-multiprocessing")
99

100
101
102
103
104
105
106
107
    try:
        compare_two_settings(MODEL_NAME, pp_args, tp_args, pp_env)
    except Exception:
        if pp_env is None:
            raise
        else:
            # Ray ADAG tests are flaky, so we don't want to fail the test
            logger.exception("Ray ADAG tests failed")