ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

模型总跑第0张卡?CUDA_VISIBLE_DEVICES不生效的排障全攻略

模型总跑第0张卡?CUDA_VISIBLE_DEVICES不生效的排障全攻略 1. 问题现象指定了显卡模型却还是跑到了第0张卡上先说结论这个问题我在生产环境里踩过不止一次而且每次的原因都不太一样。最典型的表现是你在代码里明明写了CUDA_VISIBLE_DEVICES2或者在脚本里设置了device_ids[2]但程序跑起来之后用nvidia-smi一看模型权重和中间激活全挤在了第0张显卡上。这个现象在 transformer 类模型包括 Bert、GPT、ViT 这些上出现得尤其频繁因为这类模型库往往封装得很深硬件初始化、分布式逻辑、模型并行这些环节都帮你包掉了。很多新手会把锅甩给“transformer 库有 bug”但我负责任地说绝大多数情况下是环境变量、库封装和代码优先级这三层之间打架跟模型本身关系不大。这篇文章会完整拆解我排查这个问题的全过程包括底层原因、快速验证手段、可落地的修复方案以及一些你以后大概率还会遇到的相关坑。适合自己搭过 transformer 训练或推理、被多卡环境折磨过的同学参考。注意以下所有操作都基于 Linux 环境 PyTorch HuggingFace transformers 库NVIDIA 驱动和 CUDA 已正常安装。如果你的环境是 Windows 或 AMD 显卡排查思路类似但部分命令和路径会不一样。2. 为什么“指定了”不等于“生效了”2.1 系统层面CUDA_VISIBLE_DEVICES 是第一步但不是全部很多人在代码里做文章比如torch.device(cuda:0)然后以为这就是“指定了显卡”。实际上cuda:0这个编号是逻辑编号不是物理编号。它指向的是 CUDA 运行时可见的显卡列表里的第一块而这个可见列表是由环境变量CUDA_VISIBLE_DEVICES决定的。举个例子如果机器上有4块卡物理编号是 0、1、2、3。你只设置了CUDA_VISIBLE_DEVICES2那么在这个进程里物理第2块卡会被重新映射为逻辑编号0。此时你写cuda:0实际用的就是物理第2块卡。但如果你的代码里先执行了某些操作把 CUDA 上下文提前初始化了或者你用某个上层框架启动程序时它自动设置了别的CUDA_VISIBLE_DEVICES那么你后续再import torch的时候看到的显卡编号可能就不是你以为的那个了。我遇到过最典型的场景用deepspeed或accelerate启动训练脚本它们内部会根据--num_gpus或者配置文件重新设置可见显卡列表。如果你之前在脚本里手动export CUDA_VISIBLE_DEVICES2进到框架里却被它覆盖成别的值就会发生“我明明指定了为什么没生效”的错觉。2.2 框架层面transformer 库加载模型的默认逻辑HuggingFace transformers 库在加载模型时默认行为是对齐 PyTorch 的device_map。如果你没有显式传入device_map模型权重会被加载到默认设备上——通常是 CPU然后再移动到某个 CUDA 设备。这个“某个”在很多实现里写的是cuda:0。注意这个cuda:0是逻辑编号。如果你已经正确设置了CUDA_VISIBLE_DEVICES2那么它其实去的是物理第2块卡这是没问题的。但问题恰恰在于你设置了环境变量但设置的位置不对导致整个进程根本就没接收到这个变量。另外还有一个点HuggingFace 从某个版本开始如果检测到环境里有多张显卡且你没有明确指定device_map它可能会默认启用auto模式自动把某些层分配到不同的卡上。如果你期望的是“全部放在物理第2块卡”但库给你自动分配到了别的地方表现上就是你指定了但没生效。2.3 代码层面device 指定顺序和上下文初始化更深层的问题在于 CUDA 上下文的初始化时机。PyTorch 中一旦执行了任何 CUDA 操作比如torch.zeros(1).cuda()或者加载一个已经在 GPU 上的 tensor当前进程就会绑定当前可见的显卡列表。如果在这之前你没有正确设置环境变量后续再设置也来不及了。我见过有人这样写import torch import os os.environ[CUDA_VISIBLE_DEVICES] 2 model model.cuda()这段代码从表面看没问题但如果在你import torch之前系统或某个第三方库已经执行过 CUDA 初始化那么这个环境变量就晚了。严格来说CUDA_VISIBLE_DEVICES必须在进程启动的最早期设置最好是在import torch之前甚至是在import any_cuda_related_library之前。更隐蔽的一个坑os.environ[CUDA_VISIBLE_DEVICES] 2只是在当前 Python 进程里设置了环境变量但如果你是用subprocess或multiprocessing启动的子进程这个设置不会自动继承除非你显式传给子进程。3. 快速验证三行命令定位问题在动手改代码之前先用几个命令确认当前的可见显卡和实际占用情况。3.1 查看物理显卡和当前占用nvidia-smi这个命令展示的是物理显卡的实时状态包括显存占用、温度、进程等。你可以看到物理编号和当前占用。如果某个进程占的是物理0号卡而你期望它在2号卡说明指定失败了。3.2 确认进程看到的可见显卡CUDA_VISIBLE_DEVICES2 python -c import torch; print(torch.cuda.device_count()); print(torch.cuda.current_device())如果输出是1和0说明环境变量生效了进程只看到一块卡映射为逻辑0。如果输出是4或更大的数字说明环境变量根本没传到这个进程里。3.3 确认模型参数所在设备import torch from transformers import AutoModel model AutoModel.from_pretrained(bert-base-uncased) print(model.device)如果输出是cpu说明模型还在 CPU 上后面调用.cuda()或者.to()才会移动。如果你没移动就做前向推理PyTorch 会报错“tensor must be on the same device”。如果你移动了可以通过下面的方式确认print(next(model.parameters()).device)这个输出能精确告诉你模型权重到底在哪个逻辑设备上。4. 三层修复从环境变量到代码写法4.1 第一层在进程启动前正确设置环境变量最稳妥的方式是在运行命令之前直接设置环境变量保证 Python 进程一开始就看到它export CUDA_VISIBLE_DEVICES2 python train.py如果你是在脚本里写那么必须放在所有import torch之前而且最好的做法是放在文件最开头import os os.environ[CUDA_VISIBLE_DEVICES] 2 import torch注意这里不能调换顺序。一旦import torch完成CUDA 相关的运行时可能已经被初始化尽管很多情况下是懒加载但最安全的做法是环境变量先行。如果你用 PyTorch 的torch.cuda.set_device(2)那指定的是逻辑编号不是物理编号。物理编号只能通过CUDA_VISIBLE_DEVICES来控制。这也是很多人搞混的点。4.2 第二层在 transformers 库加载时显式指定HuggingFace transformers 从 4.30 左右开始模型加载支持device_map参数。你可以直接指定from transformers import AutoModel model AutoModel.from_pretrained( bert-base-uncased, device_map{: 0} )device_map{: 0}的意思是所有没有明确映射到某层的模块代表根模块都放到逻辑设备0上。如果你的环境变量已经正确设置逻辑0就是物理目标卡。如果机器显存足够不想让模型分散到多卡还可以这样from transformers import AutoModel model AutoModel.from_pretrained( bert-base-uncased, device_mapcuda:0 )按我实际测试device_mapcuda:0在大部分情况下等效于{: 0}但前者更简洁适合 rapid prototype。如果你需要精确控制哪些层放到哪张卡可以用字典形式。另外还有一个容易忽视的地方from_pretrained加载时如果模型文件很大且本地没有缓存会先下载到磁盘再加载到内存最后移动到设备。这个过程如果耗时很长你可能误以为“卡住了”其实是正常的。4.3 第三层手动移动模型到目标设备如果环境变量和device_map都没问题但你依然发现某些子模块或 tensors 在其他设备上那是模型内部有些 buffer 没有被.to()方法一起带过去。比较少见但确实存在。这种情况下可以在 model 加载后手动把所有参数和 buffer 移动到目标设备device torch.device(cuda:0) model model.to(device) # 强制所有 buffer 也过去 for buffer in model.buffers(): buffer.data buffer.data.to(device)这个方法比较粗暴但在某些第三方自定义模型里能救命。5. 常见问题与排查技巧实录5.1 设了 CUDA_VISIBLE_DEVICES 但 nvidia-smi 看不到进程这个问题很怪但也常见。如果你设在环境变量里但进程是通过nohup或者某些任务调度系统比如 Slurm启动的可能环境变量没被传递。解决方式是在启动命令里显式使用CUDA_VISIBLE_DEVICES前缀CUDA_VISIBLE_DEVICES2 nohup python train.py train.log 21 或者在你的 Python 脚本里把 device 信息打印出来确认实际看到的编号print(torch.cuda.device_count(), torch.cuda.current_device())5.2 accelerate / deepspeed 环境下指定无效如果你用accelerate launch或deepspeed启动配置文件里的num_processes、gpu_ids会覆盖环境变量。最省心的办法是直接改配置# accelerate config num_processes: 1 gpu_ids: [2]deepspeed 的话在启动命令里加--includelocalhost:2明确指定物理卡号。5.3 用 Docker 运行时容器内编号错乱Docker 容器里看到的是NVIDIA_VISIBLE_DEVICES环境变量不是CUDA_VISIBLE_DEVICES。如果你在宿主机设置了CUDA_VISIBLE_DEVICES但没传进容器容器内以为自己看到了所有卡。正确做法docker run --gpus device2 -e NVIDIA_VISIBLE_DEVICES2 ...这样容器里只看到物理第2块卡逻辑编号也是0。5.4 模型加载到 GPU 但推理时依然报 CUDA error这种情况往往是显存不足。transformer 模型加载时会把所有权重从 CPU 拷贝到 GPU如果显存不够会在.cuda()那一步直接报CUDA out of memory。我之前在 8G 显存的卡上跑 Bert-large参数 3.4 亿大概占 1.3G但加上优化器和中间变量轻松上 4G。如果你的模型加载后 nvidia-smi 看到显存占用接近上限就别折腾指定显卡了先解决显存问题。提示用torch.cuda.memory_summary()可以打印当前显存分配情况排查哪些变量占了大头非常实用。6. 踩坑总结这些细节最容易忽略6.1 环境变量作用域问题os.environ[CUDA_VISIBLE_DEVICES] 2只影响当前 Python 进程。如果你在 Jupyter Notebook 里设置它只对当前 kernel 生效如果你在脚本 A 里调用脚本 BB 不会自动继承。这一点在多人共用服务器、分模块开发时非常致命。6.2 多进程训练时的传播性如果你用torch.multiprocessing或DataLoader的多进程模式子进程会继承父进程的环境变量但前提是父进程在fork之前已经设置好了。如果用spawn方式那环境变量不一定继承建议在主函数入口重新设置或通过参数传递。6.3 配置文件和硬编码冲突很多训练框架比如 transformers 自带的TrainingArguments有一个no_cuda参数默认 False。如果你手动设置了CUDA_VISIBLE_DEVICES但框架里还有个配置项是use_cpuTrue模型会强行跑到 CPU 上。排查时先确认框架配置里没有类似的开关。6.4 有时不指定反而是最佳选择在单进程 单卡的情况下你完全可以不指定显卡让 CUDA 默认走逻辑0。只要你确定机器上只有一块物理卡或者只有一块卡被你可见那写不写device_map没区别。我以前用 4 卡机器调试某个小模型图省事直接不指定结果模型的每一层被自动分布到4张卡上显存占用倒是均匀了但通信开销非常大推理速度反而更慢。这种问题不算“加载错卡”但同样让人摸不着头脑。7. 终极方案把设备选择抽成工具函数经过多次踩坑我现在习惯把设备选择逻辑封装成一个小工具所有项目都复用。核心思路是环境变量优先框架参数覆盖最后检查设备数。import os import torch def setup_device(target_gpu_idNone): 选择目标 GPU 并用环境变量锁定可见显卡。 Args: target_gpu_id: 物理显卡编号。如果为 None则使用逻辑0. if target_gpu_id is not None: os.environ[CUDA_VISIBLE_DEVICES] str(target_gpu_id) # 二次确认 print(fVisible GPUs: {torch.cuda.device_count()}) print(fCurrent CUDA device: {torch.cuda.current_device()}) print(fDevice name: {torch.cuda.get_device_name()}) if torch.cuda.device_count() 1: # 显式指定避免深度框架自动分布 return torch.device(cuda:0) return torch.device(cuda:0 if torch.cuda.is_available() else cpu)这个函数确保在任何import torch之前完成环境变量设置并且通过打印设备信息让你一眼看出当前进程到底看到了几张卡。调用方式device setup_device(target_gpu_id2) model.to(device)如果你用的是 HuggingFace transformers 的高层封装如Trainer直接把device对应的逻辑编号写进TrainingArguments的device_ids[0]即可因为经过环境变量过滤后逻辑0就是物理目标卡。8. 我的最终建议说实话显卡指定问题 90% 以上是环境变量没设对或者被上层框架覆盖了。解决问题的关键不是背一堆 API而是先建立“物理设备、可见设备、逻辑设备”三个概念然后从最外层环境变量开始排查逐层向内。如果你排查一天还没搞定不妨把模型换小一点先确认流程能跑通再逐步放大模型。我之前有段时间一直以为指定显卡的代码有问题结果发现是一个离线的第三方库在import时执行了 CUDA 初始化把环境变量读早了。这种问题单靠看代码很难定位需要耐心加日志。希望这篇文字能让你少踩几个坑。如果你也遇到过类似的奇怪现象并找到了不一样的查法欢迎日后在评论里聊。
返回列表