From 67d3a2ea5750d4246f2e2dd0f53dc44ae9d6425c Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Fri, 11 Sep 2026 07:26:39 +0530 Subject: [PATCH] [XPU] Make checkpoint_engine worker device-agnostic (#32382) --- .github/workflows/pr-test-xpu.yml | 5 + .../advanced_features/checkpoint_engine.mdx | 14 +- python/pyproject_xpu.toml | 1 + .../checkpoint_engine_worker.py | 23 ++- .../test_checkpoint_engine_worker.py | 141 ++++++++++++++++++ 5 files changed, 171 insertions(+), 13 deletions(-) create mode 100644 test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py diff --git a/.github/workflows/pr-test-xpu.yml b/.github/workflows/pr-test-xpu.yml index e1f114900..d432697f7 100644 --- a/.github/workflows/pr-test-xpu.yml +++ b/.github/workflows/pr-test-xpu.yml @@ -151,6 +151,11 @@ jobs: docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install --no-cache-dir --no-deps xgrammar==0.1.33 docker exec ci_sglang_xpu /bin/bash -c '/opt/venv/bin/hf auth login --token ${HF_TOKEN}' + - name: Install checkpoint-engine extra (optional; tests skip if unavailable) + continue-on-error: true + run: | + docker exec -w /sglang-checkout/python ci_sglang_xpu /opt/venv/bin/python3 -m pip install --no-cache-dir ".[checkpoint-engine]" --extra-index-url https://download.pytorch.org/whl/xpu + - name: Run tests # Old stage-a (30) + stage-b (120) = 150; +20% headroom for --enable-retry. timeout-minutes: 180 diff --git a/docs/docs/advanced_features/checkpoint_engine.mdx b/docs/docs/advanced_features/checkpoint_engine.mdx index 59472e609..6c9b1ba8e 100644 --- a/docs/docs/advanced_features/checkpoint_engine.mdx +++ b/docs/docs/advanced_features/checkpoint_engine.mdx @@ -25,7 +25,7 @@ pip install 'checkpoint-engine[p2p]' The system consists of two main components: -1. **SGLang Server**: Runs with `--wait-for-initial-weights` flag to wait for weights before becoming ready +1. **SGLang Server**: Runs with `--checkpoint-engine-wait-weights-before-ready` flag to wait for weights before becoming ready 2. **Checkpoint Engine Workers**: Separate processes (managed by torchrun) that load and distribute model weights The checkpoint engine uses a parameter server architecture with support for: @@ -43,7 +43,7 @@ python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ - --wait-for-initial-weights + --checkpoint-engine-wait-weights-before-ready ``` **Terminal 2 - Run Checkpoint Engine:** @@ -75,7 +75,7 @@ python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ - --wait-for-initial-weights \ + --checkpoint-engine-wait-weights-before-ready \ --host [IP] ``` @@ -110,7 +110,7 @@ python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ - --wait-for-initial-weights \ + --checkpoint-engine-wait-weights-before-ready \ --host [IP] ``` @@ -147,7 +147,7 @@ python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ - --wait-for-initial-weights \ + --checkpoint-engine-wait-weights-before-ready \ --host [IP] \ --dist-init-addr [IP]:9120 \ --nnodes 2 \ @@ -185,7 +185,7 @@ python -m sglang.launch_server \ --model-path Qwen/Qwen3-8B \ --tp 8 \ --load-format dummy \ - --wait-for-initial-weights \ + --checkpoint-engine-wait-weights-before-ready \ --host [IP] \ --dist-init-addr [IP]:9120 \ --nnodes 2 \ @@ -220,7 +220,7 @@ torchrun --nproc-per-node 8 \ ### SGLang Server Options - `--load-format dummy`: Use dummy format for initial loading (allows overlapping with other tasks) -- `--wait-for-initial-weights`: Wait for checkpoint engine to provide weights before becoming ready +- `--checkpoint-engine-wait-weights-before-ready`: Wait for checkpoint engine to provide weights before becoming ready - `--host`: Host address for multi-node setups - `--dist-init-addr`: Distributed initialization address for tensor parallelism diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index aea322914..6e9920528 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -73,6 +73,7 @@ dependencies = [ ] [project.optional-dependencies] +checkpoint-engine = ["checkpoint-engine @ git+https://github.com/MoonshotAI/checkpoint-engine.git@f772d1858204fa82bdb0436fdfec137f5de2972c"] diffusion = [ "addict==2.4.0", "av==16.1.0", diff --git a/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py b/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py index 6f11c7872..60dfe48bc 100644 --- a/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py +++ b/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py @@ -22,6 +22,9 @@ from typing import Callable, Dict, Optional import torch import zmq +from sglang.srt.platforms import current_platform +from sglang.srt.utils import get_device, get_device_module, is_npu + try: from checkpoint_engine.worker import update_weights_from_ipc except ImportError: @@ -100,17 +103,25 @@ class SGLangCheckpointEngineWorkerExtensionImpl(SGLangCheckpointEngineWorkerExte self.model_runner = model_runner def get_device_uuid(self) -> str: - """Get the UUID of current device.""" - # Get device UUID for current device - device_id = torch.cuda.current_device() + """Physical GPU id, matching checkpoint-engine's ParameterServer key. + + Must equal ps.py::_get_physical_gpu_id for the ZMQ handshake to resolve: + NPU uses an ``NPU-`` key, every other accelerator ``GPU-``. + NPU keeps its own branch because its key is derived from npu-smi rather + than device properties, and no NPU platform implements get_device_uuid.""" + if is_npu(): + from checkpoint_engine.device_utils import npu_generate_uuid + + return f"NPU-{npu_generate_uuid()}" + device_id = get_device_module().current_device() try: - return f"GPU-{torch.cuda.get_device_properties(device_id).uuid!s}" + return f"GPU-{current_platform.get_device_uuid(device_id)}" except AssertionError as e: raise ValueError(f"Failed to get GPU UUID for device {device_id}") from e def get_device_id(self) -> int: """Get the device ID.""" - return torch.cuda.current_device() + return get_device_module().current_device() def get_model_loader(self) -> Callable: """Get the model weight loader function.""" @@ -130,7 +141,7 @@ class SGLangCheckpointEngineWorkerExtensionImpl(SGLangCheckpointEngineWorkerExte if quant_method is not None: # Move parameters to device if needed for quantization processing target_device = torch.device( - "cuda", torch.cuda.current_device() + get_device(), get_device_module().current_device() ) with device_loading_context(module, target_device): quant_method.process_weights_after_loading(module) diff --git a/test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py b/test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py new file mode 100644 index 000000000..50ece8d34 --- /dev/null +++ b/test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py @@ -0,0 +1,141 @@ +"""Unit tests for srt/checkpoint_engine/checkpoint_engine_worker.py — no server, no model loading. + +Focus: device resolution so the ZMQ handshake key matches checkpoint-engine's +ParameterServer (ps.py::_get_physical_gpu_id) on every backend -- ``GPU-`` +for CUDA/XPU and ``NPU-`` for NPU. These paths are pure namespace routing +(``get_device`` / ``get_device_module`` / ``is_npu``) and are fully mockable on CPU. + +Skipped entirely unless the ``checkpoint-engine`` extra is installed, since the +worker module refuses to import without it. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +# Unit tests may register CPU suites only (scripts/lint/check_registered_tests.py). +# TestWorkerDeviceUuidOnXpu below still self-skips off XPU, so it stays runnable +# by hand on an Intel GPU host. +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import importlib.util +import unittest +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.utils import is_xpu +from sglang.test.test_utils import CustomTestCase + +# checkpoint-engine is an optional extra (sglang[checkpoint-engine]) that CI does +# not install, and the worker module raises ImportError at import time without it. +# Probe first so this file stays importable: an unguarded import would fail +# collection and take the whole file down rather than skipping. +_HAS_CHECKPOINT_ENGINE = importlib.util.find_spec("checkpoint_engine") is not None + +if _HAS_CHECKPOINT_ENGINE: + from sglang.srt.checkpoint_engine.checkpoint_engine_worker import ( + SGLangCheckpointEngineWorkerExtensionImpl, + ) + +_WORKER_MOD = "sglang.srt.checkpoint_engine.checkpoint_engine_worker" +_NO_CKPT_ENGINE = "requires the checkpoint-engine optional dependency" + + +@unittest.skipUnless(_HAS_CHECKPOINT_ENGINE, _NO_CKPT_ENGINE) +class TestWorkerDeviceResolution(CustomTestCase): + """get_device_uuid / get_device_id must route through the active accelerator + namespace and emit the key the ParameterServer expects.""" + + def _make_worker(self): + # model_runner is unused by the device-resolution methods under test. + return SGLangCheckpointEngineWorkerExtensionImpl(model_runner=MagicMock()) + + def _fake_device_module(self, *, current=3): + mod = MagicMock() + mod.current_device.return_value = current + return mod + + def _fake_platform(self, *, uuid="abcd-1234"): + # The uuid now comes from the platform layer (current_platform), which + # already returns str(get_device_properties(id).uuid) for cuda/xpu. + plat = MagicMock() + plat.get_device_uuid.return_value = uuid + return plat + + def test_device_uuid_cuda(self): + worker = self._make_worker() + fake = self._fake_device_module(current=0) + with ( + patch(f"{_WORKER_MOD}.is_npu", return_value=False), + patch(f"{_WORKER_MOD}.get_device_module", return_value=fake), + patch( + f"{_WORKER_MOD}.current_platform", self._fake_platform(uuid="cuda-uuid") + ), + ): + self.assertEqual(worker.get_device_uuid(), "GPU-cuda-uuid") + self.assertEqual(worker.get_device_id(), 0) + + def test_device_uuid_xpu(self): + worker = self._make_worker() + fake = self._fake_device_module(current=2) + with ( + patch(f"{_WORKER_MOD}.is_npu", return_value=False), + patch(f"{_WORKER_MOD}.get_device_module", return_value=fake), + patch( + f"{_WORKER_MOD}.current_platform", self._fake_platform(uuid="xpu-uuid") + ), + ): + # XPU shares CUDA's GPU- format; only the namespace differs. + self.assertEqual(worker.get_device_uuid(), "GPU-xpu-uuid") + self.assertEqual(worker.get_device_id(), 2) + + def test_device_uuid_npu_uses_npu_prefix(self): + # NPU must NOT be treated as CUDA: the ParameterServer keys it as + # NPU-, so a GPU- key would never resolve. + worker = self._make_worker() + with ( + patch(f"{_WORKER_MOD}.is_npu", return_value=True), + patch( + "checkpoint_engine.device_utils.npu_generate_uuid", + return_value="1.2.3.4-0", + ), + ): + self.assertEqual(worker.get_device_uuid(), "NPU-1.2.3.4-0") + + def test_device_uuid_wraps_assertion_error(self): + worker = self._make_worker() + fake = self._fake_device_module(current=1) + plat = MagicMock() + plat.get_device_uuid.side_effect = AssertionError("no uuid") + with ( + patch(f"{_WORKER_MOD}.is_npu", return_value=False), + patch(f"{_WORKER_MOD}.get_device_module", return_value=fake), + patch(f"{_WORKER_MOD}.current_platform", plat), + self.assertRaises(ValueError), + ): + worker.get_device_uuid() + + +@unittest.skipUnless(_HAS_CHECKPOINT_ENGINE, _NO_CKPT_ENGINE) +@unittest.skipUnless(is_xpu(), "requires an Intel XPU") +class TestWorkerDeviceUuidOnXpu(CustomTestCase): + """Hardware-gated: the real XPU key must match what checkpoint-engine's + ParameterServer derives, or the ZMQ handshake silently fails on XPU.""" + + def test_real_uuid_matches_parameter_server(self): + from checkpoint_engine.device_utils import DeviceManager + from checkpoint_engine.ps import _get_physical_gpu_id + + worker = SGLangCheckpointEngineWorkerExtensionImpl(model_runner=MagicMock()) + key = worker.get_device_uuid() + + self.assertTrue(key.startswith("GPU-"), key) + self.assertEqual(worker.get_device_id(), torch.xpu.current_device()) + + # Independently derived by the ParameterServer side; the two must agree. + dm = DeviceManager() + self.assertEqual(dm.device_type, "xpu") + self.assertEqual(key, _get_physical_gpu_id(dm, torch.xpu.current_device())) + + +if __name__ == "__main__": + unittest.main(verbosity=3)