90 lines
4.0 KiB
Diff
90 lines
4.0 KiB
Diff
diff --git a/frontier/config/quantization_manager.py b/frontier/config/quantization_manager.py
|
|
--- a/frontier/config/quantization_manager.py
|
|
+++ b/frontier/config/quantization_manager.py
|
|
@@ -356,9 +356,17 @@ class QuantizationManager:
|
|
target_precision = self.get_precision(op_name)
|
|
precision_match = target_precision == profiling_precision
|
|
quant_match = profiling_quant_signature == expected_quant_signature
|
|
+ # A mixed-precision profiling CSV records the model's output dtype at
|
|
+ # file level. Quantized operators in that CSV were nevertheless
|
|
+ # measured with the exact quantization scheme identified by the
|
|
+ # quant signature. Treating those samples as BF16 and scaling them
|
|
+ # again would double-apply the FP8 speedup.
|
|
+ exact_fp8_profile = quant_match and target_precision == PrecisionType.FP8
|
|
with self._lock:
|
|
- self._operation_profiling_precision[op_name] = profiling_precision
|
|
- if precision_match:
|
|
+ self._operation_profiling_precision[op_name] = (
|
|
+ target_precision if exact_fp8_profile else profiling_precision
|
|
+ )
|
|
+ if precision_match or exact_fp8_profile:
|
|
self._operation_data_sources[op_name] = "profiling"
|
|
self._operation_approximation_factors.pop(op_name, None)
|
|
self._operation_speedup_factors.pop(op_name, None)
|
|
diff --git a/tests/unit/test_quantization_profile_contract.py b/tests/unit/test_quantization_profile_contract.py
|
|
new file mode 100644
|
|
--- /dev/null
|
|
+++ b/tests/unit/test_quantization_profile_contract.py
|
|
@@ -0,0 +1,61 @@
|
|
+from unittest.mock import MagicMock
|
|
+
|
|
+import pytest
|
|
+
|
|
+from frontier.config.model_config import QuantizationConfig
|
|
+from frontier.config.precision_type import PrecisionType
|
|
+from frontier.config.quantization_manager import QuantizationManager
|
|
+
|
|
+
|
|
+def test_exact_fp8_quant_signature_does_not_rescale_mixed_profile() -> None:
|
|
+ manager = QuantizationManager()
|
|
+ manager.load_config()
|
|
+
|
|
+ quant_config = QuantizationConfig(
|
|
+ quant_method="fp8",
|
|
+ activation_scheme="dynamic",
|
|
+ is_checkpoint_fp8_serialized=True,
|
|
+ weight_block_size=(128, 128),
|
|
+ )
|
|
+ quant_signature = quant_config.get_quant_signature()
|
|
+ model_config = MagicMock()
|
|
+ model_config.get_default_precision.return_value = PrecisionType.BF16
|
|
+ model_config.get_name.return_value = "Qwen3-235B-A22B-FP8"
|
|
+ model_config.torch_dtype = "bfloat16"
|
|
+ model_config.quantization_config = quant_config
|
|
+ model_config.get_quant_signature.return_value = quant_signature
|
|
+
|
|
+ manager.configure_from_model_config(model_config)
|
|
+ manager.register_profiling_metadata(
|
|
+ operation_names=["attn_pre_proj"],
|
|
+ profiling_precision=PrecisionType.BF16,
|
|
+ profiling_quant_signature=quant_signature,
|
|
+ expected_quant_signature=quant_signature,
|
|
+ file_path="linear_op.csv",
|
|
+ )
|
|
+
|
|
+ metadata = {
|
|
+ item["operation"]: item
|
|
+ for item in manager.get_operation_precision_metadata()
|
|
+ }
|
|
+ assert metadata["attn_pre_proj"]["data_source"] == "profiling"
|
|
+ assert metadata["attn_pre_proj"]["approximation_factor"] is None
|
|
+ assert manager.has_precision_mismatch("attn_pre_proj") is False
|
|
+ assert manager.adjust_compute_time("attn_pre_proj", 1.25) == pytest.approx(1.25)
|
|
+
|
|
+
|
|
+def test_mismatched_fp8_quant_signature_still_uses_approximation() -> None:
|
|
+ manager = QuantizationManager()
|
|
+ manager.load_config()
|
|
+ manager._operation_precisions["attn_pre_proj"] = PrecisionType.FP8
|
|
+
|
|
+ manager.register_profiling_metadata(
|
|
+ operation_names=["attn_pre_proj"],
|
|
+ profiling_precision=PrecisionType.BF16,
|
|
+ profiling_quant_signature="none",
|
|
+ expected_quant_signature="method=fp8",
|
|
+ file_path="linear_op.csv",
|
|
+ )
|
|
+
|
|
+ assert manager.has_precision_mismatch("attn_pre_proj") is True
|
|
+ assert manager.adjust_compute_time("attn_pre_proj", 1.0) == pytest.approx(0.5)
|