|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | +"""Tests for NVFP4QTensor per-block FP8 scale clamping (underflow + overflow).""" |
| 17 | + |
| 18 | +from types import SimpleNamespace |
| 19 | + |
| 20 | +import torch |
| 21 | + |
| 22 | +from modelopt.torch.quantization.qtensor.nvfp4_tensor import ( |
| 23 | + NVFP4QTensor, |
| 24 | + _cast_per_block_scale_to_fp8, |
| 25 | +) |
| 26 | + |
| 27 | +_FP8_E4M3FN_MIN = 2**-9 # 0.001953125 — smallest positive FP8 E4M3FN subnormal |
| 28 | +_FP8_E4M3FN_MAX = 448.0 |
| 29 | + |
| 30 | + |
| 31 | +class TestNVFP4ScaleClamping: |
| 32 | + """Per-block weight scales outside the FP8 E4M3FN range must be clamped, not turned into 0/NaN.""" |
| 33 | + |
| 34 | + def test_no_zero_scales_for_tiny_weights(self): |
| 35 | + """Tiny per-block amax (<<FP8 min) must not underflow to zero after FP8 cast.""" |
| 36 | + block_size = 16 |
| 37 | + tiny_weight = torch.full((4, block_size), 1e-10) |
| 38 | + # wsf2=1.0 → per_block_scale = amax/(6*wsf2) ≈ 1.7e-11 << 2^-9, exercises FP8-min clamp |
| 39 | + wsf2 = torch.tensor(1.0) |
| 40 | + |
| 41 | + per_block_scale, _ = NVFP4QTensor.get_weights_scaling_factor(tiny_weight, block_size, wsf2) |
| 42 | + per_block_scale_f32 = per_block_scale.float() |
| 43 | + |
| 44 | + assert (per_block_scale_f32 > 0).all(), ( |
| 45 | + f"Zero per-block scales found after FP8 cast: {per_block_scale_f32.tolist()}. " |
| 46 | + "FP8 scale underflow clamping likely regressed." |
| 47 | + ) |
| 48 | + assert (per_block_scale_f32 >= _FP8_E4M3FN_MIN).all(), ( |
| 49 | + "Per-block scales with zero values found after FP8 cast " |
| 50 | + "(below the FP8 E4M3FN subnormal minimum — clamp would have prevented this)." |
| 51 | + ) |
| 52 | + |
| 53 | + def test_normal_weights_unaffected_by_clamp(self): |
| 54 | + """Weights with typical magnitudes must not be affected by the underflow clamp.""" |
| 55 | + block_size = 16 |
| 56 | + torch.manual_seed(42) |
| 57 | + normal_weight = torch.randn(8, block_size) |
| 58 | + |
| 59 | + per_block_scale, _ = NVFP4QTensor.get_weights_scaling_factor(normal_weight, block_size) |
| 60 | + assert (per_block_scale.float() > 0).all(), "Normal weights produced zero scales." |
| 61 | + |
| 62 | + def test_mixed_weight_no_zeros(self): |
| 63 | + """Mixed-magnitude tensor (normal + tiny blocks) must have no zero scales.""" |
| 64 | + block_size = 16 |
| 65 | + weight = torch.cat( |
| 66 | + [ |
| 67 | + torch.randn(4, block_size), |
| 68 | + torch.full((4, block_size), 1e-12), |
| 69 | + ], |
| 70 | + dim=0, |
| 71 | + ) |
| 72 | + |
| 73 | + per_block_scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, block_size) |
| 74 | + assert (per_block_scale.float() > 0).all(), ( |
| 75 | + "Zero scales in mixed-magnitude tensor after FP8 cast." |
| 76 | + ) |
| 77 | + |
| 78 | + def test_helper_clamps_overflow_to_max(self): |
| 79 | + """Values above 448 must saturate to 448, not cast to NaN (fp8_e4m3fn has no Inf).""" |
| 80 | + oversized = torch.tensor([100.0, 448.0, 1e3, 1e6]) |
| 81 | + out = _cast_per_block_scale_to_fp8(oversized).float() |
| 82 | + assert torch.isfinite(out).all(), f"FP8 cast produced non-finite values: {out.tolist()}" |
| 83 | + assert (out <= _FP8_E4M3FN_MAX).all(), f"FP8 cast values exceed 448: {out.tolist()}" |
| 84 | + |
| 85 | + def test_helper_clamps_underflow_to_min(self): |
| 86 | + """Values below the FP8 subnormal must clamp up, not collapse to 0.""" |
| 87 | + tiny = torch.tensor([0.0, 1e-12, 1e-6, _FP8_E4M3FN_MIN / 2]) |
| 88 | + out = _cast_per_block_scale_to_fp8(tiny).float() |
| 89 | + assert (out > 0).all(), f"FP8 cast produced zero scales: {out.tolist()}" |
| 90 | + |
| 91 | + def test_static_path_no_nan_when_block_amax_zero(self): |
| 92 | + """Static path: zero-amax block + small global_amax must clamp to 448, not cast to NaN.""" |
| 93 | + block_size = 16 |
| 94 | + # global_amax small enough that 1.0 * 448 / (global_amax/6) >> 448. |
| 95 | + global_amax = torch.tensor(0.01) |
| 96 | + # One block with amax=0 (triggers safety net), three normal blocks. |
| 97 | + per_block_amax = torch.tensor([[0.0, 0.005, 0.008, 0.01]]) |
| 98 | + weight = torch.randn(1, 4 * block_size) |
| 99 | + q = SimpleNamespace( |
| 100 | + global_amax=global_amax, |
| 101 | + _amax=per_block_amax, |
| 102 | + block_sizes={-1: block_size}, |
| 103 | + ) |
| 104 | + |
| 105 | + per_block_scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(q, weight) |
| 106 | + per_block_scale_f32 = per_block_scale.float() |
| 107 | + assert torch.isfinite(per_block_scale_f32).all(), ( |
| 108 | + f"NaN/Inf in exported static per-block scale: {per_block_scale_f32.tolist()}" |
| 109 | + ) |
| 110 | + assert (per_block_scale_f32 <= _FP8_E4M3FN_MAX).all(), ( |
| 111 | + f"Static per-block scale exceeds FP8 max 448: {per_block_scale_f32.tolist()}" |
| 112 | + ) |
](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1382){kind=link}
0 commit comments