Relax test_bf16_stochastic_round tolerance

[karol-brejna-i]

PR Description

TL;DR

This PR enables test/test_low_bit_optim.py in XPU CI (ci_test_xpu.sh).
It also fixes test_low_bit_optim.py::TestQuantize::test_bf16_stochastic_round_device_xpu_compile_False.

test_bf16_stochastic_round's tolerance was too tight for the stochastic-rounding estimator's natural sampling variance, causing a ~2-3% flake rate that appeared as a deterministic PVC-only failure (fixed RNG seed in the test harness). Widened the tolerance to match the actual variance and enabled test/test_low_bit_optim.py in XPU CI.

Details

• Root cause: for large x in the sampled range, the tolerance (rtol·x, up to ~3e-3) was only about 7 standard deviations above the estimator's real noise (std≈4.3e-4) — not enough margin, so it occasionally failed. Fixed seeding then turned that rare flake into a result that always repeats on the same machine.
• Evidence: Monte-Carlo failure rates on BMG / PVC 1550 / PVC 1100 and CPU all ~2-3%, bias ≈0 on all — not a hardware or version bug.
• Fix: widened atol/rtol in the test's final assert_close (no changes to the SR kernel or seeding).
• CI: added test/test_low_bit_optim.py in ci_test_xpu.sh.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures

As of commit 907671e with merge base eec7681 ():

NEW FAILURES - The following jobs have failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run 1xH100 Tests / test (H100, mt-l-x86iamx-22-225-h100, --pre torch torchvision torchaudio mslk --index-url https:/... / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t f2883a4a484718cfdd0aa7ad985a7b14bbb50f92b1e68478ca26715d580df16b /exec failed with exit code 1

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

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[karol-brejna-i]

Detailed analysis and rationale

Summary of the problem

1. The test compares the empirical mean of 100,000 stochastically-rounded BF16 samples against the original FP32 value, with given tolerance (atol/rtol=3e-5):

# test/test_low_bit_optim.py
def test_bf16_stochastic_round(self, device, compile):
    x = torch.rand(32, device=device) * 100
    x_rep = x.view(-1, 1).repeat(1, 100_000)

    func = torch.compile(_fp32_to_bf16_sr, fullgraph=True, dynamic=False, disable=not compile)
    x_rep_bf16 = func(x_rep)
    ...
    torch.testing.assert_close(x_rep_bf16.float().mean(1), x, atol=3e-5, rtol=3e-5)

2. In our runs, the test consistently passes on BMG and PVC 1550, and consistently fails on PVC 1100 — the failure always at the same index (24) with the same magnitude.

3. The test harness does not use fresh randomness per run. test/test_low_bit_optim.py pins the seed: if common_utils.SEED is None: common_utils.SEED = 1234, and PyTorch's TestCase.setUp() re-applies torch.manual_seed(1234) before every test — so x and the rounding draws are identical on every invocation.

Problem analysis

Due to seeding, the rounding draws are identical on every invocation. This makes the outcome fully deterministic per machine. That explains why we have the same results on the same machine.

On the other hand, the same seed does not produce the same random bits on every device: PyTorch's on-device RNG (counter-based, Philox-style) partitions its counter space according to the kernel launch geometry, which depends on the GPU's execution resources (tile/EU count, SIMD width). So a fixed seed lands on a slightly different sample per GPU model. This explains why the test may behave differently on PVC 1100.

We can use the following script (Monte Carlo like experiment that repeats the test logic multiple times) to measure the true statistical behavior (independent of the fixed-seed), and real error distribution on a given device:

import torch
from torchao.optim.quant_utils import _fp32_to_bf16_sr

atol = rtol = 3e-5
trials = 1000

for dev in ["cpu", "xpu"]:
    fails = 0
    all_diffs = []
    for _ in range(trials):
        x = torch.rand(32, device=dev) * 100
        xr = x.view(-1, 1).repeat(1, 100_000)
        out = _fp32_to_bf16_sr(xr).float().mean(1)
        diff = out - x
        tol = atol + rtol * x.abs()
        if (diff.abs() > tol).any():
            fails += 1
        all_diffs.append(diff)

    diffs = torch.cat(all_diffs)
    print(f"{dev}: {fails}/{trials} failures ({100*fails/trials:.2f}%) | "
          f"mean(bias)={diffs.mean().item():.6f} std={diffs.std().item():.6f} "
          f"max|diff|={diffs.abs().max().item():.6f}")

The results are:

Machine	XPU fail-rate	CPU fail-rate	Bias	std(diff) (xpu/cpu)	max|diff| (xpu/cpu)
BMG (Arc Pro B70)	2.40%	2.00%	≈0	0.000437 / 0.000430	0.002953 / 0.002823
PVC Max 1550 (uli)	2.50%	3.30%	≈0	0.000437 / 0.000434	0.002998 / 0.002747
PVC Max 1100 (sdp-pvc)	2.70%	2.50%	≈0	0.000436 / 0.000435	0.002808 / 0.002800

It shows two things:

• the test is a bit "flaky" and could fail in certain circumstances (~2-3% failure rate per test invocation).
• the observed worst-case error (max|diff|) reaches ~2.8e-3 to ~3.0e-3 on all three machines — large enough to occasionally exceed the current tolerance (hence the ~2-3% flake), while bias stays ≈0. So the operation is statistically correct (unbiased), just noisier than the current tolerance allows.

Proposed fix

I propose loosening the test tolerance to atol=rtol=5e-3:

    torch.testing.assert_close(x_rep_bf16.float().mean(1), x, atol=5e-3, rtol=5e-3)

This is justified by the measured error distribution. std(diff) across all three machines/backends is tightly clustered around 4.3e-4 to 4.4e-4, and the largest single deviation observed is ~3.0e-3 (worst case across 32,000 samples per device). Both are comfortably below 5e-3, so this value puts the tolerance floor above even the worst observed error, which means the test passes across the whole range of x. At the same time 5e-3 is still far tighter than any genuinely biased or broken SR implementation would produce (that would deviate by orders of magnitude more), so it doesn't weaken the test's ability to catch real regressions.

This makes the test pass in CI, while eliminating false positives from sampling variance alone.

[xiaowangintel]

Summary

This failure is not device-specific. After changing the test seed configuration, the same failure was reproduced on other hardware platforms as well：

Scenario	Seed	Result
XPU (original CI failure)	1234 (default)	❌
CPU, manual_seed(42)	42	❌
CUDA, manual_seed(16)	16	❌

Reproduce:

def test_bf16_stochastic_round(self, device, compile):
    torch.seed(42) # For cuda, set it to 16
    x = torch.rand(32, device=device) * 100
    x_rep = x.view(-1, 1).repeat(1, 100_000)

    func = torch.compile(_fp32_to_bf16_sr, fullgraph=True, dynamic=False, disable=not compile)
    x_rep_bf16 = func(x_rep)
    ...
    torch.testing.assert_close(x_rep_bf16.float().mean(1), x, atol=3e-5, rtol=3e-5)

LOG

CPU log

======================================================================== test session starts ========================================================================
test/test_low_bit_optim.py::TestQuantize::test_bf16_stochastic_round_device_cpu_compile_False
FAILED

============================================================================= FAILURES ==============================================================================
_________________________________________________ TestQuantize.test_bf16_stochastic_round_device_cpu_compile_False __________________________________________________

self = <test.test_low_bit_optim.TestQuantize testMethod=test_bf16_stochastic_round_device_cpu_compile_False>, device = 'cpu', compile = False

    @parametrize("device", _DEVICES)
    @parametrize("compile", [False, True])
    def test_bf16_stochastic_round(self, device, compile):
        torch.manual_seed(42)
        x = torch.rand(32, device=device) * 100
        x_rep = x.view(-1, 1).repeat(1, 100_000)

        func = torch.compile(
            _fp32_to_bf16_sr, fullgraph=True, dynamic=False, disable=not compile
        )
        x_rep_bf16 = func(x_rep)
        assert x_rep_bf16.dtype is torch.bfloat16

        # must cast BF16 tensor back to FP32 so that .mean() is accurate
>       torch.testing.assert_close(x_rep_bf16.float().mean(1), x, atol=3e-5, rtol=3e-5)
E       AssertionError: Tensor-likes are not close!
E
E       Mismatched elements: 1 / 32 (3.1%)
E       Greatest absolute difference: 0.0028076171875 at index (0,) (up to 3e-05 allowed)
E       Greatest relative difference: 3.18226775561925e-05 at index (0,) (up to 3e-05 allowed)
E
E       To execute this test, run the following from the base repo dir:
E           python test/test_low_bit_optim.py TestQuantize.test_bf16_stochastic_round_device_cpu_compile_False
E
E       This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0

test/test_low_bit_optim.py:122: AssertionError
====================================================================== short test summary info ======================================================================
FAILED test/test_low_bit_optim.py::TestQuantize::test_bf16_stochastic_round_device_cpu_compile_False - AssertionError: Tensor-likes are not close!
========================================================================= 1 failed in 2.29s =========================================================================

CUDA log

======================================================================== test session starts ========================================================================
test/test_low_bit_optim.py::TestQuantize::test_bf16_stochastic_round_device_cuda_compile_False
FAILED

============================================================================= FAILURES ==============================================================================
_________________________________________________ TestQuantize.test_bf16_stochastic_round_device_cuda_compile_False _________________________________________________

self = <test.test_low_bit_optim.TestQuantize testMethod=test_bf16_stochastic_round_device_cuda_compile_False>, device = 'cuda', compile = False

    @parametrize("device", _DEVICES)
    @parametrize("compile", [False, True])
    def test_bf16_stochastic_round(self, device, compile):
        torch.manual_seed(16)
        x = torch.rand(32, device=device) * 100
        x_rep = x.view(-1, 1).repeat(1, 100_000)

        func = torch.compile(
            _fp32_to_bf16_sr, fullgraph=True, dynamic=False, disable=not compile
        )
        x_rep_bf16 = func(x_rep)
        assert x_rep_bf16.dtype is torch.bfloat16

        # must cast BF16 tensor back to FP32 so that .mean() is accurate
>       torch.testing.assert_close(x_rep_bf16.float().mean(1), x, atol=3e-5, rtol=3e-5)
E       AssertionError: Tensor-likes are not close!
E
E       Mismatched elements: 1 / 32 (3.1%)
E       Greatest absolute difference: 0.001636505126953125 at index (0,) (up to 3e-05 allowed)
E       Greatest relative difference: 4.938915662933141e-05 at index (0,) (up to 3e-05 allowed)
E
E       To execute this test, run the following from the base repo dir:
E           python test/test_low_bit_optim.py TestQuantize.test_bf16_stochastic_round_device_cuda_compile_False
E
E       This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0

test/test_low_bit_optim.py:122: AssertionError
====================================================================== short test summary info ======================================================================
FAILED test/test_low_bit_optim.py::TestQuantize::test_bf16_stochastic_round_device_cuda_compile_False - AssertionError: Tensor-likes are not close!
========================================================================= 1 failed in 3.95s =========================================================================

[xiaowangintel]

The test_bf16_stochastic_round test failure is unrelated to the device (CPU/CUDA/XPU) and is purely a statistical design flaw in the test itself:

The test uses 32 random samples, each repeated over 100,000 stochastic rounding draws, and takes the mean to verify the unbiasedness of stochastic rounding.
However, the mean over 100,000 draws inherently carries unavoidable statistical error (standard error SE ≈ 60, corresponding to a relative error on the order of 1e-3 ~ 1e-4).
The tolerance set in the test (atol=5e-3, rtol=5e-5) happens to be very close to this statistical error magnitude.

Fix recommendation:
Relax the tolerance to cover the actual magnitude of statistical error. Alternatively, increase the sample count (e.g., from 100,000 to 500,000 / 1,000,000) to narrow the statistical error.

[xiaowangintel]

@karol-brejna-i Could you please investigate whether analogous operations exist in PyTorch, or examine relevant fp32-to-bf16 casting cases to determine their tolerance settings, and ensure that torchao is aligned with PyTorch in this regard.