已合并
fix(ttk): fix GEIR numpy import, add e8m0fnu dtype alias, and preload plugin modules #182
wangqi_ai创建于 16 天前
fix(ttk): fix GEIR numpy import, add e8m0fnu dtype alias, and preload plugin modules #182
已合并
共 3 个文件变更+58-30
| @@ -293,9 +293,7 @@ def configure_manual_data(sw, args, command): | |||
| 293 | sw.manual_data_dirs = directories | 293 | sw.manual_data_dirs = directories |
| 294 | return | 294 | return |
| 295 | 295 | ||
| 296 | - _configure_manual_data_prepare( | 296 | + _configure_manual_data_prepare(sw, command, directories, complete_prepare, is_prepare_dump) |
| 297 | - sw, command, directories, complete_prepare, is_prepare_dump | ||
| 298 | - ) | ||
| 299 | 297 | ||
| 300 | 298 | ||
| 301 | def _default_manual_data_dir(sw): | 299 | def _default_manual_data_dir(sw): |
| @@ -339,10 +337,11 @@ def _detect_framework_from_csv(input_files): | |||
| 339 | if not input_files: | 337 | if not input_files: |
| 340 | return "torch" | 338 | return "torch" |
| 341 | import csv | 339 | import csv |
| 340 | + | ||
| 342 | from ttk.core_modules.framework_api.framework_detector import detect_framework | 341 | from ttk.core_modules.framework_api.framework_detector import detect_framework |
| 343 | 342 | ||
| 344 | try: | 343 | try: |
| 345 | - with open(input_files[0], "r", newline="") as f: | 344 | + with open(input_files[0], newline="") as f: |
| 346 | reader = csv.reader(f) | 345 | reader = csv.reader(f) |
| 347 | header = next(reader, None) | 346 | header = next(reader, None) |
| 348 | if not header: | 347 | if not header: |
| @@ -375,6 +374,34 @@ def _detect_framework_from_csv(input_files): | |||
| 375 | return "torch" | 374 | return "torch" |
| 376 | 375 | ||
| 377 | 376 | ||
| 377 | +def _preload_plugin_modules(sw): | ||
| 378 | + """Import all .py modules under --plugin dirs before testcase validation. | ||
| 379 | + | ||
| 380 | + golden.py modules call FrameworkApiInfoKeeper().register() at import time | ||
| 381 | + to declare API param info for OpOverloadPacket APIs that can't be parsed | ||
| 382 | + via __annotations__. Without this preload, validation runs before the | ||
| 383 | + plugin is lazily loaded, causing INPUT_COUNT_EXCEEDED false failures. | ||
| 384 | + """ | ||
| 385 | + if not sw.plugin_path: | ||
| 386 | + return | ||
| 387 | + import importlib.util | ||
| 388 | + | ||
| 389 | + for plugin_dir in sw.plugin_path: | ||
| 390 | + plugin_path = pathlib.Path(plugin_dir) | ||
| 391 | + if not plugin_path.is_dir(): | ||
| 392 | + continue | ||
| 393 | + for py_file in plugin_path.glob("*.py"): | ||
| 394 | + if py_file.name.startswith("_"): | ||
| 395 | + continue | ||
| 396 | + mod_name = py_file.stem | ||
| 397 | + try: | ||
| 398 | + spec = importlib.util.spec_from_file_location(mod_name, str(py_file)) | ||
| 399 | + mod = importlib.util.module_from_spec(spec) | ||
| 400 | + spec.loader.exec_module(mod) | ||
| 401 | + except Exception as e: | ||
| 402 | + logging.debug(f"Preload plugin module {py_file} skipped: {e}") | ||
| 403 | + | ||
| 404 | + | ||
| 378 | def run_with_switches(sw): | 405 | def run_with_switches(sw): |
| 379 | from ttk.core_modules.tbe_logging import default_logging_config | 406 | from ttk.core_modules.tbe_logging import default_logging_config |
| 380 | from ttk.utilities import set_global_storage | 407 | from ttk.utilities import set_global_storage |
| @@ -390,6 +417,8 @@ def run_with_switches(sw): | |||
| 390 | 417 | ||
| 391 | logging.info(f"Command: ttk {sw.test_mode} -i {sw.input_files[0] if sw.input_files else ''}") | 418 | logging.info(f"Command: ttk {sw.test_mode} -i {sw.input_files[0] if sw.input_files else ''}") |
| 392 | 419 | ||
| 420 | + _preload_plugin_modules(sw) | ||
| 421 | + | ||
| 393 | if sw.test_mode == "framework-api": | 422 | if sw.test_mode == "framework-api": |
| 394 | from ttk.core_modules.framework_api.instance import FrameworkApiInstance | 423 | from ttk.core_modules.framework_api.instance import FrameworkApiInstance |
| 395 | 424 | ||
| @@ -22,12 +22,12 @@ from ttk.core_modules.tbe_logging import build_single_log_dir, default_logging_c | |||
| 22 | from ttk.core_modules.tbe_multiprocessing import DeviceLock, get_process_context | 22 | from ttk.core_modules.tbe_multiprocessing import DeviceLock, get_process_context |
| 23 | from ttk.utilities import ( | 23 | from ttk.utilities import ( |
| 24 | dump_to_file, | 24 | dump_to_file, |
| 25 | - get_global_storage, | ||
| 26 | get, | 25 | get, |
| 27 | - resolve_custom_numpy_dtypes, | 26 | + get_global_storage, |
| 28 | - waiting_for_memory, | ||
| 29 | - unpack_4bits, | ||
| 30 | is_4bit_dtype, | 27 | is_4bit_dtype, |
| 28 | + resolve_custom_numpy_dtypes, | ||
| 29 | + unpack_4bits, | ||
| 30 | + waiting_for_memory, | ||
| 31 | ) | 31 | ) |
| 32 | 32 | ||
| 33 | from .compiler import GeirCompiler | 33 | from .compiler import GeirCompiler |
| @@ -526,7 +526,7 @@ def _parse_single_output(buf, idx, output_dtypes, output_shapes, case_name): | |||
| 526 | is_complex32 = "complex32" in str(dtype_str) | 526 | is_complex32 = "complex32" in str(dtype_str) |
| 527 | is_4bit = is_4bit_dtype(dtype_str) | 527 | is_4bit = is_4bit_dtype(dtype_str) |
| 528 | if is_4bit: | 528 | if is_4bit: |
| 529 | - np_dtype = numpy.dtype("uint8") | 529 | + np_dtype = np.dtype("uint8") |
| 530 | else: | 530 | else: |
| 531 | np_dtype = np.dtype("float16") if is_complex32 else np.dtype(resolve_custom_numpy_dtypes([dtype_str])[0]) | 531 | np_dtype = np.dtype("float16") if is_complex32 else np.dtype(resolve_custom_numpy_dtypes([dtype_str])[0]) |
| 532 | try: | 532 | try: |
| @@ -11,11 +11,12 @@ | |||
| 11 | dtype utils | 11 | dtype utils |
| 12 | """ | 12 | """ |
| 13 | 13 | ||
| 14 | -import numpy | 14 | +import copy |
| 15 | -from typing import Union, Optional | ||
| 16 | import re | 15 | import re |
| 17 | import struct | 16 | import struct |
| 18 | -import copy | 17 | +from typing import Optional, Union |
| 18 | + | ||
| 19 | +import numpy | ||
| 19 | 20 | ||
| 20 | BFP16_NEEDS_FP32_FOR_NPY: Optional[bool] = None | 21 | BFP16_NEEDS_FP32_FOR_NPY: Optional[bool] = None |
| 21 | 22 | ||
| @@ -188,6 +189,7 @@ DATA_TYPE_DICT = { | |||
| 188 | "float8_e5m2": 35, | 189 | "float8_e5m2": 35, |
| 189 | "float8_e4m3fn": 36, | 190 | "float8_e4m3fn": 36, |
| 190 | "float8_e8m0": 37, | 191 | "float8_e8m0": 37, |
| 192 | + "float8_e8m0fnu": 37, | ||
| 191 | "float4_e2m1": 40, | 193 | "float4_e2m1": 40, |
| 192 | "float4_e1m2": 41, | 194 | "float4_e1m2": 41, |
| 193 | "hifloat4": 42, | 195 | "hifloat4": 42, |
| @@ -347,7 +349,9 @@ def numpy_int4(): | |||
| 347 | 349 | ||
| 348 | return int4 | 350 | return int4 |
| 349 | except ModuleNotFoundError: | 351 | except ModuleNotFoundError: |
| 350 | - raise RuntimeError("ml_dtypes is needed to support int4 dtype!!! Please install with `pip3 install ml-dtypes`") | 352 | + raise RuntimeError( |
| 353 | + "ml_dtypes is needed to support int4 dtype!!! Please install with `pip3 install ml-dtypes`" | ||
| 354 | + ) from None | ||
| 351 | 355 | ||
| 352 | 356 | ||
| 353 | def numpy_bfloat16(): | 357 | def numpy_bfloat16(): |
| @@ -365,7 +369,7 @@ def numpy_bfloat16(): | |||
| 365 | "ml-dtypes or tensorflow is needed to support bfloat16 dtype!!! " | 369 | "ml-dtypes or tensorflow is needed to support bfloat16 dtype!!! " |
| 366 | "Please install with `pip3 install ml-dtypes` " | 370 | "Please install with `pip3 install ml-dtypes` " |
| 367 | "or `pip3 install tensorflow`" | 371 | "or `pip3 install tensorflow`" |
| 368 | - ) | 372 | + ) from None |
| 369 | # some older TF version (v1.15.0) bfp16 needs to convert to fp32 to calculate in numpy | 373 | # some older TF version (v1.15.0) bfp16 needs to convert to fp32 to calculate in numpy |
| 370 | global BFP16_NEEDS_FP32_FOR_NPY | 374 | global BFP16_NEEDS_FP32_FOR_NPY |
| 371 | if BFP16_NEEDS_FP32_FOR_NPY is None: | 375 | if BFP16_NEEDS_FP32_FOR_NPY is None: |
| @@ -390,7 +394,7 @@ def numpy_float8_e5m2(): | |||
| 390 | except ModuleNotFoundError: | 394 | except ModuleNotFoundError: |
| 391 | raise RuntimeError( | 395 | raise RuntimeError( |
| 392 | "ml_dtypes is needed to support float8_e5m2 dtype!!! Please install with `pip3 install ml-dtypes`" | 396 | "ml_dtypes is needed to support float8_e5m2 dtype!!! Please install with `pip3 install ml-dtypes`" |
| 393 | - ) | 397 | + ) from None |
| 394 | 398 | ||
| 395 | 399 | ||
| 396 | def numpy_float8_e4m3fn(): | 400 | def numpy_float8_e4m3fn(): |
| @@ -402,7 +406,7 @@ def numpy_float8_e4m3fn(): | |||
| 402 | except ModuleNotFoundError: | 406 | except ModuleNotFoundError: |
| 403 | raise RuntimeError( | 407 | raise RuntimeError( |
| 404 | "ml_dtypes is needed to support float8_e4m3fn dtype!!! Please install with `pip3 install ml-dtypes`" | 408 | "ml_dtypes is needed to support float8_e4m3fn dtype!!! Please install with `pip3 install ml-dtypes`" |
| 405 | - ) | 409 | + ) from None |
| 406 | 410 | ||
| 407 | 411 | ||
| 408 | def ensure_en_dtypes_version(version): | 412 | def ensure_en_dtypes_version(version): |
| @@ -426,7 +430,7 @@ def numpy_float8_e8m0(): | |||
| 426 | except ModuleNotFoundError: | 430 | except ModuleNotFoundError: |
| 427 | raise RuntimeError( | 431 | raise RuntimeError( |
| 428 | "en_dtypes is needed to support float8_e8m0 dtype!!! Please install with `pip3 install en-dtypes`" | 432 | "en_dtypes is needed to support float8_e8m0 dtype!!! Please install with `pip3 install en-dtypes`" |
| 429 | - ) | 433 | + ) from None |
| 430 | 434 | ||
| 431 | 435 | ||
| 432 | def numpy_float4_e2m1(): | 436 | def numpy_float4_e2m1(): |
| @@ -439,7 +443,7 @@ def numpy_float4_e2m1(): | |||
| 439 | except ModuleNotFoundError: | 443 | except ModuleNotFoundError: |
| 440 | raise RuntimeError( | 444 | raise RuntimeError( |
| 441 | "en_dtypes is needed to support float4_e2m1 dtype!!! Please install with `pip3 install en-dtypes`" | 445 | "en_dtypes is needed to support float4_e2m1 dtype!!! Please install with `pip3 install en-dtypes`" |
| 442 | - ) | 446 | + ) from None |
| 443 | 447 | ||
| 444 | 448 | ||
| 445 | def numpy_float4_e1m2(): | 449 | def numpy_float4_e1m2(): |
| @@ -452,7 +456,7 @@ def numpy_float4_e1m2(): | |||
| 452 | except ModuleNotFoundError: | 456 | except ModuleNotFoundError: |
| 453 | raise RuntimeError( | 457 | raise RuntimeError( |
| 454 | "en_dtypes is needed to support float4_e1m2 dtype!!! Please install with `pip3 install en-dtypes`" | 458 | "en_dtypes is needed to support float4_e1m2 dtype!!! Please install with `pip3 install en-dtypes`" |
| 455 | - ) | 459 | + ) from None |
| 456 | 460 | ||
| 457 | 461 | ||
| 458 | def numpy_hifloat4(): | 462 | def numpy_hifloat4(): |
| @@ -465,7 +469,7 @@ def numpy_hifloat4(): | |||
| 465 | except ModuleNotFoundError: | 469 | except ModuleNotFoundError: |
| 466 | raise RuntimeError( | 470 | raise RuntimeError( |
| 467 | "en_dtypes is needed to support hifloat4 dtype!!! Please install with `pip3 install en-dtypes`" | 471 | "en_dtypes is needed to support hifloat4 dtype!!! Please install with `pip3 install en-dtypes`" |
| 468 | - ) | 472 | + ) from None |
| 469 | 473 | ||
| 470 | 474 | ||
| 471 | def IsRoundOne(sign, man, truncLen): | 475 | def IsRoundOne(sign, man, truncLen): |
| @@ -716,7 +720,6 @@ def trans_np_fp4_e1m2_tensor_to_bfloat16(in_tensor): | |||
| 716 | 720 | ||
| 717 | shape_tensor = in_tensor.shape | 721 | shape_tensor = in_tensor.shape |
| 718 | multi_shape = np.prod(shape_tensor) | 722 | multi_shape = np.prod(shape_tensor) |
| 719 | - out_tensor = np.zeros(multi_shape) | ||
| 720 | in_tensor = in_tensor.reshape(multi_shape) | 723 | in_tensor = in_tensor.reshape(multi_shape) |
| 721 | 724 | ||
| 722 | # 1个uint8包含两个fp4, 先拆成两个uint8 | 725 | # 1个uint8包含两个fp4, 先拆成两个uint8 |
| @@ -739,7 +742,6 @@ def trans_np_fp4_e2m1_tensor_to_bfloat16(in_tensor): | |||
| 739 | 742 | ||
| 740 | shape_tensor = in_tensor.shape | 743 | shape_tensor = in_tensor.shape |
| 741 | multi_shape = np.prod(shape_tensor) | 744 | multi_shape = np.prod(shape_tensor) |
| 742 | - out_tensor = np.zeros(multi_shape) | ||
| 743 | in_tensor = in_tensor.reshape(multi_shape) | 745 | in_tensor = in_tensor.reshape(multi_shape) |
| 744 | 746 | ||
| 745 | # 1个uint8包含两个fp4, 先拆成两个uint8 | 747 | # 1个uint8包含两个fp4, 先拆成两个uint8 |
| @@ -767,12 +769,12 @@ def numpy_hifloat8(): | |||
| 767 | except ModuleNotFoundError: | 769 | except ModuleNotFoundError: |
| 768 | raise RuntimeError( | 770 | raise RuntimeError( |
| 769 | "en_dtypes is needed to support hifloat8 dtype!!! Please install with `pip3 install en-dtypes`" | 771 | "en_dtypes is needed to support hifloat8 dtype!!! Please install with `pip3 install en-dtypes`" |
| 770 | - ) | 772 | + ) from None |
| 771 | except ImportError: | 773 | except ImportError: |
| 772 | raise RuntimeError( | 774 | raise RuntimeError( |
| 773 | "Please upgrade en_dtypes to v0.0.4 at least to support hifloat8 dtype!!! " | 775 | "Please upgrade en_dtypes to v0.0.4 at least to support hifloat8 dtype!!! " |
| 774 | "Command is `pip3 install --upgrade en-dtypes`" | 776 | "Command is `pip3 install --upgrade en-dtypes`" |
| 775 | - ) | 777 | + ) from None |
| 776 | 778 | ||
| 777 | 779 | ||
| 778 | def resolve_custom_numpy_dtypes(container): | 780 | def resolve_custom_numpy_dtypes(container): |
| @@ -813,7 +815,7 @@ def pack_4bits(src: numpy.ndarray): | |||
| 813 | to be continuous bytes stored in uint8 | 815 | to be continuous bytes stored in uint8 |
| 814 | """ | 816 | """ |
| 815 | if not is_4bit_dtype(src.dtype): | 817 | if not is_4bit_dtype(src.dtype): |
| 816 | - raise RuntimeError(f"Dtype of source tensor only support int4/float4") | 818 | + raise RuntimeError("Dtype of source tensor only support int4/float4") |
| 817 | pack_size = 2 | 819 | pack_size = 2 |
| 818 | shift = numpy.array([0, 4], dtype=numpy.uint8) | 820 | shift = numpy.array([0, 4], dtype=numpy.uint8) |
| 819 | array = src | 821 | array = src |
| @@ -997,10 +999,7 @@ def _mx_calculate_share_exp(fp_array: numpy.ndarray, scale_axis: int, mx_ele_dty | |||
| 997 | def _mx_calculate_share_exp_nv(fp_array: numpy.ndarray, scale_axis: int, mx_ele_dtype: str): | 999 | def _mx_calculate_share_exp_nv(fp_array: numpy.ndarray, scale_axis: int, mx_ele_dtype: str): |
| 998 | import numpy | 1000 | import numpy |
| 999 | 1001 | ||
| 1000 | - FP32_EXPONENT_BIAS = 127 | ||
| 1001 | - FP32_MIN_NORMAL = 2 ** (-FP32_EXPONENT_BIAS + 1) | ||
| 1002 | max_norm = get_dtype_range(mx_ele_dtype)[1] | 1002 | max_norm = get_dtype_range(mx_ele_dtype)[1] |
| 1003 | - ele_emax = int(numpy.log2(max_norm)) | ||
| 1004 | fp_abs_max = numpy.max(numpy.abs(fp_array), axis=scale_axis, keepdims=True).astype(numpy.float32) | 1003 | fp_abs_max = numpy.max(numpy.abs(fp_array), axis=scale_axis, keepdims=True).astype(numpy.float32) |
| 1005 | s_fp32 = fp_abs_max / max_norm | 1004 | s_fp32 = fp_abs_max / max_norm |
| 1006 | binary_ints = numpy.array(s_fp32.view(numpy.uint32)) | 1005 | binary_ints = numpy.array(s_fp32.view(numpy.uint32)) |
| @@ -1387,14 +1386,14 @@ def grouped_mx_quantize( | |||
| 1387 | return numpy.all(diff >= 0) | 1386 | return numpy.all(diff >= 0) |
| 1388 | 1387 | ||
| 1389 | if not is_non_reverse_order(group_index): | 1388 | if not is_non_reverse_order(group_index): |
| 1390 | - raise RuntimeError(f"Input tensor group_index should be non-reverse order.") | 1389 | + raise RuntimeError("Input tensor group_index should be non-reverse order.") |
| 1391 | 1390 | ||
| 1392 | axis = len(fp_array.shape) + axis if axis < 0 else axis | 1391 | axis = len(fp_array.shape) + axis if axis < 0 else axis |
| 1393 | if axis != -2 and axis != 0: | 1392 | if axis != -2 and axis != 0: |
| 1394 | raise RuntimeError(f"Not support {axis} yet!") | 1393 | raise RuntimeError(f"Not support {axis} yet!") |
| 1395 | 1394 | ||
| 1396 | if group_index[-1] != fp_array.shape[axis]: | 1395 | if group_index[-1] != fp_array.shape[axis]: |
| 1397 | - raise RuntimeError(f"The last element of group_index should match the dimension size of the input x axis.") | 1396 | + raise RuntimeError("The last element of group_index should match the dimension size of the input x axis.") |
| 1398 | 1397 | ||
| 1399 | # padding & reshape to block_size | 1398 | # padding & reshape to block_size |
| 1400 | fp_array, padded_group_index, padded_shape = _grouped_mx_reshape_to_blocks(fp_array, group_index, axis, block_size) | 1399 | fp_array, padded_group_index, padded_shape = _grouped_mx_reshape_to_blocks(fp_array, group_index, axis, block_size) |