[Bug] [Relax] BatchNorm inferred StructInfo disagrees with legalized output dtypes for FP16 inputs

Abierto
  • ¿Algo no está bien?
#20,901 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
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
48/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python
Área
compilers

Línea de trabajo

Start by running the supplied reproducer in both training modes and compare the declared and legalized return types. Then trace Relax BatchNorm type inference and legalization, along with the corresponding TOPI computation; done means inferred, legalized, and runtime output dtypes agree for FP16 inputs without breaking the FP32 controls.

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

Descripción

needs-triage type: bug
Expected behavior

BatchNorm's inferred Relax return types, legalized return types and runtime tensor dtypes should agree. With all five inputs FP16, Relax infers an FP16 tensor for each of the three outputs. Legalization should preserve that function signature.

Actual behavior

LegalizeOps changes the return dtypes even though the input types are unchanged. The runtime results agree with the changed legalized types:

Mode Inferred Relax outputs After LegalizeOps and at runtime
training=False float16, float16, float16 float32, float16, float16
training=True float16, float16, float16 float32, float32, float32

The default CPU compilation pipeline reproduces this without ConvertLayout or convolution. The input module passes check_well_formed. Both training modes reproduced in three fresh Python processes, each executing three legal inputs with positive moving variance.

The script checks the values against an independent NumPy formula before asserting the dtype contract. The final assertion detects the output type change. Direct FP32 controls preserve their inferred FP32 dtypes; explicitly casting the three outputs back to FP16 also preserves the requested FP16 interface.

The drift is already visible after LegalizeOps, before code generation. The corresponding TOPI BatchNorm computation returns the same promoted dtypes. This report concerns type consistency; the numerical values passing the script's checks do not make a changed return signature valid.

Environment
  • Apache TVM 0.26.dev0+source.8f328e8 (source snapshot)
  • Ubuntu 22.04.5, Linux 6.8.0-52-generic, x86-64
  • Python 3.11.16, LLVM 15.0.7
  • Target llvm; no GPU required
Steps to reproduce

Save as repro.py:

# Both run three inputs and fail the final dtype assertion:
python repro.py
python repro.py --training

# Controls pass, in either mode:
python repro.py --dtype float32
python repro.py --training --dtype float32
python repro.py --cast-output
python repro.py --training --cast-output

import argparse
import numpy as np
import tvm
from tvm import relax as R

p = argparse.ArgumentParser()
p.add_argument("--training", action="store_true")
p.add_argument("--dtype", choices=["float16", "float32"], default="float16")
p.add_argument("--cast-output", action="store_true")
args = p.parse_args()
params = [R.Var("x", R.TensorType([2, 2, 2, 2], args.dtype))]
params += [R.Var(name, R.TensorType([2], args.dtype))
           for name in ["gamma", "beta", "moving_mean", "moving_var"]]
bb = R.BlockBuilder()
with bb.function("main", params):
    with bb.dataflow():
        result = bb.emit(R.op.nn.batch_norm(
            *params, axis=1, training=args.training, epsilon=1e-5, momentum=0.1))
        if args.cast_output:
            fields = [bb.emit(R.op.astype(R.TupleGetItem(result, i), args.dtype))
                      for i in range(3)]
            result = R.Tuple(fields)
        out = bb.emit_output(result)
    bb.emit_func_output(out)
mod = bb.get()
assert R.analysis.check_well_formed(mod)
legalized = R.transform.LegalizeOps()(mod)
print("declared:", mod["main"].ret_ty, flush=True)
print("legalized:", legalized["main"].ret_ty, flush=True)
vm = R.VirtualMachine(tvm.compile(mod, target="llvm"), tvm.cpu())
observed = []
for seed in [42, 9001, 20261103]:
    rng = np.random.default_rng(seed)
    x = rng.uniform(-0.25, 0.25, (2, 2, 2, 2)).astype(args.dtype)
    arrays = [x, np.ones(2, args.dtype), np.zeros(2, args.dtype),
              np.zeros(2, args.dtype), np.ones(2, args.dtype)]
    actual = [v.numpy() for v in vm["main"](
        *[tvm.runtime.tensor(a, tvm.cpu()) for a in arrays])]
    data = x.astype("float64")
    if args.training:
        mean = data.mean(axis=(0, 2, 3))
        var = ((data - mean[None, :, None, None]) ** 2).mean(axis=(0, 2, 3))
        expected = [(data - mean[None, :, None, None]) /
                    np.sqrt(var[None, :, None, None] + 1e-5),
                    0.1 * mean, 0.9 + 0.1 * var]
    else:
        expected = [data / np.sqrt(1 + 1e-5), np.zeros(2), np.ones(2)]
    for value, reference in zip(actual, expected):
        np.testing.assert_allclose(value, reference, rtol=2e-3, atol=1e-3)
    observed.append([str(v.dtype) for v in actual])
print("three legal inputs checked; runtime dtypes:", observed, flush=True)
assert all(dtype == args.dtype for row in observed for dtype in row), (
    f"output dtype changed from {args.dtype} to {observed[0]}"
)
Triage
  • needs-triage
  • type: bug
Lenguaje dominante
Python
Estrellas
13.9k
Forks
4k
Merge medio
8 h 2 min
PR fusionados (30 d)
180

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.