【Bug已解决】[RT-DETRv2] MPS crash build_2d_sinusoidal_position_embedding hardcodes torch.float64, breaking Apple Silicon MPS inference 解决方案一、现象长什么样在 Apple SiliconM 系列芯片上用 MPS 后端跑 RT-DETRv2 推理初始化位置编码时直接崩# 现象 AMPS 不支持 float64 算子 RuntimeError: Input type (double) and weight type (float) should be the same # 或 RuntimeError: MPS backend does not support float64 (double) for this operation. # 现象 Bsin/cos 在 float64 张量上MPS 报不支持 # 崩溃点精确指向 build_2d_sinusoidal_position_embedding 里的 # position_embedding torch.cat([sin, cos], dim-1).to(torch.float64) # 典型触发 import torch device torch.device(mps) model RTDetrForObjectDetection.from_pretrained(rtdetr2).to(device) out model(pixel_values.to(device)) # 在构建 2D 正弦位置编码时崩溃最典型的指纹同一份代码在 CUDACPU上跑得好好的一换 MPS 就崩在位置编码——因为 MPS 对float64的支持非常有限多数算子只认 float32/bfloat16而 RT-DETRv2 把 2D 正弦位置编码硬编码成了torch.float64。二、背景RT-DETRv2 的build_2d_sinusoidal_position_embedding用来生成二维正弦位置编码因为检测任务里特征图是 2D 的需要把 (x, y) 位置编码进去。常见实现里为了 sin/cos 计算精度会把中间张量.to(torch.float64)或在arange时用dtypetorch.float64# 有 bug 的写法 grid_y, grid_x torch.meshgrid(arange(H), arange(W)) pos_x sin(grid_x * freqs).to(torch.float64) # 硬编码 float64 pos_y cos(grid_y * freqs).to(torch.float64) pos torch.cat([pos_x, pos_y], dim-1) # float64 张量CUDA 和 CPU 都良好支持 float64所以这套代码在训练机/NVIDIA 卡上从不出问题。但Apple Silicon 的 MPS 后端对 float64 支持残缺很多逐元素/规约/卷积算子不接受 double 输入于是位置编码这张 float64 张量一旦参与后续计算或与 float32 权重做运算MPS 直接RuntimeError。三、根因根因有两类位置编码硬编码 float64MPS 不支持。build_2d_sinusoidal_position_embedding内部把meshgrid/sin/cos结果转成torch.float64产出的pos是 double 张量。MPS 下这个 double 张量要么在生成 sin/cos 时就不被支持部分 MPS 算子不支持 double要么在和 float32 的 backbone 特征相加时类型不匹配 → 崩溃。float64 与模型其余部分的 dtype 不一致。 即便 MPS 勉强生成了 float64 位置编码后续pos featurefeature 是 float32会触发 Input type (double) and weight type (float) should be the same——类型不一致直接报错。本质位置编码用 float64 是为了精度但 MPS 不支持 double且和 float32 主干类型冲突。正确做法是全程用 float32 计算正弦位置编码在 float32 下精度完全够用只在必要时临时提升精度再降回。四、最小可运行复现下面用纯 Python 模拟float64 位置编码在 MPS 上与 float32 特征类型冲突from dataclasses import dataclass dataclass class FakeTensor: dtype: str # float32 / float64 def mps_add(pos, feature): MPS 下 float64 与 float32 不能相加。 if pos.dtype ! feature.dtype: raise RuntimeError( fInput type ({pos.dtype}) and weight type ({feature.dtype}) should be the same) return FakeTensor(pos.dtype) # 有 bug位置编码是 float64特征是 float32 pos FakeTensor(float64) feature FakeTensor(float32) try: mps_add(pos, feature) print(复现失败) except RuntimeError as e: print(复现成功(根因):, e) # 修正位置编码改用 float32 计算 pos_fixed FakeTensor(float32) print(修正后:, mps_add(pos_fixed, feature).dtype) # float32正常运行后float64 位置编码与 float32 特征相加触发类型冲突修正为 float32 后正常复现并修复了根因。五、解决方案第一层最小直接修复最快的止血把build_2d_sinusoidal_position_embedding里所有 float64 改成 float32或算完立即降回模型 dtypeimport torch def build_2d_sinusoidal_position_embedding(H, W, dim, device, dtypetorch.float32): 第一层修复全程 float32不再硬编码 float64。 # 1) meshgrid 用 float32不再 float64 grid_y, grid_x torch.meshgrid( torch.arange(H, dtypedtype, devicedevice), torch.arange(W, dtypedtype, devicedevice), indexingij, ) # 2) 频率向量用同样 dtype freqs torch.arange(0, dim, 2, dtypedtype, devicedevice) / dim freqs 1.0 / (10000 ** freqs) # float32 pos_x torch.sin(grid_x[..., None] * freqs) pos_y torch.sin(grid_y[..., None] * freqs) pos_x_c torch.cos(grid_x[..., None] * freqs) pos_y_c torch.cos(grid_y[..., None] * freqs) # 3) 拼成 (H, W, dim*2) 并转 (dim*2, H, W) pos torch.cat([pos_y, pos_y_c, pos_x, pos_x_c], dim-1) pos pos.permute(2, 0, 1).unsqueeze(0).to(dtype) # 保持 float32 return pos # 使用在模型 forward 里 device torch.device(mps) pos_embed build_2d_sinusoidal_position_embedding( H32, W32, dim128, devicedevice, dtypetorch.float32) # 与 float32 的 backbone 特征相加不再类型冲突第一层让用户立刻在 MPS 上跑通 RT-DETRv2位置编码全程 float32与主干一致。六、解决方案第二层结构性改进用PositionEmbedding2D把设备/ dtype 自适应做进位置编码生成器自动规避 MPS 的 float64from dataclasses import dataclass from typing import Optional dataclass class PositionEmbedding2D: 2D 正弦位置编码自动按设备选择安全 dtypeMPS 强制 float32。 dim: int 256 temperature: float 10000.0 def _safe_dtype(self, device, requested): # MPS 不支持 float64强制降到 float32 if str(device).startswith(mps) and requested torch.float64: return torch.float32 return requested def forward(self, H, W, device, dtypetorch.float32): dtype self._safe_dtype(device, dtype) grid_y, grid_x torch.meshgrid( torch.arange(H, dtypedtype, devicedevice), torch.arange(W, dtypedtype, devicedevice), indexingij, ) freqs torch.arange(0, self.dim, 2, dtypedtype, devicedevice) / self.dim freqs 1.0 / (self.temperature ** freqs) sin_y torch.sin(grid_y[..., None] * freqs) cos_y torch.cos(grid_y[..., None] * freqs) sin_x torch.sin(grid_x[..., None] * freqs) cos_x torch.cos(grid_x[..., None] * freqs) pos torch.cat([sin_y, cos_y, sin_x, cos_x], dim-1) return pos.permute(2, 0, 1).unsqueeze(0).to(dtype) # 使用模型里取代硬编码 float64 的实现 pe PositionEmbedding2D(dim128) pos pe.forward(32, 32, devicetorch.device(mps)) # 自动 float32PositionEmbedding2D的语义是位置编码的 dtype 跟着设备走MPS 上自动避开 float64从结构上根治 MPS crash。七、解决方案第三层断言 / CI 守护用 pytest 固化MPS 上位置编码为 float32、且与主干 dtype 一致import pytest import torch def test_mps_forces_float32(): from pos_emb import PositionEmbedding2D pe PositionEmbedding2D(dim128) # 即使请求 float64MPS 也应强制 float32 pos pe.forward(8, 8, devicetorch.device(mps), dtypetorch.float64) assert pos.dtype torch.float32, MPS 上位置编码必须是 float32 def test_cpu_allows_float64(): from pos_emb import PositionEmbedding2D pe PositionEmbedding2D(dim128) pos pe.forward(8, 8, devicetorch.device(cpu), dtypetorch.float64) assert pos.dtype torch.float64, CPU 上允许 float64 def test_no_dtype_mismatch_with_backbone(): from pos_emb import PositionEmbedding2D pe PositionEmbedding2D(dim128) pos pe.forward(8, 8, torch.device(mps), torch.float32) feature torch.randn(1, 128, 8, 8, dtypetorch.float32, devicemps) # 相加不应类型冲突 out pos feature assert out.dtype torch.float32CI 跑pytest tests/test_rtdetrv2_mps.py若 CI 无 MPS 设备可用按设备名模拟的单元测试覆盖_safe_dtype分支以后只要有人又把位置编码写死 float64测试立刻红灯。八、排查清单当 RT-DETRv2 在 MPS 上崩在位置编码按顺序查报错含float64/double→build_2d_sinusoidal_position_embedding硬编码了 float64改 float32。Input type (double) and weight type (float)→ 位置编码与主干 dtype 不一致统一 float32。CUDA/CPU 正常、MPS 崩 → 确认是 float64 不被 MPS 支持用PositionEmbedding2D按设备选 dtype。正弦位置编码在 float32 下精度足够没必要用 float64。长期方案位置编码生成器内置MPS 强制 float32的设备自适应而非写死精度。九、小结[RT-DETRv2] MPS crash: build_2d_sinusoidal_position_embedding hardcodes torch.float64 的根因是2D 正弦位置编码被硬编码成 float64而 Apple Silicon 的 MPS 后端对 double 支持残缺且 float64 位置编码与 float32 主干类型冲突于是 CUDA/CPU 正常、MPS 崩溃。第一层把位置编码全程改 float32meshgrid/sin/cos 同 dtype立刻在 MPS 跑通。第二层用PositionEmbedding2D按设备自动选 dtypeMPS 强制 float32结构性规避。第三层pytest 断言MPS 上位置编码为 float32、与主干一致、CPU 允许 float64防止回归。记住MPS 后端基本不支持 float64任何位置编码/归一化等中间计算都应在 float32 下完成需要高精度时临时提升再降回绝不能把 double 张量直接送进 MPS 算子或与 float32 主干混合。