Prhub

#25748 loader: yield filtered MTP weights lazily to avoid OOM hang on multi-layer EAGLE

原始 PR 作者 JustinTong0323 合并时间 2026-05-20 12:33 文件变更 1 提交数 3 评论 4 代码增减 +7 / -6

执行摘要

延迟 MTP 权重过滤修复 OOM 挂起

PR body 明确指出:多层级 EAGLE 下每个 draft runner 独立全扫描检查点,tuple(filtered_weights) 持有全部前缀重命名的张量,TP=8 时 3 个 runner 的活跃张量导致主机触发 page reclaim,所有调度线程卡在 futex_do_waitread_bytes=0,服务无法就绪。

值得精读的 bugfix 典范:一行的逻辑错误(tuple 强制物化)导致整个系统在特定配置下不可用,修复后效果显著。代码改动虽小,但不熟悉迭代器模型的人容易犯同样错误。

实现拆解

  1. 修改函数签名和返回值:将 _filter_mtp_weights 的返回值从 Tuple[Tuple[str, torch.Tensor], ...] 改为 Generator[Tuple[str, torch.Tensor], None, None],使其成为一个生成器。
  2. 移除列表并改用 yield:删除 filtered_weights = []filtered_weights.append(...),改为直接 yield (prefix + new_name, tensor),这样每个张量在被消费后才释放,不会全部积压在内存中。
  3. 保留上游行为一致性:该函数的所有调用者(位于 _get_weights_iterator 中)已经通过 return self._filter_mtp_weights(...) 将迭代器返回给上层,上层本就以迭代方式消费,因此改为生成器后完全兼容。此外,同文件的 _get_all_weights_get_weights_iterator 的非 MTP 路径已经使用生成器表达式,风格统一。
  4. 文档更新:更新了 docstring 以解释惰性 yield 避免 OOM 的原因。
文件 模块 状态 重要度
python/sglang/srt/model_loader/loader.py 模型加载器 modified 6.49

关键符号

_filter_mtp_weights _get_weights_iterator

关键源码片段

python/sglang/srt/model_loader/loader.py data-contract

唯一变更文件,包含核心修复:将 `_filter_mtp_weights` 从返回 tuple 改为生成器,避免内存爆炸。

# python/sglang/srt/model_loader/loader.py
# 关键函数:_filter_mtp_weights
# 原实现:filtered_weights = [] 后用 append 收集全部张量,最后 return tuple(...) 一次性物化。
# 新实现:使用 yield 逐项生成,让上游 buffered_multithread_safetensors_weights_iterator
# 的滑动窗口 buffer 真正生效,避免大 MoE 检查点时 CPU 内存爆炸。@classmethod
def _filter_mtp_weights(
    cls, weights_iterator, prefix: str, draft_model_idx: int
) -> Generator[Tuple[str, torch.Tensor], None, None]:
    """Filter MTP weights to keep only the specified draft model layer
    and remap it to layer 0.  Yields lazily so the upstream buffered
    iterator 's sliding window actually bounds CPU memory — eager
    materialization caused page-reclaim hangs on large MoE checkpoints
    with multi-layer EAGLE."""
    for name, tensor in weights_iterator:
        match = cls._MTP_PATTERN.match(name)
        if match is not None:
            idx = int(match.group(1))
            if idx != draft_model_idx:
                continue
            # 将 MTP 层编号重映射为 layer 0
            new_name = name.replace(match.group(), "model.mtp.layers.0.")
        else:
            new_name = name
        # 关键变更:以前是 filtered_weights.append(...),
        # 最后 return tuple(filtered_weights);
        # 现在直接 yield,每个权重被消费后即可释放。
        yield (prefix + new_name, tensor)

评论区精华

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

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

风险与影响

风险极低。

  • 回归:调用者已经以迭代方式消费结果,且同文件其他路径早已使用生成器,行为不变。
  • 性能:惰性 yield 避免了全量拷贝,理论上减少内存和延迟。
  • 兼容性:仅修改内部函数签名和实现,对外可见的接口未变。

影响范围限定在 DefaultModelLoader._filter_mtp_weights 函数。主要受益场景:使用多层级 EAGLE 且模型权重很大的用户(如 MiMo-V2.5-Pro 等大型 MoE 模型),服务启动时间从无限挂起降至数分钟,且多线程加载器可以正常工作。对于单个 draft runner 或小模型,影响微乎其微。

核心路径变更 缺少测试覆盖

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论