Enable addmm + GELU epilogue fusion via cuBLASLt

[aakhundov]

Stack from ghstack (oldest at bottom):

• -> Enable addmm + GELU epilogue fusion via cuBLASLt #103811

Summary:

Previously, addmm + GELU epilogue fusion was unconditionally disabled in ATen/native/cuda/Blas.cpp due to compilation and numerical issues in CUDA <= 11.4. This PR:

1. Enables addmm + GELU epilogue fusion for CUDA >= 11.8.

2. Restricts the usage of fused addmm epilogue to contiguous output (bugfix).

3. Extends unit tests with addmm epilogue fusion and GELU activation paths.

Test Plan:

$ python test/test_linalg.py -k test_addmm_relu -v

test_addmm_relu_cpu_bfloat16 (main.TestLinalgCPU.test_addmm_relu_cpu_bfloat16) ... ok
test_addmm_relu_cpu_float32 (main.TestLinalgCPU.test_addmm_relu_cpu_float32) ... ok
test_addmm_relu_cpu_float64 (main.TestLinalgCPU.test_addmm_relu_cpu_float64) ... ok
test_addmm_relu_cuda_bfloat16 (main.TestLinalgCUDA.test_addmm_relu_cuda_bfloat16) ... ok
test_addmm_relu_cuda_float32 (main.TestLinalgCUDA.test_addmm_relu_cuda_float32) ... ok
test_addmm_relu_cuda_float64 (main.TestLinalgCUDA.test_addmm_relu_cuda_float64) ... ok

$ python test/test_linalg.py -k test_addmm_gelu -v

test_addmm_gelu_cpu_bfloat16 (main.TestLinalgCPU.test_addmm_gelu_cpu_bfloat16) ... ok
test_addmm_gelu_cpu_float32 (main.TestLinalgCPU.test_addmm_gelu_cpu_float32) ... ok
test_addmm_gelu_cpu_float64 (main.TestLinalgCPU.test_addmm_gelu_cpu_float64) ... ok
test_addmm_gelu_cuda_bfloat16 (main.TestLinalgCUDA.test_addmm_gelu_cuda_bfloat16) ... ok
test_addmm_gelu_cuda_float32 (main.TestLinalgCUDA.test_addmm_gelu_cuda_float32) ... ok
test_addmm_gelu_cuda_float64 (main.TestLinalgCUDA.test_addmm_gelu_cuda_float64) ... ok

Reviewers: @eellison

Differential Revision: D46829884

cc @ptrblck @csarofeen @xwang233

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ No Failures

As of commit f190760:
💚 Looks good so far! There are no failures yet. 💚

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

[linux-foundation-easycla]

• ❌ The email address for the commit (0cc8142) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please submit a support request ticket.

[linux-foundation-easycla]

• ❌ The email address for the commit (0cc8142) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please submit a support request ticket.

[aakhundov]

/easycla

[aakhundov]

@aakhundov has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[IvanYashchuk]

My approval works only for the test/test_linalg.py file if merged with the GH bot.

[IvanYashchuk]

In the future, we should move these tests to the test_matmul_cuda.py file. Other than that changes look good!

[aakhundov]

@aakhundov has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[aakhundov]

@aakhundov has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[aakhundov]

@aakhundov has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[eellison]

nice !!

[eellison]

for anyone looking at this in the future: we're using 11.8 here because that's what our CI has test coverage for

[eellison]

i'm surprised the torch.half bounds are so high here, but this is a prior issue, no need to fix on your pr

[aakhundov]

My understanding is that these tests are not run on torch.half (only on float32, float64, and bfloat16 on CUDA). So maybe it's a copy / paste artifact.

[aakhundov]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team

Raised by workflow job

[aakhundov]

@pytorchbot label "topic: not user facing"

[aakhundov]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[huydhn]

cc @aakhundov I suspect that this change is causing failures on periodic Windows CUDA job https://hud.pytorch.org/pytorch/pytorch/commit/1c79003b3c13c7bc47e5796e4451d6565121f3a0. Could you help take a look into the issue? May be relax the tolerance on Windows if the failure is expected?

2023-06-22T01:54:28.9232875Z =================================== RERUNS ====================================
2023-06-22T01:54:28.9233147Z _________________ TestLinalgCUDA.test_addmm_gelu_cuda_float32 _________________
2023-06-22T01:54:28.9233410Z Traceback (most recent call last):
2023-06-22T01:54:28.9233730Z   File "C:\actions-runner\_work\pytorch\pytorch\test\test_linalg.py", line 5582, in test_addmm_gelu
2023-06-22T01:54:28.9234092Z     self._test_addmm_impl(torch._addmm_activation, "gelu", device, dtype)
2023-06-22T01:54:28.9234464Z   File "C:\actions-runner\_work\pytorch\pytorch\test\test_linalg.py", line 5526, in _test_addmm_impl
2023-06-22T01:54:28.9234830Z     self._test_addmm_addmv(func, V, m1, m2, beta=1, activation=activation)
2023-06-22T01:54:28.9235268Z   File "C:\actions-runner\_work\pytorch\pytorch\test\test_linalg.py", line 5447, in _test_addmm_addmv
2023-06-22T01:54:28.9235818Z     self.assertEqual(res1, res3)
2023-06-22T01:54:28.9236206Z   File "C:\actions-runner\_work\pytorch\pytorch\build\win_tmp\build\torch\testing\_internal\common_utils.py", line 3100, in assertEqual
2023-06-22T01:54:28.9236552Z     raise error_metas[0].to_error(
2023-06-22T01:54:28.9236787Z AssertionError: Tensor-likes are not close!
2023-06-22T01:54:28.9236938Z 
2023-06-22T01:54:28.9237027Z Mismatched elements: 67 / 250 (26.8%)
2023-06-22T01:54:28.9237294Z Greatest absolute difference: 0.00047320593148469925 at index (2, 5) (up to 0.0001 allowed)
2023-06-22T01:54:28.9237609Z Greatest relative difference: 0.3691028356552124 at index (6, 22) (up to 1.3e-06 allowed)

[aakhundov]

@huydhn Thanks for the heads-up! From the error, it seems that the GELU epilogue fusion is not available on Windows (or it doesn't use tanh approximation). In any case, we should be able to fix this by adding not IS_LINUX at this line: https://github.com/pytorch/pytorch/blob/main/test/test_linalg.py#L5419 . I'll send a PR shortly.

Btw, the failing Windows CUDA job didn't show up in the CI of the PR. Is this expected? Is there a way to trigger the Windows CUDA job manually in the PR, before merging? Thanks!

[aakhundov]

@huydhn Submitted a fix in #104031. Would be great if we could run the Windows CUDA job on this PR to make sure it remedies the issue. Thanks!

[aakhundov]

@aakhundov has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[huydhn]

Btw, the failing Windows CUDA job didn't show up in the CI of the PR. Is this expected? Is there a way to trigger the Windows CUDA job manually in the PR, before merging? Thanks!

Yes, you can add ciflow/periodic to your PR to trigger Windows CUDA jobs. As the name implies, the job only runs periodically and not on every PR by default because we only have a limited number of Windows CUDA runners. When failure like this happens, it's ok to follow up with a fix.

I have add ciflow/periodic onto #104031, so it should be running there now.