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

オープン
  • 何か問題がありますか？
#20,901 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
48/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python
領域
compilers

調査の方向性

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.

索引モデルが issue の本文から書いたものです。

説明

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
主要言語
Python
スター
13.9k
フォーク
4k
平均マージ
8時間 2分
マージ済み PR（30日）
180

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

apache/tvm のほかの issue

apache/tvm の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。