timm models accuracy regress due to triton pin bump

[shunting314]

🐛 Describe the bug

When bumping triton from old pin (989adb9a29496c22a36ef82ca69cad5dad536b9c), to new pin (45fff310c891f5a92d55445adf8cc9d29df5841e) here #123077 . We found there are accuracy test regression.

Command

time python benchmarks/dynamo/timm_models.py --backend inductor --amp --accuracy --only mixer_b16_224 --training

It pass with the old pin but fail with the new pin.

Update: The problem can also be reproed using levit_128 model.

cc @htyu @manman-ren

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang

[manman-ren]

I will take a look. In general, should we also file a github issue at triton/issues? Or that is redundant.

[shunting314]

should we also file a github issue at triton/issues? Or that is redundant.

Thanks @manman-ren . My understanding is, we can do that if we need help from OAI triton team.

[manman-ren]

I was able to reproduce this

E0401 16:49:46.154000 139907086283904 torch/_dynamo/utils.py:1405] RMSE (res-fp64): 0.00146, (ref-fp64): 0.00015 and shape=torch.Size([384]). res.dtype: torch.float32, multiplier: 3.000000, tol: 0.010000
E0401 16:49:46.155000 139907086283904 torch/_dynamo/utils.py:1319] Accuracy failed for key name blocks.2.mlp_tokens.fc1.bias
fail_accuracy

Applying the pin update patch and using "make triton" doesn't work. I updated my local Triton/llvm to the hash and was able to repro.

There is a llvm update in this pin update: from 4017f04 Jan 30 to 6f44bb7 March 22. It seems that updating the llvm caused this issue. I initially used llvm 4017f04 + new Triton, there was no accuracy issue.

How do we get the generated py file that actually triggered the issue?

[shunting314]

How do we get the generated py file that actually triggered the issue?

There are multiple ways to look at inductor generated code. I usually enable TORCHINDUCTOR_BENCHMARK_KERNEL=1 so that the path for each wrapper will be printed.

But even with the wrapper code, it's still quite hard to know for sure which piece cause the accuracy issue.

To give you more context, the test is comparing compiled model's output with eager model's output. The test fail if there 2 results does not agree with each other.

[manman-ren]

Oh got it. So every generated kernel can be the cause? How do you usually root-cause this? Is there a bisecting tool that overrides the cubin of each kernel? i.e we have two sets of cubins, one for the current pin, the other for the updated pin. We can try to bisect on the two sets by overriding the cubin file with one from the current set.

[shunting314]

Is there a bisecting tool that overrides the cubin of each kernel? i.e we have two sets of cubins, one for the current pin, the other for the updated pin. We can try to bisect on the two sets by overriding the cubin file with one from the current set.

We don't have such a tool yet. This is the first time I see accuracy regression caused by triton.

I think what you said could be a nice way to dig into this issue.

[manman-ren]

I am able to bisect it down to one kernel: triton_red_fused__to_copy_native_layer_norm_16.
Attaching the ptx before and after.
old-ptx.txt
new-ptx.txt

The difference seems to be fma.rn.f32 instead of "mul.f32" and "add.f32".

[htyu]

I am able to bisect it down to one kernel: triton_red_fused__to_copy_native_layer_norm_16. Attaching the ptx before and after. old-ptx.txt new-ptx.txt

The difference seems to be fma.rn.f32 instead of "mul.f32" and "add.f32".

Great! Yeah, that seems to be the culprint. fma should only be used in fast math mode.

BTW, curious, did you replace cubins and rerun the test during triaging?

[shunting314]

BTW, curious, did you replace cubins and rerun the test during triaging?

I'm also curious. But I would point triton cache to an alternative path

[shunting314]

Nice finding, @manman-ren ! I rerun the command in the summary with FMA manually disabled in triton, and the test can pass now. How can we properly fix this?

also cc @Chillee who is trying to disabling FMA in inductor.

[htyu]

Perhaps a fast-math flag can be added to a Triton kernel and then passed down to LLVM which already accepts fast-math flag on function level.

[Chillee]

#115435

@htyu Triton already has it.