De-flake quantize tests and break pt2_compliant_tag cascade (#5911)

Summary:
Pull Request resolved: #5911

Opcheck / test-health bookkeeping fixes for the fbgemm_dev OMH dashboard (T191384137):

1. (2c) pt2_compliant_tag cascade (test/jagged/common.py): when an op is tagged pt2_compliant_tag but has legitimately-xfailed opcheck sub-tests (recorded in failures_dict.json), the auto-generated test_pt2_compliant_tag_fbgemm_<op> test still fails with a cascading assertion. These cannot be expressed in failures_dict.json, so skip them via additional_decorators until the underlying opcheck gaps are fixed. Covers jagged_dense_elementwise_add and jagged_to_padded_dense.

2. (2e) Numeric de-flaking (test/quantize/bfloat16_test.py, mixed_dim_int8_test.py): the bf16 round-trip / numpy-reference comparisons previously used assert_close defaults (or a dynamic atol that always collapsed to 1e-2, since torch.rand inputs are in [0,1)). Replace with explicit rtol=1e-2, atol=1e-2 -- rtol already scales tolerance with magnitude, so this loosens the tight defaults without misleading dead code. mixed_dim_int8_test.py keeps the magnitude-scaled atol because its inputs come from torch.randn (magnitudes > 1), where the scaling legitimately applies.

3. (2d) Fix histogram_binning_calibration FakeTensor device bug + remove now-valid xfail: fbgemm::histogram_binning_calibration test_faketensor was marked xfail in test/sparse/failures_dict.json. The meta impl histogram_binning_calibration_abstract existed but allocated the bin_ids output on CPU instead of logit.device, so on GPU the fake output device diverged from the real kernel, failing test_faketensor (LAND_BLOCKING). Added device=logit.device to the bin_ids allocation (matching the real CPU/CUDA kernels and the sibling generic_histogram_binning_calibration_by_feature impl), then removed the now-passing xfail entry. NOTE: other failures_dicts should be re-swept with PYTORCH_OPCHECK_ACCEPT=1 on NVIDIA to catch any remaining XPASS entries.

Reviewed By: spcyppt

Differential Revision: D108686545

fbshipit-source-id: c38a8edad8ed5a8cd4989352e68520937e5fe310

Author: q10

fbgemm_gpu/fbgemm_gpu/sparse_ops.py

@@ -1191,7 +1191,11 @@ def histogram_binning_calibration_abstract(
     bin_ctr_in_use_after: int,
     bin_ctr_weight_value: float,
 ) -> tuple[Tensor, Tensor]:
-    return torch.empty_like(logit), torch.empty([logit.numel()], dtype=torch.int64)
+    # bin_ids must be allocated on logit's device to match the real kernel,
+    # otherwise the faketensor opcheck fails with a device mismatch on GPU.
+    return torch.empty_like(logit), torch.empty(
+        [logit.numel()], dtype=torch.int64, device=logit.device
+    )
 
 
 def float_to_hfp8_quantized(

fbgemm_gpu/test/jagged/common.py

@@ -9,6 +9,7 @@
 # pyre-ignore-all-errors[56]
 
 import itertools
+import unittest
 from collections.abc import Callable
 from typing import Any
 
@@ -41,7 +42,30 @@
 # e.g. "test_faketensor__test_cumsum": [unittest.expectedFailure]
 # Please avoid putting tests here, you should put operator-specific
 # skips and failures in deeplearning/fbgemm/fbgemm_gpu/test/failures_dict.json
-additional_decorators: dict[str, list[Callable[..., Any]]] = {}
+#
+# The exception is the auto-generated `test_pt2_compliant_tag_fbgemm_<op>`
+# tests. When an op is tagged `pt2_compliant_tag` but has legitimately-xfail'd
+# opcheck sub-tests (recorded in failures_dict.json), the tag test fails with a
+# cascading "failed some of the generated opcheck tests" assertion even though
+# the underlying gaps are already tracked. These cannot be expressed in
+# failures_dict.json (they are not per-op-per-input entries), so they are
+# skipped here until the underlying opcheck gaps are fixed (T191384137).
+additional_decorators: dict[str, list[Callable[..., Any]]] = {
+    # Cascades from xfail'd ElementwiseBinaryTest.test_aot_dispatch_dynamic /
+    # test_autograd_registration __test_jagged_elementwise_binary entries.
+    "test_pt2_compliant_tag_fbgemm_jagged_dense_elementwise_add": [
+        unittest.skip(
+            "pt2_compliant_tag test cascades from xfail'd jagged_dense_elementwise_add opcheck sub-tests (T191384137)"
+        ),
+    ],
+    # Cascades from xfail'd JaggedToPaddedDenseTest.test_aot_dispatch_dynamic
+    # __test_jagged_to_padded_dense entry.
+    "test_pt2_compliant_tag_fbgemm_jagged_to_padded_dense": [
+        unittest.skip(
+            "pt2_compliant_tag test cascades from xfail'd jagged_to_padded_dense opcheck sub-tests (T191384137)"
+        ),
+    ],
+}
 
 
 def lengths_to_segment_ids(lengths: torch.Tensor) -> torch.Tensor:

fbgemm_gpu/test/quantize/bfloat16_test.py

@@ -35,6 +35,9 @@ def test_dense_mlp_quantize_ops(
             input_data = torch.rand((n, k), dtype=torch.float32)
             quantized_data = torch.ops.fbgemm.FloatToBfloat16Quantized(input_data)
             dequantized_data = torch.ops.fbgemm.Bfloat16QuantizedToFloat(quantized_data)
+            # bf16 round-trip error is proportional to the value magnitude,
+            # which the relative (rtol) term of assert_close already accounts
+            # for; atol provides a floor for near-zero values.
             torch.testing.assert_close(
                 dequantized_data, input_data, rtol=1e-2, atol=1e-2
             )
@@ -100,7 +103,11 @@ def test_quantize_and_dequantize_op(self, nrows: int, ncols: int) -> None:
         ref_bfloat16 = f(input_data.numpy())
         f = np.vectorize(lambda x: bfloat_dequantize(x))
         ref_fp32 = torch.from_numpy(f(ref_bfloat16)).float()
-        torch.testing.assert_close(dequantized_data, ref_fp32)
+        # The kernel and the numpy reference may round the bf16 mantissa
+        # differently (round-to-even vs round-half-away), so loosen the tight
+        # default tolerance: rtol scales with magnitude and atol floors
+        # near-zero values.
+        torch.testing.assert_close(dequantized_data, ref_fp32, rtol=1e-2, atol=1e-2)
 
         if torch.cuda.is_available():
             input_data_gpu = input_data.cuda()
@@ -111,7 +118,9 @@ def test_quantize_and_dequantize_op(self, nrows: int, ncols: int) -> None:
                 quantized_data_gpu
             )
             # compare quantized data
-            torch.testing.assert_close(dequantized_data_gpu.cpu(), ref_fp32)
+            torch.testing.assert_close(
+                dequantized_data_gpu.cpu(), ref_fp32, rtol=1e-2, atol=1e-2
+            )
 
     @unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
     # pyre-fixme[56]: Pyre was not able to infer the type of argument
@@ -137,8 +146,11 @@ def test_quantize_and_dequantize_op_cuda_large_nrows_bf16(
             dequantized_data_gpu = torch.ops.fbgemm.Bfloat16QuantizedToFloat(
                 quantized_data_gpu
             )
-            # compare quantized data
-            torch.testing.assert_close(dequantized_data_gpu.cpu(), dequantized_data)
+            # GPU-vs-CPU bf16 rounding can differ; rtol handles magnitude
+            # scaling and atol floors near-zero values.
+            torch.testing.assert_close(
+                dequantized_data_gpu.cpu(), dequantized_data, rtol=1e-2, atol=1e-2
+            )
 
 
 if __name__ == "__main__":

fbgemm_gpu/test/quantize/mixed_dim_int8_test.py

@@ -151,7 +151,19 @@ def run_mixed_dim_8bit_dequantize_op_test(
         if output_dtype == SparseType.FP16:
             output_ref_concat = output_ref_concat.half()
 
-        torch.testing.assert_close(output_ref_concat, mixed_dim_dequant_output)
+        # Dynamic tolerance: int8 row-wise dequantization error is proportional
+        # to the per-row span (~span/255), so for inputs with larger magnitudes
+        # the absolute error grows. Scale atol by the max abs of the expected
+        # tensor instead of relying on the tight default tolerance, which makes
+        # the comparison flaky on randomly generated data.
+        atol = (
+            max(1e-2, 1e-2 * float(output_ref_concat.abs().max().float()))
+            if output_ref_concat.numel()
+            else 1e-2
+        )
+        torch.testing.assert_close(
+            output_ref_concat, mixed_dim_dequant_output, rtol=1e-2, atol=atol
+        )
 
 
 if __name__ == "__main__":

fbgemm_gpu/test/sparse/failures_dict.json

@@ -51,12 +51,7 @@
     },
     "fbgemm::generic_histogram_binning_calibration_by_feature": {},
     "fbgemm::group_index_select_dim0": {},
-    "fbgemm::histogram_binning_calibration": {
-      "HistogramBinningCalibrationTest.test_faketensor__test_histogram_binning_calibration": {
-        "comment": "",
-        "status": "xfail"
-      }
-    },
+    "fbgemm::histogram_binning_calibration": {},
     "fbgemm::histogram_binning_calibration_by_feature": {
       "HistogramBinningCalibrationTest.test_aot_dispatch_dynamic__test_histogram_binning_calibration_by_feature": {
         "comment": "",