把 Stable Diffusion 的 UNet 从 PyTorch 权重转成 TensorRT engine跑通之后推理速度能实打实上一个台阶但真正动手的人多半会在某个深夜被一行红字拦住Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select)。这个报错看起来像环境问题实际上九成以上是代码里某个张量忘了搬设备尤其集中在 TensorRT 转换脚本这条相对冷门的代码路径上。我前后在三台机器上折腾过 SD 的 TensorRT 转换从 RTX 3060 12G 到 4090踩过的坑基本覆盖了这个报错的全部变体。这篇就把整个排查思路、修复补丁、验证方法和我自己总结的速查表一次讲清楚它能帮你把一个看起来玄学的设备不一致问题拆成可复现、可定位、可修复的具体动作适合已经开始做 SD TensorRT 加速、但卡在转换或推理报错上的中阶玩家也适合想理解 PyTorch 设备语义的初学者。1. 先把报错误读懂设备不一致到底发生在哪一层1.1 报错文本的逐段拆解很多人看到这行红字第一反应是去重装 CUDA、重装 TensorRT其实完全跑偏了。这行报错的信息量非常足只要逐段读能直接砍掉一半排查时间。Expected all tensors to be on the same device是 PyTorch 的算子级约束绝大多数 CUDA 算子在执行前会做一次设备一致性检查参与的张量必须在同一个设备上要么全在 CPU要么全在同一块 GPU 上。它不会帮你自动搬也不会给你一个隐式拷贝的便利。but found at least two devices, cpu and cuda:0是实际的设备列表。这里要注意cuda:0是默认设备索引并不一定意味着你真的有两块卡很多时候就是那一块卡。看到cpu and cuda:0就说明一个张量在内存里另一个在显存里。真正最关键的是最后半句when checking argument for argument index in method wrapper_CUDA__index_select。它点名了具体算子index_select以及具体参数index。torch.index_select(input, dim, index)的语义是从 input 这张表里按 index 给出的下标把对应的行取出来。在神经网络里这个算子最典型的应用就是嵌入层Embedding也就是文本编码器里把 token id 映射成向量的那一步。所以这句话翻译成人话是查表用的权重表和张量下标一个在 CPU、一个在 GPU。报错说index参数有问题说明问题出在下标这一侧而不是权重这一侧。这个判断非常有用因为文本编码器的权重通常跟着模型一起.to(device)了剩下的可疑对象就只有 tokenizer 输出的input_ids。提示PyTorch 2.0 之前的版本报错里写的是method wrapper_index_select少一个_CUDA__前缀含义完全一致不要因为字符串不同就以为是另一个问题。1.2 为什么偏偏是 TensorRT 转换阶段爆发有意思的地方在于同一套模型在 WebUI 里跑图一切正常一进转换脚本就炸。这不是巧合而是代码路径完全不同导致的。日常推理时tokenizer 输出的input_ids确实是在 CPU 上但 WebUI 那套代码在进入文本编码器之前会做一次设备对齐或者干脆依赖某个包装层把张量搬过去。转换脚本是另一条路径它通常只关心把模型搬到 GPU然后喂 dummy input 跑一遍轨迹dummy input 往往就是随手写的torch.randint(0, 49408, (1, 77))或者torch.randn(1, 4, 64, 64)写的时候脑子里想的是形状对不对压根没考虑 device。另一个高频原因是权重加载方式。torch.load()不带map_location时权重会恢复到保存时的设备带map_locationcpu时全部落在内存里。之后如果用的是标准的model.load_state_dict(sd)PyTorch 会做一次拷贝设备是对的。但如果你写的是module.weight.data sd[weight]这种直接赋值那就等于把内存里的张量塞进了 GPU 模型的属性里后续前向一定会撞设备检查。还有一类更隐蔽的register_buffer注册的缓冲区。buffer 会跟着model.to(device)一起搬这点没问题但如果某个中间量是在forward里用纯 Python 列表或 numpy 数组现算出来的再torch.tensor(...)转成张量那它默认落在 CPU 上而且不会因为模型在 GPU 就自动跟着走。1.3 动手前先收集这三样现场信息在改任何一行代码之前先把下面三样东西收集齐。我自己的经验是缺了任何一样排查时间至少翻倍。信息类别具体内容为什么需要完整回溯栈从Traceback第一行到最后一行的全部内容包含文件名、行号、函数名最后一行只告诉你算子名倒数第几行才告诉你业务代码在哪版本矩阵torch、torchvision、CUDA、TensorRT、onnx、onnxsim、polygraphy、diffusers 或 WebUI 的版本ONNX 导出和 TRT build 对版本组合极其敏感很多设备报错其实是版本不匹配的次生现象最小复现命令能触发报错的最短命令最好带完整参数没有最小复现你永远不知道是转换逻辑的问题还是你改的那行代码的问题版本矩阵这一项常被忽略。举个真实例子onnx1.16 之后导出器对position_ids这类 buffer 的处理方式和 1.15 不同某些情况下会把它当常量折叠掉行为变化之后就会出现同一个脚本昨天能跑今天不能跑的现象。把版本记下来至少你能确认不是自己手抖改错了东西。2. 转换链路全景SD 到 TensorRT 要过几道关2.1 从 PyTorch 权重到 engine 的四段式流程SD 模型转 TensorRT 不是一步到位的操作中间要过四道关每一道都有独立的设备语义报错可能出现在任意一段。第一段是 PyTorch 前向与 ONNX 导出。这一段在你自己的 Python 进程里跑模型在cuda:0输入张量必须也在cuda:0否则就是你看到的这个报错。第二段是 ONNX 图简化与手术。这一段操作的是计算图对象不涉及实际张量计算所以理论上不会报设备错。但如果简化阶段做了常量折叠把本该在运行期传进去的输入折成了常量那么到了第三段就会因为输入对不上而报形状错很多人会误以为是设备问题。第三段是 TensorRT build engine。这一段由 TensorRT 的 builder 接管它会申请自己的显存和 PyTorch 的显存池是两套体系。这一段报的错通常长这样[TRT] Error Code 3: API Usage Error或者No supported profile was found跟设备一致性没关系。第四段是 engine 反序列化与运行时绑定。这一段需要显式指定用哪块 GPU、哪个 stream写错了会报CUDA error: invalid device ordinal之类也不是我们今天要解决的那个。明白了这四段你就能快速判断手上这个报错属于哪一类。回到正题index_select的设备不一致几乎百分之百发生在第一段也就是 PyTorch 前向或 ONNX 导出触发的那次前向。2.2 三个最容易埋雷的模块文本编码器、UNet、VAE按我的实际踩坑频率排序SD 里最容易出设备不一致的模块是这三个。文本编码器CLIP排第一因为它是index_select的主战场。结构很简单input_ids进 embedding 层底层就是F.embedding最终调用index_select。很多自定义的转换脚本会自己重新写一遍文本编码的调用逻辑比如为了拿到pooled_output而手动拆开 forward一拆就容易漏掉.to(device)。UNet 排第二它的雷点更分散。注意力模块里的 mask 张量、时间步嵌入time embedding的中间结果、以及某些自定义 attention 实现里手工构造的position_ids都可能被留在 CPU。特别是当你用了 xformers 或 sdp-attention 的替换实现时mask 的构造逻辑变了原来框架自己处理好的设备对齐可能就断了。VAE 排第三而且它的雷点很有特点分块解码tiling相关。分块推理的逻辑是把大图切成小块逐个解码再拼回去如果你看过相关实现会发现拼接时用到的索引、偏移量在大部分版本里是纯 Python 整数这本身没问题。但也有一些实现尤其是为了性能做过自定义的用torch.arange生成索引张量来做 gather/scatter这类张量如果不显式指定device默认就在 CPU立刻撞上同一个设备检查。除了这三个还有一类第三方条件注入模块值得单独提一句。比如基于人脸特征的条件控制方案它的人脸嵌入通常来自一个独立的推理库输出是 numpy 数组然后torch.from_numpy(...)转成张量。这一步出来的张量铁定在 CPU而主模型在 GPUtorch.cat一拼接就炸。这类问题有个共同特征报错发生在模型的入口附近而不是在 UNet 内部。2.3 转换脚本里 device 参数是怎么被悄悄丢掉的抛开模块层面转换脚本自身的写法有几个系统性缺陷值得单独列出来对照检查。危险写法后果正确做法torch.randn(1, 4, 64, 64)落在 CPU前向第一层就报错torch.randn(1, 4, 64, 64, devicedevice)torch.randint(0, 49408, (1, 77))同上且在文本编码器直接触发index_select加devicedevicetorch.load(path)不带map_location张量落在保存时的设备跨机器会出问题统一map_locationcpu之后走load_state_dictt torch.tensor(numpy_arr)落在 CPUtorch.as_tensor(numpy_arr, devicedevice)或torch.from_numpy(numpy_arr).to(device)w.data sd[w]直接赋值张量不搬模型属性里混进 CPU 张量改成load_state_dict或赋值后补.to(device)torch.arange(n)做索引落在 CPU加devicedevicetorch.zeros(...)做 mask落在 CPU加devicedevice或dtype同源的写法这张表几乎就是我这几年所有同类报错的来源清单。你会发现一个规律凡是在forward之外、或者在被包装过的forward里临时造张量的地方都是高风险区域。框架自带的代码通常已经处理好了出问题的永远是自己写的那几十行。还有一个小技巧值得记一下PyTorch 2.0 之后提供了torch.set_default_device(cuda)调用之后所有没显式指定 device 的张量创建都会默认落在 CUDA 上。听起来像是万能解药但它会改变进程内的全局语义包括某些第三方库内部的张量创建容易引发更奇怪的连锁问题。我的建议是只在你完全掌控的短脚本里用长期维护的代码还是老老实实显式传 device。3. 实操修复从回溯栈到具体补丁3.1 定位让报错自己告诉你哪一行出问题第一步永远是完整回溯栈。假设你拿到的栈长这样路径以你本地版本为准Traceback (most recent call last): File convert_unet.py, line 187, in module main() File convert_unet.py, line 152, in main torch.onnx.export(model, dummy_input, onnx_path, ...) File .../torch/onnx/utils.py, line 516, in export _export(...) File ldm/modules/encoders/modules.py, line 168, in forward outputs self.transformer(input_ids, attention_maskattention_mask) RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select)读法是自下往上。最后一行给算子index_select。倒数第二行给业务代码位置ldm/modules/encoders/modules.py第 168 行的self.transformer(input_ids, ...)调用。到这里基本可以锁定input_ids是那个在 CPU 上的张量而 transformer 的 embedding 权重在 GPU 上。第二步是让它错得更靠前一点。CUDA 算子默认是异步执行的报错位置有时会滞后。加一个环境变量能让执行同步化报错点会更贴近真实出错位置CUDA_LAUNCH_BLOCKING1 python convert_unet.py --model sd15 --output engines/这个变量会让性能明显下降但排查阶段非常值。等定位完记得去掉。第三步是主动插桩。如果回溯栈里全是框架内部代码看不到自己的业务逻辑就在可疑的边界上打打印def dump_device(tag, obj): if isinstance(obj, torch.Tensor): print(f[{tag}] tensor device{obj.device}, dtype{obj.dtype}, shape{tuple(obj.shape)}) elif isinstance(obj, (list, tuple)): for i, x in enumerate(obj): dump_device(f{tag}[{i}], x) elif isinstance(obj, dict): for k, v in obj.items(): dump_device(f{tag}.{k}, v) else: print(f[{tag}] {type(obj).__name__}) # 在可疑调用前后各打一次 dump_device(before_clip, input_ids) outputs self.transformer(input_ids, attention_maskattention_mask) dump_device(after_clip, outputs.last_hidden_state)attention_mask也要一起打印因为它经常是同一个问题的第二个受害者很多人把input_ids修好了下一步又在 mask 上撞一次同样的报错、同样的算子只是行号变了。3.2 修复方案A/B/C 的取舍与对比定位到位置之后修复手段有三种各有适用场景。我把它们放在一张表里对比方便你按自己的维护情况选。方案具体做法改动量覆盖范围升级风险副作用A 逐点补丁在每处出问题的代码加.to(device)小只覆盖已知位置新问题要再修高升级覆盖文件后补丁丢失无B 全局默认设备torch.set_default_device(cuda)极小覆盖所有隐式张量创建低会改变第三方库行为可能引发其他异常C 入口统一清洗在转换脚本入口写一个递归设备对齐函数对所有输入一次性处理中等覆盖所有入口张量UNet 内部问题仍需单独处理低脚本是自己维护的无且可复用我的实际选择是 C 为主、A 为辅。理由很实在B 方案虽然最省事但我在一次试验里发现它会让某个图像后处理库内部的张量创建也跑到 GPU 上结果在保存图片时多了一次没必要的显存往返还偶发过一次显存碎片导致的 OOM。C 方案的思路是在边界上做一次性对齐逻辑清晰、可测试、可复用而且不侵入第三方代码。3.3 三类典型场景的代码补丁场景一文本编码器的input_ids落在 CPU。这是最高频的一种。补丁很小但位置要对。# 修改前 batch_encoding self.tokenizer( text, truncationTrue, max_lengthself.max_length, return_lengthTrue, return_overflowing_tokensFalse, paddingmax_length, return_tensorspt, ) tokens batch_encoding[input_ids] # 修改后 batch_encoding self.tokenizer( text, truncationTrue, max_lengthself.max_length, return_lengthTrue, return_overflowing_tokensFalse, paddingmax_length, return_tensorspt, ) tokens batch_encoding[input_ids].to(self.device)这里的self.device通常是模块在初始化时记录下来的设备字符串比如torch.device(cuda, 0)。如果你的代码里没有这个属性就在__init__里加一行self.device torch.device(cuda if torch.cuda.is_available() else cpu)别用next(self.parameters()).device现算因为参数可能在之后的某个时刻被移动而input_ids用的应该是当前前向时的设备。补丁二转换脚本的 dummy input。这是 ONNX 导出失败的第二大原因。device torch.device(cuda, 0) dtype torch.float16 # 修改前 dummy_input torch.randn(1, 4, 64, 64) dummy_ids torch.randint(0, 49408, (1, 77)) # 修改后 dummy_input torch.randn(1, 4, 64, 64, devicedevice, dtypedtype) dummy_ids torch.randint(0, 49408, (1, 77), devicedevice) # 如果还要传 attention mask一起对齐 dummy_mask torch.ones_like(dummy_ids, devicedevice)注意dtype也要对齐。fp16 模型配 fp32 dummy input虽然不会直接报设备错但在某些算子尤其是 LayerNorm 和 attention 内部会因为 dtype 不一致报另一个错或者静默产生数值偏差让后面的精度对比结果莫名其妙地差。补丁三自定义条件注入与分块索引。这类问题通常出现在模型的入口和出口两端。def to_device_recursive(obj, device, dtypeNone): 把任意嵌套结构里的张量统一搬到目标设备 if isinstance(obj, torch.Tensor): out obj.to(device) if dtype is not None and out.is_floating_point(): out out.to(dtype) return out if isinstance(obj, dict): return {k: to_device_recursive(v, device, dtype) for k, v in obj.items()} if isinstance(obj, (list, tuple)): converted [to_device_recursive(v, device, dtype) for v in obj] return type(obj)(converted) if not isinstance(obj, tuple) else tuple(converted) if isinstance(obj, np.ndarray): return torch.from_numpy(obj).to(device) return obj用法是在转换脚本里把要喂给模型的整包输入过一次这个函数inputs { sample: latent, timestep: t, encoder_hidden_states: cond, face_emb: face_tensor, # 来自外部推理库的 numpy 结果 } inputs to_device_recursive(inputs, device, torch.float16)这个函数我用了很久好处是它能处理 numpy 数组把外部库输出这一类最隐蔽的来源一次性堵住。你不需要去改外部库的代码也不需要知道它内部用的什么框架。3.4 改完怎么验证真的修好了不报错不代表转对了。设备问题修完之后数值精度往往还需要一次确认因为在整个过程中你很可能同时改了 dtype。第一步是 PyTorch 与 ONNX 的输出对比。用同一份权重、同一份输入跑一次原生 PyTorch 前向再跑一次 ONNX Runtime 前向比较输出的最大绝对误差。import numpy as np import onnxruntime as ort # 原生 PyTorch 结果记为 ref with torch.no_grad(): ref model(**inputs.to(device) if hasattr(inputs, to) else inputs) ref_np ref.detach().float().cpu().numpy() # ONNX Runtime 结果 providers [(CUDAExecutionProvider, {device_id: 0})] sess ort.InferenceSession(unet.onnx, providersproviders) onnx_out sess.run(None, {k: v.detach().cpu().numpy() for k, v in inputs.items()})[0] diff np.abs(ref_np - onnx_out) print(max_abs_diff %.6f % diff.max()) print(mean_abs_diff %.6f % diff.mean())判据要按精度分档fp32 导出时max_abs_diff应该在 1e-4 量级以内fp16 导出时放宽到 1e-2 量级也可以接受但mean_abs_diff最好在 1e-3 以下。如果max_abs_diff到了 0.1 以上别急着高兴转换成功那说明某一层被错误地算成了 fp32 或者某个算子被降级到 CPU 执行了。第二步是确认没有 CPU fallback。ONNX Runtime 的 CUDA provider 在遇到不支持的算子时会自动退回 CPU这个过程通常只打一条 warning很容易被刷屏日志淹没。把日志级别调高或者直接看 session 的 provider 分配print(sess.get_providers()) # 期望输出[CUDAExecutionProvider, CPUExecutionProvider] # 如果只有 CPUExecutionProvider说明 CUDA provider 根本没加载成功我自己踩过一次这个坑转换脚本一切正常engine 也 build 出来了但推理速度只比原生 PyTorch 快了一点点。查了半天才发现是 ONNX Runtime 根本没加载到 CUDA provider全程在 CPU 上跑。那次经历之后我养成了每次转完都跑一遍速度基线对比的习惯。4. 转换成功之后engine 推理阶段的连带坑4.1 engine 反序列化与显存绑定engine 文件本身和设备是绑定的这一点必须提前说清楚。TensorRT build 出来的 engine 里包含了针对特定 GPU 架构编译的内核把在 4090 上 build 的 engine 拿到 3060 上用不一定能正常工作即使能加载也可能触发重建或性能退化。反序列化的代码看起来简单但设备索引这一环很容易漏import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # 注意这一行会隐式初始化默认设备 TRT_LOGGER trt.Logger(trt.Logger.WARNING) def load_engine(engine_path, device_id0): cuda.init() dev cuda.Device(device_id) ctx dev.make_context() # 绑定在这个上下文里创建的所有资源都归 device_id try: with open(engine_path, rb) as f, trt.Runtime(TRT_LOGGER) as runtime: engine runtime.deserialize_cuda_engine(f.read()) return engine, ctx except Exception: ctx.pop() raisepycuda.autoinit会在导入时创建一个默认上下文如果它绑定的设备和你想用的不是同一块后面就会报invalid device context或者更莫名的设备错。这也是为什么那段报错里出现cuda:0时我第一件事是确认进程里到底初始化了几个上下文。如果你同时用 PyTorch 和 PyCUDA两套运行时会有各自的显存池和上下文。PyTorch 默认会在当前设备上创建一个 primary contextPyCUDA 如果再去 make 一个就可能出现两套上下文争抢同一块显存的情况。稳妥的做法是如果主要逻辑在 PyTorch 里优先用 TensorRT 官方提供的 PyTorch 集成接口让显存管理统一只有需要精细控制时再下到 PyCUDA 层面。4.2 动态 shape 与 profile 设置SD 的 UNet 输入是 latent尺寸跟出图分辨率绑定。转换时如果没开动态 shapeengine 就只认 build 时那一个尺寸换成别的分辨率直接报No supported profile。开了动态 shape就得配 profile 的 min/opt/max 三个尺寸。出图分辨率latent 尺寸÷8建议 min建议 opt建议 max512 x 51264 x 6432 x 3264 x 6496 x 96768 x 76896 x 9648 x 4896 x 96128 x 1281024 x 1024128 x 12864 x 64128 x 128160 x 160这张表里的 min 和 max 不是随便定的。min 给到你实际会用的最小尺寸的一半左右max 给到最大尺寸的 1.25 倍这样在日常使用范围内不需要重建 engine同时 profile 区间又不会太宽导致 TensorRT 为了覆盖所有情况而选择保守的内核、拖慢速度。profile 的区间越宽engine build 时间越长运行时的内核选择也越保守。我自己的做法是按主力分辨率建一个窄 profile 的 engine再按偶尔用的分辨率建第二个切换时按需加载。代价是硬盘上多几百 MB换来的是最常用的那个尺寸跑得最快。还有一个容易忽略的点batch 维度也要进 profile。如果你平时单张出图、偶尔开四张并行那 batch 的 min/opt/max 就设成 1/1/4而不是 1/4/4。batch 维度对显存的影响是线性的opt 设大了会让 TensorRT 按大 batch 去选内核单张推理时反而不划算。4.3 精度、速度、显存的三方权衡实测下面这组数据是我在两台机器上反复测出来的量级参考用的是 SD 1.5、512x512、batch 1、20 步采样。数字会随驱动版本、TensorRT 版本、引擎配置波动所以只当作量级参考不要当成 benchmark 去对线。硬件后端配置迭代速度it/sUNet 显存占用首图额外耗时RTX 3060 12GPyTorch fp16 优化注意力约 4 到 5约 2.2 GB无RTX 3060 12GTensorRT fp16静态 shape约 8 到 9约 1.8 GB首次 build 约 6 到 12 分钟RTX 4090PyTorch fp16 优化注意力约 18 到 24约 2.4 GB无RTX 4090TensorRT fp16静态 shape约 32 到 40约 2.0 GB首次 build 约 4 到 8 分钟从这组数据能看出几个规律。加速比大概在 1.8 到 2 倍之间不会更高因为 SD 采样循环里除了 UNet 还有 VAE 解码和调度器计算这两部分没有被 TensorRT 覆盖。显存占用确实降了一点主要来自 TensorRT 更激进的内核融合减少了中间张量的驻留。build 时间这一项要有心理准备。首次转换可能占用几分钟到十几分钟而且这个过程是单线程吃 CPU 的属于典型的一次性投入。engine 文件也不小UNet 的 fp16 engine 在 1.5 GB 到 2 GB 之间加上 VAE 和文本编码器一张卡上准备 3 GB 左右的磁盘空间比较稳妥。关于 fp16 和 int8 的选择我的建议是优先 fp16。int8 需要校准数据集来做量化校准校准集选得不好会明显影响出图质量表现为细节糊、色彩偏移。如果你确实想试 int8至少准备两百张以上有代表性的图片做校准并且在转完之后用固定的种子和提示词对比出图肉眼确认没有明显退化再投入使用。5. 问题速查表与踩坑经验5.1 高频故障对照表下面这张表是我这几年攒下来的按报错特征而不是原因来索引这样你拿到报错就能直接查。报错或现象最可能的原因排查动作修复方式cpu and cuda:0 ... index_selectinput_ids或索引张量在 CPU打印 tokenizer 输出张量的 device补.to(device)同上但出现在 attention 相关层attention mask 在 CPU打印 mask 的 device 和 dtypemask 与输入同设备同 dtypefound at least two devices, cuda:0 and cuda:1多卡环境下设备索引不一致检查CUDA_VISIBLE_DEVICES与代码里的硬编码索引统一设备索引或直接限制可见卡数Expected all tensors ... cuda and cpu顺序反过来权重在 CPU输入在 GPU确认load_state_dict是否被跳过改用标准加载流程No supported profile was found输入尺寸超出 profile 区间打印实际输入 shape 与 profile 区间重建 engine放宽 profile转换成功但速度几乎没提升ONNX Runtime 回退到 CPU provider打印sess.get_providers()重装匹配版本的 onnxruntime-gpubuild engine 阶段 OOMprofile 的 max 区间过大查看 build 时显存峰值缩小 max 区间或临时关掉其他显存占用engine 加载后输出全黑或噪点精度配置与权重不匹配对比 ONNX 与 engine 输出误差统一 fp16/fp32重转invalid device ordinal设备索引超出可见范围打印torch.cuda.device_count()修正索引或调整可见设备首次转换正常第二次报文件占用engine 文件被上一次进程持有检查残留进程结束进程后重建或改用新文件名换显卡后 engine 报内部错误engine 与 GPU 架构绑定确认是否跨架构复用在新卡上重新 build转换脚本在 CPU 上跑得很慢但能跑通模型压根没搬到 GPU打印模型第一个参数的 device检查.to(device)是否被异常吞掉最后一行这个坑值得多说两句。我遇到过一次转换脚本跑了四十分钟还没结束日志一切正常没有任何报错。后来加了一行打印才发现模型的参数全在 CPU 上——model.to(device)那行写在了一个try/except里而它恰好抛了个异常被静默捕获了。这种能跑但慢得离谱的情况比直接报错更难查因为它不会主动提醒你。5.2 几条不外传的实操心得第一条心得永远在脚本的第一行固定设备语义。很多人写转换脚本时习惯用到处传 device 变量的方式结果是十个函数里有八个传对了、两个漏传。我的做法是在脚本开头就写死import os import torch os.environ.setdefault(CUDA_VISIBLE_DEVICES, 0) DEVICE torch.device(cuda, 0) DTYPE torch.float16 torch.backends.cuda.matmul.allow_tf32 True torch.backends.cudnn.allow_tf32 True torch.cuda.set_device(0)把可见设备收窄成一张卡之后所有设备索引不一致类的问题直接消失。代价是你不能用多卡并行转换但对于单卡转换这个场景这个代价可以忽略。allow_tf32这两行在 30 系和 40 系卡上能再榨出一点性能代价是极小的数值精度损失在图像生成场景里基本看不出来。第二条心得把设备断言写进自定义模块的forward里。听起来有点重但只在开发阶段加、上线前用环境变量关掉就行def forward(self, input_ids, attention_maskNone): if os.environ.get(SD_DEBUG_DEVICE) 1: assert input_ids.device self.device, \ finput_ids on {input_ids.device}, expected {self.device} if attention_mask is not None: assert attention_mask.device self.device ...这样做的好处是把运行到一半才炸变成进模块就炸报错信息还带着你自己的描述可读性比 PyTorch 原生报错高得多。我甚至在这个断言里加过assert attention_mask.dtype input_ids.dtype因为 dtype 不一致引发的问题比设备不一致更隐蔽。第三条心得转换产物要做版本归档。engine 文件和你的环境是强绑定的TensorRT 大版本升级、CUDA 大版本升级、显卡换代都可能需要重建。我的习惯是在 engine 目录旁边放一个manifest.json记录当时的 torch 版本、TensorRT 版本、CUDA 版本、显卡型号、build 命令和 profile 参数。半年后回到这个项目时这份记录能省掉一整晚的回忆时间。{ created_at: 2025-01-15T22:40:11, torch: 2.1.2cu121, tensorrt: 8.6.1, cuda: 12.1, gpu: NVIDIA GeForce RTX 3060, profile: {min: [1, 4, 32, 32], opt: [1, 4, 64, 64], max: [1, 4, 96, 96]}, precision: fp16, build_cmd: python convert_unet.py --model sd15 --fp16 --dynamic }最后再分享一个排查顺序上的经验。遇到设备不一致类报错不要上来就翻源码按这个顺序走先打印出错张量的 device 和 dtype再确认模型的next(model.parameters()).device然后检查CUDA_VISIBLE_DEVICES和torch.cuda.current_device()最后才去读源码。这三步走完八成的问题已经能定位到具体那一行了。剩下两成需要读源码的情况基本都是自定义模块或者魔改过的注意力实现那时候你已经知道该往哪看了。这套流程我最近一次用是在给一个加了人脸条件控制的 SD 流程做 TensorRT 加速报错位置在模型入口的拼接处从收到报错到定位到那行 numpy 转张量的代码全程不到十分钟。对比第一次遇到这个问题时折腾了一整个通宵的经历差别就在于有没有一套固定的排查顺序。