[Bug][Relax] FoldBatchnormToConv2D crashes on float16 BatchNorm constants

Aperta Adatta ai principianti
  • Qualcosa non va?
#20,533 1 commento 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

@Nanmur ci sta già lavorando.

Dal 3/10/2026.

  • #20537 di @Nanmur — aperta

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
86/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python
Ambito
compilers, machine-learning

Direzione di ricerca

Esegui prima la riproduzione float16 fornita, quindi esamina fold_batch_norm_to_conv2d_for_inference.py, dove viene creata la costante epsilon e aggiunta alla varianza di BatchNorm. Il lavoro è completato quando il grafo float16 non genera più un errore di dtype, i dtypes non supportati rimangono invariati se applicabile e il comportamento esistente di folding float32 continua ad avere esito positivo.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

Description

FoldBatchnormToConv2D constructs BatchNorm epsilon as float32 unconditionally.
Applying the pass to a well-formed, executable float16 graph therefore raises a
dtype error instead of producing a target module or leaving the graph unchanged.

Environment

  • TVM 0.26.dev0, executed at e269315c90e3a061c9e1c77b370ce883b1b223f4
  • Ubuntu 24.04 x86-64, LLVM 18.1.3 CPU target
  • Current upstream main checked at 043b57fbe9b84aba6c0525fc77b624ef8b7bbe2f
  • The relevant Python source is byte-identical at both revisions (SHA-256 c591c953ce9de0f8b539149ab466692572ae608978bbb7ca4875fe919f3c588e)
  • I did not have a binary built from current main, so the runtime reproduction below is from the executed revision

Minimal reproduction

import numpy as np
import tvm
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(
        x: R.Tensor((1, 1, 2, 2), "float16"),
        weight: R.Tensor((2, 1, 1, 1), "float16"),
        gamma: R.Tensor((2,), "float16"),
        beta: R.Tensor((2,), "float16"),
        mean: R.Tensor((2,), "float16"),
        variance: R.Tensor((2,), "float16"),
    ):
        conv = R.nn.conv2d(x, weight, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype="float16")
        return R.nn.batch_norm(conv, gamma, beta, mean, variance, axis=1, training=False)[0]

def tensor(values):
    return tvm.runtime.tensor(np.asarray(values, dtype="float16"))

params = {
    "weight": tensor([[[[1.5]]], [[[-2.0]]]]),
    "gamma": tensor([2.0, 3.0]),
    "beta": tensor([5.0, 7.0]),
    "mean": tensor([1.0, -1.0]),
    "variance": tensor([4.0, 9.0]),
}
source = relax.transform.BindParams("main", params)(Module)
assert relax.analysis.check_well_formed(source, check_ty=True)
relax.transform.FoldBatchnormToConv2D()(source)

Expected behavior

The pass should construct epsilon in the variance dtype, or leave unsupported dtypes unchanged without throwing.

Actual behavior

The pass raises:

TypeError: Binary operators must have the same datatype for both operands.
R.add(<float16 variance>, T.float32(1e-05)) uses float16 and float32.

The source is well formed and executes on LLVM. Channels 1 through 6 all reproduced the same error. Changing only the graph dtype to float32 gave six successful, well-formed folds whose VM outputs matched the sources.

Likely cause

fold_batch_norm_to_conv2d_for_inference.py unconditionally creates tirx.FloatImm("float32", bn_attrs["epsilon"]) and adds it to bn_variance.

Duplicate check

I searched open and closed TVM issues and PRs for this pass, float16, epsilon dtype, and the exception, and did not find the same cause. #19977/#19979 concern ONNX frontend importing, and #3260 concerns the older Relay stack.

I can send a fix PR and help with follow-up testing if this diagnosis looks right.

Triage

  • needs-triage
  • type: bug
Lingua principale
Python
Stelle
13.9k
Fork
4k
Merge medio
8h 2m
PR unite (30g)
180

Preparare l'ambiente

Questo progetto non fornisce container di sviluppo, Dockerfile né guida per i contributori, quindi l'ambiente è a tuo carico: parti dal suo README e consulta la nostra guida al primo contributo per i passaggi generali.

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di apache/tvm

Tutte le issue di apache/tvm

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.