[Reland] ROCm CI (Infra + Skips)

[petrex]

This PR to address the import error in CI + infra changes to enable ROCm CI.

This pull request introduces the skip_if_rocm decorator across various test files to skip tests that are not yet supported on ROCm. The changes ensure that tests are conditionally skipped if ROCm is detected, improving the test suite's compatibility with different environments.

Key changes include:

Cherry-pick ROCm CI infra changes from #999

Introduction of skip_if_rocm decorator:

• Added skip_if_rocm import in multiple test files to conditionally skip tests not supported on ROCm. (test/dtypes/test_affine_quantized.py, test/dtypes/test_floatx.py, test/float8/test_base.py, test/hqq/test_hqq_affine.py, test/integration/test_integration.py, test/kernel/test_galore_downproj.py, test/prototype/test_awq.py, test/prototype/test_low_bit_optim.py, test/prototype/test_splitk.py, test/quantization/test_galore_quant.py, test/quantization/test_marlin_qqq.py, test/sparsity/test_marlin.py, test/test_ops.py, test/test_s8s4_linear_cutlass.py, torchao/utils.py) [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] [15]

Application of skip_if_rocm decorator:

• Applied @skip_if_rocm("ROCm development in progress") to multiple test functions to skip them when running on ROCm. (test/dtypes/test_affine_quantized.py, test/dtypes/test_floatx.py, test/float8/test_base.py, test/hqq/test_hqq_affine.py, test/integration/test_integration.py, test/kernel/test_galore_downproj.py, test/prototype/test_awq.py, test/prototype/test_low_bit_optim.py, test/prototype/test_splitk.py, test/quantization/test_galore_quant.py, test/quantization/test_marlin_qqq.py, test/sparsity/test_marlin.py) [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] [15] [16] [17] [18] [19] [20]

Module-level skips for ROCm:

• Added module-level skips for ROCm in specific test files to skip all tests within the module if ROCm is detected. (test/test_ops.py, test/test_s8s4_linear_cutlass.py) [1] [2]

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures

As of commit 89f0bc1 with merge base ea7910e ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_quantized_tensor_subclass_int8_dyn_quant
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_quantized_tensor_subclass_int8_dyn_quant
• Run Regression Tests / test-nightly (ROCM Nightly, linux.rocm.gpu.torchao, --pre torch --index-url https://download.pyto... / linux-job (gh)
  test/prototype/test_low_bit_optim.py::TestFSDP2::test_fsdp2

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

[pytorch-bot]

Warning: Unknown label ciflow/rocm.
Currently recognized labels are

• ciflow/benchmark
• ciflow/tutorials

Please add the new label to .github/pytorch-probot.yml

[andrewor14]

Thanks @petrex. Would you also like to reopen #999 or should someone else do it?

[petrex]

Thanks @andrewor14 I will work with AMD team on that

[petrex]

@amdfaa just cherry-pick your infra changes into this PR so we can have a clearer CI signal. Please help review the changes. thx

[amdfaa]

lgtm from the infra side

[msaroufim]

The breakage on cuda doesns't seem related to you, seems like it's this test FAILED test/quantization/test_quant_api.py::TestQuantFlow::test_quantized_tensor_subclass_int8_dyn_quant - torch._inductor.exc.CppCompileError: C++ compile error in which case @jerryzh168 might need to take a look

[jithunnair-amd]

@petrex Looks like more fixes/skips are needed: https://github.com/pytorch/ao/actions/runs/12934492712/job/36084046951?pr=1581
=== 117 failed, 1493 passed, 518 skipped, 49 warnings in 3551.22s (0:59:11) ====