把 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_location='cpu'时全部落在内存里。之后如果用的是标准的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,而主模型在 GPU,torch.cat一拼接就炸。这类问题有个共同特征:报错发生在模型的"入口附近",而不是在 UNet 内部。
2.3 转换脚本里 device 参数是怎么被悄悄丢掉的
抛开模块层面,转换脚本自身的写法有几个系统性缺陷,值得单独列出来对照检查。
| 危险写法 | 后果 | 正确做法 |
|---|---|---|
torch.randn(1, 4, 64, 64) | 落在 CPU,前向第一层就报错 | torch.randn(1, 4, 64, 64, device=device) |
torch.randint(0, 49408, (1, 77)) | 同上,且在文本编码器直接触发index_select | 加device=device |
torch.load(path)不带map_location | 张量落在保存时的设备,跨机器会出问题 | 统一map_location='cpu',之后走load_state_dict |
t = torch.tensor(numpy_arr) | 落在 CPU | torch.as_tensor(numpy_arr, device=device)或torch.from_numpy(numpy_arr).to(device) |
w.data = sd['w']直接赋值 | 张量不搬,模型属性里混进 CPU 张量 | 改成load_state_dict,或赋值后补.to(device) |
torch.arange(n)做索引 | 落在 CPU | 加device=device |
torch.zeros(...)做 mask | 落在 CPU | 加device=device或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_mask=attention_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_BLOCKING=1 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_mask=attention_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, truncation=True, max_length=self.max_length, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt", ) tokens = batch_encoding["input_ids"] # 修改后 batch_encoding = self.tokenizer( text, truncation=True, max_length=self.max_length, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt", ) 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, device=device, dtype=dtype) dummy_ids = torch.randint(0, 49408, (1, 77), device=device) # 如果还要传 attention mask,一起对齐 dummy_mask = torch.ones_like(dummy_ids, device=device)注意dtype也要对齐。fp16 模型配 fp32 dummy input,虽然不会直接报设备错,但在某些算子(尤其是 LayerNorm 和 attention 内部)会因为 dtype 不一致报另一个错,或者静默产生数值偏差,让后面的精度对比结果莫名其妙地差。
补丁三:自定义条件注入与分块索引。这类问题通常出现在模型的入口和出口两端。
def to_device_recursive(obj, device, dtype=None): """把任意嵌套结构里的张量统一搬到目标设备""" 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", providers=providers) 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_id=0): 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 context,PyCUDA 如果再去 make 一个,就可能出现两套上下文争抢同一块显存的情况。稳妥的做法是:如果主要逻辑在 PyTorch 里,优先用 TensorRT 官方提供的 PyTorch 集成接口,让显存管理统一;只有需要精细控制时再下到 PyCUDA 层面。
4.2 动态 shape 与 profile 设置
SD 的 UNet 输入是 latent,尺寸跟出图分辨率绑定。转换时如果没开动态 shape,engine 就只认 build 时那一个尺寸,换成别的分辨率直接报No supported profile。开了动态 shape,就得配 profile 的 min/opt/max 三个尺寸。
| 出图分辨率 | latent 尺寸(÷8) | 建议 min | 建议 opt | 建议 max |
|---|---|---|---|---|
| 512 x 512 | 64 x 64 | 32 x 32 | 64 x 64 | 96 x 96 |
| 768 x 768 | 96 x 96 | 48 x 48 | 96 x 96 | 128 x 128 |
| 1024 x 1024 | 128 x 128 | 64 x 64 | 128 x 128 | 160 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/s) | UNet 显存占用 | 首图额外耗时 |
|---|---|---|---|---|
| RTX 3060 12G | PyTorch fp16 + 优化注意力 | 约 4 到 5 | 约 2.2 GB | 无 |
| RTX 3060 12G | TensorRT fp16,静态 shape | 约 8 到 9 | 约 1.8 GB | 首次 build 约 6 到 12 分钟 |
| RTX 4090 | PyTorch fp16 + 优化注意力 | 约 18 到 24 | 约 2.4 GB | 无 |
| RTX 4090 | TensorRT 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_select | input_ids或索引张量在 CPU | 打印 tokenizer 输出张量的 device | 补.to(device) |
| 同上但出现在 attention 相关层 | attention mask 在 CPU | 打印 mask 的 device 和 dtype | mask 与输入同设备同 dtype |
found 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-gpu |
| build engine 阶段 OOM | profile 的 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_mask=None): if os.environ.get("SD_DEBUG_DEVICE") == "1": assert input_ids.device == self.device, \ f"input_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.2+cu121", "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 转张量的代码,全程不到十分钟。对比第一次遇到这个问题时折腾了一整个通宵的经历,差别就在于有没有一套固定的排查顺序。