diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 4e4b83059..2340bbf44 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -67,7 +67,7 @@ from sglang.utils import is_in_ci try: from fastsafetensors import SafeTensorsFileLoader, SingleGroup -except ImportError as e: +except ImportError: SafeTensorsFileLoader = SingleGroup = None logger = logging.getLogger(__name__) @@ -814,10 +814,12 @@ def _prefetch_all_checkpoints( naturally adapts to any RAM size — even if the full checkpoint does not fit in page cache, the prefetch thread stays ahead of the loader. """ - import asyncio import threading import time + if num_threads < 1: + raise ValueError("weight loader prefetch num_threads must be >= 1") + # Use node-local rank so that each node independently prefetches the # full checkpoint into its own page cache. Global rank would split files # across nodes, but page cache is not shared across nodes. @@ -842,40 +844,60 @@ def _prefetch_all_checkpoints( num_threads, ) - async def _prefetch_all() -> None: - semaphore = asyncio.Semaphore(num_threads) + def _prefetch_all() -> None: completed = 0 next_log_pct = 10 - async def prefetch_one(path: str) -> None: + def record_complete() -> None: nonlocal completed, next_log_pct - try: - async with semaphore: - await asyncio.to_thread(_prefetch_checkpoint_file, path) - completed += 1 - if total_for_rank > 0 and next_log_pct <= 100: - pct = 100 * completed / total_for_rank - if pct >= next_log_pct: - logger.info( - "Rank %d: prefetching checkpoint files: %d%% (%d/%d)", - local_rank, - next_log_pct, - completed, - total_for_rank, - ) - next_log_pct += 10 - except Exception: - logger.warning( - "Failed to prefetch checkpoint file %r.", - path, - exc_info=True, - ) - await asyncio.gather(*(prefetch_one(p) for p in my_files)) + completed += 1 + if total_for_rank > 0 and next_log_pct <= 100: + pct = 100 * completed / total_for_rank + while pct >= next_log_pct and next_log_pct <= 100: + logger.info( + "Rank %d: prefetching checkpoint files: %d%% (%d/%d)", + local_rank, + next_log_pct, + completed, + total_for_rank, + ) + next_log_pct += 10 + + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + file_iter = iter(my_files) + pending: Dict[concurrent.futures.Future, str] = {} + + for path in itertools.islice(file_iter, num_threads): + pending[executor.submit(_prefetch_checkpoint_file, path)] = path + + while pending: + done, _ = concurrent.futures.wait( + pending, + return_when=concurrent.futures.FIRST_COMPLETED, + ) + for future in done: + path = pending.pop(future) + try: + future.result() + except Exception: + logger.warning( + "Failed to prefetch checkpoint file %r.", + path, + exc_info=True, + ) + finally: + record_complete() + + next_path = next(file_iter, None) + if next_path is not None: + pending[ + executor.submit(_prefetch_checkpoint_file, next_path) + ] = next_path def _run_prefetch() -> None: start = time.perf_counter() - asyncio.run(_prefetch_all()) + _prefetch_all() elapsed = time.perf_counter() - start logger.info( "Rank %d: prefetching checkpoint files into page cache " @@ -1081,7 +1103,7 @@ def buffered_multi_thread_safetensors_weights_iterator( # Seed the buffer. for st_file in itertools.islice(file_iter, buffer_size): - pending.append(executor.submit(_load_file, st_file)) + pending.append((st_file, executor.submit(_load_file, st_file))) with tqdm( total=len(hf_weights_files), @@ -1091,18 +1113,22 @@ def buffered_multi_thread_safetensors_weights_iterator( position=tqdm._get_free_pos(), ) as pbar: while pending: - future = pending.popleft() + st_file, future = pending.popleft() state_dict = future.result() del future # let GC reclaim the Future's internal result # Replenish: submit the next file to keep the buffer full. next_file = next(file_iter, None) if next_file is not None: - pending.append(executor.submit(_load_file, next_file)) + pending.append((next_file, executor.submit(_load_file, next_file))) for name in sorted(state_dict.keys()): yield name, state_dict[name] del state_dict + if drop_cache_after_load: + # DONTNEED reduces page-cache pressure after copying weights, + # but later mmap-backed tensor access may fault pages again. + _drop_file_cache_after_load(st_file) pbar.update(1) diff --git a/test/registered/unit/model_loader/test_prefetch_checkpoints.py b/test/registered/unit/model_loader/test_prefetch_checkpoints.py index e417e60b3..e4a4c861b 100644 --- a/test/registered/unit/model_loader/test_prefetch_checkpoints.py +++ b/test/registered/unit/model_loader/test_prefetch_checkpoints.py @@ -8,12 +8,15 @@ to weights loaded without prefetch. import os import tempfile import unittest +from concurrent.futures import Future from unittest.mock import patch import safetensors.torch import torch from sglang.srt.model_loader.weight_utils import ( + _prefetch_all_checkpoints, + buffered_multi_thread_safetensors_weights_iterator, safetensors_weights_iterator, ) from sglang.test.ci.ci_register import register_cpu_ci @@ -21,8 +24,40 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") -class TestPrefetchWeightsIdentical(unittest.TestCase): - """Verify that loading with prefetch yields identical weights to without.""" +class _InlineThread: + def __init__(self, target, daemon=None): + self.target = target + self.daemon = daemon + + def start(self): + self.target() + + +class _InlineExecutor: + def __init__(self, max_workers): + self.max_workers = max_workers + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def submit(self, fn, *args, **kwargs): + future = Future() + try: + future.set_result(fn(*args, **kwargs)) + except Exception as exc: + future.set_exception(exc) + return future + + +def _wait_all(fs, return_when): + return set(fs), set() + + +class TestPrefetchCheckpoints(unittest.TestCase): + """Verify coordinated checkpoint prefetch behavior.""" def _create_safetensors_files(self, tmpdir, num_shards=3): """Create real safetensors files with known tensor content.""" @@ -44,12 +79,144 @@ class TestPrefetchWeightsIdentical(unittest.TestCase): paths = self._create_safetensors_files(tmpdir) without = dict(safetensors_weights_iterator(paths, prefetch=False)) - with_pf = dict(safetensors_weights_iterator(paths, prefetch=True)) + with ( + patch("threading.Thread", _InlineThread), + patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor), + patch("concurrent.futures.wait", side_effect=_wait_all), + ): + with_pf = dict(safetensors_weights_iterator(paths, prefetch=True)) self.assertEqual(set(without.keys()), set(with_pf.keys())) for name in without: torch.testing.assert_close(without[name], with_pf[name]) + def test_prefetch_rejects_invalid_thread_count(self): + with self.assertRaisesRegex(ValueError, "num_threads"): + _prefetch_all_checkpoints(["dummy.safetensors"], num_threads=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)] + pending_sizes = [] + submitted_paths = [] + + class RecordingExecutor(_InlineExecutor): + def submit(self, fn, path): + submitted_paths.append(path) + return super().submit(fn, path) + + def record_pending_size(fs, return_when): + pending_sizes.append(len(fs)) + return _wait_all(fs, return_when) + + with ( + patch("threading.Thread", _InlineThread), + patch("concurrent.futures.ThreadPoolExecutor", RecordingExecutor), + patch("concurrent.futures.wait", side_effect=record_pending_size), + patch("sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file"), + ): + _prefetch_all_checkpoints(paths, num_threads=4) + + self.assertEqual(submitted_paths, paths) + self.assertLessEqual(max(pending_sizes), 4) + + @patch("torch.distributed.is_initialized", return_value=False) + def test_prefetch_logs_failed_futures(self, _): + paths = ["bad.safetensors"] + + def fail_prefetch(path): + raise OSError(f"failed {path}") + + with ( + patch("threading.Thread", _InlineThread), + patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor), + patch("concurrent.futures.wait", side_effect=_wait_all), + patch( + "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file", + side_effect=fail_prefetch, + ), + patch("sglang.srt.model_loader.weight_utils.logger.warning") as warning, + ): + _prefetch_all_checkpoints(paths, num_threads=1) + + warning.assert_called_once() + self.assertEqual( + warning.call_args.args[0], + "Failed to prefetch checkpoint file %r.", + ) + self.assertEqual(warning.call_args.args[1], paths[0]) + self.assertTrue(warning.call_args.kwargs["exc_info"]) + + @patch("torch.distributed.is_initialized", return_value=False) + def test_prefetch_progress_logs_all_crossed_buckets(self, _): + paths = [f"model-{i:05d}.safetensors" for i in range(3)] + + with ( + patch("threading.Thread", _InlineThread), + patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor), + patch("concurrent.futures.wait", side_effect=_wait_all), + patch("sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file"), + patch("sglang.srt.model_loader.weight_utils.logger.info") as log_info, + ): + _prefetch_all_checkpoints(paths, num_threads=1) + + progress_pcts = [ + call.args[2] + for call in log_info.call_args_list + if call.args + and call.args[0] == "Rank %d: prefetching checkpoint files: %d%% (%d/%d)" + ] + self.assertEqual(progress_pcts, list(range(10, 101, 10))) + + @patch("torch.distributed.is_initialized", return_value=True) + def test_prefetch_uses_node_local_rank_partitioning(self, _): + paths = [f"model-{i:05d}.safetensors" for i in range(10)] + loaded_paths = [] + + class FakeWorldGroup: + local_rank = 1 + local_size = 3 + world_size = 99 + + with ( + patch("threading.Thread", _InlineThread), + patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor), + patch("concurrent.futures.wait", side_effect=_wait_all), + patch( + "sglang.srt.model_loader.weight_utils.get_world_group", + return_value=FakeWorldGroup(), + ), + patch( + "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file", + side_effect=loaded_paths.append, + ), + ): + _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_buffered_loader_drops_cache_after_each_loaded_shard(self, _): + with tempfile.TemporaryDirectory() as tmpdir: + paths = self._create_safetensors_files(tmpdir, num_shards=3) + + with patch( + "sglang.srt.model_loader.weight_utils._drop_file_cache_after_load" + ) as drop_cache: + loaded = list( + buffered_multi_thread_safetensors_weights_iterator( + paths, + max_workers=2, + drop_cache_after_load=True, + ) + ) + + self.assertEqual(len(loaded), 6) + self.assertEqual( + [call.args[0] for call in drop_cache.call_args_list], + paths, + ) + if __name__ == "__main__": unittest.main()