Prhub

#30146 Disable multi-threaded load by default when prefetch is on

原始 PR 作者 mmangkad 合并时间 2026-07-09 15:28 文件变更 5 提交数 2 评论 2 代码增减 +224 / -5

执行摘要

Prefetch 开启时默认禁用多线程加载

PR body 指出,当 enable_multithread_load 默认开启(#20289)且同时启用 --weight-loader-prefetch-checkpoints 时,prefetch 线程与多线程加载器并行读取相同 shard,在 NFS/Lustre 等共享文件系统上导致 I/O 过订阅,降低加载速度。需要协调两者,让 prefetch 线程发挥预热作用而不被多线程加载干扰。

建议部署了 NFS/Lustre 的团队关注此 PR,理解 prefetch 与多线程加载的冲突。设计决策(通过额外配置 opt-in 而非完全禁止)值得学习。

讨论亮点

两位 reviewer (b8zhong, Fridge003) 均直接批准,未产生实质讨论,说明变更逻辑清晰且得到认可。

实现拆解

  1. 在 loader.py 的 _get_weights_iterator 方法中增加条件判断:当 weight_loader_prefetch_checkpoints=True、未禁用 mmap、加载格式非 FASTSAFETENSORS、当前 use_multithread 为真、且用户未在 model_loader_extra_config 中显式设置 enable_multithread_load 或 num_threads 时,将 use_multithread 设为 False 并记录警告。
  2. 更新 server_args.py 中 weight_loader_prefetch_checkpoints 参数的描述,说明默认禁用多线程加载及 opt-in 方式。
  3. 更新两处文档(model_loading.mdx 和 server_arguments.mdx)以反映这一行为。
  4. 在测试文件 test_prefetch_checkpoints.py 中新增 TestPrefetchDispatch 类,通过 mock 验证不同配置组合下是否正确分派到单线程或多线程迭代器。
文件 模块 状态 重要度
python/sglang/srt/model_loader/loader.py 加载器 modified 6.74
test/registered/unit/model_loader/test_prefetch_checkpoints.py 预取测试 modified 7.36
python/sglang/srt/server_args.py 参数配置 modified 4.32
docs_new/docs/advanced_features/model_loading.mdx 文档 modified 2.53
docs_new/docs/advanced_features/server_arguments.mdx 文档 modified 2.14

关键符号

DefaultModelLoader._get_weights_iterator TestPrefetchDispatch.test_prefetch_uses_single_thread_for_default_config TestPrefetchDispatch.test_explicit_multithread_keeps_multithread

关键源码片段

python/sglang/srt/model_loader/loader.py core-logic

核心逻辑变更,添加条件判断协调 prefetch 与 multithread。

def _get_weights_iterator(self, source):
    # 从额外配置获取多线程开关,默认 True
    use_multithread = extra_config.get("enable_multithread_load", True)
    # ... 省略准备步骤 ...
​
    server_args = get_global_server_args()
    weight_loader_disable_mmap = server_args.weight_loader_disable_mmap
    weight_loader_prefetch = server_args.weight_loader_prefetch_checkpoints
​
    # 当满足以下所有条件时,降级为单线程以避免 I/O 过订阅:
    # - 预取已启用
    # - 未禁用 mmap(即使用 mmap)
    # - 加载格式不是 FASTSAFETENSORS(该格式不走 mmap)
    # - 当前为多线程模式
    # - 用户未在额外配置中显式设置 enable_multithread_load 或 num_threads
    if (
        weight_loader_prefetch
        and not weight_loader_disable_mmap
        and self.load_config.load_format != LoadFormat.FASTSAFETENSORS
        and use_multithread
        and not ({"enable_multithread_load", "num_threads"} & extra_config.keys())
    ):
        logger.warning(
            "--weight-loader-prefetch-checkpoints is enabled; "
            "falling back to single-threaded weight loading "
            "to avoid I/O oversubscription with the prefetch threads. "
            "Set enable_multithread_load=true in "
            "--model-loader-extra-config to keep multi-threaded loading."
        )
        use_multithread = False
​
    if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
        weights_iterator = fastsafetensors_weights_iterator(hf_weights_files)
    elif use_multithread:
        weights_iterator = buffered_multi_thread_safetensors_weights_iterator(
            hf_weights_files,
            max_workers=extra_config.get("num_threads", self.DEFAULT_NUM_THREADS),
            disable_mmap=weight_loader_disable_mmap,
            prefetch=weight_loader_prefetch,
            prefetch_num_threads=prefetch_num_threads,
            drop_cache_after_load=weight_loader_drop_cache_after_load,
        )
    else:
        weights_iterator = safetensors_weights_iterator(
            hf_weights_files,
            disable_mmap=weight_loader_disable_mmap,
            prefetch=weight_loader_prefetch,
            prefetch_num_threads=prefetch_num_threads,
            drop_cache_after_load=weight_loader_drop_cache_after_load,
        )
    # ... 后续返回 weights_iterator

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

主要风险是默认行为变更可能影响依赖多线程加载速度的场景(如本地 NVMe),但用户可通过显式设置 enable_multithread_load=true 恢复。新逻辑不影响 pt、npcache 格式及 FASTSAFETENSORS 路径。测试覆盖了关键组合,降低了回归风险。

对用户:使用 --weight-loader-prefetch-checkpoints 的用户将自动获得单线程降级,改善共享文件系统加载性能;本地 NVMe 场景需手动 opt-in。对系统:减少 I/O 竞争,提高资源利用率。对团队:明确了 prefetch 与多线程加载的交互语义,降低了后续维护成本。

默认行为变更 I/O 性能回归风险 本地存储场景需 opt-in

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论