fma can cause drastically worse precision in torch.compile/Triton

[Chillee]

🐛 Describe the bug

import torch
torch.set_default_device('cuda')

scale = torch.tensor(0.180336877703666687)
x = torch.tensor(1134139801600.000000)

def f(x, scale):
    max_scaled = x * scale
    return torch.exp(max_scaled - x * scale)

print(f(x, scale))
print(torch.compile(f)(x, scale))
>>> tensor(1., device='cuda:0')
>>> tensor(inf, device='cuda:0')

The root cause here is the same as here: #121558 (comment)

I think the general structure of this can show up in many cases, but in this case, this is a Triton issue. Here, Inductor actually CSE's the x * scale call.

The generated Triton code is

@triton_heuristics.pointwise(
    size_hints=[1],
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {3: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=(3,), ids_of_folded_args=(3,), divisible_by_8=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_exp_mul_sub_0', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': 'cddd2cc4107921f6715d58b181fdb08b23055085471461101752ba6efb772ae1'},
    min_elem_per_thread=0
)
@triton.jit
def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    tmp0 = tl.load(in_ptr0 + (0))
    tmp1 = tl.broadcast_to(tmp0, [XBLOCK])
    tmp2 = tl.load(in_ptr1 + (0))
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK])
    tmp4 = tmp1 * tmp3
    tmp5 = tmp4 - tmp4
    tmp6 = tl_math.exp(tmp5)
    tl.store(out_ptr0 + (tl.full([XBLOCK], 0, tl.int32)), tmp6, None)

Here, you would imagine that tmp4 - tmp4 is 0. Unfortunately, Triton generates this PTX.

        mul.f32         %f5, %f3, %f4;
        .loc    1 30 18
        neg.f32         %f6, %f5;
        fma.rn.f32      %f7, %f3, %f4, %f6;

Which duplicates the FMA again!

Versions

N/A

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @desertfire

[Chillee]

Turning off fmas in Triton also solves this problem. #122179

[Chillee]

Here are some more repros, in increasing order of sadness.

1 vs inf. even without fma

Sadly, we can also replicate this with torch.compile's fusion even if we turn off fma.

import torch
torch.set_default_device('cuda')

scale = torch.tensor(0.180336877703666687, dtype=torch.bfloat16)
x = torch.tensor(1134139801600.000000, dtype=torch.bfloat16)

def f(x, scale):
    max_scaled = x * scale
    max_scaled = torch.ops.aten._lazy_clone(max_scaled)
    return torch.exp(max_scaled - x * scale)

print(f(x, scale))
print(torch.compile(f)(x, scale))
>>> tensor(1., device='cuda:0')
>>> tensor(inf, device='cuda:0')

Here, compile computes x * scale in fp32 and rounds to bf16 when storing. Then, the divergence with torch.exp(max_scaled - x * scale) (where the entire computation is performed in fp32) causes the inf.

nans even with normal range inputs

Here is a less synthetic example. Here, I've taken the softmax implementation and implemented it in user land (without casting to float32).

import torch
torch.set_default_device('cuda')

def _softmax(x, dim):
    # eager softmax returns a contiguous tensor. Ensure that decomp also returns
    # a contiguous tensor.
    x = x.contiguous()
    if x.numel() == 0:
        unnormalized = torch.exp(x)
    else:
        x_max = torch.amax(x, dim, keepdim=True)
        unnormalized = torch.exp(x - x_max)
    result = unnormalized / torch.sum(unnormalized, dim, keepdim=True)
    return result

def f(x, y):
    x = torch.matmul(x, x)
    x = (x / y)
    return _softmax(x, dim=-1)

torch.manual_seed(1234)
x = torch.tensor([[[[-16.1649,   5.6846,  -5.1022,  -9.1134],
          [-11.5552,  -2.2615, -12.8913,  10.6538],
          [ -7.1666,  -5.3333,   2.0776,  -9.7984],
          [  7.4469,  -2.3948,   2.7371,   0.9201]],

         [[ -8.0361, -16.3771,  22.7741,   4.4685],
          [ 20.8047,  -0.7771,  -2.4355,  -2.2299],
          [  3.8343,  -2.0914,  -2.4077,   2.2740],
          [-15.8663,  -2.7015, -12.5241,  -3.0040]],

         [[ -2.5139,  14.4393,  -3.7186,   1.2255],
          [  5.6742,  14.1842,  -8.5976,  16.8366],
          [ -9.7358,  -3.0279,  11.8164,  -4.0787],
          [ -9.0621,   8.2580,  29.9486,  -2.4107]],

         [[  7.3622,  12.5640, -20.5592,  13.6237],
          [-11.5640,   0.8832,  16.7275,  -2.5009],
          [ -2.0953, -12.2276, -26.2633,   4.5268],
          [ 15.3329, -11.7492,   6.5650,  -9.2483]]],


        [[[  7.9980,  -4.9369,   3.1508,   5.2994],
          [  3.8052,   3.9514,   8.4987, -10.5045],
          [ -2.6827,  -4.0010,  -4.0611,   6.4091],
          [-19.0318,   6.4073,   2.8923,   8.0250]],

         [[  7.1650,  -3.4585,   5.7720,  -5.0305],
          [ -0.9765,  -3.0086,  11.7114,   8.0555],
          [ -3.1027,  -3.5514,   9.6182,  -8.8526],
          [ -9.2348,  -6.0239,   6.2528,  -6.7221]],

         [[ 11.5936,  22.4139,  -0.4089,  -4.9889],
          [ 14.8217,  -2.3426, -17.6189,   3.7427],
          [  1.9546, -13.0902,   8.6293,  -7.2457],
          [ -7.6900,  -4.5796,   9.6332, -10.2631]],

         [[  0.8027,  -1.0955,  14.8404,  -0.2673],
          [  3.2143,  -1.8640,  -2.9678,   6.5165],
          [ -3.9865,   6.5230,   6.3019,  -0.4247],
          [  8.3185, -13.5076,  27.0986,  -1.6792]]]], device='cuda:0', dtype=torch.bfloat16)
y = torch.tensor([[[ 0.6331]],

        [[ 1.6358]],

        [[-0.3459]],

        [[ 1.0196]]], device='cuda:0', requires_grad=True, dtype=torch.bfloat16)
f(x, y).sum().backward()
print(y.grad)
y.grad = None
torch.compile(f)(x, y).sum().backward()
print(y.grad)

tensor([[[ 9.7656e-04]],

        [[ 7.8201e-05]],

        [[ 0.0000e+00]],

        [[-4.2725e-04]]], device='cuda:0', dtype=torch.bfloat16)

tensor([[[nan]],

        [[nan]],

        [[nan]],

        [[nan]]], device='cuda:0', dtype=torch.bfloat16)

[Chillee]

I also find this repro interesting, since it shows how things can bad in the context of optimizations like rematerialization.

There are 3 issues that go on here (this is the code: https://pastebin.com/80pbhpLt)

1. We choose to unroll sum/amax, preventing fusion with the downstream div_1.
2. We choose to rematerialize most of the pointwise ops (div, sub, and exp) (i.e. we compute it in both the first kernel and the second kernel). This is also suboptimal, since we load both inputs in two kernels for no reason.
3. In the second kernel, we are computing ((arg0_1 / arg1_1) - buf0).exp() / buf1. Here, because buf0 is supposed to be (arg0_1 / arg1_1).amax(). So, there is some element for which this should be ((arg0_1 / arg1_1) - (arg0_1 / arg1_1)).exp(), which is equal to 1. However, because of fma, this is rewritten as fma(arg0_1, (1/arg_1), -buf0), which does not equal 1.

#122179 (comment)

[lezcano]

See #115435 and the PR after.

[jansel]

I'm fine with just disabling FMA assuming the perf hit isn't bad.