minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Model Loading] Overlap checkpoint staging with CUDA graph capture during startup (#32017)

Browse Source 
Co-authored-by: Wenhui Zhu <wzhu59@asu.edu>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
Han Yu
2026-08-13 12:26:25 -07:00
committed by GitHub
co-authored by Wenhui Zhu Alex Nails
parent 0772e79ee7
commit 6b94d39f13
16 changed files with 2020 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/model_loader/test_prefetch_checkpoints.py
+201 -32
View File
@@ -7,10 +7,11 @@ to weights loaded without prefetch.
import os
import tempfile
import threading
import unittest
from concurrent.futures import Future
from types import SimpleNamespace
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import safetensors.torch
import torch
@@ -18,6 +19,7 @@ import torch
from sglang.srt.configs.load_config import LoadConfig, LoadFormat
from sglang.srt.model_loader.loader import DefaultModelLoader
from sglang.srt.model_loader.weight_utils import (
CheckpointFilePrefetchHandle,
_prefetch_all_checkpoints,
buffered_multi_thread_safetensors_weights_iterator,
fastsafetensors_weights_iterator,
@@ -37,6 +39,12 @@ class _InlineThread:
def start(self):
self.target()
def join(self, timeout=None):
pass
def is_alive(self):
return False
class _InlineExecutor:
def __init__(self, max_workers):
@@ -99,6 +107,45 @@ class TestPrefetchCheckpoints(CustomTestCase):
with self.assertRaisesRegex(ValueError, "num_threads"):
_prefetch_all_checkpoints(["dummy.safetensors"], num_threads=0)
@patch("torch.distributed.is_initialized", return_value=False)
def test_wait_returns_after_worker_thread_failure(self, _):
worker_errors = []
with (
patch(
"concurrent.futures.ThreadPoolExecutor",
side_effect=RuntimeError("worker failed"),
),
patch(
"threading.excepthook",
side_effect=lambda args: worker_errors.append(args.exc_value),
),
):
handle = _prefetch_all_checkpoints(["dummy.safetensors"], num_threads=1)
handle.wait(timeout=5)
self.assertTrue(handle.done)
self.assertTrue(handle.failed)
self.assertEqual(handle.errors, ())
self.assertEqual(len(worker_errors), 1)
self.assertIsInstance(worker_errors[0], RuntimeError)
def test_prefetch_stop_has_a_bounded_default_wait(self):
thread = MagicMock()
thread.is_alive.return_value = True
cancel_event = threading.Event()
handle = CheckpointFilePrefetchHandle(
thread=thread,
cancel_event=cancel_event,
succeeded_event=threading.Event(),
errors=[],
)
with self.assertRaisesRegex(TimeoutError, "checkpoint prefetching"):
handle.stop()
self.assertTrue(cancel_event.is_set())
thread.join.assert_called_once_with(60.0)
@patch("torch.distributed.is_initialized", return_value=False)
def test_prefetch_keeps_bounded_pending_window(self, _):
paths = [f"model-{i:05d}.safetensors" for i in range(20)]
@@ -106,9 +153,9 @@ class TestPrefetchCheckpoints(CustomTestCase):
submitted_paths = []
class RecordingExecutor(_InlineExecutor):
def submit(self, fn, path):
def submit(self, fn, path, *args):
submitted_paths.append(path)
return super().submit(fn, path)
return super().submit(fn, path, *args)
def record_pending_size(fs, return_when):
pending_sizes.append(len(fs))
@@ -129,7 +176,7 @@ class TestPrefetchCheckpoints(CustomTestCase):
def test_prefetch_logs_failed_futures(self, _):
paths = ["bad.safetensors"]
def fail_prefetch(path):
def fail_prefetch(path, cancel_event):
raise OSError(f"failed {path}")
with (
@@ -142,15 +189,18 @@ class TestPrefetchCheckpoints(CustomTestCase):
),
patch("sglang.srt.model_loader.weight_utils.logger.warning") as warning,
):
_prefetch_all_checkpoints(paths, num_threads=1)
handle = _prefetch_all_checkpoints(paths, num_threads=1)
handle.wait()
self.assertEqual(handle.errors[0][0], paths[0])
self.assertIsInstance(handle.errors[0][1], OSError)
warning.assert_called_once()
self.assertEqual(
warning.call_args.args[0],
"Failed to prefetch checkpoint file %r.",
"Failed to prefetch checkpoint file %r: %s",
)
self.assertEqual(warning.call_args.args[1], paths[0])
self.assertTrue(warning.call_args.kwargs["exc_info"])
self.assertIsInstance(warning.call_args.args[2], OSError)
@patch("torch.distributed.is_initialized", return_value=False)
def test_prefetch_progress_logs_all_crossed_buckets(self, _):
@@ -193,13 +243,40 @@ class TestPrefetchCheckpoints(CustomTestCase):
),
patch(
"sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file",
side_effect=loaded_paths.append,
side_effect=lambda path, cancel_event: loaded_paths.append(path),
),
):
_prefetch_all_checkpoints(paths, num_threads=2)
self.assertEqual(sorted(loaded_paths), sorted(paths[1::3]))
@patch("torch.distributed.is_initialized", return_value=False)
def test_prefetch_handle_cancels_before_scheduling_next_shard(self, _):
paths = [f"model-{i:05d}.safetensors" for i in range(3)]
started = threading.Event()
release = threading.Event()
loaded_paths = []
def block_first_prefetch(path, cancel_event):
loaded_paths.append(path)
started.set()
self.assertTrue(release.wait(timeout=5))
with patch(
"sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file",
side_effect=block_first_prefetch,
):
handle = _prefetch_all_checkpoints(paths, num_threads=1)
self.assertTrue(started.wait(timeout=5))
with self.assertRaisesRegex(TimeoutError, "checkpoint prefetching"):
handle.wait(timeout=0)
handle.cancel()
release.set()
handle.wait(timeout=5)
self.assertTrue(handle.cancelled)
self.assertEqual(loaded_paths, paths[:1])
@patch("torch.distributed.is_initialized", return_value=False)
def test_buffered_loader_drops_cache_after_each_loaded_shard(self, _):
with tempfile.TemporaryDirectory() as tmpdir:
@@ -319,11 +396,16 @@ class TestPrefetchDispatch(CustomTestCase):
weight_loader_drop_cache_after_load=drop_cache,
)
def _run(self, loader):
def _run(self, loader, **iterator_kwargs):
# _get_weights_iterator returns a generator wrapping the chosen
# iterator; consuming it forces the eager dispatch (the if/elif/else
# that calls the iterator factory) to execute.
list(loader._get_weights_iterator(self._make_source()))
list(
loader._get_weights_iterator(
self._make_source(),
**iterator_kwargs,
)
)
def _patch_dispatch(self, prefetch, disable_mmap=False, drop_cache=False):
return (
@@ -332,14 +414,6 @@ class TestPrefetchDispatch(CustomTestCase):
"_prepare_weights",
return_value=("/dummy", ["f.safetensors"], True),
),
patch(
"sglang.srt.model_loader.loader.get_server_args",
return_value=self._server_args(
prefetch,
disable_mmap,
drop_cache,
),
),
patch(
"sglang.srt.model_loader.loader.get_model",
return_value=self._server_args(prefetch, disable_mmap, drop_cache),
@@ -360,12 +434,11 @@ class TestPrefetchDispatch(CustomTestCase):
"""Prefetch on + no explicit multithread config -> single-threaded,
and the opt-out warning fires once."""
loader = self._make_loader({})
p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=True
)
with (
p_prep,
p_args,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
@@ -380,12 +453,11 @@ class TestPrefetchDispatch(CustomTestCase):
"""Explicit enable_multithread_load=true is the escape hatch; the
override and its warning must not fire."""
loader = self._make_loader({"enable_multithread_load": True})
p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=True
)
with (
p_prep,
p_args,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
@@ -401,12 +473,11 @@ class TestPrefetchDispatch(CustomTestCase):
default) also signals multi-thread intent, so the override must not
fire and num_threads stays live."""
loader = self._make_loader({"num_threads": 64})
p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=True
)
with (
p_prep,
p_args,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
@@ -423,12 +494,11 @@ class TestPrefetchDispatch(CustomTestCase):
"""Prefetch off -> multi-threaded iterator is used (default), no
override warning."""
loader = self._make_loader({})
p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=False
)
with (
p_prep,
p_args,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
@@ -439,16 +509,117 @@ class TestPrefetchDispatch(CustomTestCase):
mock_single.assert_not_called()
mock_warning.assert_not_called()
def test_startup_prefetch_reuses_existing_background_handle(self):
"""Startup commit reuses resolved shards and the active prefetch handle."""
loader = self._make_loader({})
source = self._make_source()
resolved_source = DefaultModelLoader.ResolvedSource(
source=source,
hf_folder="/dummy",
weight_files=("f.safetensors",),
use_safetensors=True,
)
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=False
)
with (
p_prep as mock_prepare,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
p_warn as mock_warning,
):
list(
loader._get_weights_iterator(
source,
resolved_source=resolved_source,
startup_prefetch_started=True,
startup_prefetch_active=True,
)
)
mock_prepare.assert_not_called()
mock_single.assert_called_once()
self.assertFalse(mock_single.call_args.kwargs["prefetch"])
mock_buffered.assert_not_called()
mock_warning.assert_called_once()
def test_completed_startup_prefetch_restores_multithread_loader(self):
loader = self._make_loader({})
source = self._make_source()
resolved_source = DefaultModelLoader.ResolvedSource(
source=source,
hf_folder="/dummy",
weight_files=("f.safetensors",),
use_safetensors=True,
)
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=False
)
with (
p_prep as mock_prepare,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
p_warn as mock_warning,
):
list(
loader._get_weights_iterator(
source,
resolved_source=resolved_source,
startup_prefetch_started=True,
startup_prefetch_active=False,
)
)
mock_prepare.assert_not_called()
mock_buffered.assert_called_once()
self.assertFalse(mock_buffered.call_args.kwargs["prefetch"])
mock_single.assert_not_called()
mock_warning.assert_not_called()
def test_completed_startup_prefetch_is_not_started_twice(self):
loader = self._make_loader({})
source = self._make_source()
resolved_source = DefaultModelLoader.ResolvedSource(
source=source,
hf_folder="/dummy",
weight_files=("f.safetensors",),
use_safetensors=True,
)
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=True
)
with (
p_prep,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
p_warn as mock_warning,
):
list(
loader._get_weights_iterator(
source,
resolved_source=resolved_source,
startup_prefetch_started=True,
startup_prefetch_active=False,
)
)
mock_buffered.assert_called_once()
self.assertFalse(mock_buffered.call_args.kwargs["prefetch"])
mock_single.assert_not_called()
mock_warning.assert_not_called()
def test_prefetch_does_not_override_when_mmap_disabled(self):
"""Prefetch is a no-op without mmap, so the override and its warning
must not fire."""
loader = self._make_loader({})
p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=True, disable_mmap=True
)
with (
p_prep,
p_args,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
@@ -463,7 +634,7 @@ class TestPrefetchDispatch(CustomTestCase):
"""FASTSAFETENSORS ignores both flags; override + warning must not
fire."""
loader = self._make_loader({}, load_format=LoadFormat.FASTSAFETENSORS)
p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=True
)
with (
@@ -472,7 +643,6 @@ class TestPrefetchDispatch(CustomTestCase):
return_value=iter([]),
) as mock_fast,
p_prep,
p_args,
p_model,
p_buffered as mock_buffered,
p_single as mock_single,
@@ -492,7 +662,7 @@ class TestPrefetchDispatch(CustomTestCase):
loader = self._make_loader(
{"enable_gds": False}, load_format=LoadFormat.FASTSAFETENSORS
)
p_prep, p_args, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch(
prefetch=False,
drop_cache=True,
)
@@ -502,7 +672,6 @@ class TestPrefetchDispatch(CustomTestCase):
return_value=iter([]),
) as mock_fast,
p_prep,
p_args,
p_model,
p_buffered,
p_single,