书生-S2 权重加载慢、多模态前处理报错、两段耗时分别在日志里定位

文章导读
启动慢和多模态前处理报错,通常是两条独立的链路:一条是权重从磁盘或远端读到显存的过程,另一条是图片解码、缩放、拼 tile、走 processor 的过程。放在一次启动里一起调,很容易出现改了一处、另一处照旧,判断不出改动是否生效。可行的做法是先在日志里给加载阶段打时间戳,再脱离权重加载单独跑一次前处理,把慢和报错分别变成可重复的现象,最后用同一张计时表回填确认。
📋 目录
  1. 壹 在启动日志里给加载阶段打上时间戳
  2. 贰 脱离模型加载单独跑一次图片前处理
  3. 叁 把两类报错信息分开归类
  4. 肆 用最小脚本分别复现慢与报错
  5. 伍 修复后回填同一张计时表
A A

启动慢和多模态前处理报错,通常是两条独立的链路:一条是权重从磁盘或远端读到显存的过程,另一条是图片解码、缩放、拼 tile、走 processor 的过程。放在一次启动里一起调,很容易出现改了一处、另一处照旧,判断不出改动是否生效。可行的做法是先在日志里给加载阶段打时间戳,再脱离权重加载单独跑一次前处理,把慢和报错分别变成可重复的现象,最后用同一张计时表回填确认。

加载慢看权重读取与初始化各阶段的时间戳差值,前处理报错看异常栈第一帧落在哪个模块;两者不在同一段日志里,就不要用同一个改动去验证。判断依据是同一份权重、同一张图片在改动前后耗时与异常状态是否变化,而不是启动总时间是否变短。若本地缓存的冷热状态、图片分辨率或批大小有变化,前后对比不成立,需要回到同一条件重测。

在启动日志里给加载阶段打上时间戳

目标是把权重加载从整体启动时间里单独摘出来。打点位置建议固定在这几处:进程入口第一行、读完配置之后、开始加载权重之前、每读完一个权重分片之后、processor 或 tokenizer 初始化之后、服务就绪之前。每个点都输出同一个累计计时器的当前值,而不是距上一步多少秒,否则后面没法把某一段和总时间对齐。

日志字段格式示例,数字位置只表示字段,实际以你自己的日志为准:

[boot] stage=start elapsed=0.000s
[boot] stage=config_parsed elapsed=0.412s
[boot] stage=weights_begin path=/path/to/ckpt elapsed=0.418s
[boot] stage=weights_shard shard=3/8 elapsed=9.870s
[boot] stage=weights_done elapsed=18.930s
[boot] stage=processor_init elapsed=19.102s
[boot] stage=server_ready elapsed=19.140s

打点用一个起点加一个打印函数即可,替换项是权重路径和分片列表:

import time
T0 = time.perf_counter()
def mark(stage, **kw):
    extra = ' '.join('%s=%s' % (k, v) for k, v in kw.items())
    print('[boot] stage=%s elapsed=%.3fs %s' % (
        stage, time.perf_counter() - T0, extra), flush=True)

mark('start')
mark('config_parsed')
mark('weights_begin', path=CKPT)
for i, shard in enumerate(shards, 1):
    load_shard(shard)          # 替换成你项目实际的读取或反序列化调用
    mark('weights_shard', shard='%d/%d' % (i, len(shards)))
mark('weights_done')
mark('processor_init')
mark('server_ready')

需要确认的边界:如果 weights_begin 到 weights_done 占了大头,问题在读取或反序列化这一侧;如果这一段很快而 processor_init 很长,那启动慢其实是组件初始化慢,不是读权重慢。加 flush=True 或把日志落盘,避免输出缓冲把几个时间点挤在一起看不清。

脱离模型加载单独跑一次图片前处理

这一步只回答一个问题:报错是不是前处理引起的。脚本里不加载权重、不建推理服务,只初始化 processor,读一张与线上同格式的图片,跑一次前处理。

# check_preprocess_only.py
import sys, time, traceback
from PIL import Image
from your_project import build_processor   # 替换成项目实际的 processor 入口

CKPT = '/path/to/checkpoint'      # 替换
IMG  = '/path/to/sample.jpg'      # 替换:与线上同格式、同尺寸量级的图片

t0 = time.perf_counter()
proc = build_processor(CKPT)
t1 = time.perf_counter()

img = Image.open(IMG)
t2 = time.perf_counter()
try:
    out = proc(images=img, return_tensors='pt')
except Exception:
    traceback.print_exc()
    print('PREPROCESS_FAILED')
    sys.exit(1)

print('PREPROCESS_OK init=%.3fs call=%.3fs' % (t1 - t0, time.perf_counter() - t2))
for k, v in out.items():
    if hasattr(v, 'shape'):
        print('tensor', k, tuple(v.shape))

判据:退出码为 0、末尾打印 PREPROCESS_OK、输出张量形状与预期一致,说明前处理链路本身能跑通,线上报错更可能来自调用参数、并发方式或输入内容差异。如果在这里就以同样的异常栈失败,说明问题在 processor 初始化、依赖版本或图片本身,与权重读取无关,按第 3 节的归类去查。

建议至少准备三类图片各跑一次:正常图片、带透明通道或四通道的图片、体积异常大或字段缺失的图片。这样才能区分是所有图片都报错,还是只有某类图片报错。如果项目自带 load_image 或 preprocess 函数,直接替换掉 build_processor,不要为了做这一步额外引入别的依赖。

把两类报错信息分开归类

看报错先看栈的第一帧落在哪个模块,比读报错文案更可靠。下面的归类用于快速判断属于哪一段,以及该去查什么。

书生-S2 权重加载慢、多模态前处理报错、两段耗时分别在日志里定位
归类常见关键字对应检查点
权重读取/加载FileNotFoundError、No such file or directory、PermissionError、EACCES权重路径、挂载点、运行用户权限、分片文件名是否齐全
权重读取/加载safetensors header 相关报错、EOFError、unexpected end of file、分片校验失败文件是否传输完整,大小与校验值是否对得上,远端读取是否被中断
权重读取/加载missing keys、unexpected keys、size mismatch(出现在 load_state_dict 阶段)权重与代码版本是否配套,是否漏加载了部分分片
权重读取/加载CUDA out of memory(出现在加载阶段)、NCCL timeout、connection reset显存与卡数配置、多卡加载顺序、远端存储连接是否稳定
图片前处理UnidentifiedImageError、cannot identify image file、image decode failed、truncated图片能否被解码,输入是否为损坏文件或非图片内容
图片前处理Expected 4D input、size mismatch、dimension out of range、num_frames 相关报错通道数、分辨率、批大小、tile 或帧数与 processor 期望是否一致
图片前处理TypeError: unexpected keyword argument、jinja2 模板报错、chat template 缺失调用参数名、提示词模板、processor 与模型版本的对应关系

两类关键字出现在同一条日志里是常见现象:加载阶段出问题后进程重启,恢复的第一批请求又触发前处理报错,看起来就像同一个问题。建议给两类阶段加不同前缀,例如 [boot] 与 [pre],按前缀过滤日志,不要按时间顺序猜。

用最小脚本分别复现慢与报错

慢和报错拆成两个脚本,每个脚本只保留一种变量。

复现慢:只做读取并逐片计时,替换项是按实际命名规则匹配分片文件。

# bench_load.py
import sys, time, glob, socket
T0 = time.perf_counter()
def mark(stage, **kw):
    print('[boot] stage=%s elapsed=%.3fs %s' % (
        stage, time.perf_counter() - T0,
        ' '.join('%s=%s' % (k, v) for k, v in kw.items())), flush=True)

CKPT = sys.argv[1] if len(sys.argv) > 1 else '/path/to/ckpt'
shards = sorted(glob.glob(CKPT + '/*.safetensors'))
mark('start', shards=len(shards), host=socket.gethostname())
for i, s in enumerate(shards, 1):
    t = time.perf_counter()
    open(s, 'rb').read(1 << 20)     # 替换成实际的读取或反序列化调用
    mark('shard', idx='%d/%d' % (i, len(shards)), file=s,
         cost='%.3f' % (time.perf_counter() - t))
mark('done')

复现报错:固定一组图片反复跑前处理,把异常类型和图片名一起打出来。

# bench_preprocess.py
import sys, glob, traceback
from PIL import Image
from your_project import build_processor   # 替换

proc = build_processor('/path/to/checkpoint')
ok, fail = 0, 0
for p in sorted(glob.glob(sys.argv[1] if len(sys.argv) > 1 else '/path/to/imgs/*')):
    try:
        out = proc(images=Image.open(p), return_tensors='pt')
        ok += 1
    except Exception as e:
        fail += 1
        print('FAIL file=%s error=%s' % (p, type(e).__name__))
        traceback.print_exc()
print('SUMMARY ok=%d fail=%d' % (ok, fail))

触发条件要跟着结果一起记录,否则换一台机器就复现不出来。建议记录:权重路径与是否本地盘、分片数量、机器与卡数、首次加载还是已有页缓存、图片格式与分辨率范围、批大小、并发数。把这些写进脚本参数或环境变量,和日志一起留存。复现成功的标准是同一脚本、同一参数、同一输入能稳定出现同样的慢或同样的异常类型,而不是偶尔出现一次。

修复后回填同一张计时表

改完之后不要只看启动总时间,回到第 1 节那份分段日志,把同一阶段的值填进同一张表,填写方式如下:

阶段本次改动改动前 elapsed改动后 elapsed判定
weights_done(填具体改动,如改为本地盘路径)(填改动前值)(填改动后值)下降 / 持平 / 上升
processor_init(填具体改动)(填改动前值)(填改动后值)下降 / 持平 / 上升
前处理异常数(填具体改动)(填改动前 fail 数)(填改动后 fail 数)消除 / 仍存在

结论判定上有几点要守住:一次只改一处再测,否则分不清是哪一项起作用;加载耗时对比时冷热状态必须一致,首次从远端读取和本地已有页缓存本身就是两种情形;前处理报错的结论是同一批图片从 fail 变成 ok 且输出形状符合预期才算消除,fail 数没到 0 就是未解决。同一配置重复跑几轮,看的是稳定区间,不要拿单次结果下结论。如果回填后某一阶段没有变化,说明这次改动没有作用在真正的耗时段或报错点上,回到第 3 节的归类重新判断。