From 445ded18c1184a5a44d0f41010d614adbd107ca7 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 26 May 2026 08:57:13 -0500 Subject: [PATCH] [ROCm][CI] Extend ROCm quick reduce coverage (#40990) Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 24 +- tests/distributed/test_quick_all_reduce.py | 76 +- tests/distributed/test_rocm_quick_reduce.py | 750 ++++++++++++++++++++ tests/utils.py | 57 +- 4 files changed, 869 insertions(+), 38 deletions(-) create mode 100644 tests/distributed/test_rocm_quick_reduce.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 3cadab548fb..c7ef61b719a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2703,19 +2703,35 @@ steps: optional: true working_dir: "/vllm-workspace/" source_file_dependencies: + - csrc/custom_quickreduce.cu + - csrc/ops.h + - csrc/torch_bindings.cpp - vllm/distributed/ - - vllm/v1/distributed/ + - vllm/model_executor/layers/ + - vllm/entrypoints/llm.py + - vllm/config/parallel.py - vllm/model_executor/layers/fused_moe/ + - vllm/v1/engine/ + - vllm/v1/executor/ + - vllm/v1/worker/ + - vllm/v1/distributed/ - vllm/v1/attention/backends/ - vllm/v1/attention/selector.py - - tests/distributed/test_context_parallel.py - - tests/v1/distributed/test_dbo.py - - examples/features/data_parallel/data_parallel_offline.py - vllm/_aiter_ops.py + - vllm/_custom_ops.py - vllm/platforms/rocm.py + - vllm/envs.py + - examples/offline_inference/data_parallel.py + - tests/distributed/test_context_parallel.py + - tests/distributed/test_rocm_quick_reduce.py + - tests/distributed/test_quick_all_reduce.py + - tests/v1/distributed/test_dbo.py + - tests/utils.py commands: - pytest -v -s tests/distributed/test_context_parallel.py - pytest -v -s tests/v1/distributed/test_dbo.py + - pytest -v -s tests/distributed/test_rocm_quick_reduce.py + - pytest -v -s tests/distributed/test_quick_all_reduce.py #-------------------------------------------------------- mi355 ยท entrypoints --------------------------------------------------------# diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py index 6c1ae9dda30..a9591f96a78 100644 --- a/tests/distributed/test_quick_all_reduce.py +++ b/tests/distributed/test_quick_all_reduce.py @@ -25,16 +25,34 @@ from ..utils import ( ensure_model_parallel_initialized, init_test_distributed_environment, multi_process_parallel, + set_random_seed, ) -torch.manual_seed(42) -random.seed(44) + +def on_gfx942() -> bool: + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx942 as rocm_on_gfx942 + + return rocm_on_gfx942() + return False + + +set_random_seed(42) +_test_size_rng = random.Random(44) # Size over 8MB is sufficient for custom quick allreduce. -test_sizes = [random.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8)] +test_sizes = [ + _test_size_rng.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8) +] for i, v in enumerate(test_sizes): test_sizes[i] -= v % 8 +def _assert_quickreduce(fa, inp): + assert fa is not None + assert not fa.disabled + assert fa.should_quick_allreduce(inp) + + @pytest.fixture def envs_cache_disabled(): disable_envs_cache() @@ -216,11 +234,14 @@ def graph_quickreduce( ): with monkeypatch.context() as m: m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + m.delenv("HIP_VISIBLE_DEVICES", raising=False) + m.delenv("ROCR_VISIBLE_DEVICES", raising=False) device = torch.device(f"cuda:{rank}") torch.accelerator.set_device_index(device) init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) ensure_model_parallel_initialized(tp_size, pp_size) group = get_tp_group().device_group + fa = get_tp_group().device_communicator.qr_comm # A small all_reduce for warmup. # this is needed because device communicators might be created lazily @@ -246,6 +267,8 @@ def graph_quickreduce( device_idx = torch.accelerator.current_device_index() inp1 = torch.randint(1, 23, (sz,), dtype=dtype, device=device_idx) inp2 = torch.randint(-23, 1, (sz,), dtype=dtype, device=device_idx) + _assert_quickreduce(fa, inp1) + _assert_quickreduce(fa, inp2) torch.accelerator.synchronize() graph = torch.cuda.CUDAGraph() @@ -270,6 +293,8 @@ def eager_quickreduce( ): with monkeypatch.context() as m: m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + m.delenv("HIP_VISIBLE_DEVICES", raising=False) + m.delenv("ROCR_VISIBLE_DEVICES", raising=False) device = torch.device(f"cuda:{rank}") torch.accelerator.set_device_index(device) @@ -281,12 +306,42 @@ def eager_quickreduce( inp = torch.tensor( [1.0 * ((i) % 23) for i in range(sz)], dtype=torch.float16, device=device ) + _assert_quickreduce(fa, inp) out = fa.quick_all_reduce(inp) torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1) inp = torch.tensor( [1.0 * ((i) % 23) for i in range(sz)], dtype=torch.bfloat16, device=device ) + _assert_quickreduce(fa, inp) + out = fa.quick_all_reduce(inp) + torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1) + + +@ray.remote(num_gpus=1, max_calls=1) +def bf16_cast_quickreduce( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, +): + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + m.delenv("HIP_VISIBLE_DEVICES", raising=False) + m.delenv("ROCR_VISIBLE_DEVICES", raising=False) + m.setenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1") + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + + sz = 16 * 1024 * 1024 + fa = get_tp_group().device_communicator.qr_comm + inp = torch.tensor( + [1.0 * (i % 23) for i in range(sz)], dtype=torch.bfloat16, device=device + ) + _assert_quickreduce(fa, inp) + assert fa.use_fp16_kernels out = fa.quick_all_reduce(inp) torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1) @@ -308,12 +363,27 @@ def test_custom_quick_allreduce( world_size = tp_size * pipeline_parallel_size if world_size > torch.accelerator.device_count(): pytest.skip("Not enough GPUs to run the test.") + if test_target is graph_quickreduce and on_gfx942(): + pytest.xfail( + "CUDA graph capture with quick reduce hits " + "hipErrorStreamCaptureInvalidated on gfx942" + ) monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode) multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) +@pytest.mark.skipif( + not current_platform.is_rocm(), reason="only test quick allreduce for rocm" +) +def test_custom_quick_allreduce_bf16_cast(monkeypatch: pytest.MonkeyPatch): + if torch.accelerator.device_count() < 2: + pytest.skip("Not enough GPUs to run the test.") + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "FP") + multi_process_parallel(monkeypatch, 2, 1, bf16_cast_quickreduce) + + def qr_variable_input(rank, world_size): """ When the tensor parallelism is set to 4 or 8, frequent changes diff --git a/tests/distributed/test_rocm_quick_reduce.py b/tests/distributed/test_rocm_quick_reduce.py new file mode 100644 index 00000000000..081698483a4 --- /dev/null +++ b/tests/distributed/test_rocm_quick_reduce.py @@ -0,0 +1,750 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import contextlib +import importlib +import multiprocessing as mp +import os +import queue +import traceback +from functools import lru_cache +from types import SimpleNamespace +from typing import Literal +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +from huggingface_hub import snapshot_download + +import vllm.envs as envs +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm-only quick-reduce tests", +) + +MB = 1024 * 1024 +WORLD_SIZE = 2 +QUANT_LEVELS = ["FP", "INT8", "INT6", "INT4"] + + +def _log(message: str) -> None: + print(f"[rocm_quick_reduce] {message}", flush=True) + + +def _reload_envs(): + return importlib.reload(envs) + + +def _make_quick_allreduce( + *, + disabled: bool = False, + world_size: int = 2, + quant_level: str = "FP", + use_fp16_kernels: bool = False, + qr_max_size: int = 64 * MB, +): + from vllm.distributed.device_communicators.quick_all_reduce import ( + QuickAllReduce, + QuickReduceRegime, + ) + + qar = QuickAllReduce.__new__(QuickAllReduce) + qar.disabled = disabled + qar.world_size = world_size + qar.use_fp16_kernels = use_fp16_kernels + qar.qr_quant_level = QuickReduceRegime[quant_level] + qar.qr_max_size = qr_max_size + return qar + + +def _quick_allreduce_worker( + rank: int, + port: int, + quant_level: str, + dtype_name: str, + cast_bf16: bool, +): + os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level + os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "1" if cast_bf16 else "0" + _log( + f"worker start: rank={rank} quant={quant_level} " + f"dtype={dtype_name} cast_bf16={cast_bf16}" + ) + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=WORLD_SIZE, + ) + + qar = None + try: + from vllm.distributed.device_communicators.quick_all_reduce import ( + QuickAllReduce, + ) + + qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank) + assert not qar.disabled + + num_elements = 8 * MB if dtype_name == "float16" else 4 * MB + + dtype = getattr(torch, dtype_name) + inp = torch.ones(num_elements, dtype=dtype, device=device) + + assert qar.should_quick_allreduce(inp) + if cast_bf16: + assert qar.use_fp16_kernels + + out = qar.quick_all_reduce(inp) + assert torch.allclose(out, inp * WORLD_SIZE, atol=2.5, rtol=0.1) + _log( + f"worker complete: rank={rank} quant={quant_level} " + f"dtype={dtype_name} num_elements={num_elements} " + f"use_fp16_kernels={qar.use_fp16_kernels}" + ) + finally: + if qar is not None: + qar.close() + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_two_gpu_quick_allreduce_test( + *, + quant_level: str, + dtype_name: str, + cast_bf16: bool, +): + _log( + f"launch 2-GPU case: quant={quant_level} " + f"dtype={dtype_name} cast_bf16={cast_bf16}" + ) + ctx = mp.get_context("spawn") + port = get_open_port() + procs = [] + + for rank in range(WORLD_SIZE): + proc = ctx.Process( + target=_quick_allreduce_worker, + args=(rank, port, quant_level, dtype_name, cast_bf16), + ) + proc.start() + procs.append(proc) + + for proc in procs: + proc.join(timeout=60) + assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}" + _log( + f"finished 2-GPU case: quant={quant_level} " + f"dtype={dtype_name} cast_bf16={cast_bf16}" + ) + + +MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" +E2E_PREFILL_TOKENS = 1024 +E2E_MAX_MODEL_LEN = 1536 +E2E_GPU_MEMORY_UTILIZATION = 0.3 +E2E_KV_CACHE_MEMORY_BYTES = 2 << 30 + +_BACKGROUND_LINE = ( + "Background filler: this archived operations memo repeats a routine status " + "line so the distributed test uses a realistically long prefill." +) +_BACKGROUND_BLOCK = " ".join([_BACKGROUND_LINE] * 48) + + +def _build_prompt(*, fact_block: str, question: str) -> str: + return ( + "Read the archived operations memo below. Most of the memo is filler. " + "Use only the fact block near the end when answering.\n" + f"{_BACKGROUND_BLOCK}\n" + "Fact block:\n" + f"{fact_block}\n" + f"Question: {question}\n" + "Answer in one short sentence." + ) + + +E2E_PROMPTS = [ + _build_prompt( + fact_block=( + "- Festival city: Oslo\n- Mascot animal: otter\n- Welcome drink: tea" + ), + question="Which city hosts the festival, and what animal is the mascot?", + ), + _build_prompt( + fact_block=( + "- Meeting day: Tuesday\n" + "- Planned snack: apricot cake\n" + "- Backup room: Cedar" + ), + question="What day is the meeting, and what snack is planned?", + ), +] +RECORDED_RESPONSE_TEXTS = ( + " The city hosting the festival is Oslo, and the mascot is an otter.", + " The meeting is on Tuesday and the snack planned is apricot cake.", +) +REQUIRED_WORDS = (("oslo", "otter"), ("tuesday", "apricot")) + + +def _log_prompt_summaries() -> None: + for i, prompt in enumerate(E2E_PROMPTS): + prompt_lines = prompt.splitlines() + fact_block = [line for line in prompt_lines if line.startswith("- ")] + fact_summary = "; ".join(line.removeprefix("- ") for line in fact_block) + _log(f"prompt {i} facts: {fact_summary}") + + +@lru_cache(maxsize=1) +def _get_model_path() -> str: + try: + path = snapshot_download(repo_id=MODEL_NAME, local_files_only=True) + _log(f"using cached model snapshot: {path}") + return path + except Exception: + path = snapshot_download(repo_id=MODEL_NAME) + _log(f"downloaded model snapshot: {path}") + return path + + +def _get_hidden_size(model_config) -> int: + hidden_size = getattr(model_config, "hidden_size", None) + if hidden_size is None and hasattr(model_config, "text_config"): + hidden_size = getattr(model_config.text_config, "hidden_size", None) + assert isinstance(hidden_size, int) + return hidden_size + + +def _check_tp_allreduce_uses_quick_reduce( + self, + num_tokens: int, + dtype_name: str = "float16", +) -> dict[str, int | bool]: + from vllm.distributed.communication_op import tensor_model_parallel_all_reduce + from vllm.distributed.parallel_state import get_tp_group + + assert self.device is not None + qr_comm = get_tp_group().device_communicator.qr_comm + assert qr_comm is not None + assert not qr_comm.disabled + + hidden_size = _get_hidden_size(self.model_runner.model.config) + dtype = getattr(torch, dtype_name) + sample = torch.full( + (num_tokens, hidden_size), + fill_value=float(self.rank + 1), + dtype=dtype, + device=self.device, + ) + + assert qr_comm.should_quick_allreduce(sample) + + expected = sample.clone() + reduced = tensor_model_parallel_all_reduce(sample) + dist.all_reduce(expected, group=get_tp_group().device_group) + torch.testing.assert_close(reduced, expected, atol=2.5, rtol=0.1) + + stats = { + "rank": self.rank, + "hidden_size": hidden_size, + "num_tokens": num_tokens, + "use_fp16_kernels": qr_comm.use_fp16_kernels, + } + _log( + "worker quick-reduce check: " + f"rank={self.rank} hidden_size={hidden_size} " + f"num_tokens={num_tokens} use_fp16_kernels={qr_comm.use_fp16_kernels}" + ) + return stats + + +def _check_quick_reduce_disabled(self) -> int: + from vllm.distributed.parallel_state import get_tp_group + + qr_comm = get_tp_group().device_communicator.qr_comm + assert qr_comm is not None + assert qr_comm.disabled + _log(f"worker confirmed quick reduce is disabled: rank={self.rank}") + return self.rank + + +def _collect_generations(outputs) -> list[tuple[tuple[int, ...], str]]: + return [ + (tuple(output.outputs[0].token_ids), output.outputs[0].text) + for output in outputs + ] + + +def _shutdown_llm(llm: LLM | None) -> None: + if llm is None: + cleanup_dist_env_and_memory() + return + + with contextlib.suppress(Exception): + llm.llm_engine.engine_core.shutdown() + + del llm + cleanup_dist_env_and_memory() + + +def _log_generations( + label: str, + generations: list[tuple[tuple[int, ...], str]], +) -> None: + for i, (token_ids, text) in enumerate(generations): + _log(f"{label} prompt {i} token ids: {list(token_ids)}") + _log(f"{label} prompt {i} text: {text!r}") + + +def _assert_required_words( + label: str, + generations: list[tuple[tuple[int, ...], str]], +) -> None: + for i, (_, text) in enumerate(generations): + lowered = text.lower() + missing = [word for word in REQUIRED_WORDS[i] if word not in lowered] + assert not missing, ( + f"{label} prompt {i} is missing required words {missing}. " + f"Observed text: {text!r}" + ) + + +def _collect_soft_mismatches( + baseline_generations: list[tuple[tuple[int, ...], str]], + quick_reduce_generations: list[tuple[tuple[int, ...], str]], +) -> list[str]: + mismatches = [] + + for i, (_, text) in enumerate(baseline_generations): + expected = RECORDED_RESPONSE_TEXTS[i] + if text != expected: + mismatches.append( + f"baseline prompt {i} drifted from the recorded response.\n" + f"expected={expected!r}\nactual={text!r}" + ) + + for i, (_, text) in enumerate(quick_reduce_generations): + expected = RECORDED_RESPONSE_TEXTS[i] + if text != expected: + mismatches.append( + f"quick-reduce prompt {i} drifted from the recorded response.\n" + f"expected={expected!r}\nactual={text!r}" + ) + + for i, ((_, baseline_text), (_, quick_reduce_text)) in enumerate( + zip(baseline_generations, quick_reduce_generations) + ): + if baseline_text != quick_reduce_text: + mismatches.append( + f"baseline and quick-reduce responses differ for prompt {i}.\n" + f"baseline={baseline_text!r}\nquick_reduce={quick_reduce_text!r}" + ) + + return mismatches + + +def _run_generation( + *, + backend: Literal["mp", "ray"], + quant_mode: str, + expect_quick_reduce: bool, +) -> list[tuple[tuple[int, ...], str]]: + llm = None + monkeypatch = pytest.MonkeyPatch() + with monkeypatch.context() as m: + m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode) + model_path = _get_model_path() + _log( + f"starting generation: backend={backend} quant={quant_mode} " + f"gpu_memory_utilization={E2E_GPU_MEMORY_UTILIZATION} " + f"kv_cache_bytes={E2E_KV_CACHE_MEMORY_BYTES} model={model_path}" + ) + + try: + llm = LLM( + model=model_path, + tokenizer=model_path, + tensor_parallel_size=2, + distributed_executor_backend=backend, + dtype="half", + enforce_eager=True, + max_model_len=E2E_MAX_MODEL_LEN, + max_num_seqs=len(E2E_PROMPTS), + gpu_memory_utilization=E2E_GPU_MEMORY_UTILIZATION, + kv_cache_memory_bytes=E2E_KV_CACHE_MEMORY_BYTES, + seed=0, + ) + + if not expect_quick_reduce: + assert llm.collective_rpc(_check_quick_reduce_disabled) == [0, 1] + + if expect_quick_reduce: + worker_stats = llm.collective_rpc( + _check_tp_allreduce_uses_quick_reduce, + args=(E2E_PREFILL_TOKENS,), + ) + assert [stat["rank"] for stat in worker_stats] == [0, 1] + worker_summary = "; ".join( + "rank={rank} hidden_size={hidden_size} num_tokens={num_tokens} " + "use_fp16_kernels={use_fp16_kernels}".format(**stat) + for stat in worker_stats + ) + _log(f"{backend} quick-reduce worker checks: {worker_summary}") + + outputs = llm.generate( + E2E_PROMPTS, + SamplingParams( + temperature=0.0, + max_tokens=20, + stop=["\nAnswer:", " Answer:"], + ), + use_tqdm=False, + ) + generations = _collect_generations(outputs) + assert all(text.strip() for _, text in generations) + _log_generations(f"{backend} {quant_mode}", generations) + return generations + finally: + _shutdown_llm(llm) + + +def _run_quick_reduce_llm_e2e_in_subprocess( + *, + backend: Literal["mp", "ray"], +) -> str | None: + _log(f"running LLM e2e: backend={backend}") + _log_prompt_summaries() + baseline_outputs = _run_generation( + backend=backend, + quant_mode="NONE", + expect_quick_reduce=False, + ) + quick_reduce_outputs = _run_generation( + backend=backend, + quant_mode="FP", + expect_quick_reduce=True, + ) + + _assert_required_words("baseline", baseline_outputs) + _assert_required_words("quick-reduce", quick_reduce_outputs) + + mismatches = _collect_soft_mismatches(baseline_outputs, quick_reduce_outputs) + if mismatches: + details = "\n\n".join(mismatches) + _log(f"soft response mismatch:\n{details}") + return details + + _log(f"LLM e2e backend={backend} matched the recorded responses exactly") + return None + + +def _quick_reduce_llm_e2e_worker( + result_queue: mp.Queue, + backend: Literal["mp", "ray"], +) -> None: + try: + xfail_reason = _run_quick_reduce_llm_e2e_in_subprocess(backend=backend) + except Exception: + result_queue.put({"status": "error", "reason": traceback.format_exc()}) + raise + else: + if xfail_reason is not None: + result_queue.put({"status": "xfail", "reason": xfail_reason}) + else: + result_queue.put({"status": "ok"}) + + +def run_quick_reduce_llm_e2e( + *, + backend: Literal["mp", "ray"], +) -> None: + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + proc = ctx.Process( + target=_quick_reduce_llm_e2e_worker, + args=(result_queue, backend), + ) + proc.start() + proc.join(timeout=600) + + try: + result = result_queue.get(timeout=5) + except queue.Empty as exc: + if proc.exitcode != 0: + raise AssertionError( + f"quick-reduce llm e2e subprocess failed for backend={backend} " + f"with exit code {proc.exitcode} and produced no result" + ) from exc + raise AssertionError( + f"quick-reduce llm e2e subprocess produced no result for backend={backend}" + ) from exc + + if result["status"] == "xfail": + pytest.xfail(result["reason"]) + if result["status"] == "error": + raise AssertionError( + f"quick-reduce llm e2e subprocess failed for backend={backend}:\n" + f"{result['reason']}" + ) + + assert proc.exitcode == 0, ( + f"quick-reduce llm e2e subprocess failed for backend={backend} " + f"with exit code {proc.exitcode}" + ) + + +def test_quick_reduce_regime_values(): + from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime + + assert QuickReduceRegime.FP.value == 0 + assert QuickReduceRegime.INT8.value == 1 + assert QuickReduceRegime.INT6.value == 2 + assert QuickReduceRegime.INT4.value == 3 + assert QuickReduceRegime.NONE.value == 4 + + +def test_quick_reduce_regime_names(): + from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime + + assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"} + + +@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"]) +def test_quick_reduce_quantization_env_var(monkeypatch, quant_level): + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_level) + + reloaded_envs = _reload_envs() + assert quant_level == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION + + +def test_quick_reduce_quantization_default(monkeypatch): + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", raising=False) + + reloaded_envs = _reload_envs() + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION == "NONE" + + +@pytest.mark.parametrize("cast_bf16", [True, False]) +def test_quick_reduce_cast_bf16_to_fp16_env_var(monkeypatch, cast_bf16): + monkeypatch.setenv( + "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1" if cast_bf16 else "0" + ) + + reloaded_envs = _reload_envs() + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is cast_bf16 + + +def test_quick_reduce_cast_bf16_to_fp16_default(monkeypatch): + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", raising=False) + + reloaded_envs = _reload_envs() + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is True + + +@pytest.mark.parametrize("max_mb", [128, 512, 2048, None]) +def test_quick_reduce_max_size_env_var(monkeypatch, max_mb): + if max_mb is None: + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False) + else: + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", str(max_mb)) + + reloaded_envs = _reload_envs() + assert max_mb == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB + + +def test_quick_reduce_max_size_default(monkeypatch): + monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False) + + reloaded_envs = _reload_envs() + assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is None + + +@pytest.mark.parametrize( + ("gcn_arch_name", "expected"), + [ + ("gfx942", True), + ("gfx950", True), + ("gfx90a", False), + ("", False), + ], +) +def test_quick_allreduce_rocm_arch_available(gcn_arch_name, expected): + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce + + qar = QuickAllReduce.__new__(QuickAllReduce) + qar.disabled = True + + with ( + patch( + "vllm.distributed.device_communicators.quick_all_reduce.current_platform." + "is_rocm", + return_value=True, + ), + patch( + "torch.cuda.get_device_properties", + return_value=SimpleNamespace(gcnArchName=gcn_arch_name), + ), + ): + assert qar._rocm_arch_available() is expected + + +def test_quick_allreduce_rocm_arch_available_handles_probe_failure(): + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce + + qar = QuickAllReduce.__new__(QuickAllReduce) + qar.disabled = True + + with ( + patch( + "vllm.distributed.device_communicators.quick_all_reduce.current_platform." + "is_rocm", + return_value=True, + ), + patch("torch.cuda.get_device_properties", side_effect=RuntimeError), + ): + assert qar._rocm_arch_available() is False + + +def test_quick_allreduce_rejects_disabled(): + qar = _make_quick_allreduce(disabled=True) + + inp = torch.zeros(1024, dtype=torch.float16) + assert qar.should_quick_allreduce(inp) is False + + +def test_quick_allreduce_rejects_unsupported_dtype(): + qar = _make_quick_allreduce() + + inp = torch.zeros(1024 * 1024, dtype=torch.float32) + assert qar.should_quick_allreduce(inp) is False + + +def test_quick_allreduce_rejects_non_aligned_input(): + qar = _make_quick_allreduce() + + inp = torch.zeros(5, dtype=torch.float16) + assert qar.should_quick_allreduce(inp) is False + + +def test_quick_allreduce_rejects_non_contiguous_input(): + qar = _make_quick_allreduce() + + inp = torch.zeros((1024, 1024), dtype=torch.float16)[:, ::2] + assert qar.should_quick_allreduce(inp) is False + + +def test_quick_allreduce_rejects_input_smaller_than_threshold(): + qar = _make_quick_allreduce() + + inp = torch.zeros((MB // 2) - 8, dtype=torch.float16) + assert qar.should_quick_allreduce(inp) is False + + +def test_quick_allreduce_accepts_input_at_threshold(): + qar = _make_quick_allreduce() + + inp = torch.zeros(MB // 2, dtype=torch.float16) + assert qar.should_quick_allreduce(inp) is True + + +def test_quick_allreduce_rejects_input_larger_than_max_size(): + qar = _make_quick_allreduce(qr_max_size=1 * MB) + + inp = torch.zeros(MB, dtype=torch.float16) + assert qar.should_quick_allreduce(inp) is False + + +def test_quick_allreduce_bf16_uses_fp16_threshold_when_cast_enabled(): + inp = torch.zeros(MB // 2, dtype=torch.bfloat16) + + without_cast = _make_quick_allreduce(use_fp16_kernels=False) + with_cast = _make_quick_allreduce(use_fp16_kernels=True) + + assert without_cast.should_quick_allreduce(inp) is False + assert with_cast.should_quick_allreduce(inp) is True + + +def test_quick_allreduce_supported_world_sizes(): + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce + + assert QuickAllReduce._SUPPORTED_WORLD_SIZES == [2, 4, 8] + + +def test_quick_allreduce_supported_dtypes(): + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce + + assert [torch.float16, torch.bfloat16] == QuickAllReduce._SUPPORTED_DTYPES + + +def test_quick_allreduce_min_size_table(): + from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce + + for dtype in [torch.float16, torch.bfloat16]: + for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES: + min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)] + assert len(min_sizes) == 4 + assert all(size > 0 for size in min_sizes) + + +def test_qr_max_size(): + from vllm import _custom_ops as ops + + max_size = ops.qr_max_size() + assert isinstance(max_size, int) + assert max_size > 0 + + +@pytest.mark.skipif( + current_platform.device_count() < WORLD_SIZE, + reason="requires 2 ROCm GPUs", +) +@pytest.mark.parametrize("quant_level", QUANT_LEVELS) +def test_quick_allreduce_two_gpu_correctness(quant_level): + _log(f"two-GPU correctness case: quant={quant_level}") + _run_two_gpu_quick_allreduce_test( + quant_level=quant_level, + dtype_name="float16", + cast_bf16=False, + ) + + +@pytest.mark.skipif( + current_platform.device_count() < WORLD_SIZE, + reason="requires 2 ROCm GPUs", +) +def test_quick_allreduce_bf16_cast_mode(): + _log("BF16 cast case") + _run_two_gpu_quick_allreduce_test( + quant_level="FP", + dtype_name="bfloat16", + cast_bf16=True, + ) + + +@pytest.mark.skipif( + current_platform.device_count() < WORLD_SIZE, + reason="requires 2 ROCm GPUs", +) +def test_quick_allreduce_llm_e2e(): + _log("LLM e2e case: backend=mp") + run_quick_reduce_llm_e2e(backend="mp") + + +@pytest.mark.skipif( + current_platform.device_count() < WORLD_SIZE, + reason="requires 2 ROCm GPUs", +) +def test_quick_allreduce_llm_e2e_ray(): + _log("LLM e2e case: backend=ray") + run_quick_reduce_llm_e2e(backend="ray") diff --git a/tests/utils.py b/tests/utils.py index 7fa7fd8ad85..6a32f3e2e2d 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1363,43 +1363,38 @@ def multi_process_parallel( ) -> None: import ray - # Using ray helps debugging the error when it failed - # as compared to multiprocessing. - # NOTE: We need to set working_dir for distributed tests, - # otherwise we may get import errors on ray workers - # NOTE: Force ray not to use gitignore file as excluding, otherwise - # it will not move .so files to working dir. - # So we have to manually add some of large directories - os.environ["RAY_RUNTIME_ENV_IGNORE_GITIGNORE"] = "1" + # Using ray helps debugging the error when it failed as compared to + # multiprocessing. For local Ray workers, putting the repo root on + # PYTHONPATH is enough and avoids uploading the full source tree, which + # exceeds Ray's working_dir package size limit on CI. + env_vars = { + "PYTHONPATH": os.pathsep.join( + filter(None, [str(VLLM_PATH), os.environ.get("PYTHONPATH")]) + ), + **{env_var: "1" for env_var in current_platform.ray_noset_device_env_vars}, + } ray.init( runtime_env={ - "working_dir": VLLM_PATH, - "excludes": [ - "build", - ".git", - "cmake-build-*", - "shellcheck", - "dist", - "ep_kernels_workspace", - ], + "env_vars": env_vars, } ) distributed_init_port = get_open_port() - refs = [] - for rank in range(tp_size * pp_size): - refs.append( - test_target.remote( - monkeypatch, - tp_size, - pp_size, - rank, - distributed_init_port, - ), - ) - ray.get(refs) - - ray.shutdown() + try: + refs = [] + for rank in range(tp_size * pp_size): + refs.append( + test_target.remote( + monkeypatch, + tp_size, + pp_size, + rank, + distributed_init_port, + ), + ) + ray.get(refs) + finally: + ray.shutdown() @contextmanager