mirror of
https://github.com/vllm-project/vllm.git
synced 2025-10-20 23:03:52 +08:00
328 lines
9.5 KiB
Python
328 lines
9.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
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.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Literal, NamedTuple
|
|
|
|
import pytest
|
|
|
|
from vllm.config.compilation import CompilationMode
|
|
from vllm.config.model import RunnerOption
|
|
from vllm.logger import init_logger
|
|
from vllm.utils import is_torch_equal_or_newer
|
|
|
|
from ..models.registry import HF_EXAMPLE_MODELS
|
|
from ..utils import compare_two_settings, create_new_process_for_each_test
|
|
|
|
logger = init_logger("test_sequence_parallel")
|
|
|
|
VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
|
|
|
|
|
|
class ParallelSetup(NamedTuple):
|
|
tp_size: int
|
|
pp_size: int
|
|
enable_fusion: bool
|
|
eager_mode: bool
|
|
chunked_prefill: bool
|
|
|
|
|
|
class SPTestOptions(NamedTuple):
|
|
multi_node_only: bool
|
|
load_format: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class SPTestSettings:
|
|
parallel_setups: list[ParallelSetup]
|
|
distributed_backends: list[str]
|
|
runner: RunnerOption
|
|
test_options: SPTestOptions
|
|
|
|
@staticmethod
|
|
def detailed(
|
|
*,
|
|
tp_base: int = 2,
|
|
pp_base: int = 1,
|
|
multi_node_only: bool = False,
|
|
runner: RunnerOption = "auto",
|
|
load_format: str | None = None,
|
|
):
|
|
parallel_setups = []
|
|
for eager_mode_val in [False, True]:
|
|
for pp_multiplier in [1, 2]:
|
|
for chunked_prefill_val in [False, True]:
|
|
parallel_setups.append(
|
|
ParallelSetup(
|
|
tp_size=tp_base,
|
|
pp_size=pp_multiplier * pp_base,
|
|
enable_fusion=False,
|
|
eager_mode=eager_mode_val,
|
|
chunked_prefill=chunked_prefill_val,
|
|
)
|
|
)
|
|
return SPTestSettings(
|
|
parallel_setups=parallel_setups,
|
|
distributed_backends=["mp", "ray"],
|
|
runner=runner,
|
|
test_options=SPTestOptions(
|
|
multi_node_only=multi_node_only, load_format=load_format
|
|
),
|
|
)
|
|
|
|
@staticmethod
|
|
def fast(
|
|
*,
|
|
tp_base: int = 2,
|
|
pp_base: int = 1,
|
|
runner: RunnerOption = "auto",
|
|
multi_node_only: bool = False,
|
|
load_format: str | None = None,
|
|
):
|
|
parallel_setups = []
|
|
for eager_mode_val in [False, True]:
|
|
for pp_multiplier in [1, 2]:
|
|
for chunked_prefill_val in [False, True]:
|
|
parallel_setups.append(
|
|
ParallelSetup(
|
|
tp_size=tp_base,
|
|
pp_size=pp_multiplier * pp_base,
|
|
enable_fusion=False,
|
|
eager_mode=eager_mode_val,
|
|
chunked_prefill=chunked_prefill_val,
|
|
)
|
|
)
|
|
return SPTestSettings(
|
|
parallel_setups=parallel_setups,
|
|
distributed_backends=["mp", "ray"],
|
|
runner=runner,
|
|
test_options=SPTestOptions(
|
|
multi_node_only=multi_node_only, load_format=load_format
|
|
),
|
|
)
|
|
|
|
@staticmethod
|
|
def fp8_quant(
|
|
*,
|
|
tp_base: int = 2,
|
|
pp_base: int = 1,
|
|
runner: RunnerOption = "auto",
|
|
multi_node_only: bool = False,
|
|
load_format: str | None = None,
|
|
):
|
|
parallel_setups = []
|
|
for fusion_val in [False, True]:
|
|
parallel_setups.append(
|
|
ParallelSetup(
|
|
tp_size=tp_base,
|
|
pp_size=pp_base,
|
|
enable_fusion=fusion_val,
|
|
eager_mode=True,
|
|
chunked_prefill=False,
|
|
)
|
|
)
|
|
return SPTestSettings(
|
|
parallel_setups=parallel_setups,
|
|
distributed_backends=["mp", "ray"],
|
|
runner=runner,
|
|
test_options=SPTestOptions(
|
|
multi_node_only=multi_node_only, load_format=load_format
|
|
),
|
|
)
|
|
|
|
def iter_params(self, model_id: str):
|
|
opts = self.test_options
|
|
|
|
for parallel_setup in self.parallel_setups:
|
|
for backend in self.distributed_backends:
|
|
yield (
|
|
model_id,
|
|
parallel_setup,
|
|
backend,
|
|
self.runner,
|
|
opts,
|
|
)
|
|
|
|
|
|
def _compare_sp(
|
|
model_id: str,
|
|
parallel_setup: ParallelSetup,
|
|
distributed_backend: str,
|
|
runner: RunnerOption,
|
|
test_options: SPTestOptions,
|
|
num_gpus_available: int,
|
|
use_inductor_graph_partition: bool,
|
|
*,
|
|
method: Literal["generate", "encode"],
|
|
is_multimodal: bool,
|
|
):
|
|
(
|
|
tp_size,
|
|
pp_size,
|
|
enable_fusion,
|
|
eager_mode,
|
|
chunked_prefill,
|
|
) = parallel_setup
|
|
|
|
multi_node_only, load_format = test_options
|
|
|
|
model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
|
|
model_info.check_transformers_version(on_fail="skip")
|
|
|
|
trust_remote_code = model_info.trust_remote_code
|
|
tokenizer_mode = model_info.tokenizer_mode
|
|
hf_overrides = model_info.hf_overrides
|
|
skip_tokenizer_init = model_info.skip_tokenizer_init
|
|
|
|
if load_format == "dummy":
|
|
# Avoid OOM
|
|
text_overrides = {
|
|
"num_hidden_layers": 4,
|
|
"hidden_size": 512,
|
|
"intermediate_size": 800,
|
|
"num_attention_heads": 4,
|
|
"num_key_value_heads": 1,
|
|
}
|
|
|
|
if is_multimodal:
|
|
hf_overrides.update({"text_config": text_overrides})
|
|
else:
|
|
hf_overrides.update(text_overrides)
|
|
else:
|
|
model_info.check_available_online(on_fail="skip")
|
|
|
|
if num_gpus_available < tp_size * pp_size:
|
|
pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
|
|
if VLLM_MULTI_NODE and distributed_backend == "mp":
|
|
pytest.skip(
|
|
"Skipping multi-node pipeline parallel test for "
|
|
"multiprocessing distributed backend"
|
|
)
|
|
if multi_node_only and not VLLM_MULTI_NODE:
|
|
pytest.skip("Not in multi-node setting")
|
|
|
|
common_args = [
|
|
# use half precision for speed and memory savings in CI environment
|
|
"--dtype",
|
|
"float16",
|
|
"--max-model-len",
|
|
"2048",
|
|
"--max-num-seqs",
|
|
"8",
|
|
]
|
|
if chunked_prefill:
|
|
common_args.append("--enable-chunked-prefill")
|
|
if eager_mode:
|
|
common_args.append("--enforce-eager")
|
|
if runner != "auto":
|
|
common_args.extend(["--runner", runner])
|
|
if trust_remote_code:
|
|
common_args.append("--trust-remote-code")
|
|
if tokenizer_mode:
|
|
common_args.extend(["--tokenizer-mode", tokenizer_mode])
|
|
if load_format:
|
|
common_args.extend(["--load-format", load_format])
|
|
if hf_overrides:
|
|
common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
|
|
if skip_tokenizer_init:
|
|
common_args.append("--skip-tokenizer-init")
|
|
|
|
compilation_config = {
|
|
"mode": CompilationMode.VLLM_COMPILE,
|
|
"custom_ops": ["+rms_norm"],
|
|
"compile_sizes": [4, 8],
|
|
"pass_config": {
|
|
"enable_sequence_parallelism": True,
|
|
"enable_fusion": enable_fusion,
|
|
"enable_noop": True,
|
|
},
|
|
"use_inductor_graph_partition": use_inductor_graph_partition,
|
|
}
|
|
|
|
tp_sp_args = [
|
|
*common_args,
|
|
"--tensor-parallel-size",
|
|
str(tp_size),
|
|
"--pipeline-parallel-size",
|
|
str(pp_size),
|
|
"--distributed-executor-backend",
|
|
distributed_backend,
|
|
"--compilation_config",
|
|
json.dumps(compilation_config),
|
|
]
|
|
|
|
tp_args = [
|
|
*common_args,
|
|
"--tensor-parallel-size",
|
|
str(tp_size),
|
|
"--distributed-executor-backend",
|
|
"mp",
|
|
]
|
|
|
|
compare_two_settings(model_id, tp_sp_args, tp_args, method=method)
|
|
|
|
|
|
SP_TEXT_GENERATION_MODELS = {
|
|
# [Decoder-only]
|
|
"meta-llama/Llama-3.2-1B-Instruct": SPTestSettings.fast(),
|
|
"RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8": SPTestSettings.fp8_quant(),
|
|
}
|
|
|
|
SP_TEST_MODELS = [
|
|
# TODO support other models
|
|
# [LANGUAGE GENERATION]
|
|
"meta-llama/Llama-3.2-1B-Instruct",
|
|
"RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"model_id",
|
|
"parallel_setup",
|
|
"distributed_backend",
|
|
"runner",
|
|
"test_options",
|
|
),
|
|
[
|
|
params
|
|
for model_id, settings in SP_TEXT_GENERATION_MODELS.items()
|
|
for params in settings.iter_params(model_id)
|
|
if model_id in SP_TEST_MODELS
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("use_inductor_graph_partition", [True, False])
|
|
@create_new_process_for_each_test()
|
|
def test_tp_sp_generation(
|
|
model_id: str,
|
|
parallel_setup: ParallelSetup,
|
|
distributed_backend: str,
|
|
runner: RunnerOption,
|
|
test_options: SPTestOptions,
|
|
num_gpus_available,
|
|
use_inductor_graph_partition: bool,
|
|
):
|
|
if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
|
|
pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
|
|
|
|
_compare_sp(
|
|
model_id,
|
|
parallel_setup,
|
|
distributed_backend,
|
|
runner,
|
|
test_options,
|
|
num_gpus_available,
|
|
use_inductor_graph_partition,
|
|
method="generate",
|
|
is_multimodal=False,
|
|
)
|