[ROCm][CI] Use larger atol value for INT3 in test_quick_all_reduce.py (#50450)
Signed-off-by: Dino Music <Dino.Music@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
music-dino committed
Jul 31, 2026 at 04:41 UTC
5d5f22ed7a556ecafc005bfeb34b89cadf91501b
1 file changed
+14
-4
tests/distributed/test_quick_all_reduce.py
+14
-4
@@ -38,6 +38,14 @@ def on_gfx942() -> bool:
38
return False
39
40
41
+def _quickreduce_tolerance(quant_level: QuickReduceRegime) -> dict:
42
+ # With INT3 inputs of absmax 23, per-rank error peaks at 2.75,
43
+ # doubled by the TP2 sum = 5.5.
44
+ if quant_level == QuickReduceRegime.INT3:
45
+ return {"atol": 5.5, "rtol": 0.1}
46
+ return {"atol": 2.5, "rtol": 0.1}
47
+
48
+
49
set_random_seed(42)
50
_test_size_rng = random.Random(44)
51
# Size over 8MB is sufficient for custom quick allreduce.
@@ -280,8 +288,9 @@ def graph_quickreduce(
288
out2 = tensor_model_parallel_all_reduce(inp2)
289
dist.all_reduce(inp2, group=group)
290
graph.replay()
283
- torch.testing.assert_close(out1, inp1, atol=2.5, rtol=0.1)
284
- torch.testing.assert_close(out2, inp2, atol=2.5, rtol=0.1)
291
+ tol = _quickreduce_tolerance(fa.qr_quant_level)
292
+ torch.testing.assert_close(out1, inp1, **tol)
293
+ torch.testing.assert_close(out2, inp2, **tol)
294
295
296
@ray.remote(num_gpus=1, max_calls=1)
@@ -309,14 +318,15 @@ def eager_quickreduce(
318
)
319
_assert_quickreduce(fa, inp)
320
out = fa.quick_all_reduce(inp)
312
- torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
321
+ tol = _quickreduce_tolerance(fa.qr_quant_level)
322
+ torch.testing.assert_close(out, inp * tp_size, **tol)
323
324
inp = torch.tensor(
325
[1.0 * ((i) % 23) for i in range(sz)], dtype=torch.bfloat16, device=device
326
)
327
_assert_quickreduce(fa, inp)
328
out = fa.quick_all_reduce(inp)
319
- torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
329
+ torch.testing.assert_close(out, inp * tp_size, **tol)
330
331
332
@ray.remote(num_gpus=1, max_calls=1)