| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | from unittest.mock import patch |
| 5 | |
| 6 | import pytest |
| 7 | import regex as re |
| 8 | import requests |
| 9 | import torch |
| 10 | from transformers import AutoModelForCausalLM |
| 11 | |
| 12 | from tests.entrypoints.serve.dev.rlhf.conftest import ( |
| 13 | MODEL_NAME, |
| 14 | gen, |
| 15 | health, |
| 16 | ok, |
| 17 | server, |
| 18 | ) |
| 19 | from tests.utils import multi_gpu_test |
| 20 | from vllm.distributed.weight_transfer import ( |
| 21 | HTTPVLLMWeightSyncClient, |
| 22 | ModuleSource, |
| 23 | WeightTransferTrainerFactory, |
| 24 | ) |
| 25 | from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo |
| 26 | from vllm.platforms import current_platform |
| 27 | |
| 28 | |
| 29 | def weight_checker(url: str, action: str, baseline=None) -> requests.Response: |
| 30 | body = {"action": action} |
| 31 | if baseline is not None: |
| 32 | body["baseline"] = baseline |
| 33 | return requests.post(f"{url}/weight_checker", json=body, timeout=300) |
| 34 | |
| 35 | |
| 36 | def checksums(url: str) -> dict[str, str]: |
| 37 | response = weight_checker(url, "checksum") |
| 38 | response.raise_for_status() |
| 39 | return response.json()["checksums"] |
| 40 | |
| 41 | |
| 42 | def ranks(keys) -> set[str]: |
| 43 | return {key.rsplit(":", 1)[0] for key in keys} |
| 44 | |
| 45 | |
| 46 | def reset_reload_and_compare(url: str) -> None: |
| 47 | baseline = checksums(url) |
| 48 | try: |
| 49 | weight_checker(url, "reset").raise_for_status() |
| 50 | comparison = weight_checker(url, "compare", baseline).json() |
| 51 | assert ranks(comparison["mismatches"]) == ranks(baseline) |
| 52 | finally: |
| 53 | requests.post( |
| 54 | f"{url}/collective_rpc", json={"method": "reload_weights"}, timeout=300 |
| 55 | ).raise_for_status() |
| 56 | assert weight_checker(url, "compare", baseline).json() == { |
| 57 | "match": True, |
| 58 | "mismatches": [], |
| 59 | } |
| 60 | |
| 61 | |
| 62 | @pytest.fixture(scope="class") |
| 63 | def server_url(): |
| 64 | with server() as url: |
| 65 | yield url |
| 66 | |
| 67 | |
| 68 | class TestSingleGPU: |
| 69 | @pytest.mark.parametrize( |
| 70 | "payload", |
| 71 | [{}, {"action": "frobnicate"}, {"action": "compare"}, [], None], |
| 72 | ) |
| 73 | def test_invalid_request_returns_400(self, server_url, payload): |
| 74 | response = requests.post( |
| 75 | f"{server_url}/weight_checker", json=payload, timeout=10 |
| 76 | ) |
| 77 | assert response.status_code == 400, response.text |
| 78 | assert health(server_url) == 200 |
| 79 | |
| 80 | def test_checksum_is_stable_and_compare_reports_differences(self, server_url): |
| 81 | baseline = checksums(server_url) |
| 82 | assert all( |
| 83 | re.fullmatch(r"[0-9a-f]{64}", digest) for digest in baseline.values() |
| 84 | ) |
| 85 | assert checksums(server_url) == baseline |
| 86 | |
| 87 | changed, missing, *_ = baseline |
| 88 | altered = {**baseline, changed: "0" * 64} |
| 89 | del altered[missing] |
| 90 | assert weight_checker(server_url, "compare", altered).json() == { |
| 91 | "match": False, |
| 92 | "mismatches": sorted([changed, missing]), |
| 93 | } |
| 94 | |
| 95 | def test_reset_changes_weights_and_reload_restores_them(self, server_url): |
| 96 | reset_reload_and_compare(server_url) |
| 97 | assert ok(gen(server_url)) |
| 98 | |
| 99 | |
| 100 | @pytest.mark.skipif( |
| 101 | not current_platform.is_cuda_alike(), reason="IPC weight transfer uses CUDA IPC." |
| 102 | ) |
| 103 | def test_ipc_weight_transfer_restores_reset_weights(): |
| 104 | """The documented RL flow: reset, then transfer the checkpoint back.""" |
| 105 | if current_platform.is_xpu(): |
| 106 | pytest.skip("IPC weight transfer backend uses CUDA IPC handles") |
| 107 | args = ["--weight-transfer-config", '{"backend": "ipc"}'] |
| 108 | with ( |
| 109 | patch.dict(os.environ, {"VLLM_ALLOW_INSECURE_SERIALIZATION": "1"}), |
| 110 | server(extra_args=args, port=8772) as url, |
| 111 | ): |
| 112 | baseline = checksums(url) |
| 113 | weight_checker(url, "reset").raise_for_status() |
| 114 | trainer_model = AutoModelForCausalLM.from_pretrained( |
| 115 | MODEL_NAME, torch_dtype=torch.bfloat16 |
| 116 | ).cuda() |
| 117 | WeightTransferTrainerFactory.trainer_init( |
| 118 | IPCTrainerInitInfo(rank=0), |
| 119 | client=HTTPVLLMWeightSyncClient(url), |
| 120 | source=ModuleSource(trainer_model), |
| 121 | ).send_weights() |
| 122 | assert weight_checker(url, "compare", baseline).json() == { |
| 123 | "match": True, |
| 124 | "mismatches": [], |
| 125 | } |
| 126 | |
| 127 | |
| 128 | @multi_gpu_test(num_gpus=2) |
| 129 | @pytest.mark.parametrize( |
| 130 | "args, expected", |
| 131 | [ |
| 132 | (["--data-parallel-size", "2"], {"dp0:pp0:pcp0:tp0", "dp1:pp0:pcp0:tp0"}), |
| 133 | (["--tensor-parallel-size", "2"], {"dp0:pp0:pcp0:tp0", "dp0:pp0:pcp0:tp1"}), |
| 134 | ], |
| 135 | ) |
| 136 | def test_checksum_and_reset_cover_every_worker(args, expected): |
| 137 | with server(extra_args=args, port=8771, timeout=600) as url: |
| 138 | assert ranks(checksums(url)) == expected |
| 139 | reset_reload_and_compare(url) |