首页 / 资讯中心 / 文章详情

CPU环境跑BERT模型超时?先别怪CPU,问题可能卡在权重下载

CPU环境跑BERT模型超时?先别怪CPU,问题可能卡在权重下载 ★ FEATURED ARTICLE
1. 超时不背锅你的环境大概率没问题权重根本没下载下来1.1 A同学的现场装得很顺卡在from_pretrained前阵子某开发者跑来找我说他在一台只有CPU的服务器上搭Transformer环境conda创建虚拟环境、装CPU版PyTorch、装transformers库整个过程都很顺利。结果一执行from_pretrained(bert-base-chinese)终端就卡住不动短则两三分钟长则五分钟然后抛出一串ConnectionError或者ReadTimeoutError。他连着试了七八次每次都是不同的报错有时候是RemoteDisconnected有时候是SSLError一度怀疑是不是CPU太弱扛不住BERT模型。我先让他别急着改代码先做两个判断第一跑一个CPU密集型的简单任务比如让PyTorch跑一遍矩阵乘法看看CPU算力是否正常第二盯着终端看卡住的那段时间里CPU占用率到底是多少。结果很有意思卡住的时候CPU占用率基本是0。这说明它根本不是算不动而是程序一直等在某个外部输入上。CPU环境里搭Transformer用TensorFlow还是PyTorch、是几核几线程在下载权重这个阶段完全不重要因为模型文件还在硬盘之外压根还没轮到CPU上场。如果你也遇到类似情况可以先对照一下报错类型大概率是下面这类网络层问题。常见报错真实原因排查方向ConnectionErrorTCP连接建立失败服务器不可达网络链路、端点配置ReadTimeoutError连接建立了但数据很久没传输完带宽、文件体积、超时阈值RemoteDisconnected连接被服务端主动断开请求频率、代理链路、机房策略SSLErrorTLS握手失败或证书中断环境变量中的证书、系统时间OSError: file not found下载未完成缓存里只有部分文件清缓存重新下载1.2 一次from_pretrained调用到底做了多少事很多人对from_pretrained的理解是一句魔法给个模型名字模型就被加载出来了。实际上它内部做的事情比你想象得多。第一步解析模型标识符。bert-base-chinese这种名字会被解析成官方模型仓库里的一个仓库ID。第二步检查本地缓存。transformers库和huggingface_hub会把下载过的文件缓存在本地默认路径是~/.cache/huggingface/hub。如果本地已有完整文件它就直接跳过下载。第三步如果本地没有它就会向官方仓库发起HTTP请求先请求一个带版本信息的etag文件再对比本地缓存判断有哪些文件缺失。第四步逐个下载缺失文件常用的包括config.json、tokenizer.json、tokenizer_config.json、vocab.txt、model.safetensors或pytorch_model.bin。第五步校验文件SHA256哈希确认完整之后才真正加载到内存里建模。这一步棋的关键在于第四步和第五步是网络操作的集中区。以bert-base-chinese为例光权重文件pytorch_model.bin就有接近400MB加上tokenizer相关的几个文件总共要下载约400MB的数据。如果网络丢包率偏高、带宽只有几十KB/s那么一次下载很可能持续好几十分钟。而默认下载请求是有超时控制的一旦卡太久就会抛出ReadTimeoutError或者RemoteDisconnected。我见过最夸张的一种情况某开发者把模型文件下载了一半缓存目录里躺着一个300多MB的.incomplete文件每次代码启动都会去校验校验失败就从头下载于是每次都是卡住-超时-重来的死循环。1.3 换GPU也一样别在超时时怀疑CPU这一点我想单独拎出来说因为太多人把问题归错因。如果在CPU环境下加载BERT超时换到GPU机器上只要网络、缓存、磁盘这些外部条件不变照样会超时。模型加载本身是一个IO密集加内存分配的过程GPU只影响后续的前向计算不参与模型文件的下载和反序列化。所以遇到超时第一反应不应该是我的机器性能不行而应该是我的文件到底有没有完整出现在本地。这一步排查其实也简单。在你执行代码之前先去缓存目录看一眼文件大小。以bert-base-chinese为例一个完整缓存的模型目录下pytorch_model.bin或者model.safetensors应该有大约380MB到420MB大小。如果小于这个范围或者目录里只有.incomplete前缀的文件那不用看了就是下载环节的问题。2. CPU环境里被低估的三个因素网络、磁盘、内存峰值2.1 网络来自官方仓库的请求经常半路断开如果你是在办公网络或者云服务器上做这件事大概率会碰到一个很现实的问题默认的模型仓库端点对很多网络环境并不友好直连的稳定性很差。这不是配置问题也不是Python环境问题单纯是网络链路的质量问题。表现就是请求时通时不通下载到一半断掉重试几次还是断。解决思路不是把超时时间调到无限大而是换一条更稳定的路。transformers库和它底层的huggingface_hub都支持通过环境变量HF_ENDPOINT来指定模型仓库的访问地址。你可以把它指向一个自己测试过延迟低的镜像服务或者团队内部搭建的模型缓存服务。配置方式很简单在启动脚本里加一行export HF_ENDPOINThttps://你的镜像加速地址这条环境变量设置之后后续所有模型的下载、etag检查都会走这个地址。设置完成后先用一个小模型验证连通性比如prajjwal1/bert-tiny这个模型只有几十MB下载快适合用来验证端点是否生效。千万别一上来就直接下载400MB的大模型万一端点配错了你又得等半天超时。2.2 磁盘缓存目录与权限问题会让你每次都从零开始网络问题解决之后第二个容易被忽视的坑是磁盘和缓存目录。transformers的缓存路径在不同版本之间变化过老版本是~/.cache/huggingface/transformers新版本是~/.cache/huggingface/hub。如果多个项目共用一台服务器不同用户跑代码可能会遇到权限问题一个用户下载了权重另一个用户没有读取权限于是在另一个用户看来就是缓存不存在重新下载然后再次超时。更隐蔽的问题是磁盘空间不足。400MB的模型下载加上下载时的临时文件至少需要1GB左右的空闲空间。如果~/.cache所在分区空间不足下载会神不知鬼不觉地中断。你从日志里看到的是网络超时实际上却是磁盘写不进去导致的连带失败。我的建议是手动把缓存目录和数据下载目录统一管理不要放任它写到系统盘根目录下。用环境变量固定一个你确认过空间和权限的位置export HF_HOME/data/hf_home然后清理一次默认缓存避免旧的半成品文件干扰判断rm -rf ~/.cache/huggingface/hub清理完再重新下载。这一步尤其重要如果你之前已经下载过几次都超时了缓存目录里大概率躺着多个.incomplete文件这些半成品会让后续每一次下载都陷入校验失败、重新下载、再次超时的死循环不清掉没法根治。2.3 内存加载大权重时的峰值可能让你的服务器假死下载问题解决之后加载阶段还有一个容易和超时混淆的坑内存峰值。CPU服务器往往配置不算高如果内存只有8GB甚至4GB加载BERT这种亿级参数的模型时要特别小心。BertModel的参数量在1亿出头单精度FP32权重就占约400MB但from_pretrained在加载过程中会有额外的临时缓冲区。在某些情况下内存峰值可以达到模型权重的两三倍也就是接近1GB以上。如果服务器还跑着其他服务内存不够时操作系统会疯狂用swap表现出来就是程序不报错但整个机器卡得像死机一样任务管理器里看不出进度很像卡住超时。解决办法是显式开启低内存加载模式from transformers import BertModel, BertTokenizer model_path bert-base-chinese tokenizer BertTokenizer.from_pretrained(model_path) model BertModel.from_pretrained(model_path, low_cpu_mem_usageTrue)low_cpu_mem_usageTrue会让加载过程避免一些不必要的中间拷贝峰值内存会明显下降。配合设置PyTorch的线程数也能减少内存开销和无效争抢import torch torch.set_num_threads(4) model.to(cpu)如果你的机器内存实在紧张建议加载之前用free -h看一眼剩余内存如果内存低于2GB就不要用bert-base级别了直接换小模型会更实际。3. 一套可以照抄的加载流程从镜像加速到完全离线3.1 用镜像把文件完整预下载到本地缓存我现在的习惯是不让代码运行时去下载模型而是先单独执行一次预下载。这样模型加载和网络环境彻底解耦问题定位也会清晰很多。最简单的方式是用huggingface_hub的snapshot_download把整个仓库快照拉下来。import os os.environ[HF_ENDPOINT] https://你的镜像加速地址 from huggingface_hub import snapshot_download snapshot_download(repo_idbert-base-chinese, local_dir/data/model_cache/bert-base-chinese)repo_id指定模型仓库IDlocal_dir指定下载到哪个目录。下载完成后去看一下目标目录里的文件正常情况下应该包含以下这些/data/model_cache/bert-base-chinese/ ├── config.json ├── model.safetensors ├── tokenizer.json ├── tokenizer_config.json └── vocab.txt如果你的transformers版本比较老下载下来的权重文件可能是pytorch_model.bin而不是model.safetensors这都没关系关键是config.json和tokenizer相关文件必须齐全。只下载权重文件、不下载config和vocab后面照样会触发联网请求。有开发者问过我为什么不直接改代码里的下载超时时间我的看法是超时时间只是兜底不该把它当主策略。如果网络质量本身不行把超时从30秒调到300秒只是让失败来得更慢并不会让下载成功。正确思路是先把文件落到本地这一步用镜像加速验完文件大小再进下一步。3.2 直接用本地路径加载不再触发在线检查预下载完成之后代码里的加载方式就要改掉了。不要再传模型名称改成传本地路径from transformers import BertModel, BertTokenizer model_path /data/model_cache/bert-base-chinese tokenizer BertTokenizer.from_pretrained(model_path) model BertModel.from_pretrained(model_path, low_cpu_mem_usageTrue)传本地路径时transformers会优先读取目录下的文件不会因为本地文件不完整而临时发起在线请求。这也意味着只要你目录里的文件是完整的网络再怎么抖都不会影响加载。这一步做完超时问题基本上就只剩一个源头本地文件不完整、或者缓存目录权限不对。补充一个细节如果你以前从官方仓库成功下载过完整的模型并且缓存目录没有被清理过直接使用本地路径时可能会发现加载速度依然很慢。这时候可以检查一下是不是目录中还有.cache目录或者存在一些旧的序列化文件。实际上模型文件只要在本地加载通常只需要几十秒。3.3 断网环境下必开的离线模式还有一种场景是服务器本身没有外网权限或者只有内网访问。这种情况下除了预下载还要显式开启离线模式。transformers提供两个环境变量来控制export TRANSFORMERS_OFFLINE1 export HF_HUB_OFFLINE1TRANSFORMERS_OFFLINE是老版本transformers使用的环境变量HF_HUB_OFFLINE是新版huggingface_hub使用的变量。两个都写上兼容性最好。也可以在Python脚本开头设置import os os.environ[TRANSFORMERS_OFFLINE] 1 os.environ[HF_HUB_OFFLINE] 1这样设置之后模型加载时如果发现本地没有缓存会直接报错而不会傻傻地尝试联网避免长时间悬挂。对于生产环境我建议直接把这两个环境变量写进部署脚本让程序明确知道这个环境不该联网请给我用本地文件。离线模式还有一个附带好处可以防止代码里面不小心用了新的模型标识符却没提前下载导致线上运行时突然卡住。宁可让它立刻报文件缺失也不要让它卡五分钟超时这在自动化任务里非常重要。4. 加载不超时之后还有三个新坑在排队等你4.1 版本依赖冲突最容易伪装成网络超时的假卡死你以为超时全是网络问题我刚开始也这么想。直到有一次在某开发者的环境里模型文件早就下载好了本地路径加载也试过了程序还是会卡在读取tokenizer和vocab那一步。那次卡住之后CPU占用率只有个位数但系统日志里什么都看不出来看起来特别像在网络等待。后来折腾了半天发现是安装transformers的时候顺带装了一个版本不兼容的tokenizers库导致分词器读文件时死锁。这类问题在CPU环境里尤其常见因为很多时候大家用的是预编译的包依赖关系并不完全可控。最容易踩坑的是这几个包tokenizers、protobuf、safetensors、torch。如果你还同时装了TensorFlowprotobuf版本冲突很可能就是罪魁祸首。我的建议是专门建一个干净的虚拟环境把关键依赖版本锁死。下面这组是我在类似场景里用起来比较稳定的组合依赖建议版本说明python3.9兼容性好不易踩编译坑torch2.x CPU版从CPU only安装源获取transformers4.36.x稳定且接口完整tokenizers0.15.x与transformers 4.36匹配safetensors0.4.x如果下载的是safetensors格式需要protobuf3.20.x如果环境影响TensorFlow选3.x较稳装完之后记得执行一次依赖完整性检查pip check这个命令会列出当前环境里所有版本冲突。很多看起来像网络超时的问题其实是某个库在初始化时卡住了。先跑一次pip check再怀疑网络能省下大量排查时间。4.2 CPU上推理慢的取舍ONNX量化与更小的模型等加载问题彻底解决你会发现下一个有点痛的现实在CPU上跑BERT推理确实不快。尤其在文本长度较长、batch size稍大时一次前向可能要几百毫秒甚至几秒。如果你要做的是批量离线任务这个速度勉强能接受如果要做实时服务那基本撑不住。两个方向可以缓解。第一个方向是用ONNX Runtime加载模型把模型导出为ONNX格式再用CPU专用的优化执行引擎来跑。优点是推理速度通常能提升一倍左右内存占用也相对可控。操作步骤是先用optimum库导出pip install optimum onnxruntime optimum-cli export onnx -m /data/model_cache/bert-base-chinese /data/model_onnx/导出之后用ONNX Runtime加载并进行推理不再直接用transformers。这个方案的代价是如果后续要做模型的继续训练或者微调ONNX格式就不太方便了它更适合纯推理场景。第二个方向是直接换更小的模型。如果你只是想要语义向量不需要12层的大模型那么换一个层数更少、参数量更小的压缩版模型推理速度会有数量级的提升。方案加载方式推理速度适用场景原版BERTtransformers直接加载慢需要精确特征或微调ONNX Runtime导出后加载较快只做推理追求稳定延迟小模型/压缩版transformers直接加载明显更快对精度要求不极致4.3 把超时防御写进代码重试、超时阈值与日志最后分享一个我在实际项目里落地的做法不要指望环境永远一帆风顺把防御逻辑直接写进代码。huggingface_hub支持通过环境变量控制网络请求的超时阈值我一般会设置成这样export HF_HUB_DOWNLOAD_TIMEOUT30 export HF_HUB_ETAG_TIMEOUT10HF_HUB_DOWNLOAD_TIMEOUT控制下载单个文件的超时时间HF_HUB_ETAG_TIMEOUT控制版本检查请求的超时时间。设置之后下载环节最多等30秒就会失败不会动不动悬挂五分钟。注意这两个值不能设得太小否则网络稍微慢一点就会误杀也不能设太大不然又回到卡顿等待的老路。同时我写加载代码时会包一层带退避重试的逻辑import time import logging from transformers import BertModel, BertTokenizer logging.basicConfig(levellogging.INFO) logger logging.getLogger(model_loader) def load_bert_with_retry(model_source, max_retry3): for attempt in range(max_retry): try: tokenizer BertTokenizer.from_pretrained(model_source) model BertModel.from_pretrained(model_source, low_cpu_mem_usageTrue) return tokenizer, model except Exception as exc: logger.warning(第%s次加载失败: %s, attempt 1, exc) if attempt max_retry - 1: time.sleep(2 ** attempt) raise RuntimeError(模型加载连续多次失败请检查网络与本地缓存)这里用指数退避第一次失败等1秒第二次等2秒第三次直接抛出异常不无限重试。配合前面的预下载策略这个函数在正常情况下基本不会触发重试。回到题目的问题在CPU环境下搭Transformer环境调取BERT超时先别怀疑CPU也别急着重装环境。打开日志确认模型文件有没有完整落地再把模型下载和模型加载这两件事拆开。我用这个思路帮A同学处理完之后前后大概花了半小时问题根因就是第一次下载没走镜像缓存里留了一堆半成品。清理缓存、镜像预下载、本地路径加载三步走之后之后再也没出现过超时。我现在自己的习惯是无论新机器还是老机器只要部署涉及transformers一定先跑预下载脚本把模型文件固定到数据盘然后启动脚本里直接加上TRANSFORMERS_OFFLINE1和HF_HUB_OFFLINE1。这样后续网络不管怎么抖动都不会再出现这种八竿子打不着的超时问题。
阅读完成 · 觉得有帮助?
咨询建站