先说一个我记忆中很典型的画面去年底我帮朋友排查一个推理服务他在启动脚本里明明写了CUDA_VISIBLE_DEVICES2程序跑起来之后nvidia-smi一看2号卡纹丝不动0号卡显存飙了 20 多 GB。他在群里发了一句“transformer 指定显卡后模型依然加载到默认第 1 张显卡上怎么解决”评论区瞬间炸出一堆同样遭遇的人。这类问题在 Hugging Face transformers、PyTorch 默认行为、CUDA 环境变量这三者交叉的地方非常高频而且“看起来像指定了实际上哪儿都没指定”的情况远比想象中多。这篇文章我会把这个问题拆开揉碎从 CUDA_VISIBLE_DEVICES 到底改了什么到from_pretrained加载模型时的默认设备分配逻辑再到分布式训练和device_mapauto大模型推理场景下的特殊表现最后给出一套可以直接照着操作的排查链路和指定方案。不管你是刚入门 transformer 的小白还是被多卡环境折腾过的老手这篇文章都能帮你少走几次弯路。1. 先分清“显存占用”和“模型计算设备”问题定位的第一步很多人一看到nvidia-smi里 0 号卡显存涨了就认定“模型加载到了默认第 1 张显卡上”。但这里藏着一个非常容易混淆的细节显存占用高不代表模型的计算逻辑一定在跑这块卡模型权重被加载到某块卡上也不代表程序眼里“当前设备”就是那块卡。如果不先把这两个概念拆开后面所有排查都会跑偏。1.1 一次典型的“指定失败”现场复现先说一个最常见的复现场景。你有 8 张卡物理编号 0 到 7你想把模型放到物理第 5 张卡上编号 4跑推理于是你写了这样一段代码import os os.environ[CUDA_VISIBLE_DEVICES] 4 from transformers import AutoModel model AutoModel.from_pretrained(bert-base-uncased)然后你运行nvidia-smi发现物理编号 4 的卡没有显存占用物理编号 0 的卡却涨了一块。这时候你的第一反应大概率是“transformers 把模型加载到默认第 1 张显卡上了”。但真实情况往往是AutoModel.from_pretrained在这个过程中确实只负责加载权重它默认把参数放进 CPU 内存并不会主动把模型推到CUDA_VISIBLE_DEVICES指定的卡上。真正把模型占用到 0 号物理卡上的是你后面某一步操作比如model model.to(cuda) # 或者 model.cuda() # 或 outputs model(**inputs) # 某些 wrapper 内部默认用了 cuda:0当CUDA_VISIBLE_DEVICES4生效后程序里的“cuda:0”已经不是物理 0 号卡了而是“当前可见卡中的第一张”也就是物理 4 号卡。如果模型最终占用了物理 0 号卡说明环境变量根本没生效或者你的代码在设置环境变量之前就已经初始化了 CUDA 上下文。1.2 两个层面的设备概念物理卡号与 CUDA 逻辑卡号要彻底理解这个问题你必须先接受一个事实在 PyTorch / CUDA 的程序里你口中说的“第 1 张卡”和物理主板上插着的“第 1 张卡”不一定是一回事。物理设备编号由系统和驱动决定nvidia-smi里显示的 GPU 编号就是物理编号。CUDA 可见设备编号通过CUDA_VISIBLE_DEVICES这个环境变量可以过滤并重映射设备编号。程序里看到的cuda:0、cuda:1是对应可见设备列表中的顺序不是物理编号。举个例子物理卡有 0、1、2、3 四张你设置CUDA_VISIBLE_DEVICES2,0那么程序里cuda:0对应物理 2 号卡cuda:1对应物理 0 号卡物理 1 号和 3 号卡对程序完全不可见也就是说“默认第 1 张显卡”这个说法本身就值得推敲。如果环境变量生效了程序默认的cuda:0就是你指定列表里的第一张物理卡如果环境变量没生效程序默认的cuda:0才会真的落到物理 0 号卡上。很多人在代码里写了model.to(cuda)或AutoModel.from_pretrained(...).to(cuda)看到显存占用在 0 号物理卡就误以为 transformers 忽略了环境变量其实问题往往出在环境变量设置的位置和方式上。1.3 用三行代码确认模型到底落在哪张卡与其靠nvidia-smi猜不如直接在程序里打印模型参数所在的设备。在模型加载完成后加这么几行import torch # 查看所有模型参数所在的设备集合 device_set {param.device for param in model.parameters()} print(模型参数所在设备集合:, device_set) # 查看当前默认设备 print(当前默认CUDA设备:, torch.cuda.current_device()) # 查看可用的CUDA设备数量 print(可见CUDA设备数量:, torch.cuda.device_count())这三个输出交叉验证能帮你快速判断device_set里的设备号是逻辑设备号。如果集合里是{cuda:0}说明模型所有参数都在“可见设备列表的第一张卡”上这个cuda:0到底对应哪张物理卡再看nvidia-smi的占用情况。torch.cuda.current_device()返回程序当前默认使用的逻辑设备通常和model.to(cuda)落到的卡一致。torch.cuda.device_count()返回可见卡数量。如果你设置了CUDA_VISIBLE_DEVICES4它应该返回 1如果返回 8说明环境变量根本没有真正传入程序。把这几个值打出来你基本就能定位到问题是在环境变量传递环节还是在模型的设备映射环节而不是一上来就盯着nvidia-smi的显存占用瞎猜。2. 最经典的翻车现场环境变量设了但模型根本没用上在这么多年的排查经验里“指定显卡失败”案例中占比最高的其实是这类看上去有点“低级”的问题代码确实设置了CUDA_VISIBLE_DEVICES但模型加载和后续计算走的依然是默认设备逻辑。这个坑在 Hugging Face transformers 生态里尤其常见因为from_pretrained的设计哲学是“先加载权重不负责具体搬卡”。2.1 CUDA_VISIBLE_DEVICES 解决的是“可见性”不是“放置位置”我把这个结论放在最前面CUDA_VISIBLE_DEVICES只解决“程序能看见哪些显卡”的问题它不会自动帮你把模型权重搬到任何一张可见的卡上。可以这样理解它相当于给 CUDA 运行时的“视野”做了裁剪。在程序看来没有出现在这个环境变量里的物理卡就像不存在一样。但模型参数在初始加载时即使一张卡都不可见也可能先放在 CPU 上模型推理时要跑在 GPU 上必须显式调用.to(cuda)、.cuda()或者通过device_map参数让框架自动分配设备。所以下面的写法是无效的import os os.environ[CUDA_VISIBLE_DEVICES] 4 from transformers import AutoModel, AutoTokenizer model AutoModel.from_pretrained(bert-base-uncased) # 没有 .to(cuda)也没有 device_map # 模型默认留在 CPU 上后续哪张卡都不会占用这种写法没有占显存但如果后面你直接这样调用inputs tokenizer(测试文本, return_tensorspt) outputs model(**inputs)大概率会报错说输入是 CUDA tensor 而模型参数在 CPU 上如果你提前把 inputs 转成了 CUDA tensor而模型还在 CPU 上一样报错。更隐蔽的是有些代码里你会看到model AutoModel.from_pretrained(...)之后没有显式.to(cuda)但运行nvidia-smi的时候 0 号卡却有显存占用。这种情况通常是因为后面某个环节比如模型内部某个子模块或者某个 wrapper内部执行了.cuda()而默认逻辑号cuda:0恰好映射到了物理 0 号卡上。2.2 正确写法把“指定可见卡”和“指定放置位置”两步都做了我不想只说错误写法直接给一段经过验证的推荐写法import os # 在导入任何 CUDA/torch 相关库之前设置环境变量 os.environ[CUDA_VISIBLE_DEVICES] 4 import torch from transformers import AutoModel, AutoTokenizer # 程序里逻辑设备号 cuda:0 此时对应物理第 5 张卡物理编号 4 device torch.device(cuda:0) if torch.cuda.is_available() else torch.device(cpu) model AutoModel.from_pretrained(bert-base-uncased) model.to(device) tokenizer AutoTokenizer.from_pretrained(bert-base-uncased) inputs tokenizer(测试文本, return_tensorspt).to(device) outputs model(**inputs)这里有三个关键点任何一个漏掉都可能导致模型落在默认卡上os.environ[CUDA_VISIBLE_DEVICES] 4必须在import torch之前执行。因为 PyTorch 在导入时就会检查环境变量如果已经初始化 CUDA 上下文后面再改环境变量通常不生效。device要显式指定为cuda:0。在设置了环境变量后cuda:0已经不是物理 0 号卡了而是可见设备列表里的第一张卡。不要因为代码里写的是cuda:0就觉得它还是默认第一张物理卡。输入张量inputs也要.to(device)。模型参数和设备全部搬过去之后输入还在 CPU 上会直接报错如果输入在别的卡上则可能触发设备不一致的异常。还有一个小技巧如果你的卡比较新显存够用建议在模型加载之后顺手把torch.cuda.set_device(device)也写上确保后面如果有临时张量被创建默认落在你想要的那张卡上而不是程序自动选的 0 号逻辑卡。torch.cuda.set_device设置的是程序后续默认创建 CUDA 张量的设备跟环境变量的作用互补。2.3 pipeline 场景下的隐性默认设备Hugging Face transformers 的pipelineAPI 封装得太好反而掩盖了很多设备分配的细节。很多人这样写from transformers import pipeline # 想指定到物理2号卡 classifier pipeline(sentiment-analysis, device2)如果你没有设置CUDA_VISIBLE_DEVICESdevice2一般会直接映射到物理 2 号卡。但如果你设置了CUDA_VISIBLE_DEVICES4物理 2 号卡对程序是不可见的这行代码会报错。正确写法有二选一不设置CUDA_VISIBLE_DEVICES直接device2因为此时逻辑编号等于物理编号。设置CUDA_VISIBLE_DEVICES4然后device0因为可见设备列表里只有一张卡逻辑编号从 0 开始。从device2改成device0对不熟悉 CUDA 逻辑编号的人来说非常反直觉。我在实际项目中更推荐的做法是除非有特别复杂的多进程多卡交互需求否则尽量不跟device参数较劲而是统一用CUDA_VISIBLE_DEVICES控制“我只让程序看见一张卡”然后用cuda:0或device0来指代它。这样代码里的编号永远不会超过你真正想用的物理卡范围。3. Hugging Face transformers 在大模型场景下的显卡分配逻辑与容易踩的坑上面聊的是中小模型from_pretrained加载后手动.to(cuda)基本就能解决问题。但如果模型大到单卡放不下比如 LLM 动辄 7B、13B、70B你需要用到device_mapauto或者accelerate的自动设备映射这时候“指定显卡”的玩法又完全不一样了。相关的热词里不少人都搜过“加载本地模型”“transformer模型详解”这背后其实是一个共同需求怎么让本地加载的大模型稳定地落在自己指定的那张卡上。3.1 device_mapauto 的工作原理当你调用from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( meta-llama/Llama-2-7b-chat-hf, device_mapauto, torch_dtypetorch.float16, )transformers 底层会借助accelerate库来推断每一层的放置策略。它的默认逻辑是查看你当前有几张可见显卡。根据每张卡的显存大小把模型的 embedding 层、各 transformer block、norm 层、lm_head 层尽量均匀地切分到所有可见显卡上。如果所有显卡的显存加起来还不够再把多余的层放到 CPU 内存甚至磁盘 offload。这里最容易让人困惑的点是device_mapauto会对所有可见显卡做“雨露均沾”而不是默认只加载到cuda:0。所以如果你设置了CUDA_VISIBLE_DEVICES4然后使用device_mapauto模型确实只会加载到物理 4 号卡上因为程序眼里只有这一张可见卡。看起来好像没毛病。但问题往往出在另一种写法上model AutoModelForCausalLM.from_pretrained( meta-llama/Llama-2-7b-chat-hf, device_map{: cuda:0}, # 或者 device_mapcuda:0 )如果你没有设置CUDA_VISIBLE_DEVICES这个cuda:0就是物理 0 号卡如果你设置了CUDA_VISIBLE_DEVICES4这个cuda:0才是物理 4 号卡。很多人把外部环境变量和device_map里的逻辑编号搞混设置device_map{: cuda:4}然后CUDA_VISIBLE_DEVICES4只让程序看见一张卡结果直接报CUDA device 4 does not exist或Invalid device id因为程序眼里根本没有 4 号逻辑卡。3.2 device_map 中的编号是“逻辑号”不是“物理号”这里我想强调一个规则所有框架层面的cuda:n编号都是逻辑编号不是物理编号。逻辑编号的映射关系由CUDA_VISIBLE_DEVICES决定。我用一张表来做个对应关系方便理解环境变量设置程序里的 cuda:0程序里的 cuda:1物理不可见卡不设置物理 0 号物理 1 号无CUDA_VISIBLE_DEVICES4物理 4 号无0,1,2,3,5,6,7CUDA_VISIBLE_DEVICES2,0物理 2 号物理 0 号1,3,4,5,6,7CUDA_VISIBLE_DEVICES0,2,4物理 0 号物理 2 号1,3,5,6,7所以当你写device_map{: cuda:0}的时候你不是在说“给我物理 0 号卡”而是在说“给我所有可见卡里编号为 0 的那张”。如果这个时刻可见卡列表的第一张恰好是物理 4 号卡那么一切正常但如果你的环境变量设置失效了或者它在导入 torch 之后才设置那么cuda:0就真的对应物理 0 号卡模型也就真的加载到了“默认第一张显卡”上。这也是为什么我建议在大模型加载时尽量通过启动命令而不是写在代码中央来设置CUDA_VISIBLE_DEVICES。比如CUDA_VISIBLE_DEVICES4 python my_script.py或者在 shell 脚本里export CUDA_VISIBLE_DEVICES4 python my_script.py这样环境变量在 Python 解释器启动前就已经生效能最大程度避免“import torch 之后再去改环境变量”造成的失效。3.3 加载本地模型时的特殊注意点从用户提供的热词看很多人还在搜“加载本地模型”这个场景下的“指定显卡失败”也有一些特殊性。当你从from_pretrained加载本地模型目录比如./models/llama-2-7b-chat时模型的权重读取走的是本地文件系统跟显卡无关。显卡分配发生在权重读取之后、模型初始化过程中。所以 “本地模型加载到默认第一张卡” 这个现象跟远端模型还是本地模型没有本质区别原因还是出在设备映射逻辑。但本地模型场景有一个额外容易踩的坑offload_folder和max_memory参数。如果你显存不够用想强制模型部分层留在 CPU你可能会这样写model AutoModelForCausalLM.from_pretrained( ./models/llama-2-7b-chat, device_mapauto, max_memory{0: 10GiB, 1: 10GiB}, offload_folderoffload, )这里max_memory的键0、1也是逻辑设备编号。如果当前程序只看得见一张物理卡你写max_memory{0: 10GiB, 1: 10GiB}就会因为设备数量不足而报错。很多人在本地加载 7B 模型时明明只想用一张 80G 的卡却参考网上代码写了max_memory{0: 20GiB}发现模型很慢或者有几层被放到了 CPU。这不是“指定显卡失败”而是device_mapauto发现单卡显存不够主动把部分层分配到 CPU 去了。想确认这一点可以打印model.hf_device_mapprint(model.hf_device_map)这个字典会显示每一层被放到了什么设备上。如果你看到model.layers.0: 0、model.layers.5: cpu或者disk说明设备映射里混入了 CPU/磁盘程序就不会严格按照你以为的“指定到第几号卡”去执行。3.4 多模型叠加加载时的显存溢出假象另一个让人误以为是“指定显卡失败”的场景是你在同一个进程里加载了多个模型。比如先加载一个 bert-base再加载一个 Llama-2-7B两个模型都把device_map或.to(cuda)写成了同一个逻辑设备。此时nvidia-smi显示 0 号卡显存爆炸而 1 号卡很空你会以为是第二个模型没有按照指定加载到 1 号卡上。但更准确的解释是两个模型确实都加载到了各自的逻辑cuda:0上。如果第一个模型没有指定物理 1 号卡第二个模型也只是“逻辑 cuda:0”那么它们当然都落在同一张物理卡上。这不是 transformers 的 bug而是你对“逻辑编号”和“物理编号”的映射关系还没建立起直觉。我的建议是如果真有一个进程要同时加载多个模型并且希望它们分别落在不同的物理卡上最稳妥的方式是给每个模型单独设置CUDA_VISIBLE_DEVICES并用子进程隔离或者使用device_map直接指定不同的逻辑编号。比如物理卡 0、1 都可见时模型 A 放cuda:0模型 B 放cuda:1。但如果你只让程序看见一张物理卡比如CUDA_VISIBLE_DEVICES0那么就算你写了cuda:1也会报错因为程序眼里根本没有 1 号卡。4. 分布式训练和科学计算场景下“指定显卡”为什么会静默失效前面讲的更多是单进程推理、单卡加载的场景。真正让人头大的是分布式训练或科学计算场景里的“指定显卡”问题。在这个场景下模型占用哪张显存往往不是由from_pretrained决定的而是被 PyTorch 的分布式通信组、torchrun的进程分配逻辑、甚至 NVIDIA NCCL 的默认行为左右。很多看起来像是“transformers 不听话”的问题根因其实在框架和运行时层面。4.1 torchrun 对显卡编号语义的改写用torchrun启动训练的时候它的分配逻辑是这样的torchrun --nproc_per_node2 --nnodes1 train.py它会在同一台机器上启动 2 个进程每个进程被分配的LOCAL_RANK分别是 0 和 1。如果此时代码里写的是local_rank int(os.environ[LOCAL_RANK]) device fcuda:{local_rank} model.to(device)那么进程 0 会把自己放到cuda:0进程 1 会把自己放到cuda:1。逻辑上看起来没问题。但只要细想一步LOCAL_RANK是 torchrun 分配的“进程本地序号”它不等于物理显卡编号也不等于CUDA_VISIBLE_DEVICES过滤后的逻辑编号。如果外部还设置了CUDA_VISIBLE_DEVICES2,3那么进程 0 看到的cuda:0是物理 2 号卡进程 1 看到的cuda:1是物理 3 号卡这个过程比较丝滑不太会出错。但如果你在代码里又加了一句os.environ[CUDA_VISIBLE_DEVICES] 2想“保险起见”把进程限制到物理 2 号卡上那后面就乱了。因为 torchrun 已经初始化了分布式环境你再改CUDA_VISIBLE_DEVICES通常不会影响已经建立的 CUDA 上下文倒是可能让某些子进程看到的设备数量变成 1导致 NCCL 通信初始化失败。更常见的“静默失效”场景是这样的你手动设置了CUDA_VISIBLE_DEVICES3,4,5,6然后用torchrun --nproc_per_node4启动训练。按理说 4 个进程应该分别落在物理 3、4、5、6 号卡上。但是如果你代码里的device不是从LOCAL_RANK推导而是直接写死了cuda:0那么 4 个进程会全部挤到物理 3 号卡上另外三张卡空空如也。这不是“模型默认第一张卡”的问题而是“所有进程都指认逻辑 0 号卡”的问题。4.2 显存分配器与多进程的“先到先得”效应即便你正确使用了LOCAL_RANK在分布式训练里还有一个不太起眼但很折磨人的问题CUDA 缓存分配器caching allocator的显存预留行为。PyTorch 第一次在某张卡上执行任何 CUDA 运算时会初始化这块卡的缓存分配器并预留一部分显存作为缓存块。如果在训练脚本的开头所有进程都先执行了一些公共代码比如数据预处理或者某个默认devicecuda的操作就有可能出现多个进程在LOCAL_RANK的设备映射执行之前已经把显存注册到了默认卡上。这种情况在nvidia-smi里的表现非常迷惑你明明指定了进程 0 去 3 号卡进程 1 去 4 号卡但它们的显存占用可能同时出现在物理 0 号卡上。原因可能是某个算子没有指定设备默认用了torch.cuda.current_device()。某个 DataLoader worker 在 fork 之前就持有 CUDA 上下文。某个第三方库在 import 时自动调用了torch.cuda.set_device(0)。我实际排查下来最容易出问题的是 Hugging Face 的datasets、tokenizers库在某些情况下会初始化 CUDA 上下文。它们本身不一定用 GPU但一旦初始化了就可能导致后续的设备分配出现“先到先得”的错位。解决思路也很简单在显式设定当前设备之前不要做任何可能触发 CUDA 初始化的操作。比如from transformers import ...尽量集中在文件顶部CUDA_VISIBLE_DEVICES在进程启动时由 shell 注入而不是在 Python 代码里临时设置。4.3 “指定”的层级物理层、运行时层、框架层可能互相打架我需要再往上抽象一层所谓“指定显卡”实际上是三个层级共同作用的结果。物理层nvidia-smi看到的 GPU 编号由硬件和驱动决定。运行时层CUDA Runtime 通过CUDA_VISIBLE_DEVICES过滤后的可见设备列表。所有 CUDA 调用都基于这个列表来解析cuda:n。框架层PyTorch、Hugging Face transformers、accelerate 等库在运行时层之上再维护一套“当前设备”和“设备映射”的逻辑。这三层只要有一层没对齐就会出现标题里那种“指定了显卡模型还是加载到默认第一张卡上”的现象。最常见的错位是你在框架层写了devicecuda:4但运行时层的可见设备列表只有 2 张卡于是报错或回退。你在运行时层设置了CUDA_VISIBLE_DEVICES4但框架层的device_map字典写的是{: cuda:4}于是报无效设备。你在框架层没有写任何设备映射完全依赖运行时层但运行时层的环境变量在 Python 启动前没设置成功于是模型默认落到物理 0 号卡。所以每次遇到“指定失败”不要只在一个层面找原因。优先确认三层是否一致这比瞎试device参数管用十倍。5. 一套完整的排查链路和最终可行的指定方案这些年我帮不少人处理过这类问题最后沉淀下来的排查流程非常固定。你不需要把上面的原理全部记住但要记得“先看现象、再定位层、最后下方案”这个顺序。5.1 五步定位法从现象确定问题出在哪一层第一步检查环境变量是否真的传入了程序。在代码最开头加一行import os print(CUDA_VISIBLE_DEVICES:, os.environ.get(CUDA_VISIBLE_DEVICES, 未设置))如果是通过 shell 启动的这一步基本能排除“环境变量没设置”的低级问题。第二步检查可见显卡数量。在import torch之后执行print(可见显卡数量:, torch.cuda.device_count())如果你设置的是CUDA_VISIBLE_DEVICES4这里应该输出 1。如果输出 8说明环境变量没有在 import torch 之前生效。第三步检查模型参数的设备集合device_set {p.device for p in model.parameters()} print(模型参数设备集合:, device_set)这一步能确认模型所有参数在哪个逻辑设备上。第四步检查逻辑设备到物理设备的映射。用nvidia-smi看哪张卡有进程占用把进程号跟当前程序的 PID 对应一下nvidia-smi --query-compute-appspid,used_memory,gpu_uuid --formatcsv第五步检查是否有model.hf_device_map。如果加载的是大模型打印出来确认有没有层落到了 CPU/磁盘print(getattr(model, hf_device_map, 无 device_map可能未使用自动映射))这五步走完问题出在哪一层基本一目了然。5.2 单卡推理场景的正确指定方式如果你只是单卡跑一个中小模型推荐用这种组合简单而且可复现CUDA_VISIBLE_DEVICES4 python inference.pyinference.py内部import torch from transformers import AutoModel, AutoTokenizer # 环境变量已经在启动命令行中注入 device torch.device(cuda:0 if torch.cuda.is_available() else cpu) model AutoModel.from_pretrained(./model_dir) model.to(device) tokenizer AutoTokenizer.from_pretrained(./model_dir) inputs tokenizer(测试文本, return_tensorspt).to(device) with torch.no_grad(): outputs model(**inputs)注意这里的cuda:0指的是“可见卡列表里的第一张”也就是物理 4 号卡。不要被这个数字吓到。5.3 大模型推理场景的正确指定方式大模型加载推荐用device_mapauto配合CUDA_VISIBLE_DEVICES来控制卡的范围CUDA_VISIBLE_DEVICES4 python run_llm.pyrun_llm.py内部import torch from transformers import AutoModelForCausalLM, AutoTokenizer model AutoModelForCausalLM.from_pretrained( ./models/llama-2-7b-chat, device_mapauto, torch_dtypetorch.float16, ) tokenizer AutoTokenizer.from_pretrained(./models/llama-2-7b-chat) inputs tokenizer(你好, return_tensorspt).to(cuda) outputs model.generate(**inputs, max_new_tokens100) print(tokenizer.decode(outputs[0], skip_special_tokensTrue))如果不想用device_mapauto也可以手动指定device_map {: cuda:0}在这种场景下cuda:0依然代表可见设备列表的第一张卡。配合CUDA_VISIBLE_DEVICES4启动它就是物理 4 号卡。如果模型很大需要拆分到多张卡比如物理 4、5 号卡都可见那么可以CUDA_VISIBLE_DEVICES4,5 python run_llm_multi.py代码里用model AutoModelForCausalLM.from_pretrained( ./models/llama-2-13b-chat, device_mapauto, torch_dtypetorch.float16, )这样显存充足时layers 会被自动分配到cuda:0和cuda:1上分别对应物理 4、5 号卡。想更精细控制每张卡的显存上限可以加max_memory{0: 40GiB, 1: 40GiB}这里面的 0、1 同样是逻辑编号。5.4 分布式训练场景的正确指定方式分布式训练推荐完全交给torchrun管设备代码里只依据LOCAL_RANK推导设备不要在代码里写死物理卡号CUDA_VISIBLE_DEVICES2,3 torchrun --nproc_per_node2 train.py训练脚本内部import os import torch import torch.distributed as dist from transformers import AutoModelForSequenceClassification local_rank int(os.environ[LOCAL_RANK]) dist.init_process_group(backendnccl) torch.cuda.set_device(local_rank) model AutoModelForSequenceClassification.from_pretrained(./model_dir) model.to(fcuda:{local_rank}) model torch.nn.parallel.DistributedDataParallel(model, device_ids[local_rank])这里有一个细节容易忽略from_pretrained在分布式场景下如果指定device_mapauto有时会跟DistributedDataParallel冲突因为device_mapauto可能把不同层分配到不同设备上而 DDP 要求模型在单一设备上复制。所以我个人建议分布式训练中小模型不要用device_map直接在from_pretrained之后用.to(fcuda:{local_rank})大模型训练要走张量并行或流水线并行的专门框架而不是简单 DDP 硬扛。5.5 一些补充排查命令和技巧最后分享几个平时排查特别管用的命令都是我在实际工作里反复用到的# 实时查看每张卡的显存占用和进程 watch -n 1 nvidia-smi # 只看占用显存的进程 nvidia-smi --query-compute-appspid,used_gpu_memory,process_name --formatcsv # 确认当前进程对每张卡的可见性 python -c import os; print(os.environ.get(CUDA_VISIBLE_DEVICES)) python -c import torch; print(torch.cuda.device_count(), torch.cuda.current_device())还有一个我自己常用的技巧在代码里输出torch.cuda.memory_summary()它会列出每张可见卡上当前的内存分配情况包括缓存块、活跃块、保留块。结合这个输出你可以判断某张卡上到底是被模型参数占用了还是被缓存分配器保留了排查起来更精准。个人在实际操作中最大的体会是这类“指定显卡失败”的坑绝大多数不是某个库的 bug而是设备编号背后的映射关系没理顺。只要把CUDA_VISIBLE_DEVICES当成一个过滤器把框架里的cuda:n都理解为可见设备列表的下标并且记住“在 import torch 之前设置环境变量”这条铁律80% 的问题都能当场解决。剩下的 20%基本都能通过hf_device_map和nvidia-smi交叉验证找到答案。如果你手头已经在用device_mapauto加载大模型建议下一次遇到类似问题时先别急着改代码把model.hf_device_map打出来看看。我曾经排查过一个 70B 模型推理任务用户一直抱怨“明明指定了 2 号卡但显存占用出现在 0 号卡”最后发现是因为device_mapauto检测到 0 号卡显存更大自动把 embedding 层和第一段 layers 放在了cuda:0而CUDA_VISIBLE_DEVICES根本没设置于是物理 0 号卡就成了逻辑 0 号卡。这根本不是“指定失败”而是“自动分配策略把模型放到了显存更大的卡上”。你越理解这些背后的行为越不会在排查时被表面现象带偏。