[Bug] FoldBatchnormToConv2D turns a finite float32 output into NaN

Abierto Apto para principiantes
  • ¿Algo no está bien?
#20,584 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
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
85/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python
Área
compilers

Línea de trabajo

El error está en python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py; empieza por la lógica que calcula wt = scale / sqrt(variance + epsilon) y lo multiplica por los pesos de convolución. Ajusta el escalado para que la multiplicación se realice con mayor precisión (por ejemplo, float64) y así evitar el desbordamiento de float32; después, vuelve a convertir el resultado a float32. Ejecuta el script de reproducción proporcionado para confirmar que la salida plegada sigue siendo finita y coincide con la línea base anterior al plegado.

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

Descripción

needs-triage type: bug

Thanks for participating in the TVM community! We use https://discuss.tvm.apache.org/ for any general usage questions and discussions. The issue tracker is used for actionable items such as feature proposals discussion, roadmaps, and bug tracking. You are always welcomed to post on the forum first 😸

Issues that are inactive for a period of time may get closed. We adopt this policy so that we won't lose track of actionable issues that may fall at the bottom of the pile. Feel free to reopen a new one if you feel there is an additional problem that needs attention when an old one gets closed.

Expected behavior

For this valid inference graph, folding BatchNorm into Conv2D should preserve a finite output. The unfused graph returns 50331648.0 on the input below. This report concerns the finite-to-NaN change, not exact bitwise equality under floating-point reassociation.

Actual behavior

With the same input and llvm target, FoldBatchnormToConv2D changes the IR and the result becomes NaN:

  • Input after float32 rounding: 9999999827968.0.
  • Each original Conv2D product magnitude: 9.999999655936003e+25, below the float32 maximum 3.4028234663852886e+38.
  • Before folding: 50331648.0 (finite); after folding: NaN.
  • Mathematical result evaluated in float64: about 49999624.14.

The source rewrite computes wt = scale / sqrt(variance + epsilon), then wt_conv = conv_weight * wt and executes Conv2D with wt_conv ([pass source](https://github.com/apache/tvm/blob/main/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py)). Here wt is about 1e13. The folded weights remain finite at about 1e26, but their products with the 1e13 input are about 1e39, beyond the float32 limit. Opposite-signed terms can therefore become +Inf and -Inf, and their sum becomes NaN. Before folding, the two Conv2D products are about +1e26 and -1e26 (both representable); they cancel before BatchNorm applies its scale.

Only the folding pass causes this discrepancy in a separate pass isolation check: parameter binding alone and layout conversion alone both preserve the baseline output. This is a numerical regression caused by finite-precision overflow after an algebraically valid transformation. It requires unusually large values; this report does not claim that ordinary inference inputs are affected.

Environment
  • TVM 0.26.dev0, locally built from commit 5d4dc8721307cf9677becacf0395c19be130e206.
  • Python 3.10.20; Linux x86_64 (WSL2, kernel 6.6.114.1-microsoft-standard-WSL2).
  • float32, Relax VM, CPU target llvm; Conv2D NCHW / OIHW; inference BatchNorm with axis=1, center=True, scale=True, training=False.
  • Reproduced in two separate fresh processes. The executable reproduction was run on the local commit above; current upstream source was inspected, but a current-main binary was not run.
Steps to reproduce

Run this script with a TVM source build that provides Relax and LLVM:

import numpy as np
import tvm
import tvm_ffi
from tvm import relax
from tvm.script import relax as R

v = np.float32(1e13)
x_data = np.full((1, 2, 1, 1), v, dtype="float32")
weights = np.zeros((1, 2, 3, 3), dtype="float32")
weights[0, 0, 1, 1] = v
weights[0, 1, 1, 1] = -v
params = {
    "weight": weights,
    "scale": np.array([v], dtype="float32"),
    "bias": np.array([v], dtype="float32"),
    "mean": np.array([1.0], dtype="float32"),
    "variance": np.array([1.0], dtype="float32"),
}

@tvm.script.ir_module
class Case:
    @R.function
    def main(x: R.Tensor((1, 2, 1, 1), "float32"),
             weight: R.Tensor((1, 2, 3, 3), "float32"),
             scale: R.Tensor((1,), "float32"),
             bias: R.Tensor((1,), "float32"),
             mean: R.Tensor((1,), "float32"),
             variance: R.Tensor((1,), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
        conv = R.nn.conv2d(x, weight, strides=(1, 1), padding=1,
                           dilation=(1, 1), groups=1,
                           data_layout="NCHW", kernel_layout="OIHW")
        bn: R.Tuple(R.Tensor((1, 1, 1, 1), "float32"),
                    R.Tensor((1,), "float32"),
                    R.Tensor((1,), "float32")) = R.nn.batch_norm(
            conv, scale, bias, mean, variance, axis=1, epsilon=1e-5,
            center=True, scale=True, training=False)
        return R.TupleGetItem(bn, 0)

def run(mod):
    exe = relax.build(relax.pipeline.zero_pipeline()(mod), target="llvm")
    vm = relax.VirtualMachine(exe, tvm.cpu())
    return vm["main"](tvm.runtime.tensor(x_data)).numpy().item()

base = relax.transform.ConvertToDataflow(min_size=1)(Case)
bound = relax.transform.BindParams(
    "main", {name: tvm.runtime.tensor(value) for name, value in params.items()}
)(base)
folded = relax.transform.FoldBatchnormToConv2D()(bound)
print("IR changed:", not tvm_ffi.structural_equal(bound, folded))
print("before:", run(bound))
print("after:", run(folded))

Observed in both fresh processes:

IR changed: True
before: 50331648.0
after: nan

The operands and BatchNorm parameters are finite, variance=1 is valid, and the original Conv2D products remain within float32 range. The only difference between the two compiled modules is the explicit FoldBatchnormToConv2D transform; both use the same build pipeline and target.

Related reports [#20530](https://github.com/apache/tvm/issues/20530)–[[#20534](https://github.com/apache/tvm/issues/20534)](https://github.com/apache/tvm/issues/20534) cover disabled scale/center, non-channel axis, training mode, float16 constants, and IOHW layout respectively; the associated open PRs [#20535](https://github.com/apache/tvm/pull/20535)–[[#20537](https://github.com/apache/tvm/pull/20537)](https://github.com/apache/tvm/pull/20537) address those cases. None of these eight items reports this float32 finite-to-NaN overflow case. This is based on the issues and PRs inspected on 2026-10-07, not a claim that every upstream report has been ruled out.

Triage

Please refer to the list of label tags here to find the relevant tags and add them below in a bullet format (example below).

  • 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.