| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import ray |
| 6 | import torch |
| 7 | import torch.distributed as dist |
| 8 | |
| 9 | from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops |
| 10 | from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa |
| 11 | from vllm.distributed.parallel_state import get_dp_group, get_tp_group, graph_capture |
| 12 | from vllm.envs import disable_envs_cache |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | from ..utils import ( |
| 16 | assert_rocm_custom_allreduce_backend_state, |
| 17 | ensure_model_parallel_initialized, |
| 18 | init_test_distributed_environment, |
| 19 | multi_gpu_test, |
| 20 | multi_process_parallel, |
| 21 | ) |
| 22 | |
| 23 | pytestmark = pytest.mark.skipif( |
| 24 | not current_platform.is_rocm(), |
| 25 | reason="ROCm-only AITER custom allreduce tests", |
| 26 | ) |
| 27 | |
| 28 | test_cases = [ |
| 29 | ((2, 7168), torch.float16), |
| 30 | ((2, 7168), torch.bfloat16), |
| 31 | ((128, 8192), torch.float16), |
| 32 | ((128, 8192), torch.bfloat16), |
| 33 | ] |
| 34 | |
| 35 | |
| 36 | def _configure_aiter_custom_ar_env(monkeypatch: pytest.MonkeyPatch) -> None: |
| 37 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 38 | monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) |
| 39 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 40 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1") |
| 41 | monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE") |
| 42 | disable_envs_cache() |
| 43 | rocm_aiter_ops.refresh_env_variables() |
| 44 | |
| 45 | |
| 46 | def _assert_aiter_handles_input(inp: torch.Tensor) -> None: |
| 47 | aiter_ar_comm = get_tp_group().device_communicator.aiter_ar_comm |
| 48 | assert aiter_ar_comm is not None |
| 49 | assert aiter_ar_comm.should_custom_ar(inp), ( |
| 50 | f"AITER CustomAllreduce does not support input shape {inp.shape}." |
| 51 | ) |
| 52 | |
| 53 | |
| 54 | @ray.remote(num_gpus=1, max_calls=1) |
| 55 | def graph_allreduce( |
| 56 | monkeypatch: pytest.MonkeyPatch, |
| 57 | tp_size, |
| 58 | pp_size, |
| 59 | rank, |
| 60 | distributed_init_port, |
| 61 | ) -> None: |
| 62 | with monkeypatch.context() as m: |
| 63 | _configure_aiter_custom_ar_env(m) |
| 64 | |
| 65 | device = torch.device(f"cuda:{rank}") |
| 66 | torch.accelerator.set_device_index(device) |
| 67 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 68 | ensure_model_parallel_initialized(tp_size, pp_size) |
| 69 | assert_rocm_custom_allreduce_backend_state(True, "NONE") |
| 70 | group = get_tp_group().device_group |
| 71 | |
| 72 | # A small all_reduce for warmup. |
| 73 | # this is needed because device communicators might be created lazily |
| 74 | # (e.g. NCCL). This will ensure that the communicator is initialized |
| 75 | # before any communication happens, so that this group can be used for |
| 76 | # graph capture immediately. |
| 77 | data = torch.zeros(1) |
| 78 | data = data.to(device=device) |
| 79 | dist.all_reduce(data, group=group) |
| 80 | torch.accelerator.synchronize() |
| 81 | del data |
| 82 | |
| 83 | for shape, dtype in test_cases: |
| 84 | with graph_capture(device=device) as graph_capture_context: |
| 85 | inp = torch.ones(shape, dtype=dtype, device=device) |
| 86 | _assert_aiter_handles_input(inp) |
| 87 | expected = inp * tp_size |
| 88 | |
| 89 | torch.accelerator.synchronize() |
| 90 | graph = torch.cuda.CUDAGraph() |
| 91 | with torch.cuda.graph(graph, stream=graph_capture_context.stream): |
| 92 | out = tensor_model_parallel_all_reduce(inp) |
| 93 | |
| 94 | graph.replay() |
| 95 | torch.testing.assert_close(out, expected) |
| 96 | |
| 97 | |
| 98 | @ray.remote(num_gpus=1, max_calls=1) |
| 99 | def eager_allreduce( |
| 100 | monkeypatch: pytest.MonkeyPatch, |
| 101 | tp_size, |
| 102 | pp_size, |
| 103 | rank, |
| 104 | distributed_init_port, |
| 105 | ) -> None: |
| 106 | with monkeypatch.context() as m: |
| 107 | _configure_aiter_custom_ar_env(m) |
| 108 | |
| 109 | device = torch.device(f"cuda:{rank}") |
| 110 | torch.accelerator.set_device_index(device) |
| 111 | init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) |
| 112 | ensure_model_parallel_initialized(tp_size, pp_size) |
| 113 | assert_rocm_custom_allreduce_backend_state(True, "NONE") |
| 114 | |
| 115 | for shape, dtype in test_cases: |
| 116 | inp = torch.ones(shape, dtype=dtype, device=device) |
| 117 | _assert_aiter_handles_input(inp) |
| 118 | expected = inp * tp_size |
| 119 | out = tensor_model_parallel_all_reduce(inp) |
| 120 | torch.testing.assert_close(out, expected) |
| 121 | |
| 122 | |
| 123 | @pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed") |
| 124 | @multi_gpu_test(num_gpus=2) |
| 125 | @pytest.mark.parametrize("tp_size", [2]) |
| 126 | @pytest.mark.parametrize("pipeline_parallel_size", [1]) |
| 127 | @pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) |
| 128 | def test_rocm_aiter_custom_allreduce( |
| 129 | monkeypatch: pytest.MonkeyPatch, |
| 130 | tp_size, |
| 131 | pipeline_parallel_size, |
| 132 | test_target, |
| 133 | ): |
| 134 | multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) |
| 135 | |
| 136 | |
| 137 | def _get_aiter_ag_rs_comm(): |
| 138 | device_communicator = get_dp_group().device_communicator |
| 139 | assert device_communicator.use_aiter_ag_rs, ( |
| 140 | "AITER custom AG/RS was not enabled on the DP group." |
| 141 | ) |
| 142 | aiter_comm = device_communicator.aiter_ar_comm |
| 143 | assert aiter_comm is not None, "AITER custom AG/RS was not initialized." |
| 144 | assert not aiter_comm.disabled, "AITER custom AG/RS is disabled." |
| 145 | return aiter_comm |
| 146 | |
| 147 | |
| 148 | def _assert_aiter_handles_ag(aiter_comm, inp: torch.Tensor) -> None: |
| 149 | assert aiter_comm.should_custom_ag(inp), ( |
| 150 | f"AITER custom all-gather does not support input shape {inp.shape}." |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | def _assert_aiter_handles_rs(aiter_comm, inp: torch.Tensor) -> None: |
| 155 | assert aiter_comm.should_custom_rs(inp, dim=0), ( |
| 156 | f"AITER custom reduce-scatter does not support input shape {inp.shape}." |
| 157 | ) |
| 158 | |
| 159 | |
| 160 | @ray.remote(num_gpus=1, max_calls=1) |
| 161 | def eager_ag_rs( |
| 162 | monkeypatch: pytest.MonkeyPatch, |
| 163 | tp_size, |
| 164 | pp_size, |
| 165 | rank, |
| 166 | distributed_init_port, |
| 167 | data_parallel_size, |
| 168 | data_parallel_master_port, |
| 169 | ) -> None: |
| 170 | with monkeypatch.context() as m: |
| 171 | _configure_aiter_custom_ar_env(m) |
| 172 | |
| 173 | device = torch.device(f"cuda:{rank}") |
| 174 | torch.accelerator.set_device_index(device) |
| 175 | init_test_distributed_environment( |
| 176 | tp_size, |
| 177 | pp_size, |
| 178 | rank, |
| 179 | distributed_init_port, |
| 180 | data_parallel_size=data_parallel_size, |
| 181 | data_parallel_master_port=data_parallel_master_port, |
| 182 | ) |
| 183 | |
| 184 | dp_group = get_dp_group() |
| 185 | group = dp_group.device_group |
| 186 | dp_world = dp_group.world_size |
| 187 | aiter_comm = _get_aiter_ag_rs_comm() |
| 188 | |
| 189 | for shape, dtype in test_cases: |
| 190 | num_tokens, hidden = shape |
| 191 | |
| 192 | # all-gather: each rank contributes (num_tokens, hidden). |
| 193 | inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1) |
| 194 | _assert_aiter_handles_ag(aiter_comm, inp) |
| 195 | expected = torch.empty( |
| 196 | (num_tokens * dp_world, hidden), dtype=dtype, device=device |
| 197 | ) |
| 198 | dist.all_gather_into_tensor(expected, inp, group=group) |
| 199 | out = aiter_comm.custom_all_gather(inp, dim=0) |
| 200 | assert out is not None |
| 201 | torch.testing.assert_close(out, expected) |
| 202 | |
| 203 | # reduce-scatter: each rank contributes (num_tokens * dp, hidden). |
| 204 | rs_in = torch.ones( |
| 205 | (num_tokens * dp_world, hidden), dtype=dtype, device=device |
| 206 | ) * (rank + 1) |
| 207 | _assert_aiter_handles_rs(aiter_comm, rs_in) |
| 208 | rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device) |
| 209 | dist.reduce_scatter_tensor(rs_expected, rs_in, group=group) |
| 210 | rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device) |
| 211 | aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0) |
| 212 | torch.testing.assert_close(rs_out, rs_expected) |
| 213 | |
| 214 | |
| 215 | @ray.remote(num_gpus=1, max_calls=1) |
| 216 | def graph_ag_rs( |
| 217 | monkeypatch: pytest.MonkeyPatch, |
| 218 | tp_size, |
| 219 | pp_size, |
| 220 | rank, |
| 221 | distributed_init_port, |
| 222 | data_parallel_size, |
| 223 | data_parallel_master_port, |
| 224 | ) -> None: |
| 225 | with monkeypatch.context() as m: |
| 226 | _configure_aiter_custom_ar_env(m) |
| 227 | |
| 228 | device = torch.device(f"cuda:{rank}") |
| 229 | torch.accelerator.set_device_index(device) |
| 230 | init_test_distributed_environment( |
| 231 | tp_size, |
| 232 | pp_size, |
| 233 | rank, |
| 234 | distributed_init_port, |
| 235 | data_parallel_size=data_parallel_size, |
| 236 | data_parallel_master_port=data_parallel_master_port, |
| 237 | ) |
| 238 | |
| 239 | dp_group = get_dp_group() |
| 240 | group = dp_group.device_group |
| 241 | dp_world = dp_group.world_size |
| 242 | aiter_comm = _get_aiter_ag_rs_comm() |
| 243 | |
| 244 | # Warmup so DP comms is initialized before graph capture |
| 245 | data = torch.zeros(1, device=device) |
| 246 | dist.all_reduce(data, group=group) |
| 247 | torch.accelerator.synchronize() |
| 248 | del data |
| 249 | |
| 250 | for shape, dtype in test_cases: |
| 251 | num_tokens, hidden = shape |
| 252 | |
| 253 | # all-gather under graph capture. |
| 254 | inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1) |
| 255 | _assert_aiter_handles_ag(aiter_comm, inp) |
| 256 | ag_expected = torch.empty( |
| 257 | (num_tokens * dp_world, hidden), dtype=dtype, device=device |
| 258 | ) |
| 259 | dist.all_gather_into_tensor(ag_expected, inp, group=group) |
| 260 | with graph_capture(device=device) as graph_capture_context: |
| 261 | torch.accelerator.synchronize() |
| 262 | graph = torch.cuda.CUDAGraph() |
| 263 | with torch.cuda.graph(graph, stream=graph_capture_context.stream): |
| 264 | ag_out = aiter_comm.custom_all_gather(inp, dim=0) |
| 265 | graph.replay() |
| 266 | torch.testing.assert_close(ag_out, ag_expected) |
| 267 | |
| 268 | # reduce-scatter under graph capture. |
| 269 | rs_in = torch.ones( |
| 270 | (num_tokens * dp_world, hidden), dtype=dtype, device=device |
| 271 | ) * (rank + 1) |
| 272 | _assert_aiter_handles_rs(aiter_comm, rs_in) |
| 273 | rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device) |
| 274 | dist.reduce_scatter_tensor(rs_expected, rs_in, group=group) |
| 275 | rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device) |
| 276 | with graph_capture(device=device) as graph_capture_context: |
| 277 | torch.accelerator.synchronize() |
| 278 | graph = torch.cuda.CUDAGraph() |
| 279 | with torch.cuda.graph(graph, stream=graph_capture_context.stream): |
| 280 | aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0) |
| 281 | graph.replay() |
| 282 | torch.testing.assert_close(rs_out, rs_expected) |
| 283 | |
| 284 | |
| 285 | @pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed") |
| 286 | @multi_gpu_test(num_gpus=2) |
| 287 | @pytest.mark.parametrize("tp_size", [1]) |
| 288 | @pytest.mark.parametrize("pipeline_parallel_size", [1]) |
| 289 | @pytest.mark.parametrize("data_parallel_size", [2]) |
| 290 | @pytest.mark.parametrize("test_target", [eager_ag_rs, graph_ag_rs]) |
| 291 | def test_rocm_aiter_custom_ag_rs( |
| 292 | monkeypatch: pytest.MonkeyPatch, |
| 293 | tp_size, |
| 294 | pipeline_parallel_size, |
| 295 | data_parallel_size, |
| 296 | test_target, |
| 297 | ): |
| 298 | multi_process_parallel( |
| 299 | monkeypatch, |
| 300 | tp_size, |
| 301 | pipeline_parallel_size, |
| 302 | test_target, |
| 303 | data_parallel_size=data_parallel_size, |
| 304 | ) |