HistogramObserver: score the full range before trimming in _non_linear_param_search (#4975)

[pquochuy]

Summary:
X-link: pytorch/pytorch#199981

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/4976

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 1ba7cce with merge base cff77b4 ():
💚 Looks good so far! There are no failures yet. 💚

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.