在 /vllm/model_executor/weight_utils.py 文件中,
elif use_safetensors:
for st_file in hf_weights_files:
with safe_open(st_file, framework="pt") as f:
for name in f.keys(): # noqa: SIM118
param = f.get_tensor(name)
yield name, param
safe_open 的默认设备是 'cpu',这在某些情况下会严重降低权重加载速度。例如,在我们的情况下,加载一个 llama2-7b 模型需要 194 秒。
为了解决这个问题,我们稍微修改了一下实现:
elif use_safetensors:
for st_file in hf_weights_files:
with safe_open(st_file, framework="pt", **device=device**) as f:
for name in f.keys(): # noqa: SIM118
param = f.get_tensor(name)
yield name, param
将加载时间从 194 秒提高到了 6 秒,提高了 32 倍。
4条答案
按热度按时间sg2wtvxw1#
请提交一个PR!这很有意义。
oewdyzsn2#
感谢您分享这些信息!我也对优化模型加载时间感兴趣。快速问题:我观察到模型的初始加载时间明显比后续加载要长得多。vLLM是否包含一个缓存机制,以加速后续加载?
bxpogfeg3#
不,我相信你观察到这个现象的原因是操作系统中的文件系统内存缓存。
kqhtkvqz4#
谢谢!这很有道理。