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

Đang mở
  • Có gì không ổn?
#20,901 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
48/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python
Lĩnh vực
compilers

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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
Ngôn ngữ chính
Python
Star
13.9k
Fork
4k
Merge trung bình
8 giờ 2 phút
Pull request đã merge (30 ngày)
180

Chuẩn bị môi trường

Dự án này không cung cấp dev container, Dockerfile hay hướng dẫn đóng góp, nên bạn cần tự thiết lập môi trường: hãy bắt đầu từ README và xem hướng dẫn đóng góp lần đầu của chúng tôi để biết các bước chung.

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của apache/tvm

Tất cả issue của apache/tvm

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.