diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 2aa02340018..26e6de2a9c8 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -53,6 +53,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwiset python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_current_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_current_scaling_exact.py || test_fail "test_float8_current_scaling_exact.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_tiny_amax_scale_host.xml $TE_PATH/tests/cpp/operator/test_tiny_amax_scale_host.py || test_fail "test_tiny_amax_scale_host.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH/tests/pytorch/test_gqa.py || test_fail "test_gqa.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_qk_norm.xml $TE_PATH/tests/pytorch/test_qk_norm.py || test_fail "test_qk_norm.py" diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index aa3ec536712..9cd74cfeb6b 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -67,3 +67,7 @@ if("$ENV{NVTE_TEST_DEBUG_CPP_UT_IN_PARALLEL}" STREQUAL "1") else() add_test(NAME test_operator COMMAND test_operator) endif() + +add_executable(test_tiny_amax_scale_host test_tiny_amax_scale_host.cpp) +target_compile_options(test_tiny_amax_scale_host PRIVATE -O2 -Wall -Wextra) +add_test(NAME test_tiny_amax_scale_host COMMAND test_tiny_amax_scale_host) diff --git a/tests/cpp/operator/test_cast_current_scaling.cu b/tests/cpp/operator/test_cast_current_scaling.cu index 7cca0d72e0f..86aa79ad880 100644 --- a/tests/cpp/operator/test_cast_current_scaling.cu +++ b/tests/cpp/operator/test_cast_current_scaling.cu @@ -4,6 +4,9 @@ * See LICENSE for license information. ************************************************************************/ +#include +#include +#include #include #include #include @@ -58,12 +61,17 @@ void compute_amax_scale_ref(const InputType *data, float scale = 1.f; float scale_inv = 1.f; - if (isinf(clamp_amax) || clamp_amax == 0.f) { + if (isinf(clamp_amax) || isnan(clamp_amax) || clamp_amax == 0.f) { *scale_ptr = scale; *scale_inv_ptr = scale_inv; return; } + const float amax_floor = max_fp8 * FLT_MIN; + if (clamp_amax < amax_floor) { + clamp_amax = amax_floor; + } + // use ieee_div in CPU scale = max_fp8 / clamp_amax; @@ -213,3 +221,65 @@ INSTANTIATE_TEST_SUITE_P( } return name; }); + +namespace { + +using namespace test; + +uint32_t float_bits(float value) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +// Call the device scale kernel, then quantize so scale_inv is stored. +void check_device_scale(DType otype, float amax, uint32_t expect_scale, uint32_t expect_inv, + bool expect_nonzero_code) { + const std::vector shape = {128, 128}; + const size_t n = 128 * 128; + Tensor input("tiny_input", shape, DType::kFloat32); + Tensor output("tiny_output", shape, otype, true, false); + float *data = input.rowwise_cpu_dptr(); + for (size_t i = 0; i < n; ++i) { + data[i] = amax; + } + input.from_cpu(); + output.set_amax(amax); + + QuantizationConfigWrapper config; + nvte_compute_scale_from_amax(output.data(), config, 0); + nvte_quantize(input.data(), output.data(), 0); + cudaDeviceSynchronize(); + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError()); + EXPECT_EQ(float_bits(output.scale()), expect_scale); + EXPECT_EQ(float_bits(output.rowwise_scale_inv()), expect_inv); + + output.to_cpu(); + bool nonzero = false; + if (otype == DType::kFloat8E4M3) { + const auto *codes = output.rowwise_cpu_dptr(); + for (size_t i = 0; i < n; ++i) { + uint8_t raw = 0; + std::memcpy(&raw, codes + i, sizeof(raw)); + nonzero = nonzero || raw != 0; + } + } else { + const auto *codes = output.rowwise_cpu_dptr(); + for (size_t i = 0; i < n; ++i) { + uint8_t raw = 0; + std::memcpy(&raw, codes + i, sizeof(raw)); + nonzero = nonzero || raw != 0; + } + } + EXPECT_EQ(nonzero, expect_nonzero_code); +} + +} // namespace + +TEST(CastCSTinyAmax, DeviceScaleInvIsFltMin) { + const float tiny = 7.7e-37f; + for (DType otype : {DType::kFloat8E4M3, DType::kFloat8E5M2}) { + check_device_scale(otype, tiny, 0x7e800000u, 0x00800000u, true); + check_device_scale(otype, 0.f, 0x3f800000u, 0x3f800000u, false); + } +} diff --git a/tests/cpp/operator/test_cast_current_scaling_grouped.cu b/tests/cpp/operator/test_cast_current_scaling_grouped.cu index 3076a9e0bc4..7dad10d6919 100644 --- a/tests/cpp/operator/test_cast_current_scaling_grouped.cu +++ b/tests/cpp/operator/test_cast_current_scaling_grouped.cu @@ -5,6 +5,7 @@ ************************************************************************/ #include +#include #include #include #include @@ -57,12 +58,17 @@ void compute_amax_scale_ref(const InputType *data, const size_t size, float *ama float scale = 1.0f; float scale_inv = 1.0f; - if (std::isinf(clamp_amax) || clamp_amax == 0.0f || std::isnan(clamp_amax)) { + if (std::isinf(clamp_amax) || std::isnan(clamp_amax) || clamp_amax == 0.0f) { *scale_ptr = scale; *scale_inv_ptr = scale_inv; return; } + const float amax_floor = max_fp8 * FLT_MIN; + if (clamp_amax < amax_floor) { + clamp_amax = amax_floor; + } + scale = max_fp8 / clamp_amax; if (std::isinf(scale)) { scale = std::numeric_limits::max(); diff --git a/tests/cpp/operator/test_cast_float8blockwise.cu b/tests/cpp/operator/test_cast_float8blockwise.cu index d50589ca437..6f726a65b00 100644 --- a/tests/cpp/operator/test_cast_float8blockwise.cu +++ b/tests/cpp/operator/test_cast_float8blockwise.cu @@ -4,6 +4,7 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include #include @@ -51,15 +52,22 @@ void scales_from_amax(float amax, const QuantizationOptions& opts, float* qscale float quant_type_max_val = Quantized_Limits::max(); float eps = opts.amax_epsilon; amax = std::max(amax, eps); - float qscale = quant_type_max_val / amax; - if (std::isinf(qscale)) { - qscale = input_type_max_val; - } - if (std::isnan(qscale) || amax == 0) { - qscale = 1.0; + const float guard = quant_type_max_val * FLT_MIN; + float qscale = 1.0f; + if (!(std::isinf(amax) || std::isnan(amax) || amax == 0.0f)) { + if (amax < guard) { + amax = guard; + } + qscale = quant_type_max_val / amax; + if (std::isinf(qscale)) { + qscale = input_type_max_val; + } + if (std::isnan(qscale)) { + qscale = 1.0f; + } } - if (opts.force_pow_2_scales && qscale != 0.0) { + if (opts.force_pow_2_scales && qscale != 0.0f) { uint32_t scale_bits = *reinterpret_cast(&qscale); // Scale must be positive, shift it uint8_t exp = scale_bits >> 23; diff --git a/tests/cpp/operator/test_cast_transpose_current_scaling.cu b/tests/cpp/operator/test_cast_transpose_current_scaling.cu index 225d24317a3..75335727d90 100644 --- a/tests/cpp/operator/test_cast_transpose_current_scaling.cu +++ b/tests/cpp/operator/test_cast_transpose_current_scaling.cu @@ -4,6 +4,8 @@ * See LICENSE for license information. ************************************************************************/ +#include +#include #include #include #include @@ -62,12 +64,17 @@ void compute_amax_scale_ref(const InputType *data, float scale = 1.f; float scale_inv = 1.f; - if (isinf(clamp_amax) || clamp_amax == 0.f) { + if (isinf(clamp_amax) || isnan(clamp_amax) || clamp_amax == 0.f) { *scale_ptr = scale; *scale_inv_ptr = scale_inv; return; } + const float amax_floor = max_fp8 * FLT_MIN; + if (clamp_amax < amax_floor) { + clamp_amax = amax_floor; + } + // use ieee_div in CPU scale = max_fp8 / clamp_amax; diff --git a/tests/cpp/operator/test_tiny_amax_scale_host.cpp b/tests/cpp/operator/test_tiny_amax_scale_host.cpp new file mode 100644 index 00000000000..ce820c91258 --- /dev/null +++ b/tests/cpp/operator/test_tiny_amax_scale_host.cpp @@ -0,0 +1,179 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Host IEEE-754 model of compute_scale_from_amax without fast-math. + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +float bits_to_f32(uint32_t bits) { + float value; + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +uint32_t f32_to_bits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +bool is_subnormal(float value) { + const uint32_t bits = f32_to_bits(value); + return (bits & 0x7f800000u) == 0u && (bits & 0x007fffffu) != 0u; +} + +bool is_finite_normal(float value) { + return std::isfinite(value) && !is_subnormal(value) && value != 0.f; +} + +struct Scales { + float scale; + float scale_inv; +}; + +Scales compute_scale(float amax, float max_fp8, bool force_pow_2, bool floor_tiny, + float epsilon) { + if (amax < epsilon) { + amax = epsilon; + } + float scale = 1.f; + const bool early = std::isinf(amax) || std::isnan(amax) || amax == 0.f; + if (!early) { + if (floor_tiny) { + const float amax_floor = max_fp8 * FLT_MIN; + if (amax < amax_floor) { + amax = amax_floor; + } + } + scale = max_fp8 / amax; + if (std::isinf(scale)) { + scale = std::numeric_limits::max(); + } + if (force_pow_2) { + uint32_t scale_bits = f32_to_bits(scale); + scale_bits &= 0xFF800000u; + scale = bits_to_f32(scale_bits); + } + } + Scales out; + out.scale = scale; + out.scale_inv = 1.0f / scale; + return out; +} + +void require(bool condition, const std::string &message) { + if (!condition) { + std::fprintf(stderr, "FAIL: %s\n", message.c_str()); + std::exit(1); + } +} + +void emit(const char *fmt, int pow2, const char *label, float amax, const Scales &fixed, + const Scales &legacy) { + std::printf("%s %d %s %08x %08x %08x %08x %08x\n", fmt, pow2, label, f32_to_bits(amax), + f32_to_bits(fixed.scale), f32_to_bits(fixed.scale_inv), f32_to_bits(legacy.scale), + f32_to_bits(legacy.scale_inv)); +} + +void check_format(const char *name, float max_fp8, uint32_t expected_threshold_bits) { + const float threshold = max_fp8 * FLT_MIN; + require(f32_to_bits(threshold) == expected_threshold_bits, std::string(name) + " threshold bits"); + require(static_cast(threshold) / static_cast(max_fp8) == + static_cast(FLT_MIN), + std::string(name) + " exact quotient"); + + const float below = bits_to_f32(expected_threshold_bits - 1u); + const float above = bits_to_f32(expected_threshold_bits + 1u); + require(static_cast(below) / static_cast(max_fp8) < static_cast(FLT_MIN), + std::string(name) + " -1 ULP is mathematically below FLT_MIN"); + require(static_cast(above) / static_cast(max_fp8) > static_cast(FLT_MIN), + std::string(name) + " +1 ULP is mathematically above FLT_MIN"); + + const float tiny = 7.7e-37f; + require(tiny > 0.f && tiny < threshold, std::string(name) + " 7.7e-37 is below the guard"); + + const float samples[] = {0.f, tiny, below, threshold, above}; + const char *labels[] = {"zero", "tiny_7_7e_37", "minus_1ulp", "exact", "plus_1ulp"}; + + for (int pow2 = 0; pow2 <= 1; ++pow2) { + for (int i = 0; i < 5; ++i) { + const Scales fixed = compute_scale(samples[i], max_fp8, pow2, true, 0.f); + const Scales legacy = compute_scale(samples[i], max_fp8, pow2, false, 0.f); + emit(name, pow2, labels[i], samples[i], fixed, legacy); + + if (i == 0) { + require(fixed.scale == 1.f && fixed.scale_inv == 1.f, std::string(name) + " zero scale"); + require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv, + std::string(name) + " zero matches legacy"); + } else if (i < 3) { + // Below the floor, scale is 2^126 and scale_inv is FLT_MIN. + require(f32_to_bits(fixed.scale) == 0x7e800000u, std::string(name) + " floored scale 2^126"); + require(f32_to_bits(fixed.scale_inv) == 0x00800000u, std::string(name) + " floored inv FLT_MIN"); + require(samples[i] * fixed.scale > 1.f, std::string(name) + " tiny value stays above 1"); + // Power-of-two rounding already collapses the -1 ULP sample onto 2^126. + if (pow2 && i == 2) { + require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv, + std::string(name) + " pow2 -1 ULP matches legacy"); + } else { + require(is_subnormal(legacy.scale_inv), std::string(name) + " legacy inv subnormal"); + require(f32_to_bits(fixed.scale_inv) != f32_to_bits(legacy.scale_inv), + std::string(name) + " floor changes the inverse"); + } + } else if (i == 3) { + require(f32_to_bits(fixed.scale) == 0x7e800000u, std::string(name) + " exact scale 2^126"); + require(f32_to_bits(fixed.scale_inv) == 0x00800000u, std::string(name) + " exact inv FLT_MIN"); + require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv, + std::string(name) + " exact matches legacy"); + } else { + require(is_finite_normal(fixed.scale) && is_finite_normal(fixed.scale_inv), + std::string(name) + " +1 ULP normal"); + require(fixed.scale_inv >= FLT_MIN, std::string(name) + " +1 ULP inv >= FLT_MIN"); + require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv, + std::string(name) + " +1 ULP matches legacy"); + } + } + } + + // Epsilon above the guard keeps the historical floor behavior. + const float epsilon = 1.0e-6f; + const Scales lifted = compute_scale(tiny, max_fp8, false, true, epsilon); + require(lifted.scale == max_fp8 / epsilon, "epsilon floor still applies before the guard"); + + // Normal magnitudes, zeros, inf, and NaN stay bitwise-identical. + const float normals[] = {1.0e-20f, 1.0e-5f, 0.125f, 1.f, 448.f, 1.0e10f, + std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN(), -0.f, 0.f}; + for (float amax : normals) { + for (int pow2 = 0; pow2 <= 1; ++pow2) { + const Scales fixed = compute_scale(amax, max_fp8, pow2, true, 0.f); + const Scales legacy = compute_scale(amax, max_fp8, pow2, false, 0.f); + require(f32_to_bits(fixed.scale) == f32_to_bits(legacy.scale), "normal scale bits"); + require(f32_to_bits(fixed.scale_inv) == f32_to_bits(legacy.scale_inv), "normal inv bits"); + } + } +} + +} // namespace + +int main() { + std::printf("THRESHOLD e4m3 %08x\n", f32_to_bits(448.0f * FLT_MIN)); + std::printf("THRESHOLD e5m2 %08x\n", f32_to_bits(57344.0f * FLT_MIN)); + std::printf("FLT_MIN %08x\n", f32_to_bits(FLT_MIN)); + std::printf("TINY %08x\n", f32_to_bits(7.7e-37f)); + check_format("e4m3", 448.0f, 0x04e00000u); + check_format("e5m2", 57344.0f, 0x08600000u); + std::printf("HOST_OK\n"); + return 0; +} diff --git a/tests/cpp/operator/test_tiny_amax_scale_host.py b/tests/cpp/operator/test_tiny_amax_scale_host.py new file mode 100644 index 00000000000..0e409236473 --- /dev/null +++ b/tests/cpp/operator/test_tiny_amax_scale_host.py @@ -0,0 +1,589 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Host checks for flooring a tiny positive FP8 amax at max_fp8 * FLT_MIN.""" + +from __future__ import annotations + +import json +import math +import shutil +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parents[3] +HOST_SRC = Path(__file__).resolve().parent / "test_tiny_amax_scale_host.cpp" +RECIPE_CUH = ROOT / "transformer_engine" / "common" / "recipe" / "recipe_common.cuh" +REFERENCE_PY = ROOT / "tests" / "pytorch" / "references" / "quantize_scale_calc.py" +CUSTOM_REFERENCE_PY = ( + ROOT / "transformer_engine" / "pytorch" / "custom_recipes" / "reference_current_scaling.py" +) +JAX_QUANTIZER_PY = ROOT / "transformer_engine" / "jax" / "quantize" / "quantizer.py" + +E4M3_MAX = np.float32(448.0) +E5M2_MAX = np.float32(57344.0) +FLT_MIN = np.float32(np.finfo(np.float32).tiny) +TINY = np.float32(7.7e-37) + + +def bits(value) -> int: + return int(np.asarray(value, dtype=np.float32).view(np.uint32)) + + +def from_bits(value: int) -> np.float32: + return np.asarray(np.uint32(value), dtype=np.uint32).view(np.float32) + + +def is_subnormal(value) -> np.ndarray: + raw = np.asarray(value, dtype=np.float32).view(np.uint32) + exponent = raw & np.uint32(0x7F800000) + mantissa = raw & np.uint32(0x007FFFFF) + return (exponent == 0) & (mantissa != 0) + + +def scales_from_amax(amax, max_fp8, pow2: bool, guarded: bool, epsilon: float = 0.0): + """Host copy of compute_scale_from_amax, vectorized over amax.""" + amax = np.asarray(amax, dtype=np.float32).copy() + epsilon_f = np.float32(epsilon) + amax = np.where(amax < epsilon_f, epsilon_f, amax) + max_fp8 = np.float32(max_fp8) + guard = np.float32(max_fp8 * FLT_MIN) + early = (~np.isfinite(amax)) | (amax == np.float32(0)) + usable = amax + if guarded: + usable = np.where(amax < guard, guard, amax) + usable = np.where(early, np.float32(1), usable) + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + scale = np.divide(max_fp8, usable) + scale = np.where(np.isinf(scale), np.float32(np.finfo(np.float32).max), scale).astype( + np.float32 + ) + if pow2: + masked = scale.view(np.uint32) & np.uint32(0xFF800000) + scale = masked.view(np.float32) + scale = np.where(early, np.float32(1), scale).astype(np.float32) + scale_inv = np.divide(np.float32(1), scale).astype(np.float32) + return scale, scale_inv + + +def fp8_grid(fmt: str): + if fmt == "e4m3": + bias, man_bits, max_exp, skip = 7, 3, 15, {(15, 7)} + elif fmt == "e5m2": + bias, man_bits, max_exp, skip = 15, 2, 30, set() + else: + raise ValueError(fmt) + values = [np.float32(0)] + for exp in range(0, max_exp + 1): + for man in range(1 << man_bits): + if exp == 0 and man == 0: + continue + if (exp, man) in skip: + continue + if exp == 0: + value = np.float32((man / (1 << man_bits)) * 2.0 ** (1 - bias)) + else: + value = np.float32((1.0 + man / (1 << man_bits)) * 2.0 ** (exp - bias)) + values.append(value) + return np.asarray(values, dtype=np.float32) + + +GRIDS = {"e4m3": fp8_grid("e4m3"), "e5m2": fp8_grid("e5m2")} + + +def quantize_fp8(scaled, fmt: str) -> np.ndarray: + """Round scaled binary32 values to the nearest finite FP8 code and return them.""" + scaled = np.asarray(scaled, dtype=np.float32) + grid = GRIDS[fmt] + sign = np.sign(scaled).astype(np.float32) + sign = np.where(sign == 0, np.float32(1), sign) + magnitude = np.abs(scaled) + # Values below half the smallest subnormal quantize to zero. + index = np.argmin(np.abs(magnitude[..., None] - grid), axis=-1) + return (sign * grid[index]).astype(np.float32) + + +def quant_dequant(x, amax, scale, scale_inv, fmt: str): + del amax + quantized = quantize_fp8(np.asarray(x, dtype=np.float32) * scale, fmt) + return quantized, (quantized * scale_inv).astype(np.float32) + + +def block_views_1d(x, block: int): + matrix = np.asarray(x, dtype=np.float32) + rows, cols = matrix.shape + if cols % block != 0: + raise ValueError(f"K={cols} is not divisible by block {block}") + view = matrix.reshape(rows, cols // block, block) + amax = np.max(np.abs(view), axis=-1) + return view, amax + + +def apply_1d(x, block, max_fp8, pow2, guarded, fmt): + view, amax = block_views_1d(x, block) + scale, scale_inv = scales_from_amax(amax, max_fp8, pow2, guarded) + quantized = quantize_fp8(view * scale[..., None], fmt) + dequant = (quantized * scale_inv[..., None]).astype(np.float32) + return ( + dequant.reshape(x.shape), + quantized.reshape(x.shape), + scale, + scale_inv, + ) + + +def apply_2d(weight, block, max_fp8, pow2, guarded, fmt): + matrix = np.asarray(weight, dtype=np.float32) + rows, cols = matrix.shape + if rows % block != 0 or cols % block != 0: + raise ValueError(f"shape {matrix.shape} is not divisible by {block}") + view = matrix.reshape(rows // block, block, cols // block, block).transpose(0, 2, 1, 3) + amax = np.max(np.abs(view), axis=(-1, -2)) + scale, scale_inv = scales_from_amax(amax, max_fp8, pow2, guarded) + quantized = quantize_fp8(view * scale[..., None, None], fmt) + dequant = (quantized * scale_inv[..., None, None]).astype(np.float32) + restored = dequant.transpose(0, 2, 1, 3).reshape(matrix.shape) + return restored, scale, scale_inv + + +def run_host() -> list[dict]: + with tempfile.TemporaryDirectory() as tmp: + binary = Path(tmp) / "fp8_scale_guard_host" + compile_cmd = [ + "g++", + "-O2", + "-std=c++17", + "-Wall", + "-Wextra", + str(HOST_SRC), + "-o", + str(binary), + ] + compiled = subprocess.run(compile_cmd, check=True, capture_output=True, text=True) + if compiled.stderr.strip(): + print(compiled.stderr, file=sys.stderr) + completed = subprocess.run([str(binary)], check=True, capture_output=True, text=True) + rows = [] + for line in completed.stdout.splitlines(): + parts = line.split() + if parts[0] in {"e4m3", "e5m2"}: + rows.append( + { + "format": parts[0], + "pow2": bool(int(parts[1])), + "label": parts[2], + "amax": parts[3], + "scale": parts[4], + "scale_inv": parts[5], + "legacy_scale": parts[6], + "legacy_scale_inv": parts[7], + } + ) + elif parts[0] == "HOST_OK": + continue + else: + print(line, file=sys.stderr) + if "HOST_OK" not in completed.stdout: + raise AssertionError("host checker did not report success") + return rows + + +def cross_check_numpy(rows: list[dict]) -> None: + maxima = {"e4m3": E4M3_MAX, "e5m2": E5M2_MAX} + for row in rows: + amax = from_bits(int(row["amax"], 16)) + for guarded, scale_key, inv_key in ( + (True, "scale", "scale_inv"), + (False, "legacy_scale", "legacy_scale_inv"), + ): + scale, scale_inv = scales_from_amax(amax, maxima[row["format"]], row["pow2"], guarded) + if bits(scale) != int(row[scale_key], 16) or bits(scale_inv) != int(row[inv_key], 16): + raise AssertionError(f"numpy/host mismatch for {row} guarded={guarded}") + + +def check_sources() -> None: + recipe = RECIPE_CUH.read_text() + if ( + "amax < amax_floor" not in recipe + or "amax == 0.f" not in recipe + or "#include " not in recipe + ): + raise AssertionError("recipe_common.cuh does not floor a positive tiny amax") + reference = REFERENCE_PY.read_text() + if ( + "torch.finfo(torch.float32).tiny" not in reference + or "amax < guard" not in reference + or "amax == 0" not in reference + ): + raise AssertionError("quantize_scale_calc.py does not mirror the floor") + custom = CUSTOM_REFERENCE_PY.read_text() + if "amax < guard" not in custom or "amax == 0" not in custom: + raise AssertionError("reference_current_scaling.py does not mirror the floor") + jax = JAX_QUANTIZER_PY.read_text() + if ( + "jnp.finfo(jnp.float32).tiny" not in jax + or "amax < guard" not in jax + or "amax > 0" not in jax + ): + raise AssertionError("jax quantizer.py does not mirror the floor") + + +def assert_random_bitwise(fmt: str, max_fp8, pow2: bool) -> dict: + rng = np.random.default_rng(20261028) + tensor = rng.standard_normal((64, 256)).astype(np.float32) + new_dq, new_q, new_scale, new_inv = apply_1d(tensor, 128, max_fp8, pow2, True, fmt) + old_dq, old_q, old_scale, old_inv = apply_1d(tensor, 128, max_fp8, pow2, False, fmt) + if not np.array_equal(new_q.view(np.uint32), old_q.view(np.uint32)): + raise AssertionError(f"random quantized bits changed for {fmt} pow2={pow2}") + if not np.array_equal(new_scale.view(np.uint32), old_scale.view(np.uint32)): + raise AssertionError(f"random scale bits changed for {fmt} pow2={pow2}") + if not np.array_equal(new_inv.view(np.uint32), old_inv.view(np.uint32)): + raise AssertionError(f"random scale_inv bits changed for {fmt} pow2={pow2}") + if not np.array_equal(new_dq.view(np.uint32), old_dq.view(np.uint32)): + raise AssertionError(f"random dequant bits changed for {fmt} pow2={pow2}") + return { + "format": fmt, + "pow2": pow2, + "elements": int(tensor.size), + "blocks": int(new_scale.size), + "bitwise_identical": True, + } + + +def _finite_normal_invs(invs) -> None: + for inv in invs: + array = np.asarray(inv, dtype=np.float32) + if not np.isfinite(array).all() or np.any(is_subnormal(array)) or np.any(array < FLT_MIN): + raise AssertionError("scale_inv is non-finite or smaller than FLT_MIN") + + +def linear_case(recipe: str) -> dict: + """Column-shaped GEMM with one tiny block, for blockwise and tensorwise scales.""" + rng = np.random.default_rng(7) + rows, in_features, out_features = 64, 256, 128 + activation = rng.standard_normal((rows, in_features)).astype(np.float32) * np.float32(0.2) + weight = rng.standard_normal((out_features, in_features)).astype(np.float32) * np.float32(0.05) + grad_out = rng.standard_normal((rows, out_features)).astype(np.float32) * np.float32(0.2) + activation[0, 0:128] = TINY + grad_out[0, 0:128] = TINY + pow2 = recipe == "blockwise" + if recipe == "blockwise": + weight[0:128, 0:128] = TINY + + def quantize_weight(guarded: bool): + if recipe == "blockwise": + return apply_2d(weight, 128, E4M3_MAX, True, guarded, "e4m3") + scale, scale_inv = scales_from_amax(np.max(np.abs(weight)), E4M3_MAX, False, guarded) + quantized = quantize_fp8(weight * scale, "e4m3") + return (quantized * scale_inv).astype(np.float32), scale, scale_inv + + def quantize_matrix(matrix, guarded: bool): + if recipe == "blockwise": + return apply_1d(matrix, 128, E4M3_MAX, True, guarded, "e4m3") + scale, scale_inv = scales_from_amax(np.max(np.abs(matrix)), E4M3_MAX, False, guarded) + quantized = quantize_fp8(matrix * scale, "e4m3") + dequant = (quantized * scale_inv).astype(np.float32) + return dequant, quantized, scale, scale_inv + + def forward_backward(guarded: bool, x_value, dy_value): + x_dq, x_q, x_scale, x_inv = quantize_matrix(x_value, guarded) + w_dq, w_scale, w_inv = quantize_weight(guarded) + dy_dq, dy_q, dy_scale, dy_inv = quantize_matrix(dy_value, guarded) + output = (x_dq @ w_dq.T).astype(np.float32) + dgrad = (dy_dq @ w_dq).astype(np.float32) + wgrad = (dy_dq.T @ x_dq).astype(np.float32) + return output, dgrad, wgrad, (x_inv, w_inv, dy_inv, x_q, dy_q) + + output, dgrad, wgrad, meta = forward_backward(True, activation, grad_out) + for name, tensor in (("y", output), ("dgrad", dgrad), ("wgrad", wgrad)): + if not np.isfinite(tensor).all(): + raise AssertionError(f"{recipe} {name} contains a non-finite value") + _finite_normal_invs(meta[:3]) + legacy = forward_backward(False, activation, grad_out) + legacy_subnormal = any(bool(np.any(is_subnormal(inv))) for inv in legacy[3][:3]) + result = { + "recipe": recipe, + "pow2_scales": pow2, + "y_finite": True, + "dgrad_finite": True, + "wgrad_finite": True, + "scale_inv_ge_flt_min": True, + "y_max_abs": float(np.max(np.abs(output))), + "dgrad_max_abs": float(np.max(np.abs(dgrad))), + "wgrad_max_abs": float(np.max(np.abs(wgrad))), + "legacy_subnormal_on_partial_blocks": legacy_subnormal, + } + if recipe == "blockwise": + if not legacy_subnormal: + raise AssertionError("blockwise legacy path did not reproduce a subnormal scale_inv") + injected_codes = meta[3][0, 0:128] + if np.all(injected_codes == 0): + raise AssertionError("injected 1x128 activation was flushed to zero") + if bits(meta[0][0, 0]) != 0x00800000: + raise AssertionError("injected block scale_inv is not FLT_MIN") + result["injected_block_kept"] = True + result["injected_block_scale_inv_bits"] = f"{bits(meta[0][0, 0]):08x}" + result["injected_block_max_abs_code"] = float(np.max(np.abs(injected_codes))) + return result + if legacy_subnormal: + raise AssertionError("partial injection unexpectedly changed the tensorwise amax") + if not np.array_equal(output.view(np.uint32), legacy[0].view(np.uint32)): + raise AssertionError("tensorwise partial injection changed the guarded result") + tiny_activation = np.full_like(activation, TINY) + tiny_out, tiny_dgrad, tiny_wgrad, tiny_meta = forward_backward(True, tiny_activation, grad_out) + legacy_tiny = forward_backward(False, tiny_activation, grad_out) + if not ( + np.isfinite(tiny_out).all() + and np.isfinite(tiny_dgrad).all() + and np.isfinite(tiny_wgrad).all() + ): + raise AssertionError("tiny tensorwise activation produced a non-finite gradient") + _finite_normal_invs(tiny_meta[:3]) + if not np.any(is_subnormal(legacy_tiny[3][0])): + raise AssertionError("legacy tensorwise path did not leave a subnormal scale_inv") + if np.all(tiny_meta[3] == 0): + raise AssertionError("fully tiny tensorwise activation was flushed to zero") + if bits(tiny_meta[0]) != 0x00800000: + raise AssertionError("tiny tensorwise scale_inv is not FLT_MIN") + result["partial_block_bitwise_identical"] = True + result["full_tensor_tiny_kept"] = True + result["full_tensor_tiny_y_max_abs"] = float(np.max(np.abs(tiny_out))) + result["full_tensor_tiny_scale_inv_bits"] = f"{bits(tiny_meta[0]):08x}" + result["legacy_subnormal_on_full_tensor"] = True + return result + + +def grouped_linear_case() -> dict: + rng = np.random.default_rng(11) + experts, tokens, in_features, out_features = 2, 32, 128, 128 + activation = rng.standard_normal((experts, tokens, in_features)).astype(np.float32) + weight = rng.standard_normal((experts, out_features, in_features)).astype( + np.float32 + ) * np.float32(0.05) + grad_out = rng.standard_normal((experts, tokens, out_features)).astype(np.float32) + activation[1].fill(TINY) + grad_out[1].fill(TINY) + outputs = [] + for expert in range(experts): + x_dq, _, _, x_inv = apply_1d(activation[expert], 128, E4M3_MAX, True, True, "e4m3") + w_dq, _, w_inv = apply_2d(weight[expert], 128, E4M3_MAX, True, True, "e4m3") + dy_dq, _, _, dy_inv = apply_1d(grad_out[expert], 128, E4M3_MAX, True, True, "e4m3") + output = x_dq @ w_dq.T + dgrad = dy_dq @ w_dq + wgrad = dy_dq.T @ x_dq + for tensor in (output, dgrad, wgrad): + if not np.isfinite(tensor).all(): + raise AssertionError("grouped linear produced a non-finite value") + for tensor in (x_inv, w_inv, dy_inv): + if np.any(is_subnormal(tensor)) or np.any( + np.asarray(tensor, dtype=np.float32) < FLT_MIN + ): + raise AssertionError("grouped linear stored a subnormal scale_inv") + outputs.append(output) + if expert == 1: + if np.all(x_dq == 0) or np.all(output == 0): + raise AssertionError("tiny expert activation was flushed to zero") + if not np.all(x_inv == FLT_MIN): + raise AssertionError("tiny expert scale_inv was not floored at FLT_MIN") + return { + "experts": experts, + "tiny_expert_output_kept": True, + "tiny_expert_output_max_abs": float(np.max(np.abs(outputs[1]))), + "normal_expert_output_max_abs": float(np.max(np.abs(outputs[0]))), + } + + +def train_proxy(steps: int = 15) -> dict: + rng = np.random.default_rng(19) + tokens, in_features, hidden, out_features = 32, 128, 128, 128 + data_x = rng.standard_normal((steps, tokens, in_features)).astype(np.float32) * np.float32(0.3) + data_y = rng.standard_normal((steps, tokens, out_features)).astype(np.float32) + + start_w1 = rng.standard_normal((hidden, in_features)).astype(np.float32) * np.float32(0.05) + start_w2 = rng.standard_normal((out_features, hidden)).astype(np.float32) * np.float32(0.05) + + def execute(guarded: bool, inject: bool): + w1 = start_w1.copy() + w2 = start_w2.copy() + losses = [] + grad_norms = [] + subnormal_inv = False + started = time.perf_counter() + for step in range(steps): + activation = data_x[step].copy() + if inject: + activation[0, :] = TINY + x_dq, _, _, x_inv = apply_1d(activation, 128, E4M3_MAX, True, guarded, "e4m3") + if inject and guarded and np.all(x_dq[0] == 0): + raise AssertionError("injected activation row was flushed to zero") + subnormal_inv = subnormal_inv or bool(np.any(is_subnormal(x_inv))) + hidden_pre = (x_dq @ w1.T).astype(np.float32) + h_dq, _, _, h_inv = apply_1d(hidden_pre, 128, E4M3_MAX, True, guarded, "e4m3") + subnormal_inv = subnormal_inv or bool(np.any(is_subnormal(h_inv))) + pred = (h_dq @ w2.T).astype(np.float32) + residual = (pred - data_y[step]).astype(np.float32) + loss = np.float32(np.mean(np.square(residual))) + dy = (np.float32(2.0) * residual / np.float32(residual.size)).astype(np.float32) + if inject: + dy[0, :] = TINY + dy_dq, _, _, dy_inv = apply_1d(dy, 128, E4M3_MAX, True, guarded, "e4m3") + subnormal_inv = subnormal_inv or bool(np.any(is_subnormal(dy_inv))) + dw2 = (dy_dq.T @ h_dq).astype(np.float32) + dh = (dy_dq @ w2).astype(np.float32) + dh_dq, _, _, dh_inv = apply_1d(dh, 128, E4M3_MAX, True, guarded, "e4m3") + subnormal_inv = subnormal_inv or bool(np.any(is_subnormal(dh_inv))) + dw1 = (dh_dq.T @ x_dq).astype(np.float32) + if not np.isfinite(loss) or not np.isfinite(dw1).all() or not np.isfinite(dw2).all(): + raise AssertionError(f"non-finite training state guarded={guarded} inject={inject}") + grad_norm = np.float32( + math.sqrt(float(np.sum(np.square(dw1)) + np.sum(np.square(dw2)))) + ) + w1 -= np.float32(0.05) * dw1 + w2 -= np.float32(0.05) * dw2 + losses.append(float(loss)) + grad_norms.append(float(grad_norm)) + return { + "losses": losses, + "grad_norms": grad_norms, + "elapsed_sec": time.perf_counter() - started, + "subnormal_scale_inv": subnormal_inv, + } + + normal_new = execute(True, False) + normal_old = execute(False, False) + injected_new = execute(True, True) + injected_old = execute(False, True) + if ( + normal_new["losses"] != normal_old["losses"] + or normal_new["grad_norms"] != normal_old["grad_norms"] + ): + raise AssertionError("normal-data loss or grad norm changed") + if injected_new["subnormal_scale_inv"] or not injected_old["subnormal_scale_inv"]: + raise AssertionError("injected smoke did not isolate the legacy subnormal") + return { + "steps": steps, + "normal_losses": normal_new["losses"], + "normal_grad_norms": normal_new["grad_norms"], + "normal_elapsed_new_sec": normal_new["elapsed_sec"], + "normal_elapsed_old_sec": normal_old["elapsed_sec"], + "bitwise_loss_and_grad_norm": True, + "injected_losses_new": injected_new["losses"], + "injected_grad_norms_new": injected_new["grad_norms"], + "injected_losses_old": injected_old["losses"], + "injected_elapsed_new_sec": injected_new["elapsed_sec"], + "injected_elapsed_old_sec": injected_old["elapsed_sec"], + "injected_new_subnormal": injected_new["subnormal_scale_inv"], + "injected_old_subnormal": injected_old["subnormal_scale_inv"], + } + + +def probe_gpu_stack() -> dict: + torch_status = "not installed" + try: + import torch # type: ignore + + torch_status = f"{torch.__version__} cuda={torch.cuda.is_available()}" + except Exception as exc: # noqa: BLE001 - the absence itself is the result. + torch_status = f"unavailable ({type(exc).__name__}: {exc})" + te_status = "not installed" + try: + import transformer_engine # type: ignore + + te_status = getattr(transformer_engine, "__file__", "imported") + except Exception as exc: # noqa: BLE001 + te_status = f"unavailable ({type(exc).__name__}: {exc})" + return { + "nvidia_smi": shutil.which("nvidia-smi") or "not found", + "nvcc": shutil.which("nvcc") or "not found", + "torch": torch_status, + "transformer_engine_import": te_status, + "cuda_operator_tests": "not executed; test_operator requires nvcc, CUDA, and a GPU", + "pytorch_quantization_tests": ( + "not executed; no PyTorch or CUDA runtime in this environment" + ), + } + + +def test_recipe_sources_have_amax_floor() -> None: + check_sources() + + +def test_host_binary_matches_numpy() -> None: + import pytest + + if shutil.which("g++") is None: + pytest.skip("g++ is not available") + cross_check_numpy(run_host()) + + +def test_random_blocks_match_legacy() -> None: + for fmt, max_fp8 in (("e4m3", E4M3_MAX), ("e5m2", E5M2_MAX)): + for pow2 in (False, True): + assert_random_bitwise(fmt, max_fp8, pow2) + rng = np.random.default_rng(3) + weight = rng.standard_normal((256, 256)).astype(np.float32) + for pow2 in (False, True): + new = apply_2d(weight, 128, E4M3_MAX, pow2, True, "e4m3") + old = apply_2d(weight, 128, E4M3_MAX, pow2, False, "e4m3") + if any(not np.array_equal(a.view(np.uint32), b.view(np.uint32)) for a, b in zip(new, old)): + raise AssertionError("2D random block quantization changed") + + +def test_tiny_block_stays_representable() -> None: + linear_case("blockwise") + linear_case("tensorwise") + + +def test_grouped_tiny_expert_stays_representable() -> None: + grouped_linear_case() + + +def test_training_proxy_normal_data_matches_legacy() -> None: + train_proxy() + + +def main() -> None: + check_sources() + rows = run_host() + cross_check_numpy(rows) + random_rows = [ + assert_random_bitwise("e4m3", E4M3_MAX, False), + assert_random_bitwise("e4m3", E4M3_MAX, True), + assert_random_bitwise("e5m2", E5M2_MAX, False), + assert_random_bitwise("e5m2", E5M2_MAX, True), + ] + # 2D 128x128 random tiles must also be unchanged. + rng = np.random.default_rng(3) + weight = rng.standard_normal((256, 256)).astype(np.float32) + for pow2 in (False, True): + new = apply_2d(weight, 128, E4M3_MAX, pow2, True, "e4m3") + old = apply_2d(weight, 128, E4M3_MAX, pow2, False, "e4m3") + if any(not np.array_equal(a.view(np.uint32), b.view(np.uint32)) for a, b in zip(new, old)): + raise AssertionError("2D random block quantization changed") + linear = [linear_case("blockwise"), linear_case("tensorwise")] + grouped = grouped_linear_case() + training = train_proxy() + environment = probe_gpu_stack() + # Confirm the FP8 grid contains the format maxima used by the guard. + if E4M3_MAX not in GRIDS["e4m3"] or E5M2_MAX not in GRIDS["e5m2"]: + raise AssertionError("FP8 grid is missing its maximum finite value") + summary = { + "host_rows": rows, + "random": random_rows, + "linear": linear, + "grouped": grouped, + "training_proxy": training, + "environment": environment, + } + print(json.dumps(summary)) + + +if __name__ == "__main__": + main() diff --git a/tests/pytorch/references/quantize_scale_calc.py b/tests/pytorch/references/quantize_scale_calc.py index a6ff425133f..77758dd8c51 100644 --- a/tests/pytorch/references/quantize_scale_calc.py +++ b/tests/pytorch/references/quantize_scale_calc.py @@ -24,12 +24,20 @@ def scale_from_amax_tensor( - amax: Amax tensor with updates made for extrema values. """ assert amax.dtype == torch.float, "amax must be a float tensor." - fp8_max = torch.finfo(quant_dtype).max + fp8_max = torch.tensor(torch.finfo(quant_dtype).max, dtype=torch.float32, device=amax.device) # Clamping amax to avoid division by small numbers - amax = torch.max(amax, torch.tensor(eps)) + eps_tensor = torch.tensor(eps, dtype=amax.dtype, device=amax.device) + # `<` is false for NaN, matching the device function's epsilon floor. + amax = torch.where(amax < eps_tensor, eps_tensor, amax) + + # Floor a positive amax at max_fp8 * FLT_MIN so scale_inv is at least FLT_MIN. + guard = fp8_max * torch.finfo(torch.float32).tiny + invalid = torch.isinf(amax) | torch.isnan(amax) | (amax == 0) + floored = torch.where(amax < guard, guard, amax) + safe_amax = torch.where(invalid, torch.ones_like(amax), floored) # Compute scale factor - scale = torch.div(fp8_max, amax) + scale = torch.div(fp8_max, safe_amax) # Note frexp doesn't give back inf for exponent with an inf input # We take care of inf before pow_2_scales scale = torch.where(scale == torch.inf, torch.finfo(x_dtype).max, scale) @@ -40,19 +48,14 @@ def scale_from_amax_tensor( # a mantissa in [0.5, 1.0). Because a normal float has a mantissa with # hidden bit in [1.0, 2.0), the exponent will be off by exactly one because # of the shift. Subnormal and zero cases need not be considered because - # the smallest possible result of fp8_max / amax is still normal. + # a floored amax keeps the smallest fp8_max / amax at 2^126. exp = exp - 1 # No subnormals and zero. assert (exp > -127).all() unity = torch.tensor([1.0], device=exp.device) torch.ldexp(unity, exp, out=scale) - # Case where amax is inf. The frexp, ldexp logic changes 0.0 scales - # Return 0.0 for 0.0 scale for consistency with non-pow2 scale - # calculation. - scale = torch.where(amax == float("inf"), 0.0, scale) - # Handle overflow cases for amax zero causing NaN - scale = torch.where(amax == 0, 1.0, scale) + scale = torch.where(invalid, torch.ones_like(scale), scale) # Compute scale_inv scale_inv = torch.reciprocal(scale) diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index 1d45d1919af..8028dfe8fe1 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -808,6 +808,24 @@ def test_current_scaling_quantization_versus_reference(self, M, N, dtype): torch.testing.assert_close(qx_native, qx_ref, atol=0.0, rtol=0.0) torch.testing.assert_close(sx_native, sx_ref, atol=0.0, rtol=0.0) + def test_tiny_positive_amax_scale_inv_is_flt_min(self): + # Public quantizer: a positive tiny amax stores scale_inv FLT_MIN, not zero codes. + x = torch.full((128, 128), 7.7e-37, dtype=torch.float32, device="cuda") + te_quant, ref_quant = self._make_quantizers(rowwise=True, columnwise=False) + x_te = te_quant(x) + x_ref = ref_quant.quantize(x) + qx_native = x_te._data.view(dtype=torch.float8_e4m3fn) + torch.testing.assert_close(qx_native, x_ref.data, atol=0.0, rtol=0.0) + torch.testing.assert_close(x_te._scale_inv, x_ref.scale, atol=0.0, rtol=0.0) + flt_min = torch.full_like(x_te._scale_inv, torch.finfo(torch.float32).tiny) + torch.testing.assert_close(x_te._scale_inv, flt_min, atol=0.0, rtol=0.0) + assert not torch.all(qx_native.to(dtype=torch.float32) == 0) + + zeros = torch.zeros_like(x) + z_te = te_quant(zeros) + ones = torch.ones_like(z_te._scale_inv, dtype=torch.float32) + torch.testing.assert_close(z_te._scale_inv.float(), ones, atol=0.0, rtol=0.0) + @pytest.mark.parametrize( "M, K, N, out_dtype, accumulate", [ diff --git a/transformer_engine/common/recipe/recipe_common.cuh b/transformer_engine/common/recipe/recipe_common.cuh index 07839407a3f..296b345a52c 100644 --- a/transformer_engine/common/recipe/recipe_common.cuh +++ b/transformer_engine/common/recipe/recipe_common.cuh @@ -7,6 +7,8 @@ #ifndef TRANSFORMER_ENGINE_RECIPE_RECIPE_COMMON_CUH_ #define TRANSFORMER_ENGINE_RECIPE_RECIPE_COMMON_CUH_ +#include + #include "common/common.h" namespace transformer_engine { @@ -21,10 +23,16 @@ __device__ __forceinline__ float compute_scale_from_amax(float amax, float max_f float scale = 1.f; - if (isinf(amax) || amax == 0.f || isnan(amax)) { + if (isinf(amax) || isnan(amax) || amax == 0.f) { return scale; } + // Floor a positive amax at max_fp8 * FLT_MIN so scale_inv is at least FLT_MIN. + const float amax_floor = max_fp8 * FLT_MIN; + if (amax < amax_floor) { + amax = amax_floor; + } + // Here we don't use "scale = max_fp8 / amax" because it has different results with/without // "--use_fast_math". // "__fdiv_rn" has the same behavior with "max_fp8 / amax" when not using fast math. diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index db56db935dc..2a6caa2a42f 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -64,9 +64,13 @@ def compute_scale_from_amax( fp8_max = jnp.astype(jnp.finfo(q_dtype).max, jnp.float32) if scale is None: scale = jnp.ones((1,)) - sf = (fp8_max / amax) / (2**margin) - sf = jnp.where(amax > 0.0, sf, scale) - sf = jnp.where(jnp.isfinite(amax), sf, scale) + # Floor a positive amax at fp8_max * FLT_MIN so scale_inv is at least FLT_MIN. + guard = fp8_max * jnp.finfo(jnp.float32).tiny + usable = jnp.isfinite(amax) & (amax > 0) + floored = jnp.where(amax < guard, guard, amax) + safe_amax = jnp.where(usable, floored, jnp.ones_like(amax)) + sf = (fp8_max / safe_amax) / (2**margin) + sf = jnp.where(usable, sf, scale) assert sf.shape == (1,), f"Expected sf.shape == (1,), but got {sf.shape}" return sf diff --git a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index 52f192e4340..bd3a3f0eb8e 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -171,13 +171,21 @@ def _scale_from_amax_tensor( - amax: Amax tensor with updates made for extrema values. """ assert amax.dtype == torch.float, "amax must be a float tensor." - fp8_max = torch.finfo(quant_dtype).max + fp8_max = torch.tensor(torch.finfo(quant_dtype).max, dtype=torch.float32, device=amax.device) # Clamping amax to avoid division by small numbers - amax = torch.max(amax, torch.tensor(eps)) + eps_tensor = torch.tensor(eps, dtype=amax.dtype, device=amax.device) + # `<` is false for NaN, matching the device function's epsilon floor. + amax = torch.where(amax < eps_tensor, eps_tensor, amax) + + # Floor a positive amax at max_fp8 * FLT_MIN so scale_inv is at least FLT_MIN. + guard = fp8_max * torch.finfo(torch.float32).tiny + invalid = torch.isinf(amax) | torch.isnan(amax) | (amax == 0) + floored = torch.where(amax < guard, guard, amax) + safe_amax = torch.where(invalid, torch.ones_like(amax), floored) # Compute scale factor - scale = torch.div(fp8_max, amax) + scale = torch.div(fp8_max, safe_amax) # Take care of inf before pow_2_scales scale = torch.where(scale == torch.inf, torch.finfo(x_dtype).max, scale) @@ -188,10 +196,8 @@ def _scale_from_amax_tensor( assert (exp > -127).all() unity = torch.tensor([1.0], device=exp.device) torch.ldexp(unity, exp, out=scale) - scale = torch.where(amax == float("inf"), 0.0, scale) - # Handle overflow cases for amax zero causing NaN - scale = torch.where(amax == 0, 1.0, scale) + scale = torch.where(invalid, torch.ones_like(scale), scale) # Compute scale_inv scale_inv = torch.reciprocal(scale)