ROCm Support : Tile_Layout kernel by petrex · Pull Request #1201 · pytorch/ao

petrex · 2024-10-30T19:07:52Z

This pull request introduces support for ROCm (Radeon Open Compute) in addition to CUDA for GPU acceleration. The changes primarily focus on enabling the build and execution of ROCm-specific code paths alongside existing CUDA paths.

In this PR, I use tensor_core_tiled_layout as proof of concept, but it generalizes to other kernels (for example, fp6_llm or sparse_marlin) with minimum effort. Feedback are welcome

co-author : @lcskrishna

Features:

ROCm Support Integration:

setup.py: Added detection for ROCm and adjusted the logic for compiling GPU extensions based on the availability of CUDA or ROCm.

Conditional Compilation for ROCm:

torchao/csrc/cuda/tensor_core_tiled_layout/tensor_core_tiled_layout.cu: Introduced conditional compilation directives to include ROCm-specific headers and adjust constants and operations for ROCm.

These changes ensure that the codebase can compile and run efficiently on both CUDA and ROCm platforms, leveraging the best available GPU acceleration technology.

Usage

With ROCm pytorch nightly docker , simply run PYTORCH_ROCM_ARCH=gfx942 python setup.py install

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 4 New Failures, 1 Pending

As of commit 2272a3f with merge base 9bcd73b ():

NEW FAILURES - The following jobs have failed:

Build M1 Wheels / pytorch/ao / upload / wheel-py3_9-cpu (gh)
The process '/usr/bin/git' failed with exit code 1
Code Analysis with Ruff / build (3.9) (gh)
Process completed with exit code 1.
Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
torchao/experimental/tests/test_int8_dynamic_activation_intx_weight.py::TestInt8DynamicActivationIntxWeight::test_export_compile_aoti_PackedLinearInt8DynamicActivationIntxWeightLayout
Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
AttributeError: '_OpNamespace' 'torchao' object has no attribute '_linear_fp_act_1bit_weight'

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

msaroufim

Thanks! Left some prelim feedback more of a n00b question but is this the main way you've seen ROCm supported in new codebases? The other approach would be to have some code duplication and avoid a boolean flag like USE_ROCM in the .cu file

msaroufim · 2024-11-04T05:22:38Z

        print("PyTorch GPU support is not available. Skipping compilation of CUDA extensions")
-    if CUDA_HOME is None and torch.cuda.is_available():
-        print("CUDA toolkit is not available. Skipping compilation of CUDA extensions")
+    if CUDA_HOME is None or not IS_ROCM and torch.cuda.is_available():


some of these if conditions might cause our binary builds to break - check CI, since we run binary builds when people modify setup.py

Until you get your first PR merged I'll have to manually approve every CI run so might be easiest if you wanna make some simple doc level PR and I can merge that fast

thanks for approving the CI flow, we will look into this

petrex · 2024-11-13T17:37:44Z

@msaroufim can we have the CI flow to test the change? thanks

msaroufim · 2024-11-13T17:42:09Z

I gave you triage access @petrex - what might help is also merging some simple PR since after that you shouldn't need me to trigger CI for you

…ytorch#1201) Co-authored-by: Martin Yuan <myuan@meta.com>

Fixes builds for non-rocm.

petrex · 2025-01-06T21:48:40Z

@msaroufim happy new year! The 2 failing CI check seems irrelevant to this PR? Any tips?

  /usr/bin/git -c protocol.version=2 fetch --no-tags --prune --no-recurse-submodules --depth=1 origin +refs/heads/rocm_enablement_staging*:refs/remotes/origin/rocm_enablement_staging* +refs/tags/rocm_enablement_staging*:refs/tags/rocm_enablement_staging*
  Error: The process '/usr/bin/git' failed with exit code 1

jerryzh168 · 2025-01-09T22:52:08Z

thanks! this is the dequantize and unpack_dequantize op btw, the packing op and int4mm op lives in pytorch right now: https://github.com/pytorch/pytorch/blob/b1c2c3967adf035356abfcb90d0b39ff7ee63707/aten/src/ATen/native/native_functions.yaml#L4162-L4170

but this can be used to test build for ROCm

jcaip

This looks good, but there are no tests? We should be able to turn on the int4wo tests to run on ROCm to verify this is correct right? Or is there additional work to be done to turn on the test?

Minor variable renaming for clarity in extension directory path resolution

petrex · 2025-02-28T18:08:48Z

This looks good, but there are no tests? We should be able to turn on the int4wo tests to run on ROCm to verify this is correct right? Or is there additional work to be done to turn on the test?

Thanks , This PR introduced optimized kernel for MI300x. We are working on allocating Mi300x nodes for torchao CI.
Will enable mi300x specific tests in following PRs.

jcaip

@petrex OK, that makes sense, approving for now.

We should also run some updated benchmarks now that we have int4+ marlin support. I'll try to find some time next week to do so.

jcaip · 2025-02-28T19:07:54Z

cc @petrex I'm assuming you want to merge in this first and then merge #1206 in? In that case I think it should be okay so long as CI passes for 1206 since you already have approval on that PR.

@lcskrishna

This pull request introduces support for ROCm (Radeon Open Compute) in addition to CUDA for GPU acceleration. The changes primarily focus on enabling the build and execution of ROCm-specific code paths alongside existing CUDA paths. In this PR, I use `tensor_core_tiled_layout` as proof of concept, but it generalizes to other kernels (for example, fp6_llm or sparse_marlin) with minimum effort. Feedback are welcome co-author : @lcskrishna ## Features: # ROCm Support Integration: * [`setup.py`](diffhunk://#diff-60f61ab7a8d1910d86d9fda2261620314edcae5894d5aaa236b821c7256badd7R49-R53): Added detection for ROCm and adjusted the logic for compiling GPU extensions based on the availability of CUDA or ROCm. # Conditional Compilation for ROCm: * [`torchao/csrc/cuda/tensor_core_tiled_layout/tensor_core_tiled_layout.cu`](diffhunk://#diff-29bb1a2fd9317c74c807a7f558f5de755af0def91b9a49c81c409f8e76f736ddL1-R1): Introduced conditional compilation directives to include ROCm-specific headers and adjust constants and operations for ROCm. These changes ensure that the codebase can compile and run efficiently on both CUDA and ROCm platforms, leveraging the best available GPU acceleration technology. ## Usage With ROCm pytorch nightly docker , simply run `PYTORCH_ROCM_ARCH=gfx942 python setup.py install ` ## Next - [ ] AMD specific unit tests (tensor_core_tiled_layout) - [ ] workload and platform specific optimization (tensor_core_tiled_layout)

lcskrishna and others added 10 commits October 16, 2024 05:19

enable build for rocm for fp6_llm

6d92e40

Merge pull request #1 from lcskrishna/cl/rocm-enablement

14b3fce

ROCm build infrastructure

enable tiled layout extension

f1a22cf

fix build error related to option

0bef6ca

require rocm 6.2

893ae03

enable tensor tiled layout extension with successful compilation

a0d3788

enable successful build

e4e654d

clean-up

3e2c6a1

Merge pull request #3 from lcskrishna/csrikris_enable_tensor_tile

c86880e

[ROCm] Enable Tiled layout extension and minor changes to setup

fix potential memory access issue

91d3c75

pytorch-bot Bot added the device: rocm label Oct 30, 2024

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 30, 2024

petrex mentioned this pull request Oct 31, 2024

ROCm Sparse Marlin Kernels #1206

Closed

1 task

msaroufim self-requested a review November 2, 2024 22:50

msaroufim reviewed Nov 4, 2024

View reviewed changes

Comment thread setup.py

Comment thread setup.py Outdated

Comment thread torchao/csrc/cuda/tensor_core_tiled_layout/tensor_core_tiled_layout.cu

Comment thread torchao/csrc/cuda/tensor_core_tiled_layout/tensor_core_tiled_layout.cu

msaroufim reviewed Nov 4, 2024

View reviewed changes

jcaip mentioned this pull request Nov 11, 2024

AMD integration tracker #1260

Open

3 tasks

Peter Yeh added 2 commits November 12, 2024 16:17

fix __nv_bfloat162 init

38b7d1c

add comment for MI300x isa

279f4b3

petrex changed the title ~~ROCm Support Integration~~ ROCm Support : HIP kernel Generator Nov 13, 2024

petrex mentioned this pull request Nov 13, 2024

[1/2] Wean off of PYBIND in favor of torch.ops.load_library #1276

Merged

Merge branch 'main' into rocm_enablement_staging

612ad14

yanbing-j pushed a commit to yanbing-j/ao that referenced this pull request Dec 9, 2024

[Flamingo] Fix the memory usage of 2x checkpoint size after loading (p…

f0a03a7

…ytorch#1201) Co-authored-by: Martin Yuan <myuan@meta.com>

lcskrishna and others added 3 commits January 6, 2025 06:39

fix build for non-rocm

bbf5a72

Merge pull request #4 from lcskrishna/rocm_enablement

735570e

Fixes builds for non-rocm.

Merge branch 'main' into rocm_enablement_staging

253c188

petrex added the topic: new feature Use this tag if this PR adds a new feature label Jan 6, 2025

petrex self-assigned this Jan 6, 2025

petrex requested review from lcskrishna and msaroufim January 8, 2025 15:59

Peter Y. Yeh added 2 commits January 9, 2025 16:39

better naming

26fa19c

lint

8f02096

petrex changed the title ~~ROCm Support : HIP kernel Generator~~ ROCm Support : Tile_Layout Lernel Jan 9, 2025

resolve merge conflict

ced9647

petrex changed the title ~~ROCm Support : Tile_Layout Lernel~~ ROCm Support : Tile_Layout kernel Jan 9, 2025

Merge branch 'main' into rocm_enablement_staging

452fa2a

renich mentioned this pull request Jan 14, 2025

ROCm support xhinker/sd_embed#25

Closed

petrex and others added 3 commits January 15, 2025 16:08

Merge branch 'main' into rocm_enablement_staging

2ec95f9

lint

89bae5f

Merge branch 'main' into rocm_enablement_staging

135c2b2

petrex requested review from jcaip and removed request for lcskrishna February 27, 2025 20:12

jcaip reviewed Feb 27, 2025

View reviewed changes

Comment thread setup.py Outdated

Refactor: Rename this_dir to curdir in setup.py

4148828

Minor variable renaming for clarity in extension directory path resolution

jcaip approved these changes Feb 28, 2025

View reviewed changes

Merge branch 'main' into rocm_enablement_staging

2272a3f

jcaip merged commit d4be9e4 into pytorch:main Mar 4, 2025

jcaip mentioned this pull request Mar 4, 2025

ROCm Sparse Marlin Kernels #1206 #1834

Merged

1 task

Conversation

petrex commented Oct 30, 2024

Features:

ROCm Support Integration:

Conditional Compilation for ROCm:

Usage

Next

Uh oh!

pytorch-bot Bot commented Oct 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 4 New Failures, 1 Pending

Uh oh!

msaroufim left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

msaroufim Nov 4, 2024

Choose a reason for hiding this comment

Uh oh!

petrex Nov 5, 2024

Choose a reason for hiding this comment

Uh oh!

petrex commented Nov 13, 2024

Uh oh!

msaroufim commented Nov 13, 2024

Uh oh!

petrex commented Jan 6, 2025

Uh oh!

jerryzh168 commented Jan 9, 2025

Uh oh!

jcaip left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

petrex commented Feb 28, 2025

Uh oh!

jcaip left a comment

Choose a reason for hiding this comment

Uh oh!

jcaip commented Feb 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

pytorch-bot Bot commented Oct 30, 2024 •

edited

Loading

jcaip commented Feb 28, 2025 •

edited

Loading