diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c45d80253863..7ecf37801bea 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -64,6 +64,8 @@ Changelog - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. +- Fix AutoCast turning FP64 initializers and ``Constant`` values that exceed the FP32 range into infinities. They are now clamped to the largest FP32 value and the affected tensor is named in a warning. +- Fix AutoCast flipping the sign of an initializer value too small for the target precision; such values now keep their sign. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. 0.48.0 (2026-10-xx) diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py index 58bd678159c5..708a80d77233 100644 --- a/modelopt/onnx/autocast/graphsanitizer.py +++ b/modelopt/onnx/autocast/graphsanitizer.py @@ -517,6 +517,23 @@ def _get_initializer_value(self, name: str, return_array: bool = False) -> np.nd return value if return_array else value.item() return None + @staticmethod + def _clamp_to_fp32_range(data: np.ndarray, name: str) -> np.ndarray: + """Clamp finite values outside the FP32 range, leaving existing +/-inf untouched. + + Casting a finite FP64 value larger than FP32's max straight to FP32 silently turns it + into infinity. Genuine infinities already in the data are not "too large for FP32" and + must be preserved as-is. + """ + fp32_max = np.finfo(np.float32).max + out_of_range = np.isfinite(data) & (np.abs(data) > fp32_max) + if np.any(out_of_range): + logger.warning( + f"{name} contains FP64 values outside the FP32 range; clamping to {fp32_max}." + ) + data = np.where(out_of_range, np.copysign(fp32_max, data), data) + return data + def _convert_fp64_initializers(self) -> bool: """Convert FP64 initializers to FP32. @@ -529,6 +546,7 @@ def _convert_fp64_initializers(self) -> bool: if initializer.data_type == onnx.TensorProto.DOUBLE: # Convert the data to FP32 fp64_data = numpy_helper.to_array(initializer, base_dir=self.external_data_dir) + fp64_data = self._clamp_to_fp32_range(fp64_data, initializer.name) fp32_data = fp64_data.astype(np.float32) # Create new initializer with FP32 data @@ -588,6 +606,9 @@ def _convert_fp64_nodes(self) -> bool: if attr.name == "value" and attr.t.data_type == onnx.TensorProto.DOUBLE: # Convert the tensor value to FP32 fp64_data = numpy_helper.to_array(attr.t, base_dir=self.external_data_dir) + fp64_data = self._clamp_to_fp32_range( + fp64_data, f"{node.op_type} node {node.name}" + ) fp32_data = fp64_data.astype(np.float32) new_tensor = numpy_helper.from_array(fp32_data) attr.t.CopyFrom(new_tensor) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index a78249fdca4c..32e5c758eaeb 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -1240,7 +1240,7 @@ def _convert_initializer_data( self._warned_values_clamp_min = True np_array = np.where( (np_array != 0.0) & (np.abs(np_array) < data_lowest), - data_lowest, + np.copysign(data_lowest, np_array), np_array, ) new_array = np_array.astype(to_type.numpy_type) diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py index f42324a6338b..2670bb1cab42 100644 --- a/tests/unit/onnx/autocast/test_graphsanitizer.py +++ b/tests/unit/onnx/autocast/test_graphsanitizer.py @@ -235,6 +235,104 @@ def test_convert_fp64_initializers(): assert init.data_type == TensorProto.FLOAT +def test_convert_fp64_initializers_clamps_out_of_range_values(): + """Values FP32 cannot represent must be clamped, not cast to infinity.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2]) + fp64_max = np.finfo(np.float64).max + + graph = helper.make_graph( + nodes=[ + helper.make_node("Clip", ["X", "lower", "upper"], ["Y"], name="clip"), + ], + name="fp64_out_of_range", + inputs=[x], + outputs=[y], + initializer=[ + numpy_helper.from_array(np.array([-fp64_max], np.float64), name="lower"), + numpy_helper.from_array(np.array([fp64_max], np.float64), name="upper"), + ], + ) + sanitizer = GraphSanitizer(helper.make_model(graph)) + + assert sanitizer._convert_fp64_initializers() is True + + fp32_max = np.finfo(np.float32).max + converted = { + init.name: numpy_helper.to_array(init) for init in sanitizer.model.graph.initializer + } + for name in ("lower", "upper"): + assert np.all(np.isfinite(converted[name])), f"{name} was cast to infinity" + assert converted["lower"] == -fp32_max + assert converted["upper"] == fp32_max + + +def test_convert_fp64_initializers_preserves_infinity(): + """An infinity the model already held is not a value FP32 cannot represent.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3]) + # An oversized finite value alongside the infinities must not drag them down with it. + values = np.array([np.inf, -np.inf, np.finfo(np.float64).max], dtype=np.float64) + + graph = helper.make_graph( + nodes=[helper.make_node("Mul", ["X", "mixed"], ["Y"], name="mul")], + name="fp64_with_infinity", + inputs=[x], + outputs=[y], + initializer=[numpy_helper.from_array(values, name="mixed")], + ) + sanitizer = GraphSanitizer(helper.make_model(graph)) + + assert sanitizer._convert_fp64_initializers() is True + + converted = numpy_helper.to_array( + next(init for init in sanitizer.model.graph.initializer if init.name == "mixed") + ) + assert converted[0] == np.inf + assert converted[1] == -np.inf + assert converted[2] == np.finfo(np.float32).max + + +@pytest.mark.parametrize("op_type", ["Constant", "ConstantOfShape"]) +def test_convert_fp64_nodes_clamps_out_of_range_values(op_type): + """A node's FP64 value attribute must be clamped the same way an initializer is.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2]) + value = numpy_helper.from_array( + np.array([np.finfo(np.float64).max], np.float64), name="big_value" + ) + + initializer = [] + if op_type == "ConstantOfShape": + producer = helper.make_node( + "ConstantOfShape", ["shape"], ["big"], name="big_node", value=value + ) + initializer.append(numpy_helper.from_array(np.array([2], dtype=np.int64), name="shape")) + else: + producer = helper.make_node("Constant", [], ["big"], name="big_node", value=value) + + graph = helper.make_graph( + nodes=[producer, helper.make_node("Mul", ["X", "big"], ["Y"], name="mul")], + name="fp64_node_out_of_range", + inputs=[x], + outputs=[y], + initializer=initializer, + ) + sanitizer = GraphSanitizer(helper.make_model(graph)) + + assert sanitizer._convert_fp64_nodes() is True + + converted = next( + numpy_helper.to_array(attr.t) + for node in sanitizer.model.graph.node + if node.name == "big_node" + for attr in node.attribute + if attr.name == "value" + ) + assert np.all(np.isfinite(converted)), f"{op_type} value was cast to infinity" + assert converted == np.finfo(np.float32).max + + def test_convert_fp64_io_types(): """Test conversion of FP64 input/output types to FP32.""" # Create inputs and outputs with FP64 types diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index d3e0a786a65a..d3c672a18b9d 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -624,6 +624,38 @@ def test_clamping_fp16_initializers_out_of_range( assert np.all(add_init_fp32_array == add_init_out_of_range) +@pytest.mark.parametrize("use_standalone_type_inference", [True, False]) +def test_underflowing_fp16_initializers_keep_their_sign( + model_with_multiple_consumers, use_standalone_type_inference +): + """A value too small for FP16 is replaced by the smallest subnormal of the same sign.""" + model, value_info_map, initializer_map, node_to_init_map = model_with_multiple_consumers + + smallest = np.finfo(np.float16).smallest_subnormal + add_init_underflow = np.array([[-1e-10, 1e-10]], dtype=np.float32) + model.graph.initializer[1].CopyFrom( + numpy_helper.from_array(add_init_underflow, name="add_init") + ) + + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + use_standalone_type_inference=use_standalone_type_inference, + ) + converter._convert_initializers(low_precision_nodes=["add1", "add2"], high_precision_nodes=[]) + + converted = next( + numpy_helper.to_array(init) + for init in converter.model.graph.initializer + if init.name == "add_init" + ) + assert converted.dtype == np.float16 + assert converted[0, 0] == -smallest, "a negative value must not come back positive" + assert converted[0, 1] == smallest + + @pytest.mark.parametrize("use_standalone_type_inference", [True, False]) def test_bf16_no_clamping_initializers_out_of_range( model_with_multiple_consumers, use_standalone_type_inference