Clean logging under --weight-loader-prefetch-checkpoints (#33930)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
This commit is contained in:
Brayden Zhong
2026-09-04 20:05:53 -07:00
committed by GitHub
co-authored by Brayden Zhong Mohammad Miadh Angkad
parent 0645398a32
commit 92a4d8b5ee
19 changed files with 111 additions and 78 deletions
@@ -95,10 +95,6 @@ import { GLM5Deployment } from '/src/snippets/autoregressive/glm-5-deployment.js
- For other configuration tips (MTP, DSA kernel, Context Parallel, HiSparse, NVFP4, Index Cache), see the [DeepSeek-V3.2 cookbook page](../DeepSeek/DeepSeek-V3_2). GLM-5 and DeepSeek-V3.2 share the same model structure, so the optimization techniques are common. - For other configuration tips (MTP, DSA kernel, Context Parallel, HiSparse, NVFP4, Index Cache), see the [DeepSeek-V3.2 cookbook page](../DeepSeek/DeepSeek-V3_2). GLM-5 and DeepSeek-V3.2 share the same model structure, so the optimization techniques are common.
- Use `--json-model-override-args '{"index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"}'` for GLM-5-FP8 if you want to enable the [IndexCache](https://github.com/THUDM/IndexCache) method. This feature is supported through [this PR](https://github.com/sgl-project/sglang/pull/21405) and introduces only a small accuracy loss. However, if you are running rigorous accuracy evaluations, it is not recommended to enable this feature. - Use `--json-model-override-args '{"index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"}'` for GLM-5-FP8 if you want to enable the [IndexCache](https://github.com/THUDM/IndexCache) method. This feature is supported through [this PR](https://github.com/sgl-project/sglang/pull/21405) and introduces only a small accuracy loss. However, if you are running rigorous accuracy evaluations, it is not recommended to enable this feature.
<Warning>
**FP8 KV Cache**: `--kv-cache-dtype fp8_e4m3` quantizes the KV cache to FP8 at runtime. Since these FP8 model checkpoints do not include pre-calibrated KV cache scaling factors, SGLang defaults to a scale of 1.0, which may cause noticeable accuracy degradation on reasoning-heavy tasks. It is not included in the generated commands above; add it manually only if memory constraints require the trade-off.
</Warning>
## 4. Model Invocation ## 4. Model Invocation
Deploy GLM-5 with the following command (FP8 on H200, all features enabled): Deploy GLM-5 with the following command (FP8 on H200, all features enabled):
@@ -322,10 +322,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
super().__init__(exclude_modules, packed_modules_mapping) super().__init__(exclude_modules, packed_modules_mapping)
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
if is_checkpoint_nvfp4_serialized: if is_checkpoint_nvfp4_serialized:
logger.warning( logger.info("Detected nvfp4 checkpoint.")
"Detected nvfp4 checkpoint. Please note that the "
"format is experimental and subject to change."
)
self.group_size = group_size self.group_size = group_size
self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv
self.swap_weight_nibbles = swap_weight_nibbles self.swap_weight_nibbles = swap_weight_nibbles
@@ -60,7 +60,6 @@ class CpDecodeAttnTpContext:
else: else:
self.decode_tp_rank = None self.decode_tp_rank = None
self.decode_tp_size = None self.decode_tp_size = None
logger.info("Disable CP decode attention TP")
self.use_decode_attn_tp = False self.use_decode_attn_tp = False
self._slice_cache: Dict = {} self._slice_cache: Dict = {}
@@ -83,7 +83,6 @@ from sglang.srt.utils import (
is_cpu, is_cpu,
is_hip, is_hip,
is_npu, is_npu,
print_info_once,
round_up, round_up,
) )
from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.custom_op import register_custom_op
@@ -474,7 +473,7 @@ class FusedMoE(torch.nn.Module):
global _deferred_finalize_info_logged global _deferred_finalize_info_logged
if not _deferred_finalize_info_logged: if not _deferred_finalize_info_logged:
_deferred_finalize_info_logged = True _deferred_finalize_info_logged = True
logging.getLogger(__name__).info( logging.getLogger(__name__).debug(
"FlashInfer TRTLLM MoE deferred finalize is " "FlashInfer TRTLLM MoE deferred finalize is "
f"{'enabled' if self.supports_deferred_finalize else 'disabled'} " f"{'enabled' if self.supports_deferred_finalize else 'disabled'} "
f"(moe_runner_backend={get_exec().moe.moe_runner_backend}, " f"(moe_runner_backend={get_exec().moe.moe_runner_backend}, "
@@ -516,10 +515,6 @@ class FusedMoE(torch.nn.Module):
get_moe_runner_backend().is_flashinfer_trtllm_routed() get_moe_runner_backend().is_flashinfer_trtllm_routed()
or get_moe_runner_backend().is_flashinfer_trtllm() or get_moe_runner_backend().is_flashinfer_trtllm()
): ):
if self.moe_runner_config.inplace:
print_info_once(
"Setting inplace to False for FlashInfer TRTLLM MoE backend."
)
self.moe_runner_config.inplace = False self.moe_runner_config.inplace = False
self.should_fuse_routed_scaling_factor_in_topk = ( self.should_fuse_routed_scaling_factor_in_topk = (
@@ -1429,10 +1429,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping) super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping)
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
if is_checkpoint_nvfp4_serialized: if is_checkpoint_nvfp4_serialized:
logger.warning( logger.info("Detected nvfp4 checkpoint.")
"Detected nvfp4 checkpoint. Please note that the "
"format is experimental and subject to change."
)
self.is_awq = is_awq self.is_awq = is_awq
self.is_w4a16 = False self.is_w4a16 = False
self.group_size = group_size self.group_size = group_size
@@ -45,10 +45,7 @@ class PetitNvFp4Config(QuantizationConfig):
) -> None: ) -> None:
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
if is_checkpoint_nvfp4_serialized: if is_checkpoint_nvfp4_serialized:
logger.warning( logger.info("Detected nvfp4 checkpoint.")
"Detected nvfp4 checkpoint. Please note that the "
"format is experimental and subject to change."
)
self.group_size = group_size self.group_size = group_size
self.kv_cache_quant_algo = kv_cache_quant_algo self.kv_cache_quant_algo = kv_cache_quant_algo
self.exclude_modules = exclude_modules self.exclude_modules = exclude_modules
@@ -287,6 +287,7 @@ def capture_prefill_graph(
"""Initialize a prefill graph and return its startup resource usage.""" """Initialize a prefill graph and return its startup resource usage."""
memory_phase = "draft_prefill" if model_runner.is_draft_worker else "prefill" memory_phase = "draft_prefill" if model_runner.is_draft_worker else "prefill"
role = "draft" if model_runner.is_draft_worker else "target"
def result( def result(
runner: Optional[BaseRunner], runner: Optional[BaseRunner],
@@ -430,8 +431,10 @@ def capture_prefill_graph(
layer_model = layer_model.model layer_model = layer_model.model
if not hasattr(layer_model, "layers"): if not hasattr(layer_model, "layers"):
logger.warning( log_info_on_rank0(
"Disable prefill CUDA graph because the model does not have a 'layers' attribute" logger,
f"Disable {role} prefill CUDA graph because the {role} model does "
"not have a 'layers' attribute",
) )
return result(None) return result(None)
@@ -467,7 +470,6 @@ def capture_prefill_graph(
tic = time.perf_counter() tic = time.perf_counter()
before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id) before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id)
role = "draft" if model_runner.is_draft_worker else "target"
capture_name = f"{role} prefill" capture_name = f"{role} prefill"
logger.info( logger.info(
f"Capture {capture_name} CUDA graph begin. " f"Capture {capture_name} CUDA graph begin. "
@@ -132,8 +132,7 @@ def load_kv_cache_scales(*, model, kv_cache_dtype: str) -> None:
else: else:
logger.warning( logger.warning(
"Using FP8 KV cache but no scaling factors " "Using FP8 KV cache but no scaling factors "
"provided. Defaulting to scaling factors of 1.0. " "provided. Defaulting to scaling factors of 1.0."
"This may lead to less accurate results!"
) )
@@ -23,6 +23,7 @@ buffers to keep break-point tensors at stable addresses.
""" """
import threading import threading
import warnings
from contextvars import ContextVar from contextvars import ContextVar
from typing import Any, Callable, Optional from typing import Any, Callable, Optional
@@ -394,7 +395,12 @@ class BreakableCUDAGraphCapture:
forked.clear() forked.clear()
graph = self._current_graph graph = self._current_graph
assert graph is not None assert graph is not None
graph.capture_end() # A segment that enqueued no kernels (back-to-back breaks, or a segment
# whose ops all ran eagerly) captures an empty graph, which replays as a
# no-op. Torch warns about it on every such capture_end; expected here.
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message="The CUDA Graph is empty")
graph.capture_end()
self.cuda_graph._append_segment(graph, self._current_graph_needs_instantiate) self.cuda_graph._append_segment(graph, self._current_graph_needs_instantiate)
self._current_graph = None self._current_graph = None
self._current_graph_needs_instantiate = False self._current_graph_needs_instantiate = False
+1 -1
View File
@@ -640,7 +640,7 @@ class DefaultModelLoader(BaseModelLoader):
{"enable_multithread_load", "num_threads"} & extra_config.keys() {"enable_multithread_load", "num_threads"} & extra_config.keys()
) )
): ):
logger.warning( logger.debug(
"Checkpoint prefetching is active; falling " "Checkpoint prefetching is active; falling "
"back to single-threaded weight loading to avoid I/O " "back to single-threaded weight loading to avoid I/O "
"oversubscription with the prefetch threads. Set " "oversubscription with the prefetch threads. Set "
@@ -1021,7 +1021,7 @@ def _prefetch_all_checkpoints(
succeeded_event = threading.Event() succeeded_event = threading.Event()
errors: List[Tuple[str, Exception]] = [] errors: List[Tuple[str, Exception]] = []
logger.info( logger.debug(
"Rank %d: prefetching %d/%d checkpoint shards into page cache " "Rank %d: prefetching %d/%d checkpoint shards into page cache "
"(background, %d local ranks sharing the work, %d threads per rank)...", "(background, %d local ranks sharing the work, %d threads per rank)...",
local_rank, local_rank,
@@ -1042,7 +1042,7 @@ def _prefetch_all_checkpoints(
if total_for_rank > 0 and next_log_pct <= 100: if total_for_rank > 0 and next_log_pct <= 100:
pct = 100 * completed / total_for_rank pct = 100 * completed / total_for_rank
while pct >= next_log_pct and next_log_pct <= 100: while pct >= next_log_pct and next_log_pct <= 100:
logger.info( logger.debug(
"Rank %d: prefetching checkpoint files: %d%% (%d/%d)", "Rank %d: prefetching checkpoint files: %d%% (%d/%d)",
local_rank, local_rank,
next_log_pct, next_log_pct,
@@ -1093,7 +1093,7 @@ def _prefetch_all_checkpoints(
start = time.perf_counter() start = time.perf_counter()
_prefetch_all() _prefetch_all()
succeeded_event.set() succeeded_event.set()
logger.info( logger.debug(
"Rank %d: prefetching checkpoint files into page cache finished in %.2fs", "Rank %d: prefetching checkpoint files into page cache finished in %.2fs",
local_rank, local_rank,
time.perf_counter() - start, time.perf_counter() - start,
@@ -80,7 +80,7 @@ class BailingMoEModelNextN(nn.Module):
config.for_nextn_model = True config.for_nextn_model = True
if quant_config is not None and quant_config.get_name() == "modelopt_fp4": if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning( logger.debug(
"Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model." "Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model."
) )
quant_config = None quant_config = None
+1 -1
View File
@@ -116,7 +116,7 @@ class DeepseekModelNextN(nn.Module):
moe_quant_config_override = None moe_quant_config_override = None
if quant_config is not None and quant_config.get_name() == "modelopt_fp4": if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning( logger.debug(
"Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model." "Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model."
) )
quant_config = None quant_config = None
@@ -50,7 +50,7 @@ class Glm4MoeLiteModelNextN(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
if quant_config is not None and quant_config.get_name() == "modelopt_fp4": if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning( logger.debug(
"Overriding Glm4MoeLiteForCausalLMNextN quant config for modelopt_fp4 " "Overriding Glm4MoeLiteForCausalLMNextN quant config for modelopt_fp4 "
"GLM-4.7-Flash model." "GLM-4.7-Flash model."
) )
+1 -1
View File
@@ -47,7 +47,7 @@ class Glm4MoeModelNextN(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
if quant_config is not None and quant_config.get_name() == "modelopt_fp4": if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning( logger.debug(
"Overriding Glm4MoeForCausalLMNextN quant config for modelopt_fp4 GLM-4.5 / GLM-4.6 / GLM-4.7 model." "Overriding Glm4MoeForCausalLMNextN quant config for modelopt_fp4 GLM-4.5 / GLM-4.6 / GLM-4.7 model."
) )
quant_config = None quant_config = None
+1 -1
View File
@@ -49,7 +49,7 @@ class GlmOcrModelNextN(nn.Module):
) -> None: ) -> None:
super().__init__() super().__init__()
if quant_config is not None and quant_config.get_name() == "modelopt_fp4": if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning( logger.debug(
"Overriding GlmOcrModelNextN quant config for modelopt_fp4 GLM-OCR model." "Overriding GlmOcrModelNextN quant config for modelopt_fp4 GLM-OCR model."
) )
quant_config = None quant_config = None
+10 -1
View File
@@ -58,6 +58,7 @@ from importlib.metadata import PackageNotFoundError, version
from importlib.util import find_spec from importlib.util import find_spec
from io import BytesIO from io import BytesIO
from json import JSONDecodeError from json import JSONDecodeError
from multiprocessing import parent_process
from multiprocessing.reduction import ForkingPickler from multiprocessing.reduction import ForkingPickler
from pathlib import Path from pathlib import Path
from typing import ( from typing import (
@@ -2409,6 +2410,14 @@ def configure_logger(server_args, prefix: str = ""):
for name in ("httpx", "httpcore"): for name in ("httpx", "httpcore"):
logging.getLogger(name).setLevel(logging.WARNING) logging.getLogger(name).setLevel(logging.WARNING)
# Server-sent hub warnings (e.g. the unauthenticated-request / HF_TOKEN
# hint) are deduplicated per process, so a TP-N launch repeats each one N
# times. Keep them only in the launching process -- every worker (scheduler,
# detokenizer, DP controller, ...) is spawned via multiprocessing, whether
# or not it passes a log prefix -- where they are printed exactly once.
if parent_process() is not None:
logging.getLogger("huggingface_hub.utils._http").setLevel(logging.ERROR)
if is_flashinfer_available(): if is_flashinfer_available():
from flashinfer.jit.core import logger as flashinfer_logger from flashinfer.jit.core import logger as flashinfer_logger
@@ -4106,7 +4115,7 @@ def freeze_gc(context: str):
g0_before, g1_before, g2_before = gc_object_counts() g0_before, g1_before, g2_before = gc_object_counts()
gc.freeze() gc.freeze()
g0_after, g1_after, g2_after = gc_object_counts() g0_after, g1_after, g2_after = gc_object_counts()
logger.info( logger.debug(
f"Freezing GC in {context} process. " f"Freezing GC in {context} process. "
f"gen0: {g0_before}->{g0_after}, " f"gen0: {g0_before}->{g0_after}, "
f"gen1: {g1_before}->{g1_after}, " f"gen1: {g1_before}->{g1_after}, "
@@ -53,6 +53,8 @@ def apply_all():
return return
_applied = True _applied = True
_mute_diffusers_torchao_probe()
# v5.4 patches # v5.4 patches
_patch_flash_attn_availability() _patch_flash_attn_availability()
_patch_rope_parameters_validation() _patch_rope_parameters_validation()
@@ -71,6 +73,23 @@ def apply_all():
logger.debug("transformers compatibility patches applied") logger.debug("transformers compatibility patches applied")
def _mute_diffusers_torchao_probe():
"""Silence diffusers' torchao-Tensor-subclass probe warning.
diffusers lazily imports its torchao quantizer and warns when the installed
torchao has moved the optional Tensor subclasses it probes for. It only
affects loading torchao-serialized diffusers checkpoints, which no sglang
path does. Set here rather than in ``configure_logger`` because the import
can land before logging is configured, and the level sticks whenever the
lazy import happens.
"""
import logging
logging.getLogger("diffusers.quantizers.torchao.torchao_quantizer").setLevel(
logging.ERROR
)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Public API: on-demand helpers (called explicitly by other modules) # Public API: on-demand helpers (called explicitly by other modules)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -203,15 +222,22 @@ def _patch_removed_symbols():
# Importing modeling_llama triggers a deep import chain: # Importing modeling_llama triggers a deep import chain:
# modeling_llama -> modeling_utils -> quantizers -> torchao # modeling_llama -> modeling_utils -> quantizers -> torchao
# torchao emits a noisy warning about incompatible torch versions # torchao emits a noisy warning about incompatible torch versions, and
# that is irrelevant here — suppress it during this import. # its register_as_pytree_constant() calls on Enum types make
_torchao_logger = logging.getLogger("torchao") # torch.utils._pytree log a deprecation warning once per Enum and per
_prev_level = _torchao_logger.level # rank. Neither is actionable here — suppress both during this import.
_torchao_logger.setLevel(logging.ERROR) _muted = [
logging.getLogger("torchao"),
logging.getLogger("torch.utils._pytree"),
]
_prev_levels = [lg.level for lg in _muted]
for lg in _muted:
lg.setLevel(logging.ERROR)
try: try:
from transformers.models.llama import modeling_llama from transformers.models.llama import modeling_llama
finally: finally:
_torchao_logger.setLevel(_prev_level) for lg, level in zip(_muted, _prev_levels):
lg.setLevel(level)
if not hasattr(modeling_llama, "LlamaFlashAttention2"): if not hasattr(modeling_llama, "LlamaFlashAttention2"):
if hasattr(modeling_llama, "LlamaAttention"): if hasattr(modeling_llama, "LlamaAttention"):
@@ -211,13 +211,13 @@ class TestPrefetchCheckpoints(CustomTestCase):
patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor), patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor),
patch("concurrent.futures.wait", side_effect=_wait_all), 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._prefetch_checkpoint_file"),
patch("sglang.srt.model_loader.weight_utils.logger.info") as log_info, patch("sglang.srt.model_loader.weight_utils.logger.debug") as log_debug,
): ):
_prefetch_all_checkpoints(paths, num_threads=1) _prefetch_all_checkpoints(paths, num_threads=1)
progress_pcts = [ progress_pcts = [
call.args[2] call.args[2]
for call in log_info.call_args_list for call in log_debug.call_args_list
if call.args if call.args
and call.args[0] == "Rank %d: prefetching checkpoint files: %d%% (%d/%d)" and call.args[0] == "Rank %d: prefetching checkpoint files: %d%% (%d/%d)"
] ]
@@ -426,14 +426,24 @@ class TestPrefetchDispatch(CustomTestCase):
"sglang.srt.model_loader.loader.safetensors_weights_iterator", "sglang.srt.model_loader.loader.safetensors_weights_iterator",
return_value=iter([]), return_value=iter([]),
), ),
patch("sglang.srt.model_loader.loader.logger.warning"), patch("sglang.srt.model_loader.loader.logger.debug"),
) )
@staticmethod
def _override_notices(mock_log):
"""The single-thread override notice among the captured log calls."""
return [
call
for call in mock_log.call_args_list
if call.args
and "falling back to single-threaded weight loading" in call.args[0]
]
def test_prefetch_uses_single_thread_for_default_config(self): def test_prefetch_uses_single_thread_for_default_config(self):
"""Prefetch on + no explicit multithread config -> single-threaded, """Prefetch on + no explicit multithread config -> single-threaded,
and the opt-out warning fires once.""" and the opt-out notice fires once."""
loader = self._make_loader({}) loader = self._make_loader({})
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=True prefetch=True
) )
with ( with (
@@ -441,18 +451,18 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
self._run(loader) self._run(loader)
mock_single.assert_called_once() mock_single.assert_called_once()
mock_buffered.assert_not_called() mock_buffered.assert_not_called()
mock_warning.assert_called_once() self.assertEqual(len(self._override_notices(mock_log)), 1)
def test_explicit_enable_multithread_keeps_buffered_with_prefetch(self): def test_explicit_enable_multithread_keeps_buffered_with_prefetch(self):
"""Explicit enable_multithread_load=true is the escape hatch; the """Explicit enable_multithread_load=true is the escape hatch; the
override and its warning must not fire.""" override and its warning must not fire."""
loader = self._make_loader({"enable_multithread_load": True}) loader = self._make_loader({"enable_multithread_load": True})
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=True prefetch=True
) )
with ( with (
@@ -460,19 +470,19 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
self._run(loader) self._run(loader)
mock_buffered.assert_called_once() mock_buffered.assert_called_once()
mock_single.assert_not_called() mock_single.assert_not_called()
mock_warning.assert_not_called() self.assertEqual(self._override_notices(mock_log), [])
def test_num_threads_only_keeps_buffered_with_prefetch(self): def test_num_threads_only_keeps_buffered_with_prefetch(self):
"""num_threads alone (relying on the enable_multithread_load=True """num_threads alone (relying on the enable_multithread_load=True
default) also signals multi-thread intent, so the override must not default) also signals multi-thread intent, so the override must not
fire and num_threads stays live.""" fire and num_threads stays live."""
loader = self._make_loader({"num_threads": 64}) loader = self._make_loader({"num_threads": 64})
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=True prefetch=True
) )
with ( with (
@@ -480,20 +490,20 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
self._run(loader) self._run(loader)
mock_buffered.assert_called_once() mock_buffered.assert_called_once()
# num_threads is forwarded as max_workers to the buffered iterator. # num_threads is forwarded as max_workers to the buffered iterator.
self.assertEqual(mock_buffered.call_args.kwargs["max_workers"], 64) self.assertEqual(mock_buffered.call_args.kwargs["max_workers"], 64)
mock_single.assert_not_called() mock_single.assert_not_called()
mock_warning.assert_not_called() self.assertEqual(self._override_notices(mock_log), [])
def test_no_prefetch_uses_multithread(self): def test_no_prefetch_uses_multithread(self):
"""Prefetch off -> multi-threaded iterator is used (default), no """Prefetch off -> multi-threaded iterator is used (default), no
override warning.""" override warning."""
loader = self._make_loader({}) loader = self._make_loader({})
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=False prefetch=False
) )
with ( with (
@@ -501,12 +511,12 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
self._run(loader) self._run(loader)
mock_buffered.assert_called_once() mock_buffered.assert_called_once()
mock_single.assert_not_called() mock_single.assert_not_called()
mock_warning.assert_not_called() self.assertEqual(self._override_notices(mock_log), [])
def test_startup_prefetch_reuses_existing_background_handle(self): def test_startup_prefetch_reuses_existing_background_handle(self):
"""Startup commit reuses resolved shards and the active prefetch handle.""" """Startup commit reuses resolved shards and the active prefetch handle."""
@@ -518,7 +528,7 @@ class TestPrefetchDispatch(CustomTestCase):
weight_files=("f.safetensors",), weight_files=("f.safetensors",),
use_safetensors=True, use_safetensors=True,
) )
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=False prefetch=False
) )
with ( with (
@@ -526,7 +536,7 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
list( list(
loader._get_weights_iterator( loader._get_weights_iterator(
@@ -541,7 +551,7 @@ class TestPrefetchDispatch(CustomTestCase):
mock_single.assert_called_once() mock_single.assert_called_once()
self.assertFalse(mock_single.call_args.kwargs["prefetch"]) self.assertFalse(mock_single.call_args.kwargs["prefetch"])
mock_buffered.assert_not_called() mock_buffered.assert_not_called()
mock_warning.assert_called_once() self.assertEqual(len(self._override_notices(mock_log)), 1)
def test_completed_startup_prefetch_restores_multithread_loader(self): def test_completed_startup_prefetch_restores_multithread_loader(self):
loader = self._make_loader({}) loader = self._make_loader({})
@@ -552,7 +562,7 @@ class TestPrefetchDispatch(CustomTestCase):
weight_files=("f.safetensors",), weight_files=("f.safetensors",),
use_safetensors=True, use_safetensors=True,
) )
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=False prefetch=False
) )
with ( with (
@@ -560,7 +570,7 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
list( list(
loader._get_weights_iterator( loader._get_weights_iterator(
@@ -575,7 +585,7 @@ class TestPrefetchDispatch(CustomTestCase):
mock_buffered.assert_called_once() mock_buffered.assert_called_once()
self.assertFalse(mock_buffered.call_args.kwargs["prefetch"]) self.assertFalse(mock_buffered.call_args.kwargs["prefetch"])
mock_single.assert_not_called() mock_single.assert_not_called()
mock_warning.assert_not_called() self.assertEqual(self._override_notices(mock_log), [])
def test_completed_startup_prefetch_is_not_started_twice(self): def test_completed_startup_prefetch_is_not_started_twice(self):
loader = self._make_loader({}) loader = self._make_loader({})
@@ -586,7 +596,7 @@ class TestPrefetchDispatch(CustomTestCase):
weight_files=("f.safetensors",), weight_files=("f.safetensors",),
use_safetensors=True, use_safetensors=True,
) )
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=True prefetch=True
) )
with ( with (
@@ -594,7 +604,7 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
list( list(
loader._get_weights_iterator( loader._get_weights_iterator(
@@ -608,13 +618,13 @@ class TestPrefetchDispatch(CustomTestCase):
mock_buffered.assert_called_once() mock_buffered.assert_called_once()
self.assertFalse(mock_buffered.call_args.kwargs["prefetch"]) self.assertFalse(mock_buffered.call_args.kwargs["prefetch"])
mock_single.assert_not_called() mock_single.assert_not_called()
mock_warning.assert_not_called() self.assertEqual(self._override_notices(mock_log), [])
def test_prefetch_does_not_override_when_mmap_disabled(self): def test_prefetch_does_not_override_when_mmap_disabled(self):
"""Prefetch is a no-op without mmap, so the override and its warning """Prefetch is a no-op without mmap, so the override and its warning
must not fire.""" must not fire."""
loader = self._make_loader({}) loader = self._make_loader({})
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=True, disable_mmap=True prefetch=True, disable_mmap=True
) )
with ( with (
@@ -622,18 +632,18 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
self._run(loader) self._run(loader)
mock_buffered.assert_called_once() mock_buffered.assert_called_once()
mock_single.assert_not_called() mock_single.assert_not_called()
mock_warning.assert_not_called() self.assertEqual(self._override_notices(mock_log), [])
def test_prefetch_does_not_override_for_fastsafetensors(self): def test_prefetch_does_not_override_for_fastsafetensors(self):
"""FASTSAFETENSORS ignores both flags; override + warning must not """FASTSAFETENSORS ignores both flags; override + warning must not
fire.""" fire."""
loader = self._make_loader({}, load_format=LoadFormat.FASTSAFETENSORS) loader = self._make_loader({}, load_format=LoadFormat.FASTSAFETENSORS)
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=True prefetch=True
) )
with ( with (
@@ -645,7 +655,7 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered as mock_buffered, p_buffered as mock_buffered,
p_single as mock_single, p_single as mock_single,
p_warn as mock_warning, p_log as mock_log,
): ):
self._run(loader) self._run(loader)
mock_fast.assert_called_once_with( mock_fast.assert_called_once_with(
@@ -655,13 +665,13 @@ class TestPrefetchDispatch(CustomTestCase):
) )
mock_buffered.assert_not_called() mock_buffered.assert_not_called()
mock_single.assert_not_called() mock_single.assert_not_called()
mock_warning.assert_not_called() self.assertEqual(self._override_notices(mock_log), [])
def test_fastsafetensors_gds_can_be_disabled(self): def test_fastsafetensors_gds_can_be_disabled(self):
loader = self._make_loader( loader = self._make_loader(
{"enable_gds": False}, load_format=LoadFormat.FASTSAFETENSORS {"enable_gds": False}, load_format=LoadFormat.FASTSAFETENSORS
) )
p_prep, p_model, p_buffered, p_single, p_warn = self._patch_dispatch( p_prep, p_model, p_buffered, p_single, p_log = self._patch_dispatch(
prefetch=False, prefetch=False,
drop_cache=True, drop_cache=True,
) )
@@ -674,7 +684,7 @@ class TestPrefetchDispatch(CustomTestCase):
p_model, p_model,
p_buffered, p_buffered,
p_single, p_single,
p_warn, p_log,
): ):
self._run(loader) self._run(loader)
mock_fast.assert_called_once_with( mock_fast.assert_called_once_with(