Skip to content

Commit

Permalink
Support histogram_binning_calibration for export
Browse files Browse the repository at this point in the history
Summary: Add fake tensor for `histogram_binning_calibration` which is needed to export old PA.

Differential Revision: D69089371
  • Loading branch information
trieuat authored and facebook-github-bot committed Feb 4, 2025
1 parent bdcce9c commit 422f0f5
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 4 deletions.
25 changes: 25 additions & 0 deletions fbgemm_gpu/fbgemm_gpu/sparse_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -1100,6 +1100,27 @@ def fused_8_bit_rowwise_quantized_to_half(
return torch.empty(output_shape, dtype=torch.float16, device=input_t.device)


def generic_histogram_binning_calibration_by_feature(
logit: Tensor,
segment_value: Tensor,
segment_lengths: Tensor,
num_segments: int,
bin_num_examples: Tensor,
bin_num_positives: Tensor,
bin_boundaries: Tensor,
positive_weight: float,
bin_ctr_in_use_after: int,
bin_ctr_weight_value: float,
) -> Tuple[Tensor, Tensor]:
torch._check(bin_num_examples.numel() == bin_num_positives.numel())
torch._check(
bin_num_examples.numel() == (num_segments + 1) * (bin_boundaries.numel() + 1)
)
return torch.empty_like(logit), torch.empty(
[logit.numel()], dtype=torch.int64, device=logit.device
)


def _setup() -> None:
# pyre-ignore[16]
_setup.done = getattr(_setup, "done", False)
Expand Down Expand Up @@ -1233,6 +1254,10 @@ def impl_autograd(op_name, fn, setup_context: Optional[Callable] = None) -> None
"fbgemm::histogram_binning_calibration",
histogram_binning_calibration_abstract,
)
impl_abstract(
"fbgemm::generic_histogram_binning_calibration_by_feature",
generic_histogram_binning_calibration_by_feature,
)
impl_abstract(
"fbgemm::FloatToHFP8Quantized",
float_to_hfp8_quantized,
Expand Down
8 changes: 4 additions & 4 deletions fbgemm_gpu/test/sparse/failures_dict.json
Original file line number Diff line number Diff line change
Expand Up @@ -141,19 +141,19 @@
"fbgemm::generic_histogram_binning_calibration_by_feature": {
"HistogramBinningCalibrationTest.test_aot_dispatch_dynamic__test_generic_histogram_binning_calibration_by_feature": {
"comment": "",
"status": "xfail"
"status": "xsuccess"
},
"HistogramBinningCalibrationTest.test_aot_dispatch_dynamic__test_generic_histogram_binning_calibration_by_feature_cpu_gpu": {
"comment": "",
"status": "xfail"
"status": "xsuccess"
},
"HistogramBinningCalibrationTest.test_faketensor__test_generic_histogram_binning_calibration_by_feature": {
"comment": "",
"status": "xfail"
"status": "xsuccess"
},
"HistogramBinningCalibrationTest.test_faketensor__test_generic_histogram_binning_calibration_by_feature_cpu_gpu": {
"comment": "",
"status": "xfail"
"status": "xsuccess"
}
},
"fbgemm::group_index_select_dim0": {
Expand Down

0 comments on commit 422f0f5

Please sign in to comment.