| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm AsyncTP numerical coverage kept outside the shared distributed suite.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from tests.compile.passes.distributed.test_async_tp import ( |
| 9 | TestAGMMModel, |
| 10 | TestMMRSModel, |
| 11 | async_tp_pass_on_test_model, |
| 12 | ) |
| 13 | from tests.utils import multi_gpu_test |
| 14 | from vllm.distributed import ( |
| 15 | tensor_model_parallel_all_gather, |
| 16 | tensor_model_parallel_reduce_scatter, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.utils.network_utils import get_file_store_init_method |
| 20 | |
| 21 | |
| 22 | class TestRocmMMRSModel(TestMMRSModel): |
| 23 | def forward(self, hidden_states): |
| 24 | view = hidden_states.reshape(-1, self.hidden_size) |
| 25 | mm = torch.ops.vllm.rocm_unquantized_gemm(view, self.gate_proj, None) |
| 26 | return tensor_model_parallel_reduce_scatter(mm, dim=0) |
| 27 | |
| 28 | def ops_in_model_before(self): |
| 29 | return super().ops_in_model_before() + [ |
| 30 | torch.ops.vllm.rocm_unquantized_gemm.default |
| 31 | ] |
| 32 | |
| 33 | |
| 34 | class TestRocmAGMMModel(TestAGMMModel): |
| 35 | def __init__(self, hidden_size=16, dtype=torch.bfloat16): |
| 36 | super().__init__(hidden_size, dtype) |
| 37 | self.weight = torch.nn.Parameter( |
| 38 | torch.empty((hidden_size * 2, hidden_size)), requires_grad=False |
| 39 | ) |
| 40 | torch.nn.init.normal_(self.weight, std=0.02) |
| 41 | |
| 42 | def forward(self, hidden_states): |
| 43 | view = hidden_states.reshape(-1, self.hidden_size) |
| 44 | gathered = tensor_model_parallel_all_gather(view, dim=0) |
| 45 | return torch.ops.vllm.rocm_unquantized_gemm(gathered, self.weight, None) |
| 46 | |
| 47 | def ops_in_model_before(self): |
| 48 | return super().ops_in_model_before() + [ |
| 49 | torch.ops.vllm.rocm_unquantized_gemm.default |
| 50 | ] |
| 51 | |
| 52 | |
| 53 | @multi_gpu_test(num_gpus=2) |
| 54 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm backend contract") |
| 55 | @pytest.mark.parametrize( |
| 56 | "test_model", [TestMMRSModel, TestAGMMModel, TestRocmMMRSModel, TestRocmAGMMModel] |
| 57 | ) |
| 58 | @pytest.mark.parametrize("dynamic", [False, True]) |
| 59 | def test_rocm_async_tp_bf16_rewrite_and_correctness(test_model, dynamic): |
| 60 | """Real symmetric-memory AG/GEMM and GEMM/RS must replace the collectives.""" |
| 61 | torch.multiprocessing.spawn( |
| 62 | async_tp_pass_on_test_model, |
| 63 | args=( |
| 64 | 2, |
| 65 | test_model, |
| 66 | 8, |
| 67 | 16, |
| 68 | 128, |
| 69 | torch.bfloat16, |
| 70 | dynamic, |
| 71 | get_file_store_init_method(), |
| 72 | True, |
| 73 | ), |
| 74 | nprocs=2, |
| 75 | ) |