[Bug已解决] 用 CPU int16 索引 CUDA 张量报错（FakeTensor 实际张量传播，DISABLED test）解决方案

[Bug已解决] 用 CPU int16 索引 CUDA 张量报错FakeTensor 实际张量传播DISABLED test解决方案一、现象长什么样PyTorch CI 里有一条DISABLED test_index_cuda_with_cpu_int16_propagate_real_tensors (__main__.PropagateRealTensorsFakeTensorTest)它测试的是当一个CUDA 张量被一个CPU 上的 int16 索引张量索引时在FakeTensor模式编译期用「假张量」推导形状 / 设备常见于torch.compile/torch.export下设备信息能否正确传播。被DISABLED说明CUDA 张量 CPU int16 索引 这个组合在 FakeTensor 的设备传播逻辑里有 bug官方暂时禁用了测试。你如果在torch.compile里写了「CUDA 张量被 CPU 索引」的代码可能遇到编译期报错或奇怪的设备断言。本文讲清楚 FakeTensor 是什么、设备传播为何会错、以及如何规避。二、FakeTensor 是什么torch.compile/torch.export在编译期并不真正跑数据而是用FakeTensor假张量来「模拟」计算它只记录形状、dtype、设备不分配真实显存。这样编译器就能在「没有真实数据」的情况下推导整张图的形状 / 设备 / 内存。PropagateRealTensors是 FakeTensor 模式的一个增强当图里混入了**真实张量real tensor有数据的**和假张量时要把真实张量的元信息设备 / 形状正确「传播」到推论结果上。三、为什么「CUDA 张量 CPU int16 索引」会出问题正常索引要求索引张量和被索引张量设备一致都在 CUDA或都在 CPU。但FakeTensor模式在推导时对「索引张量的设备」处理有 bug被索引的是 CUDA 张量 → 结果应该也在 CUDA但索引是 CPU int16 → FakeTensor 的设备传播逻辑可能「误把结果设备推导成 CPU」或反过来导致后续真实执行时设备对不上或者int16 这种「非默认索引 dtype」在 FakeTensor 下没被正确处理触发断言。于是 CI 用「结果设备应与被索引张量一致」的断言守护发现推导错就DISABLED。四、可运行复现「CUDA 张量被 CPU 索引」下面脚本演示真实的「设备不匹配索引」会怎样无 GPU 时只展示逻辑import torch def demo_cpu_index_on_cuda(): if not torch.cuda.is_available(): print(无 GPU仅说明CPU 索引去索引 CUDA 张量会报设备不匹配) return data torch.randn(10, 4, devicecuda) idx torch.tensor([0, 2, 4], dtypetorch.int16, devicecpu) # CPU int16 try: out data[idx] # CUDA 张量被 CPU 索引 print(结果设备, out.device) except Exception as e: print(直接执行报错, type(e).__name__, -, e) if __name__ __main__: demo_cpu_index_on_cuda()真实执行时PyTorch 通常会报「Expected all tensors to be on the same device」或自动把索引搬到 CUDA取决于版本。而在torch.compile的 FakeTensor 模式下这个「自动搬设备」的推导可能没做对于是编译期就错。五、解决方案一索引前把索引搬到被索引张量的设备最根本的修复保证索引张量和被索引张量设备一致。在进torch.compile之前显式把索引.to(data.device)import torch def safe_index(data, idx): # 确保索引和被索引张量同设备、同 dtype if idx.device ! data.device: idx idx.to(data.device) if idx.dtype ! torch.long: idx idx.long() # 索引通常用 long 最稳 return data[idx] # 使用 data torch.randn(10, 4, devicecuda) idx torch.tensor([0, 2, 4], dtypetorch.int16) # CPU int16 out safe_index(data, idx) print(结果设备, out.device) # cuda这样 FakeTensor 模式推导时索引和设备一致设备传播正确不会触发 bug。六、解决方案二索引用 int64long而非 int16虽然 int16 索引在部分场景可用但int64long是 PyTorch 索引的「黄金 dtype」FakeTensor 对其支持最完整。把 int16 转 longidx idx.long() # int16 - int64 out data[idx]这既规避了「int16 在 FakeTensor 下支持不全」的 bug也更符合 PyTorch 索引的通用约定。除非你刻意省内存否则索引一律用 long 最稳。七、解决方案三用 fullgraphFalse 容忍编译期设备推导问题如果不确定哪处索引踩了 FakeTensor 的设备传播 bug可以fullgraphFalse让编译器在出错处退回 eagercompiled torch.compile(model, fullgraphFalse)代价是图被切断性能略降但索引这种「小操作」走 eagerFakeTensor 不再为难它。八、解决方案四避免在 compile 图里混入 CPU 张量torch.compile的图理应尽量「纯 CUDA」或「纯 CPU」。如果你把 CPU 索引张量尤其是作为捕获的输入喂进被编译的函数FakeTensor 就要处理「跨设备传播」正是 bug 温床。做法把索引张量在图外.cuda()后再作为编译函数的输入或把索引计算也放进图里用 CUDA 上的操作生成索引。# 图外准备好同设备索引再进 compile idx_cuda idx.to(cuda).long() compiled_model torch.compile(model) out compiled_model(data, idx_cuda)九、解决方案五升级 PyTorchPropagateRealTensorsFakeTensorTest的 disabled 测试是 FakeTensor 设备传播持续完善的 Known Issue。新版本可能已修复 int16 / 跨设备索引的传播。查看版本import torch print(PyTorch, torch.__version__)十、如何判断你踩的是同一条你用了torch.compile/torch.export触发 FakeTensor 模式代码里用索引且索引是CPU 张量或int16报错 / 断言发生在「编译期」而非真实执行期关于设备不一致索引转 long 同设备后编译通过。命中即说明踩中该 disabled 测试覆盖的 FakeTensor 设备传播问题。十一、小结DISABLED test_index_cuda_with_cpu_int16_propagate_real_tensors揭示CUDA 张量被 CPU int16 索引时FakeTensor 的设备传播有 bug。应对索引前.to(data.device)保证设备一致第五节索引一律用longint64而非 int16规避 FakeTensor 支持缺口第六节不确定就fullgraphFalse让索引走 eager第七节避免在编译图里混入 CPU 张量索引先.cuda()第八节升级到修复 FakeTensor 设备传播的 PyTorch 版本。FakeTensor 是编译的「影子世界」——它只推形状和设备不碰真实数据。一旦你在影子里混入「跨设备 / 冷门 dtype」的张量影子的推导就会出错进而拖垮真实执行。索引保持「同设备 long」是让影子世界不翻车的最简法则。