把 Stable Diffusion 模型往 TensorRT 引擎转的过程里最让人血压升高的一类报错不是显存爆了而是那一行红字Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument ...)。它不像 OOM 那样干脆利落地告诉你显存不够也不像算子不支持那样明确指认某个节点它只是冷冷地告诉你有两个张量不在同一个地方你自己找吧。很多人第一反应是把 batch size 调小、把分辨率降低、重启一下 webui甚至怀疑是不是 TensorRT 装错了——这些操作基本都没用因为锅根本不在 TensorRT 那边。这篇内容围绕 Stable Diffusion 转 TensorRT 时这个经典的设备错位报错展开把报错文本逐字拆开、把常见触发场景列清楚、给出一套五分钟能定位到具体层的排查手法再把从.safetensors加载一路到 ONNX 导出、TRT engine 构建的整条链路上该在哪里钉设备讲透。适合已经在跑 SD 本地部署、正在尝试 TensorRT 加速的玩家也适合写导出脚本时被这行报错卡住的开发者。不管你是用整合包还是自己搭环境下面这些方法都能直接用。1. 先看清楚这行红字到底在说什么1.1method wrapper_CUDA__这个前缀的来历这行报错的措辞是 PyTorch 2.0 以后才有的。在此之前老版本给的是RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 mat2或者Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!。新版本把 dispatcher 的内部结构暴露到了错误信息里于是你看到了wrapper_CUDA__linear这种看起来像天书的前缀。拆开读就很好理解wrapper是 dispatcher 给算子套的壳CUDA表示这次调用已经按 CUDA 后端分发下去了__后面跟的是算子的真实名字。所以wrapper_CUDA__linear翻译成人话就是——PyTorch 已经认定这次linear调用应该走 CUDA 实现结果在参数检查阶段发现里面混进了一个 CPU 张量。顺着这个思路不同后缀指向的位置完全不同这也是定位问题的第一条线索报错中的算子后缀大概率出问题的位置SD 里的具体含义wrapper_CUDA__linear/wrapper_CUDA__addmm全连接层权重或输入CLIP 文本编码器的投影层、UNet 注意力里的 QKV 线性层wrapper_CUDA__embedding嵌入表或索引文本编码器第一层的 token embeddingwrapper_CUDA__native_layer_normLayerNorm 的 weight/bias注意力块前后的归一化层wrapper_CUDA__index_select索引张量embedding 查表时的输入 idwrapper_CUDA__conv2d卷积核或特征图UNet 下采样/上采样卷积记住一句话报错里的算子名就是地图。看到embedding就别去翻 UNet 的卷积层直接去查文本编码器。1.2 argument for argument 后面那个词才是真凶完整报错里通常还会跟一句(when checking argument for argument weight ...)或者... for argument mat2 ...。这个argument后面的名字指的是第一个被检查出来不在同一设备上的张量参数名。这非常关键因为它能帮你分清两种截然不同的故障模式如果指的是weight、bias、running_mean这类名字说明模型参数没搬全是模型侧的问题。如果指的是input、mat1、mat2、other、index、src这类名字说明调用方传进来的张量在 CPU 上是数据侧的问题最常见的就是导出脚本里的 dummy input 创建时忘了写device。这两类的修法完全不一样。前者要遍历整个nn.Module把漏掉的层搬到 GPU后者只需要在构造输入张量的地方补一个devicecuda。很多人上来就无脑加.cuda()结果加错地方报错从weight变成mat2白白浪费一晚上。1.3 为什么 TensorRT 的转换日志里会蹦出 PyTorch 的异常这是最容易被误解的一点。整套转换流程其实分成五段从.safetensors/.ckpt加载权重到 PyTorch 模型把 UNet、VAE、CLIP 等子模块拼装成可导出的结构挂上 LoRA、ControlNet、自定义 VAE用torch.onnx.export做一次 trace生成 ONNX 计算图对 ONNX 图做常量折叠、算子简化交给 TensorRT 的 builder 编译成.plan/.engine。第 1 到第 4 步全程跑在 PyTorch 和 ONNX 里TensorRT 只参与最后一步。设备错位这个异常只可能来自前三步因为 TRT builder 拿到的是已经序列化好的图它压根不关心你原来的张量在哪个设备上。所以当你在转换脚本日志里看到这行红字时正确的心态是暂时把 TensorRT 忘掉把它当成一个纯粹的 PyTorch 模型导出问题来处理。我见过太多人把 TensorRT 卸了重装三遍问题还在原地——因为问题从来就不在 TensorRT 上。2. 三类最容易出事的设备错位场景2.1 权重加载时张量留在 CPU 上这是占比最高的一类。典型写法长这样import torch ckpt torch.load(model.safetensors) model.load_state_dict(ckpt[state_dict]) # 忘了 model model.to(device)或者只搬了一部分torch.load不带map_location时会把权重还原到保存时所在的设备。如果你的 checkpoint 是在 CUDA 上保存的而当前机器的 CUDA 编号对不上或者压根没检测到 CUDAPyTorch 会静默地落到 CPU 上。更隐蔽的情况是checkpoint 保存时在cpu你以为它会自动跟模型走其实不会——load_state_dict只是把数据拷进已存在的参数里不会改变参数原本的设备。还有一种只搬了一半的写法在 SD 里特别常见model.diffusion_model.to(device) # UNet 搬了 model.cond_stage_model.to(device) # 忘了搬 CLIP model.first_stage_model.to(device) # 忘了搬 VAESD 的结构是由好几个独立子模型拼起来的diffusion_modelUNet、cond_stage_modelCLIP 文本编码器、first_stage_modelVAE各有各的参数。只对最外层调一次.to()通常能递归下去但如果中间有人用nn.ModuleList之外的方式持有子模块比如塞进了一个普通 list、dict或者挂在 Python attribute 上.to()就递归不到那部分就会永远停在 CPU。2.2 示例输入忘了指定 device这个坑踩的人最多也最好修。导出 ONNX 必须给模型喂一组 dummy input 来 trace而很多脚本里是这么写的sample torch.randn(2, 4, 64, 64) timestep torch.tensor([1]) encoder_hidden_states torch.randn(2, 77, 768) torch.onnx.export(unet, (sample, timestep, encoder_hidden_states), unet.onnx)torch.randn不带device参数默认就是 CPU。而 UNet 已经被搬到cuda:0了。第一个卷积层一执行报错立刻出现。正确的写法是让输入跟模型保持绝对一致不要靠记忆去写cuda:0而是从模型里读device next(unet.parameters()).device dtype next(unet.parameters()).dtype sample torch.randn(2, 4, 64, 64, devicedevice, dtypedtype) timestep torch.tensor([1], devicedevice) encoder_hidden_states torch.randn(2, 77, 768, devicedevice, dtypedtype)用next(model.parameters()).device去拿设备好处是无论你后面改成cuda:1、改成 CPU 调试、还是改半精度这段代码都不用动。这个小习惯我建议从第一天写导出脚本时就养成能省掉后面无数次排查。2.3 没注册成 buffer 的野生张量这一类最阴因为它不会在model.named_parameters()里出现你用常规方法打印设备分布也查不到。典型的长这样class PosEmbed(nn.Module): def __init__(self, dim): super().__init__() self.pe torch.zeros(1, 197, dim) # 普通 attribute不是 buffer这里的self.pe只是一个挂在模块上的普通 Python 属性。model.to(cuda)会把parameters()和buffers()逐个搬运但不会碰self.pe。于是模型在 GPU 上跑self.pe还在 CPU一相加就炸。两种修法选后者更稳# 写法一注册成 buffer随 .to() 自动迁移 self.register_buffer(pe, torch.zeros(1, 197, dim)) # 写法二干脆不存在 forward 里按输入设备现场生成 pe torch.zeros(1, x.size(1), x.size(-1), devicex.device, dtypex.dtype)还有一个更隐蔽的变体safetensors.torch.load_file()和load_model()返回的张量永远是 CPU 张量不管你当前模型在哪儿。如果你用 safetensors 加载 LoRA 权重然后直接拿去做F.linear(x, lora_up)x在 cuda、lora_up在 cpu报错必现。这一点后面第 4 章还会细讲。把这三类整理成一张对照表排查时可以按图索骥场景报错里的 argument 名最快的验证方式最小修复权重没搬全weight/bias打印各子模块设备分布对根模块调.to(device)或递归搬运dummy input 在 CPUmat2/other/input检查导出脚本里张量创建处补devicedevice野生张量 / LoRA 权重mat2/other用 forward hook 打设备集合注册 buffer 或显式.to()3. 五分钟定位法把出问题的层揪出来3.1 打印全模型的设备分布排查第一步永远是先看一眼现场。写一个十几行的审计函数比瞎猜快一百倍import torch from collections import Counter def audit_devices(model, tagmodel): stat Counter() for _, p in model.named_parameters(): stat[str(p.device)] 1 for _, b in model.named_buffers(): stat[str(b.device)] 1 print(f[{tag}] device stats: {dict(stat)}) # 把不在目标设备上的层名打印出来最多 30 条 target cuda:0 leaked [(n, str(p.device)) for n, p in model.named_parameters() if str(p.device) ! target] leaked [(n, str(b.device)) for n, b in model.named_buffers() if str(b.device) ! target] for n, d in leaked[:30]: print(f leaked: {d:10} {n}) print(f[{tag}] leaked total: {len(leaked)}) return leaked这个函数一跑如果device stats里同时出现cuda:0和cpu基本就锁定是权重没搬全了。而且leaked列表会直接把层名和它所在的设备打出来前三十条足够你看清规律——是集中在某个子模块还是散落各处。我自己的习惯是在模型拼装完之后、torch.onnx.export之前各调一次audit_devices。两次结果一对比如果第二次多出泄漏项说明中间某段代码又创建了新张量。3.2 用 forward pre-hook 抓第一现场设备分布看着没问题但一跑还是崩那问题多半出在 forward 内部动态生成的张量上。这时候named_parameters()那套就不够用了得用 hook 在每一层执行前检查。def install_device_guard(model, targetcuda:0): def make_pre(name): def pre_hook(module, args): devs {str(t.device) for t in args if torch.is_tensor(t)} devs | {str(p.device) for p in module.parameters(recurseFalse)} devs | {str(b.device) for b in module.buffers(recurseFalse)} if len(devs) 1: print(f[MISMATCH] {name} - {sorted(devs)}) return pre_hook for n, m in model.named_modules(): m.register_forward_pre_hook(make_pre(n))挂上去之后跑一次前向导出脚本里那次 trace 就够了日志里第一个[MISMATCH]就是第一现场。因为报错是按执行顺序触发的第一个不匹配的层往往就是真正的病灶后面跟着刷屏的那些多半只是它连带的受害者。用recurseFalse是有意为之只看这一层自己持有的参数和 buffer不看子模块的。否则父模块会把子模块的参数全算进去永远判定为不匹配日志就没法看了。3.3 打开 CUDA_LAUNCH_BLOCKING 和 C 栈PyTorch 的报错栈有时候会指向一个跟你写的代码完全无关的位置比如报在torch/nn/modules/module.py的_call_impl里。这是因为 CUDA 是异步执行的Python 侧早就跑到后面去了错误才在 GPU 上爆发出来。加两个环境变量能大幅改善这一点CUDA_LAUNCH_BLOCKING1 TORCH_SHOW_CPP_STACKTRACES1 python export_onnx.pyCUDA_LAUNCH_BLOCKING1强制每个 CUDA 调用同步报错会精确落在真正执行它的那一行。TORCH_SHOW_CPP_STACKTRACES1会把 C 侧的调用栈也打出来能直接看到是at::native::linear还是at::native::embedding触发的。注意CUDA_LAUNCH_BLOCKING1会明显拖慢速度只在排查阶段开定位到了立刻关掉。3.4 二分法排除按子模块单独跑如果上面三招还是没头绪就用最笨但最有效的办法——把模型拆开一个个跑。with torch.inference_mode(): # 先单独测 CLIP ids torch.randint(0, 49408, (2, 77)).cuda() clip_out sd_model.cond_stage_model(ids) print(clip ok, clip_out.device) # 再单独测 UNet x torch.randn(2, 4, 64, 64).cuda() t torch.tensor([1]).cuda() unet_out sd_model.model.diffusion_model(x, t, contextclip_out) print(unet ok, unet_out.device) # 最后测 VAE latent unet_out.cuda() img sd_model.decode_first_stage(latent) print(vae ok, img.device)哪一段先崩问题就在那一段里。这个方法的另一个好处是——它顺手帮你确认了每个子模块的接口签名对后面写 ONNX 导出脚本很有帮助。4. 从 .safetensors 到 ONNX整条链路逐段对齐设备4.1 权重加载阶段map_location 的正确写法先明确一个原则加载阶段统一落到 CPU之后显式搬到目标设备。不要图省事直接map_locationcuda。# 推荐写法 ckpt torch.load(ckpt_path, map_locationcpu) model.load_state_dict(ckpt[state_dict], strictFalse) model model.to(devicetarget_device, dtypetarget_dtype) audit_devices(model, after-load)为什么先落到 CPU 再搬三个理由。第一如果 checkpoint 是在 8 卡机器上存的、你现在只有 1 张卡直接map_locationcuda会因为设备编号对不上直接抛异常。第二CPU 加载对显存的瞬时压力更小低显存机器上不容易在加载阶段就 OOM。第三先 CPU 后搬的过程里你有机会在搬之前做一些张量级的预处理比如 dtype 转换、权重融合。半精度的顺序也有讲究先.to(device)再.half()或者干脆一次写成.to(devicedevice, dtypetorch.float16)。顺序反了的话在 CPU 上做半精度转换不仅慢某些算子还会因为 CPU 缺少 FP16 实现而报另一种错。对于 safetensors 格式加载方式略有不同但落点一样from safetensors.torch import load_file state load_file(model.safetensors, devicecpu) # 需要显式搬运 state {k: v.to(devicetarget_device, dtypetarget_dtype) for k, v in state.items()}safetensors.torch.load_file从 0.4 版本起支持device参数但默认值是cpu。这就是前面提到的第一类隐蔽坑模型已经在 GPU 上你以为加载出来的权重也会在 GPU 上实际上不会。4.2 模块拼装阶段补齐 LoRA / ControlNet / VAE 的搬运SD 转 TRT 时真正麻烦的是后挂载的组件。LoRA 是权重差分、ControlNet 是额外 UNet 分支、自定义 VAE 是替换掉first_stage_model——这三样都有各自独立的加载路径而且都容易在搬运上漏掉。LoRA 融合的典型形态是lora_up load_file(lora.safetensors)[lora_up.weight] # CPU 张量 w base_weight alpha * (lora_down lora_up) # base_weight 在 GPU炸修法很直接把加载出来的张量统一下沉到目标设备和 dtypedef to_target(t, device, dtype): return t.to(devicedevice, dtypedtype, non_blockingTrue) lora_up to_target(load_file(path)[lora_up.weight], device, dtype)non_blockingTrue在固定内存配合下能加速拷贝但前提是源张量锁页不锁页时加了也无害只是不生效。ControlNet 的情况是它本身是一个完整的网络加载完之后必须整体.to(device)同时它的 conditioning 输入比如边缘图、深度图也要跟它同设备。我见过有人把 ControlNet 权重搬了但 hint 图还是 CPU 上来的 numpy 转换结果一样炸。自定义 VAE 替换更容易被忽略因为它经常是在 webui 启动参数里指定的比如--no-half-vae、外挂的 VAE 文件替换逻辑散落在不同位置。替换完之后一定要再audit_devices一遍。4.3 导出阶段dummy input、dynamic_axes 与 dtype 的一致性torch.onnx.export阶段要同时保证三件事一致设备一致、dtype 一致、shape 一致。设备不一致就是本文这个报错dtype 不一致会报expected scalar type Half but found Floatshape 不一致会在 TRT 构建阶段报维度冲突。用torch.set_default_device可以从根上减少第一类的发生概率import torch torch.set_default_device(cuda:0)这行之后所有torch.randn、torch.zeros、torch.tensor不带device参数时都会默认建在cuda:0上。PyTorch 2.0 开始支持还有对应的上下文管理器写法with torch.device(cuda:0): sample torch.randn(2, 4, 64, 64)提醒set_default_device是个全局开关某些第三方库内部依赖默认在 CPU这个前提打开之后可能引入新的问题。用在导出脚本这种短生命周期的场景里最安全别在整个 webui 进程里乱开。一份能直接用的导出骨架大致是这样import torch unet sd_model.model.diffusion_model device next(unet.parameters()).device dtype next(unet.parameters()).dtype unet.eval() audit_devices(unet, before-export) with torch.inference_mode(): sample torch.randn(2, 4, 64, 64, devicedevice, dtypedtype) timestep torch.tensor([1], devicedevice) ctx torch.randn(2, 77, 768, devicedevice, dtypedtype) dummy (sample, timestep, ctx) torch.onnx.export( unet, dummy, unet_fp16.onnx, input_names[sample, timestep, context], output_names[out_sample], dynamic_axes{ sample: {0: batch, 2: height, 3: width}, timestep: {0: batch}, context: {0: batch}, out_sample: {0: batch, 2: height, 3: width}, }, opset_version17, do_constant_foldingTrue, )dynamic_axes里把 batch 和 H/W 标成动态是为了后面能一次编译出支持多种分辨率的引擎。但注意动态轴越多TRT builder 优化空间越小而且更容易在 shape 推导阶段触发隐藏的算子不支持。如果只是自己用先把分辨率固定成 512x512 或 768x768 编译一个版本跑通了再上动态轴这个顺序能省很多时间。4.4 SDXL 双文本编码器与多 GPU 的特殊处理SDXL 有两个文本编码器一个 CLIP-L、一个 CLIP-bigG它们的输出在通道维度拼接之后才喂给 UNet。这里有两个高频问题一是两个编码器只搬了一个二是两个编码器的 dtype 不一致一个 FP16、一个 FP32。稳妥的做法是统一处理text_encoders [sd_model.cond_stage_model, sd_model.cond_stage_model_2] for te in text_encoders: te.to(devicedevice, dtypedtype).eval()多 GPU 场景下还有一层复杂度如果你用device_map或者accelerate把模型分散到多张卡上next(model.parameters()).device拿到的只是第一张卡用它去构造 dummy input 会导致后续所有在其它卡上的层全部不匹配。这种情况下要么显式指定每段输入落哪张卡要么老老实实把整段模型收敛到单卡上再导出。转 TRT 这个场景里单卡是默认推荐——TRT engine 本身就是按单设备编译的多卡分散反而增加不确定性。5. 修好之后怎么验证不是假绿5.1 ONNX 侧onnxruntime 与 PyTorch 的数值对齐脚本不报错了、ONNX 文件也生成了不代表导出是对的。ONNX 导出最常见的失败模式是图能跑数值不对尤其是 FP16 导出时。所以必须做一次数值对齐import numpy as np import onnxruntime as ort import torch with torch.inference_mode(): ref unet(*dummy).cpu().numpy() sess ort.InferenceSession(unet_fp16.onnx, providers[CUDAExecutionProvider]) inputs { sample: dummy[0].cpu().numpy(), timestep: dummy[1].cpu().numpy(), context: dummy[2].cpu().numpy(), } out sess.run([out_sample], inputs)[0] diff np.abs(ref - out) print(max abs diff:, diff.max(), mean:, diff.mean()) np.testing.assert_allclose(ref, out, rtol1e-2, atol1e-2)容差定在1e-2是针对 FP16 的经验值。FP32 导出应该收紧到1e-4左右。如果 max diff 在 1 以上说明图里有算子被错误替换了或者某个节点被常量折叠折坏了这时候要回去检查do_constant_folding和 opset 版本。提示ONNX Runtime 和 PyTorch 的卷积实现在不同硬件上会有微小差异即使 FP32 也不可能逐位相同。看到 1e-5 级别的差异不用担心。5.2 TRT 侧engine 构建与两次推理比对引擎构建分两条路。命令行用trtexec最省事trtexec --onnxunet_fp16.onnx \ --saveEngineunet_fp16.plan \ --fp16 \ --minShapessample:1x4x64x64 \ --optShapessample:2x4x64x64 \ --maxShapessample:4x4x64x64 \ --workspace4096如果模型有多个输入--minShapes/--optShapes/--maxShapes必须把所有动态输入都列全漏一个就会报 shape 不匹配。构建完之后用 Python 侧的 TensorRT runtime 跑一次推理跟 ONNX Runtime 的输出比对这一次的容差可以放宽到2e-2——因为 FP16 的累加顺序在 TRT 里可能跟 ONNX Runtime 不一样。两次比对都通过才能说这次转换是干净的。5.3 端到端出图对比与常见精度容差最后一步还是得看图。用固定 seed、固定 prompt、固定步数分别用原始 PyTorch 路径和 TRT 引擎路径各出一张图然后做像素级对比import numpy as np from PIL import Image a np.asarray(Image.open(out_torch.png), dtypenp.float32) b np.asarray(Image.open(out_trt.png), dtypenp.float32) print(mean abs diff:, np.abs(a - b).mean())正常情况下的差异应该在个位数级别0-255 的尺度上肉眼几乎看不出区别。如果画面上出现明显的色块、噪点或者结构错乱说明某个阶段的精度损失已经越界了这时候的排查方向是先看是不是 VAE 用 FP16 导的VAE 对精度特别敏感再看 UNet 的 attention 层有没有被错误简化。我自己的习惯是固定一组验收用例——一张 512x512、一张 768x768、一张带 ControlNet 的三张图对着看。只要这三张稳定基本可以放心用了。6. 环境与版本里的隐性坑6.1 PyTorch / CUDA / TensorRT 版本组合这个报错的措辞本身就能反推版本wrapper_CUDA__是 PyTorch 2.0 的格式。如果你的环境里 PyTorch 还是 1.13看到的是另一套措辞说明你可能在照着旧教程操作新版本或者反过来。一个被反复验证过的组合是这样的组件推荐版本说明PyTorch2.1.x / 2.2.x支持set_default_deviceONNX 导出对 opset 17 支持稳定CUDA11.8 / 12.111.8 兼容性最广12.1 需要对应驱动TensorRT8.6 / 9.x8.6 生态最成熟9.x 移除了部分旧 APIONNX1.15低于 1.14 对 opset 17 的部分算子支持不全onnxruntime-gpu1.16与 CUDA 版本必须严格对应版本错配最典型的症状不是本文这个报错而是构建 engine 时提示某个算子不受支持、或者精度莫名其妙地崩。但有一种情况要注意TensorRT 9.x 在某些版本下不再支持动态 shape 的某些组合会迫使你回头改dynamic_axes改的过程中如果顺手动了输入构造代码就可能引入设备错位。所以改环境之前先把脚本备份一份。6.2 onnxruntime 与 onnxruntime-gpu 的共存冲突这是个经典陷阱pip install onnxruntime和pip install onnxruntime-gpu装的是同一个 import 名onnxruntime后装的会覆盖先装的。如果你验证时发现自己明明装了 GPU 版跑起来却是 CPU 执行用这一行确认import onnxruntime as ort print(ort.get_available_providers())输出里没有CUDAExecutionProvider就说明当前生效的是 CPU 版。清理方式是先pip uninstall onnxruntime onnxruntime-gpu -y再只装需要的那一个。这事儿跟本文的报错有间接关系如果你的导出脚本里恰好用 onnxruntime 做中间验证而它跑在 CPU 上你在 CPU 侧看到的张量当然都在 CPU——这会掩盖真实的设备问题让你在错误的方向上排查。6.3 半精度导出的边界与最小显存配置FP16 导出能省一半显存但有几个硬边界要知道。VAE 的解码器在 FP16 下容易出现 NaN 或者色偏所以很多人的做法是 UNet 和 CLIP 用 FP16、VAE 用 FP32。这样做的时候两个不同精度之间的衔接张量必须显式转换否则会从设备错位变成 dtype 错位latent unet_out.to(torch.float32) # 从 FP16 的 UNet 输出转回 FP32 喂给 VAE img vae.decode(latent)显存方面用trtexec构建 512x512 的 SD 1.5 UNet 引擎--workspace给到 4096MB 一般够用SDXL 建议直接给 8192MB。工作空间不足时 builder 会退化到更保守的策略编译出来的引擎反而更慢所以这个值别抠。还有一点--fp16只是允许用 FP16 做计算不代表所有层都会用。构建完可以用trtexec --loadEnginexx.plan --dumpLayerInfo看一下各层的精度确认关键层没有被意外降精度。7. 几个我踩过的边界情况7.1 只在某个分辨率才炸非 64 倍数有一类特别迷惑的现象512x512 导出得好好的改成 520x520 就报设备错位。这跟设备其实没关系——真正的原因是 520 不是 64 的倍数UNet 的下采样在某一层产生奇数尺寸触发了某个不支持该尺寸的分支而在那个分支里恰好有一段代码动态创建了 CPU 张量。所以遇到换个分辨率就炸的情况先别怀疑设备代码先把分辨率调回 64 的倍数验证一下。SD 1.5 常用 512、576、640、704、768SDXL 常用 1024、1152。如果你的dynamic_axes允许任意尺寸最好在输入校验里加一道硬性检查assert h % 64 0 and w % 64 0, fresolution must be multiple of 64, got {h}x{w}7.2 换了显卡之后重新出现换卡之后重装环境是常规操作但有个细节容易被忽略如果你之前生成的某些中间文件比如已经转换好的 LoRA 缓存、预处理的 embedding 缓存是用旧卡生成的里面可能存着cuda:0的张量。新机器上用torch.load加载时会尝试还原到cuda:0如果新卡的编号体系不一样比如从单卡变双卡或者 CUDA_VISIBLE_DEVICES 设过就会落到 CPU 上然后引发错位。修法是统一在加载时加map_locationcpu这一条对任何缓存文件都适用值得写进你的工具函数里。7.3 webui 启动参数带来的副作用跑在 webui 整合包里的转换脚本会继承主进程的启动参数这里面有几个会直接影响设备行为--precision full会禁用半精度跟你在脚本里写的.half()冲突导致部分层 FP32、部分层 FP16--no-half-vae会让 VAE 保持 FP32前面提到的衔接转换就必须加--medvram/--lowvram会在运行期动态搬运模块转换脚本执行到一半某些层可能被挪走了--skip-torch-cuda-test只是跳过检测不会修复任何设备问题别指望它。所以我的建议是做 TensorRT 转换时用独立的 Python 脚本不要在 webui 进程里跑。模型文件、VAE、LoRA 都从磁盘直接加载环境干净、设备状态可控、报错栈也短。整合包适合玩不适合做这种对状态一致性要求高的操作。最后分享一个我自己一直在用的小技巧把设备审计函数封装成一个装饰器加在任何导出函数头上函数进来先审计、出去再审计。看着有点笨但它帮我抓到过三四次只在特定参数组合下才出现的隐性泄漏比事后追查便宜太多了。
Stable Diffusion 转 TensorRT 设备错位报错排查与修复
把 Stable Diffusion 模型往 TensorRT 引擎转的过程里最让人血压升高的一类报错不是显存爆了而是那一行红字Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument ...)。它不像 OOM 那样干脆利落地告诉你显存不够也不像算子不支持那样明确指认某个节点它只是冷冷地告诉你有两个张量不在同一个地方你自己找吧。很多人第一反应是把 batch size 调小、把分辨率降低、重启一下 webui甚至怀疑是不是 TensorRT 装错了——这些操作基本都没用因为锅根本不在 TensorRT 那边。这篇内容围绕 Stable Diffusion 转 TensorRT 时这个经典的设备错位报错展开把报错文本逐字拆开、把常见触发场景列清楚、给出一套五分钟能定位到具体层的排查手法再把从.safetensors加载一路到 ONNX 导出、TRT engine 构建的整条链路上该在哪里钉设备讲透。适合已经在跑 SD 本地部署、正在尝试 TensorRT 加速的玩家也适合写导出脚本时被这行报错卡住的开发者。不管你是用整合包还是自己搭环境下面这些方法都能直接用。1. 先看清楚这行红字到底在说什么1.1method wrapper_CUDA__这个前缀的来历这行报错的措辞是 PyTorch 2.0 以后才有的。在此之前老版本给的是RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 mat2或者Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!。新版本把 dispatcher 的内部结构暴露到了错误信息里于是你看到了wrapper_CUDA__linear这种看起来像天书的前缀。拆开读就很好理解wrapper是 dispatcher 给算子套的壳CUDA表示这次调用已经按 CUDA 后端分发下去了__后面跟的是算子的真实名字。所以wrapper_CUDA__linear翻译成人话就是——PyTorch 已经认定这次linear调用应该走 CUDA 实现结果在参数检查阶段发现里面混进了一个 CPU 张量。顺着这个思路不同后缀指向的位置完全不同这也是定位问题的第一条线索报错中的算子后缀大概率出问题的位置SD 里的具体含义wrapper_CUDA__linear/wrapper_CUDA__addmm全连接层权重或输入CLIP 文本编码器的投影层、UNet 注意力里的 QKV 线性层wrapper_CUDA__embedding嵌入表或索引文本编码器第一层的 token embeddingwrapper_CUDA__native_layer_normLayerNorm 的 weight/bias注意力块前后的归一化层wrapper_CUDA__index_select索引张量embedding 查表时的输入 idwrapper_CUDA__conv2d卷积核或特征图UNet 下采样/上采样卷积记住一句话报错里的算子名就是地图。看到embedding就别去翻 UNet 的卷积层直接去查文本编码器。1.2 argument for argument 后面那个词才是真凶完整报错里通常还会跟一句(when checking argument for argument weight ...)或者... for argument mat2 ...。这个argument后面的名字指的是第一个被检查出来不在同一设备上的张量参数名。这非常关键因为它能帮你分清两种截然不同的故障模式如果指的是weight、bias、running_mean这类名字说明模型参数没搬全是模型侧的问题。如果指的是input、mat1、mat2、other、index、src这类名字说明调用方传进来的张量在 CPU 上是数据侧的问题最常见的就是导出脚本里的 dummy input 创建时忘了写device。这两类的修法完全不一样。前者要遍历整个nn.Module把漏掉的层搬到 GPU后者只需要在构造输入张量的地方补一个devicecuda。很多人上来就无脑加.cuda()结果加错地方报错从weight变成mat2白白浪费一晚上。1.3 为什么 TensorRT 的转换日志里会蹦出 PyTorch 的异常这是最容易被误解的一点。整套转换流程其实分成五段从.safetensors/.ckpt加载权重到 PyTorch 模型把 UNet、VAE、CLIP 等子模块拼装成可导出的结构挂上 LoRA、ControlNet、自定义 VAE用torch.onnx.export做一次 trace生成 ONNX 计算图对 ONNX 图做常量折叠、算子简化交给 TensorRT 的 builder 编译成.plan/.engine。第 1 到第 4 步全程跑在 PyTorch 和 ONNX 里TensorRT 只参与最后一步。设备错位这个异常只可能来自前三步因为 TRT builder 拿到的是已经序列化好的图它压根不关心你原来的张量在哪个设备上。所以当你在转换脚本日志里看到这行红字时正确的心态是暂时把 TensorRT 忘掉把它当成一个纯粹的 PyTorch 模型导出问题来处理。我见过太多人把 TensorRT 卸了重装三遍问题还在原地——因为问题从来就不在 TensorRT 上。2. 三类最容易出事的设备错位场景2.1 权重加载时张量留在 CPU 上这是占比最高的一类。典型写法长这样import torch ckpt torch.load(model.safetensors) model.load_state_dict(ckpt[state_dict]) # 忘了 model model.to(device)或者只搬了一部分torch.load不带map_location时会把权重还原到保存时所在的设备。如果你的 checkpoint 是在 CUDA 上保存的而当前机器的 CUDA 编号对不上或者压根没检测到 CUDAPyTorch 会静默地落到 CPU 上。更隐蔽的情况是checkpoint 保存时在cpu你以为它会自动跟模型走其实不会——load_state_dict只是把数据拷进已存在的参数里不会改变参数原本的设备。还有一种只搬了一半的写法在 SD 里特别常见model.diffusion_model.to(device) # UNet 搬了 model.cond_stage_model.to(device) # 忘了搬 CLIP model.first_stage_model.to(device) # 忘了搬 VAESD 的结构是由好几个独立子模型拼起来的diffusion_modelUNet、cond_stage_modelCLIP 文本编码器、first_stage_modelVAE各有各的参数。只对最外层调一次.to()通常能递归下去但如果中间有人用nn.ModuleList之外的方式持有子模块比如塞进了一个普通 list、dict或者挂在 Python attribute 上.to()就递归不到那部分就会永远停在 CPU。2.2 示例输入忘了指定 device这个坑踩的人最多也最好修。导出 ONNX 必须给模型喂一组 dummy input 来 trace而很多脚本里是这么写的sample torch.randn(2, 4, 64, 64) timestep torch.tensor([1]) encoder_hidden_states torch.randn(2, 77, 768) torch.onnx.export(unet, (sample, timestep, encoder_hidden_states), unet.onnx)torch.randn不带device参数默认就是 CPU。而 UNet 已经被搬到cuda:0了。第一个卷积层一执行报错立刻出现。正确的写法是让输入跟模型保持绝对一致不要靠记忆去写cuda:0而是从模型里读device next(unet.parameters()).device dtype next(unet.parameters()).dtype sample torch.randn(2, 4, 64, 64, devicedevice, dtypedtype) timestep torch.tensor([1], devicedevice) encoder_hidden_states torch.randn(2, 77, 768, devicedevice, dtypedtype)用next(model.parameters()).device去拿设备好处是无论你后面改成cuda:1、改成 CPU 调试、还是改半精度这段代码都不用动。这个小习惯我建议从第一天写导出脚本时就养成能省掉后面无数次排查。2.3 没注册成 buffer 的野生张量这一类最阴因为它不会在model.named_parameters()里出现你用常规方法打印设备分布也查不到。典型的长这样class PosEmbed(nn.Module): def __init__(self, dim): super().__init__() self.pe torch.zeros(1, 197, dim) # 普通 attribute不是 buffer这里的self.pe只是一个挂在模块上的普通 Python 属性。model.to(cuda)会把parameters()和buffers()逐个搬运但不会碰self.pe。于是模型在 GPU 上跑self.pe还在 CPU一相加就炸。两种修法选后者更稳# 写法一注册成 buffer随 .to() 自动迁移 self.register_buffer(pe, torch.zeros(1, 197, dim)) # 写法二干脆不存在 forward 里按输入设备现场生成 pe torch.zeros(1, x.size(1), x.size(-1), devicex.device, dtypex.dtype)还有一个更隐蔽的变体safetensors.torch.load_file()和load_model()返回的张量永远是 CPU 张量不管你当前模型在哪儿。如果你用 safetensors 加载 LoRA 权重然后直接拿去做F.linear(x, lora_up)x在 cuda、lora_up在 cpu报错必现。这一点后面第 4 章还会细讲。把这三类整理成一张对照表排查时可以按图索骥场景报错里的 argument 名最快的验证方式最小修复权重没搬全weight/bias打印各子模块设备分布对根模块调.to(device)或递归搬运dummy input 在 CPUmat2/other/input检查导出脚本里张量创建处补devicedevice野生张量 / LoRA 权重mat2/other用 forward hook 打设备集合注册 buffer 或显式.to()3. 五分钟定位法把出问题的层揪出来3.1 打印全模型的设备分布排查第一步永远是先看一眼现场。写一个十几行的审计函数比瞎猜快一百倍import torch from collections import Counter def audit_devices(model, tagmodel): stat Counter() for _, p in model.named_parameters(): stat[str(p.device)] 1 for _, b in model.named_buffers(): stat[str(b.device)] 1 print(f[{tag}] device stats: {dict(stat)}) # 把不在目标设备上的层名打印出来最多 30 条 target cuda:0 leaked [(n, str(p.device)) for n, p in model.named_parameters() if str(p.device) ! target] leaked [(n, str(b.device)) for n, b in model.named_buffers() if str(b.device) ! target] for n, d in leaked[:30]: print(f leaked: {d:10} {n}) print(f[{tag}] leaked total: {len(leaked)}) return leaked这个函数一跑如果device stats里同时出现cuda:0和cpu基本就锁定是权重没搬全了。而且leaked列表会直接把层名和它所在的设备打出来前三十条足够你看清规律——是集中在某个子模块还是散落各处。我自己的习惯是在模型拼装完之后、torch.onnx.export之前各调一次audit_devices。两次结果一对比如果第二次多出泄漏项说明中间某段代码又创建了新张量。3.2 用 forward pre-hook 抓第一现场设备分布看着没问题但一跑还是崩那问题多半出在 forward 内部动态生成的张量上。这时候named_parameters()那套就不够用了得用 hook 在每一层执行前检查。def install_device_guard(model, targetcuda:0): def make_pre(name): def pre_hook(module, args): devs {str(t.device) for t in args if torch.is_tensor(t)} devs | {str(p.device) for p in module.parameters(recurseFalse)} devs | {str(b.device) for b in module.buffers(recurseFalse)} if len(devs) 1: print(f[MISMATCH] {name} - {sorted(devs)}) return pre_hook for n, m in model.named_modules(): m.register_forward_pre_hook(make_pre(n))挂上去之后跑一次前向导出脚本里那次 trace 就够了日志里第一个[MISMATCH]就是第一现场。因为报错是按执行顺序触发的第一个不匹配的层往往就是真正的病灶后面跟着刷屏的那些多半只是它连带的受害者。用recurseFalse是有意为之只看这一层自己持有的参数和 buffer不看子模块的。否则父模块会把子模块的参数全算进去永远判定为不匹配日志就没法看了。3.3 打开 CUDA_LAUNCH_BLOCKING 和 C 栈PyTorch 的报错栈有时候会指向一个跟你写的代码完全无关的位置比如报在torch/nn/modules/module.py的_call_impl里。这是因为 CUDA 是异步执行的Python 侧早就跑到后面去了错误才在 GPU 上爆发出来。加两个环境变量能大幅改善这一点CUDA_LAUNCH_BLOCKING1 TORCH_SHOW_CPP_STACKTRACES1 python export_onnx.pyCUDA_LAUNCH_BLOCKING1强制每个 CUDA 调用同步报错会精确落在真正执行它的那一行。TORCH_SHOW_CPP_STACKTRACES1会把 C 侧的调用栈也打出来能直接看到是at::native::linear还是at::native::embedding触发的。注意CUDA_LAUNCH_BLOCKING1会明显拖慢速度只在排查阶段开定位到了立刻关掉。3.4 二分法排除按子模块单独跑如果上面三招还是没头绪就用最笨但最有效的办法——把模型拆开一个个跑。with torch.inference_mode(): # 先单独测 CLIP ids torch.randint(0, 49408, (2, 77)).cuda() clip_out sd_model.cond_stage_model(ids) print(clip ok, clip_out.device) # 再单独测 UNet x torch.randn(2, 4, 64, 64).cuda() t torch.tensor([1]).cuda() unet_out sd_model.model.diffusion_model(x, t, contextclip_out) print(unet ok, unet_out.device) # 最后测 VAE latent unet_out.cuda() img sd_model.decode_first_stage(latent) print(vae ok, img.device)哪一段先崩问题就在那一段里。这个方法的另一个好处是——它顺手帮你确认了每个子模块的接口签名对后面写 ONNX 导出脚本很有帮助。4. 从 .safetensors 到 ONNX整条链路逐段对齐设备4.1 权重加载阶段map_location 的正确写法先明确一个原则加载阶段统一落到 CPU之后显式搬到目标设备。不要图省事直接map_locationcuda。# 推荐写法 ckpt torch.load(ckpt_path, map_locationcpu) model.load_state_dict(ckpt[state_dict], strictFalse) model model.to(devicetarget_device, dtypetarget_dtype) audit_devices(model, after-load)为什么先落到 CPU 再搬三个理由。第一如果 checkpoint 是在 8 卡机器上存的、你现在只有 1 张卡直接map_locationcuda会因为设备编号对不上直接抛异常。第二CPU 加载对显存的瞬时压力更小低显存机器上不容易在加载阶段就 OOM。第三先 CPU 后搬的过程里你有机会在搬之前做一些张量级的预处理比如 dtype 转换、权重融合。半精度的顺序也有讲究先.to(device)再.half()或者干脆一次写成.to(devicedevice, dtypetorch.float16)。顺序反了的话在 CPU 上做半精度转换不仅慢某些算子还会因为 CPU 缺少 FP16 实现而报另一种错。对于 safetensors 格式加载方式略有不同但落点一样from safetensors.torch import load_file state load_file(model.safetensors, devicecpu) # 需要显式搬运 state {k: v.to(devicetarget_device, dtypetarget_dtype) for k, v in state.items()}safetensors.torch.load_file从 0.4 版本起支持device参数但默认值是cpu。这就是前面提到的第一类隐蔽坑模型已经在 GPU 上你以为加载出来的权重也会在 GPU 上实际上不会。4.2 模块拼装阶段补齐 LoRA / ControlNet / VAE 的搬运SD 转 TRT 时真正麻烦的是后挂载的组件。LoRA 是权重差分、ControlNet 是额外 UNet 分支、自定义 VAE 是替换掉first_stage_model——这三样都有各自独立的加载路径而且都容易在搬运上漏掉。LoRA 融合的典型形态是lora_up load_file(lora.safetensors)[lora_up.weight] # CPU 张量 w base_weight alpha * (lora_down lora_up) # base_weight 在 GPU炸修法很直接把加载出来的张量统一下沉到目标设备和 dtypedef to_target(t, device, dtype): return t.to(devicedevice, dtypedtype, non_blockingTrue) lora_up to_target(load_file(path)[lora_up.weight], device, dtype)non_blockingTrue在固定内存配合下能加速拷贝但前提是源张量锁页不锁页时加了也无害只是不生效。ControlNet 的情况是它本身是一个完整的网络加载完之后必须整体.to(device)同时它的 conditioning 输入比如边缘图、深度图也要跟它同设备。我见过有人把 ControlNet 权重搬了但 hint 图还是 CPU 上来的 numpy 转换结果一样炸。自定义 VAE 替换更容易被忽略因为它经常是在 webui 启动参数里指定的比如--no-half-vae、外挂的 VAE 文件替换逻辑散落在不同位置。替换完之后一定要再audit_devices一遍。4.3 导出阶段dummy input、dynamic_axes 与 dtype 的一致性torch.onnx.export阶段要同时保证三件事一致设备一致、dtype 一致、shape 一致。设备不一致就是本文这个报错dtype 不一致会报expected scalar type Half but found Floatshape 不一致会在 TRT 构建阶段报维度冲突。用torch.set_default_device可以从根上减少第一类的发生概率import torch torch.set_default_device(cuda:0)这行之后所有torch.randn、torch.zeros、torch.tensor不带device参数时都会默认建在cuda:0上。PyTorch 2.0 开始支持还有对应的上下文管理器写法with torch.device(cuda:0): sample torch.randn(2, 4, 64, 64)提醒set_default_device是个全局开关某些第三方库内部依赖默认在 CPU这个前提打开之后可能引入新的问题。用在导出脚本这种短生命周期的场景里最安全别在整个 webui 进程里乱开。一份能直接用的导出骨架大致是这样import torch unet sd_model.model.diffusion_model device next(unet.parameters()).device dtype next(unet.parameters()).dtype unet.eval() audit_devices(unet, before-export) with torch.inference_mode(): sample torch.randn(2, 4, 64, 64, devicedevice, dtypedtype) timestep torch.tensor([1], devicedevice) ctx torch.randn(2, 77, 768, devicedevice, dtypedtype) dummy (sample, timestep, ctx) torch.onnx.export( unet, dummy, unet_fp16.onnx, input_names[sample, timestep, context], output_names[out_sample], dynamic_axes{ sample: {0: batch, 2: height, 3: width}, timestep: {0: batch}, context: {0: batch}, out_sample: {0: batch, 2: height, 3: width}, }, opset_version17, do_constant_foldingTrue, )dynamic_axes里把 batch 和 H/W 标成动态是为了后面能一次编译出支持多种分辨率的引擎。但注意动态轴越多TRT builder 优化空间越小而且更容易在 shape 推导阶段触发隐藏的算子不支持。如果只是自己用先把分辨率固定成 512x512 或 768x768 编译一个版本跑通了再上动态轴这个顺序能省很多时间。4.4 SDXL 双文本编码器与多 GPU 的特殊处理SDXL 有两个文本编码器一个 CLIP-L、一个 CLIP-bigG它们的输出在通道维度拼接之后才喂给 UNet。这里有两个高频问题一是两个编码器只搬了一个二是两个编码器的 dtype 不一致一个 FP16、一个 FP32。稳妥的做法是统一处理text_encoders [sd_model.cond_stage_model, sd_model.cond_stage_model_2] for te in text_encoders: te.to(devicedevice, dtypedtype).eval()多 GPU 场景下还有一层复杂度如果你用device_map或者accelerate把模型分散到多张卡上next(model.parameters()).device拿到的只是第一张卡用它去构造 dummy input 会导致后续所有在其它卡上的层全部不匹配。这种情况下要么显式指定每段输入落哪张卡要么老老实实把整段模型收敛到单卡上再导出。转 TRT 这个场景里单卡是默认推荐——TRT engine 本身就是按单设备编译的多卡分散反而增加不确定性。5. 修好之后怎么验证不是假绿5.1 ONNX 侧onnxruntime 与 PyTorch 的数值对齐脚本不报错了、ONNX 文件也生成了不代表导出是对的。ONNX 导出最常见的失败模式是图能跑数值不对尤其是 FP16 导出时。所以必须做一次数值对齐import numpy as np import onnxruntime as ort import torch with torch.inference_mode(): ref unet(*dummy).cpu().numpy() sess ort.InferenceSession(unet_fp16.onnx, providers[CUDAExecutionProvider]) inputs { sample: dummy[0].cpu().numpy(), timestep: dummy[1].cpu().numpy(), context: dummy[2].cpu().numpy(), } out sess.run([out_sample], inputs)[0] diff np.abs(ref - out) print(max abs diff:, diff.max(), mean:, diff.mean()) np.testing.assert_allclose(ref, out, rtol1e-2, atol1e-2)容差定在1e-2是针对 FP16 的经验值。FP32 导出应该收紧到1e-4左右。如果 max diff 在 1 以上说明图里有算子被错误替换了或者某个节点被常量折叠折坏了这时候要回去检查do_constant_folding和 opset 版本。提示ONNX Runtime 和 PyTorch 的卷积实现在不同硬件上会有微小差异即使 FP32 也不可能逐位相同。看到 1e-5 级别的差异不用担心。5.2 TRT 侧engine 构建与两次推理比对引擎构建分两条路。命令行用trtexec最省事trtexec --onnxunet_fp16.onnx \ --saveEngineunet_fp16.plan \ --fp16 \ --minShapessample:1x4x64x64 \ --optShapessample:2x4x64x64 \ --maxShapessample:4x4x64x64 \ --workspace4096如果模型有多个输入--minShapes/--optShapes/--maxShapes必须把所有动态输入都列全漏一个就会报 shape 不匹配。构建完之后用 Python 侧的 TensorRT runtime 跑一次推理跟 ONNX Runtime 的输出比对这一次的容差可以放宽到2e-2——因为 FP16 的累加顺序在 TRT 里可能跟 ONNX Runtime 不一样。两次比对都通过才能说这次转换是干净的。5.3 端到端出图对比与常见精度容差最后一步还是得看图。用固定 seed、固定 prompt、固定步数分别用原始 PyTorch 路径和 TRT 引擎路径各出一张图然后做像素级对比import numpy as np from PIL import Image a np.asarray(Image.open(out_torch.png), dtypenp.float32) b np.asarray(Image.open(out_trt.png), dtypenp.float32) print(mean abs diff:, np.abs(a - b).mean())正常情况下的差异应该在个位数级别0-255 的尺度上肉眼几乎看不出区别。如果画面上出现明显的色块、噪点或者结构错乱说明某个阶段的精度损失已经越界了这时候的排查方向是先看是不是 VAE 用 FP16 导的VAE 对精度特别敏感再看 UNet 的 attention 层有没有被错误简化。我自己的习惯是固定一组验收用例——一张 512x512、一张 768x768、一张带 ControlNet 的三张图对着看。只要这三张稳定基本可以放心用了。6. 环境与版本里的隐性坑6.1 PyTorch / CUDA / TensorRT 版本组合这个报错的措辞本身就能反推版本wrapper_CUDA__是 PyTorch 2.0 的格式。如果你的环境里 PyTorch 还是 1.13看到的是另一套措辞说明你可能在照着旧教程操作新版本或者反过来。一个被反复验证过的组合是这样的组件推荐版本说明PyTorch2.1.x / 2.2.x支持set_default_deviceONNX 导出对 opset 17 支持稳定CUDA11.8 / 12.111.8 兼容性最广12.1 需要对应驱动TensorRT8.6 / 9.x8.6 生态最成熟9.x 移除了部分旧 APIONNX1.15低于 1.14 对 opset 17 的部分算子支持不全onnxruntime-gpu1.16与 CUDA 版本必须严格对应版本错配最典型的症状不是本文这个报错而是构建 engine 时提示某个算子不受支持、或者精度莫名其妙地崩。但有一种情况要注意TensorRT 9.x 在某些版本下不再支持动态 shape 的某些组合会迫使你回头改dynamic_axes改的过程中如果顺手动了输入构造代码就可能引入设备错位。所以改环境之前先把脚本备份一份。6.2 onnxruntime 与 onnxruntime-gpu 的共存冲突这是个经典陷阱pip install onnxruntime和pip install onnxruntime-gpu装的是同一个 import 名onnxruntime后装的会覆盖先装的。如果你验证时发现自己明明装了 GPU 版跑起来却是 CPU 执行用这一行确认import onnxruntime as ort print(ort.get_available_providers())输出里没有CUDAExecutionProvider就说明当前生效的是 CPU 版。清理方式是先pip uninstall onnxruntime onnxruntime-gpu -y再只装需要的那一个。这事儿跟本文的报错有间接关系如果你的导出脚本里恰好用 onnxruntime 做中间验证而它跑在 CPU 上你在 CPU 侧看到的张量当然都在 CPU——这会掩盖真实的设备问题让你在错误的方向上排查。6.3 半精度导出的边界与最小显存配置FP16 导出能省一半显存但有几个硬边界要知道。VAE 的解码器在 FP16 下容易出现 NaN 或者色偏所以很多人的做法是 UNet 和 CLIP 用 FP16、VAE 用 FP32。这样做的时候两个不同精度之间的衔接张量必须显式转换否则会从设备错位变成 dtype 错位latent unet_out.to(torch.float32) # 从 FP16 的 UNet 输出转回 FP32 喂给 VAE img vae.decode(latent)显存方面用trtexec构建 512x512 的 SD 1.5 UNet 引擎--workspace给到 4096MB 一般够用SDXL 建议直接给 8192MB。工作空间不足时 builder 会退化到更保守的策略编译出来的引擎反而更慢所以这个值别抠。还有一点--fp16只是允许用 FP16 做计算不代表所有层都会用。构建完可以用trtexec --loadEnginexx.plan --dumpLayerInfo看一下各层的精度确认关键层没有被意外降精度。7. 几个我踩过的边界情况7.1 只在某个分辨率才炸非 64 倍数有一类特别迷惑的现象512x512 导出得好好的改成 520x520 就报设备错位。这跟设备其实没关系——真正的原因是 520 不是 64 的倍数UNet 的下采样在某一层产生奇数尺寸触发了某个不支持该尺寸的分支而在那个分支里恰好有一段代码动态创建了 CPU 张量。所以遇到换个分辨率就炸的情况先别怀疑设备代码先把分辨率调回 64 的倍数验证一下。SD 1.5 常用 512、576、640、704、768SDXL 常用 1024、1152。如果你的dynamic_axes允许任意尺寸最好在输入校验里加一道硬性检查assert h % 64 0 and w % 64 0, fresolution must be multiple of 64, got {h}x{w}7.2 换了显卡之后重新出现换卡之后重装环境是常规操作但有个细节容易被忽略如果你之前生成的某些中间文件比如已经转换好的 LoRA 缓存、预处理的 embedding 缓存是用旧卡生成的里面可能存着cuda:0的张量。新机器上用torch.load加载时会尝试还原到cuda:0如果新卡的编号体系不一样比如从单卡变双卡或者 CUDA_VISIBLE_DEVICES 设过就会落到 CPU 上然后引发错位。修法是统一在加载时加map_locationcpu这一条对任何缓存文件都适用值得写进你的工具函数里。7.3 webui 启动参数带来的副作用跑在 webui 整合包里的转换脚本会继承主进程的启动参数这里面有几个会直接影响设备行为--precision full会禁用半精度跟你在脚本里写的.half()冲突导致部分层 FP32、部分层 FP16--no-half-vae会让 VAE 保持 FP32前面提到的衔接转换就必须加--medvram/--lowvram会在运行期动态搬运模块转换脚本执行到一半某些层可能被挪走了--skip-torch-cuda-test只是跳过检测不会修复任何设备问题别指望它。所以我的建议是做 TensorRT 转换时用独立的 Python 脚本不要在 webui 进程里跑。模型文件、VAE、LoRA 都从磁盘直接加载环境干净、设备状态可控、报错栈也短。整合包适合玩不适合做这种对状态一致性要求高的操作。最后分享一个我自己一直在用的小技巧把设备审计函数封装成一个装饰器加在任何导出函数头上函数进来先审计、出去再审计。看着有点笨但它帮我抓到过三四次只在特定参数组合下才出现的隐性泄漏比事后追查便宜太多了。
