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

Aperta
  • Qualcosa non va?
#20,901 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
48/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python
Ambito
compilers

Direzione di ricerca

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.

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

Descrizione

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
Lingua principale
Python
Stelle
13.9k
Fork
4k
Merge medio
8h 5m
PR unite (30g)
170

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.