[Bug][Relax][PyTorch] float64 `matmul` / `mm` / `conv` results are downcast to float32 by `from_exported_program`

Abierto
  • ¿Algo no está bien?
#20,883 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 1 día

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
3/5
Tiempo estimado
1-2 días
Aptitud para principiantes
42/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
machine-learning

Línea de trabajo

The float32 cast is likely hard-coded in the matmul, mm, addmm and conv converters of the PyTorch frontend under python/tvm/relax/frontend/torch, so start by searching there for out_dtype="float32". Reproduce with the matmul snippet in the issue and the Linear-with-bias case, which currently raises a TypeError. Done means float64 and bfloat16 inputs keep their dtype, and the bias case converts and matches eager output.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

needs-triage type: bug
Expected behavior

torch.mm on float64 inputs returns float64. The model converted with from_exported_program should keep the dtype and precision.

Actual behavior

The result is silently float32, so the precision is lost. The same happens for torch.matmul, torch.einsum, Conv1d, Conv2d, ConvTranspose2d and Linear without bias, and for bfloat16 inputs. With a float64 bias (Linear, Conv2d), the conversion fails instead:

eager: torch.float64 [[5.0000000002]]
tvm:   float32 [[5.0]]
TypeError: Binary operators must have the same datatype for both operands.

Possibly out_dtype="float32" is hard-coded in the matmul, mm, addmm and conv converters.

Environment
  • TVM: main @ 2bf0532 (0.28.dev0), built from source, target llvm
  • PyTorch: 2.12.0+cu130
  • Python: 3.12, Ubuntu x86_64
Steps to reproduce
import torch
import tvm
import tvm_tirx
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program


class M(torch.nn.Module):
    def forward(self, x):
        return torch.mm(x, x.T)


x = torch.tensor([[1.0 + 1e-10, 2.0]], dtype=torch.float64)
e = M()(x)
print("eager:", e.dtype, e.tolist())

ep = torch.export.export(M(), (x,))
mod = from_exported_program(ep)
ex = tvm.compile(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm_tirx.cpu())
out = vm["main"](tvm_tirx.runtime.tensor(x.numpy()))[0].numpy()
print("tvm:  ", out.dtype, out.tolist())

# With a bias, the conversion itself fails
m = torch.nn.Linear(2, 2, dtype=torch.float64)
ep = torch.export.export(m, (x,))
from_exported_program(ep)
Lenguaje dominante
Python
Estrellas
13.9k
Forks
4k
Merge medio
7 h 59 min
PR fusionados (30 d)
159

Preparar el entorno

Este proyecto no incluye contenedor de desarrollo, Dockerfile ni guía de contribución, así que la configuración corre por tu cuenta: empieza por su README y consulta nuestra guía para la primera contribución para los pasos generales.

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de apache/tvm

Todos los issues de apache/tvm

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.