HistogramObserver: score the full range before trimming in _non_linear_param_search

[pquochuy]

Summary:
HistogramObserver._non_linear_param_search initializes norm_min = float("inf") and only scores ranges after one quantile trim, so the untrimmed range (start_bin=0, end_bin=bins-1) is never a candidate and the first trim (1e-5 of the element count, taken from the heavier side) is always accepted. A sparse far tail of <= 1e-5 of the elements is therefore always clipped, even when the observer's own _compute_quantization_error rates that orders of magnitude worse than keeping the full range.

Fix: seed norm_min with _compute_quantization_error(0, bins - 1). The loop is unchanged; a trim is accepted only if it does not increase the error relative to the current best, including the untrimmed range. Ties still accept the trim (norm > norm_min is kept), which preserves the old behavior in degenerate cases where every range scores 0 (e.g. dst_bin_width == 0).

The same change is applied to torch/ao/quantization/observer.py and to the torchao copy in torchao/quantization/pt2e/observer.py (used by the PT2E quantizers).

Repro (synthetic): 2M values, 64% zeros, exponential bulk plus 20 tail values in [39, 136], int16 per-tensor symmetric, bins=2048.

• Before: chosen range 38.93, observer error 9.0e4, reconstruction SSE 9.0e4.
• After: chosen range 136.28, observer error 0.72, reconstruction SSE 1.04 (same as MinMaxObserver).
• A plain randn(2M) at int16 is also clipped today (range 4.24 vs max 5.22; observer error 1.69 vs 3.7e-3 at full range). At qint8 the clipped range is still chosen after the fix (error 168 vs 240 at full range).

Test changes in test/quantization/core/test_workflow_module.py:

• _ReferenceHistogramObserver mirrors the same seed so test_histogram_observer_against_reference keeps comparing like with like.
• test_histogram_observer hard-coded the old result: a bins=3 histogram [2, 3, 3] over [2, 8], where the old search clipped the top bin (3 of 8 values) at error 4.05 vs 3.7e-4 for the full range. Expected scales are updated to the full range (affine 8/255, symmetric 8/127.5; zero points unchanged).
• New test_histogram_observer_keeps_sparse_far_tail (int16 symmetric, 1M values with a 1e-5 far tail): the chosen range equals the data range and its error is <= the full-range error.
• New test_histogram_observer_still_clips_when_it_lowers_error (qint8 symmetric, randn): the range is still trimmed on both sides and the trimmed error is < the full-range error.

Differential Revision: D123598050

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4975

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit d6680b4 with merge base e31a5c9 ():

NEW FAILURE - The following job has failed:

• Run PT2E GPU Tests / test-rocm (ROCM Nightly, linux.rocm.gpu.ecosystem.mi350.1, --pre torchvision torch --index-url ht... / linux-job (gh)
  Could not load credentials from any providers

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@pquochuy has exported this pull request. If you are a Meta employee, you can view the originating Diff in D123598050.

[pquochuy]

Superseded by #4976, re-exported from D123598050 to an upstream branch to resolve the fork-related ROCm OIDC permission issue. Please continue review and CI tracking on #4976.