pytorch - 💡(How to fix) Fix torch.compile/Inductor produces different result for float16 cast before elementwise add [1 comments, 2 participants]

import torch

print(torch.__version__)

x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) / 10
c = (torch.arange(24) % 5 == 0).reshape(2, 3, 4)

def f(x, c):
    z = torch.where(c, torch.full_like(x, 0.5), torch.full_like(x, -0.5))
    y = z.to(torch.float16) + x.to(torch.float16)
    return y, y.float().sum()

eager = f(x, c)
compiled = torch.compile(f, backend="inductor", fullgraph=True)(x, c)

print("eager sum:   ", eager[1])
print("compiled sum:", compiled[1])
print("max diff:    ", (eager[0].float() - compiled[0].float()).abs().max())

torch.testing.assert_close(compiled[0], eager[0], rtol=0, atol=0)
torch.testing.assert_close(compiled[1], eager[1], rtol=0, atol=0)

---

eager sum:    tensor(20.5996)
compiled sum: tensor(20.6000)
max diff:     tensor(0.0010)
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
[/tmp/ipykernel_37381/3887225747.py](https://localhost:8080/#) in <cell line: 0>()
     18 print("max diff:    ", (eager[0].float() - compiled[0].float()).abs().max())
     19 
---> 20 torch.testing.assert_close(compiled[0], eager[0], rtol=0, atol=0)
     21 torch.testing.assert_close(compiled[1], eager[1], rtol=0, atol=0)

[/usr/local/lib/python3.12/dist-packages/torch/testing/_comparison.py](https://localhost:8080/#) in assert_close(actual, expected, allow_subclasses, rtol, atol, equal_nan, check_device, check_dtype, check_layout, check_stride, msg)
   1629     if error_metas:
   1630         # TODO: compose all metas into one AssertionError
-> 1631         raise error_metas[0].to_error(msg)
   1632 
   1633 

AssertionError: Tensor-likes are not equal!

Mismatched elements: 8 / 24 (33.3%)
Greatest absolute difference: 0.0009765625 at index (1, 2, 2)
Greatest relative difference: 0.0012197494506835938 at index (0, 1, 0)