SLURM集群中DeepSpeed加载MobileVLM预训练模型遇文件锁问题

问题描述 #

基于MobileVLM模型开发，使用Hugging Face Transformers库加载预训练模型，在SLURM集群运行推理脚本时触发如下错误：

Exception ignored in atexit callback: <function matmul_ext_update_autotune_table at 0x7fbb99fa4ca0>
Traceback (most recent call last):
  File "/public/home/swun-caiy2/miniconda3/envs/mobilevlm/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py", line 444, in matmul_ext_update_autotune_table
    fp16_matmul._update_autotune_table()
  File "/public/home/swun-caiy2/miniconda3/envs/mobilevlm/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py", line 421, in _update_autotune_table
    TritonMatmul._update_autotune_table(__class__.__name__ + "_2d_kernel", __class__._2d_kernel)
  File "/public/home/swun-caiy2/miniconda3/envs/mobilevlm/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py", line 150, in _update_autotune_table
    cache_manager.put(autotune_table)
  File "/public/home/swun-caiy2/miniconda3/envs/mobilevlm/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py", line 66, in put
    with FileLock(self.lock_path):
  File "/public/home/swun-caiy2/miniconda3/envs/mobilevlm/lib/python3.10/site-packages/filelock/_api.py", line 297, in __enter__
    self.acquire()
  File "/public/home/swun-caiy2/miniconda3/envs/mobilevlm/lib/python3.10/site-packages/filelock/_api.py", line 255, in acquire
    self._acquire()
  File "/public/home/swun-caiy2/miniconda3/envs/mobilevlm/lib/python3.10/site-packages/filelock/_unix.py", line 48, in _acquire
    raise NotImplementedError(msg) from exception
NotImplementedError: FileSystem does not appear to support flock; user SoftFileLock instead

推理脚本核心代码：

from scripts.inference import inference_once
import torch

model_path = "/public/home/swun-caiy2/wensm/mobilevlm-v2/MobileVLM-main/mtgv"
image_file = "assets/samples/my_book.jpg"
prompt_str = "What is the title of this book?"

torch.cuda.set_device(0)

args = type('Args', (), {
    "model_path": model_path,
    "image_file": image_file,
    "prompt": prompt_str,
    "conv_mode": "v1",
    "temperature": 0, 
    "top_p": None,
    "num_beams": 1,
    "max_new_tokens": 512,
    "load_8bit": False,
    "load_4bit": False,
})()

inference_once(args)

解决方案 #

1. 强制使用软文件锁 #

通过设置环境变量让filelock库使用SoftFileLock替代默认的硬锁，避免依赖flock机制：

• 在SLURM脚本开头添加：

  export FILELOCK_BACKEND=soft
  

• 或者在Python代码最顶部导入模块前设置：

  import os
  os.environ["FILELOCK_BACKEND"] = "soft"
  

2. 转移DeepSpeed缓存到本地磁盘 #

集群共享存储通常不支持flock，将DeepSpeed的Triton自动调优缓存目录指向节点本地临时目录（如/tmp）：

export DS_TRITON_CACHE_PATH=/tmp/deepspeed_triton_cache_$USER

这样文件锁操作会在本地磁盘执行，规避共享存储的限制。

3. 禁用Triton自动调优 #

如果不需要自动调优带来的性能提升，可以直接关闭该功能，跳过缓存更新和文件锁操作：

export DS_TRITON_AUTOTUNE=0

内容的提问来源于stack exchange，提问作者user25698720