Update CI to use TheRock by VeeraRajasekhar · Pull Request #602 · ROCm/TransformerEngine

VeeraRajasekhar · 2026-05-29T15:33:31Z

No description provided.

Copilot

Pull request overview

Updates the ROCm CI configuration to use TheRock-based container images (with per-architecture image selection) and refreshes related tooling/tests to work with the new environment and profiler behavior.

Changes:

Switch default ROCm CI images to TheRock tags, selected per runner GPU architecture (mi30x vs mi35x).
Update the CI deps Docker image build to install ROCm via TheRock tarballs and parameterize GPU arch at build time.
Adjust ROCm attention benchmarking to use rocprofv3 output, and add targeted test skips/workarounds for HIP-specific constraints.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
`tests/pytorch/test_fusible_ops.py`	Adds a HIP/arch-specific skip for an unsupported hipBLASLt configuration.
`tests/pytorch/distributed/run_fsdp2_fused_adam.py`	Updates async checkpoint save behavior for newer Torch versions using a blocking stager.
`ci/README.md`	Documents per-runner-architecture default image selection and tags.
`ci/ci_config.json`	Changes `default` image entry from a single tag to a per-arch mapping.
`benchmarks/attention/benchmark_attention_rocm.py`	Migrates profiling from `rocprof` to `rocprofv3` and adapts kernel-stats parsing.
`.github/workflows/rocm-ci.yml`	Selects and uses per-arch Docker images in ROCm CI jobs.
`.github/workflows/ci-deps-docker-publish.yml`	Adds `gpu_arch` input and passes it into the CI deps Docker build.
`.github/scripts/Dockerfile.ci.deps`	Rebuilds CI deps image on Ubuntu 24.04 using TheRock ROCm tarball installation and updated Python deps.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

ipanfilo · 2026-06-30T12:59:35Z

+            gfx942|gfx950) ;;
+            *) echo "gpu_arch must be gfx942 or gfx950" >&2; exit 1 ;;
+          esac
+          echo "IMAGE_TAG=${IMAGE_TAG}" >> "${GITHUB_ENV}"


Why IMAGE_TAG and only it goes to GITHUB_ENV?

ipanfilo · 2026-06-30T13:39:52Z

+        future = dcp.async_save(
+            save_state,
+            checkpoint_id=checkpoint_dir,
+            async_stager=BlockingAsyncStager(),


That should be optimization parameter, not functional one. It is related to test timeouts?

VeeraRajasekhar · 2026-07-16T20:47:16Z

Rebase the changes onto dev.

ipanfilo

Make sure it loads on non-ROCm platform.

Move _rocm_init out of the installable package and generate it at build time via BuildPy. CUDA wheels no longer ship the module, so import skips rocm_sdk init via ImportError instead of relying on rocm_sdk being absent at runtime.

ipanfilo

One remaining comment about handling of rocm-smi issue in CI tests.
And please review Copilot comments

- setup.py: write _rocm_init.py into the source tree ahead of build_py, so the checkout and the wheel carry the same file; drop a stale copy on non-ROCm builds. Gitignored. - test_checkpoint.py, distributed/run_numerics.py: append their own directory to sys.path before sibling imports -- both run as scripts, where PYTHONSAFEPATH suppresses the implicit entry. - amdsmi: dropped from the CI deps image, since a copy that fails to enumerate makes torch.cuda.device_count() report 0 and every test module errors at collection. test_sanity_import.py installs it for test_lazy_init and removes it after.

ipanfilo · 2026-07-27T14:29:51Z

 # to bound each child below this outer limit -- hence the exports below.
 # All are overridable from the environment.
 export PYTHONFAULTHANDLER=1
+# ROCm wheels ship _rocm_init.py (generated at build time) but not the source tree.


Importing of source_tree TE vs installed one is more global problem than just _rocm_init. So it is a good solution for running CI from TE root but update comments.
Also, make sure it works with editable install

VeeraRajasekhar · 2026-07-27T18:56:39Z

[failed] test_distributed_layernorm_mlp.TestDistributedLayernormMLP::test_layernorm_mlp_grad[with_jax_gemm_True-Float8CurrentScaling-use_bias_True-dtype_float16-activation_type_gelu-input_shape_4_64_256-fsdp2_tpsp2] (in test_distributed_layernorm_mlp.auto) -- AssertionError:

Checking this test failure on CI

…orm MLP

VeeraRajasekhar · 2026-07-27T21:57:26Z

[failed] test_distributed_layernorm_mlp.TestDistributedLayernormMLP::test_layernorm_mlp_grad[with_jax_gemm_True-Float8CurrentScaling-use_bias_True-dtype_float16-activation_type_gelu-input_shape_4_64_256-fsdp2_tpsp2] (in test_distributed_layernorm_mlp.auto) -- AssertionError:

Checking this test failure on CI

Fixed it.

VeeraRajasekhar · 2026-07-27T21:58:37Z

CI Run: https://github.com/ROCm/TransformerEngine/actions/runs/30308954763

The ragged/THD and SBHD-padding paths route CK's log-sum-exp through the scratch buffer devPtrSoftmaxLSEWithoutPadding, then copy it into softmax_aux with pad_remap_lse<Add>. That copy walks the full padded range, but CK only writes the slots covered by cu_seqlens, so padded rows carry uninitialized workspace memory and clobber the 0xF0 sentinel the framework layers write. JAX striped ring attention reads those rows when merging LSE across CP steps via sigmoid(aux_step - aux), turning stale NaN/inf bytes into NaN output. Finite garbage passes silently, so the failing configs shifted per environment. Prefill the buffer with 0xF0 before the CK call, matching the sentinel in jax/csrc/extensions/attention.cpp. Fixes 8 NaN failures in tests/jax/test_distributed_fused_attn.py.

ipanfilo · 2026-07-29T04:31:30Z

+AMDSMI_SRC = "/opt/rocm/share/amd_smi"
+
+
+def _install_amdsmi():


Why need installing it? If it is known where it is located in ROCm dir, can it be just added to pythonpath?

ipanfilo · 2026-07-29T04:33:50Z


        if not async_save:
            dcp.save(save_state, checkpoint_id=checkpoint_dir)
+        elif te.torch_version() >= (2, 9, 0):


check for IS_HIP_EXTENSION?

* Update CI to use TheRock (#602) * Update CI image to use TheRock * Fixed sgpu tests * Fixed gfx950 test failures * Fixed test_kv_cache failures on gfx942 * Addressed reviews * Fix ROCm 7.13 import abort by preloading torch before TE core lib On ROCm 7.13, libtransformer_engine.so pulls in libamd_comgr/LLVM at load time. If that happens before torch/rocm_sdk initializes COMGR, LLVM 23 aborts with duplicate spirv-expand-step registration. Preload torch on Linux ROCm builds immediately before loading the core shared library. * Upgrade CI to ROCm 7.13 and align TE with rocm-sdk - Bump the CI deps image to TheRock 7.13, build JAX 0.10.2 from source against rocm-sdk devel, and drop global ROCM_PATH/LD_LIBRARY_PATH so runtime libs come from ROCK wheels instead of the system tarball. - Auto-detect rocm-sdk devel for TE builds in utils.py and rocm_utils.cmake, and preload rocm-sdk-core libraries before loading libtransformer_engine.so to fix the ROCm 7.13 LLVM/COMGR double-init abort without import-order hacks. * Initialize rocm_sdk before TE native libs and fix cpp test ROCm loading. Move ROCm SDK preload to _rocm_init so it runs before transformer_engine.common loads libtransformer_engine.so, and set ROCM_PATH from rocm_sdk when unset. Link ROCm libs directly on cpp test binaries with RPATH so gtest_discover_tests can resolve transitive deps without LD_LIBRARY_PATH. * Addressed reviews * Addressed reviews * Added cmake prefix path for aotriton * Updated preload libraries list for rocm_sdk * Generate _rocm_init.py only for ROCm builds Move _rocm_init out of the installable package and generate it at build time via BuildPy. CUDA wheels no longer ship the module, so import skips rocm_sdk init via ImportError instead of relying on rocm_sdk being absent at runtime. * Addressed reviews * Set PYTHONSAFEPATH in CI to avoid source tree shadowing ROCm wheels * Fix _rocm_init.py packaging and test environment - setup.py: write _rocm_init.py into the source tree ahead of build_py, so the checkout and the wheel carry the same file; drop a stale copy on non-ROCm builds. Gitignored. - test_checkpoint.py, distributed/run_numerics.py: append their own directory to sys.path before sibling imports -- both run as scripts, where PYTHONSAFEPATH suppresses the implicit entry. - amdsmi: dropped from the CI deps image, since a copy that fails to enumerate makes torch.cuda.device_count() report 0 and every test module errors at collection. test_sanity_import.py installs it for test_lazy_init and removes it after. * Updated comments * Work around JAX 0.10.2 Triton FP8 GEMM failures in distributed layernorm MLP * Initialize padded LSE workspace in CK fused attention forward The ragged/THD and SBHD-padding paths route CK's log-sum-exp through the scratch buffer devPtrSoftmaxLSEWithoutPadding, then copy it into softmax_aux with pad_remap_lse<Add>. That copy walks the full padded range, but CK only writes the slots covered by cu_seqlens, so padded rows carry uninitialized workspace memory and clobber the 0xF0 sentinel the framework layers write. JAX striped ring attention reads those rows when merging LSE across CP steps via sigmoid(aux_step - aux), turning stale NaN/inf bytes into NaN output. Finite garbage passes silently, so the failing configs shifted per environment. Prefill the buffer with 0xF0 before the CK call, matching the sentinel in jax/csrc/extensions/attention.cpp. Fixes 8 NaN failures in tests/jax/test_distributed_fused_attn.py. * Addressed reviews * Bump ck_jit submodule * Added yaml to ci prerequisites

_rocm_init (added by #602 for the TheRock CI migration) set ROCM_PATH to the rocm-sdk devel wheel (get_devel_root()) when unset. FlyDSL 0.3.0's MLIR AMDGPU serialization resolves ld.lld relative to ROCM_PATH and cannot drive the devel-wheel layout, so fresh kernel compiles fail with "lld invocation failed" unless the user manually exports ROCM_PATH. Prefer the system /opt/rocm tree when present (the CI container provides it, and FlyDSL's linker resolves it), falling back to the devel wheel only for wheel-only environments with no system tree. This matches the Dockerfile, which builds with ROCM_PATH=/opt/rocm. Verified: after a clean rebuild with the flydsl cache cleared, the MXFP8 GEMM test links and passes with no ROCM_PATH export. Co-Authored-By: Claude <noreply@anthropic.com>

VeeraRajasekhar self-assigned this May 29, 2026

VeeraRajasekhar added ci-level 3 CI test level 3 and removed ci-level 3 CI test level 3 labels May 29, 2026

Copilot started work on behalf of VeeraRajasekhar June 3, 2026 16:56 View session

Copilot AI changed the title ~~upgrade ci to TheRock~~ Fix CI artifact download failures in sGPU/mGPU test jobs Jun 3, 2026

Copilot finished work on behalf of VeeraRajasekhar June 3, 2026 17:13

VeeraRajasekhar force-pushed the veergopu/upgrade_ci_to_rock branch 2 times, most recently from a6a47d9 to abfd1d7 Compare June 5, 2026 16:31

VeeraRajasekhar force-pushed the veergopu/upgrade_ci_to_rock branch from e804052 to 6130f86 Compare June 26, 2026 16:22

VeeraRajasekhar changed the title ~~Fix CI artifact download failures in sGPU/mGPU test jobs~~ Update CI to use TheRock Jun 26, 2026

VeeraRajasekhar requested a review from Copilot June 26, 2026 18:27

Copilot started reviewing on behalf of VeeraRajasekhar June 26, 2026 18:28 View session

Copilot AI reviewed Jun 26, 2026

View reviewed changes

Comment thread .github/scripts/Dockerfile.ci.deps Outdated

Comment thread .github/workflows/ci-deps-docker-publish.yml Outdated

Comment thread benchmarks/attention/benchmark_attention_rocm.py Outdated

Comment thread tests/pytorch/test_fusible_ops.py Outdated

VeeraRajasekhar force-pushed the veergopu/upgrade_ci_to_rock branch 5 times, most recently from 1a01245 to 66a3c8f Compare June 29, 2026 04:16

VeeraRajasekhar marked this pull request as ready for review June 29, 2026 07:10

VeeraRajasekhar requested review from ipanfilo, wangye805 and wenchenvincent as code owners June 29, 2026 07:10

VeeraRajasekhar requested a review from Copilot June 29, 2026 07:17

Copilot started reviewing on behalf of VeeraRajasekhar June 29, 2026 07:18 View session

Copilot AI reviewed Jun 29, 2026

View reviewed changes

Comment thread .github/workflows/rocm-ci.yml Outdated

Comment thread .github/scripts/Dockerfile.ci.deps Outdated

ipanfilo reviewed Jun 30, 2026

View reviewed changes

VeeraRajasekhar force-pushed the veergopu/upgrade_ci_to_rock branch 3 times, most recently from 2ee2eaf to 7347cc6 Compare July 16, 2026 20:46

Added cmake prefix path for aotriton

413e16f

ipanfilo reviewed Jul 24, 2026

View reviewed changes

Comment thread build_tools/templates/_rocm_init.py

Comment thread transformer_engine/__init__.py

Updated preload libraries list for rocm_sdk

6ced6d3

VeeraRajasekhar requested a review from ipanfilo July 24, 2026 13:49

ipanfilo reviewed Jul 24, 2026

View reviewed changes

Generate _rocm_init.py only for ROCm builds

f4f6ddd

Move _rocm_init out of the installable package and generate it at build time via BuildPy. CUDA wheels no longer ship the module, so import skips rocm_sdk init via ImportError instead of relying on rocm_sdk being absent at runtime.

VeeraRajasekhar requested a review from ipanfilo July 24, 2026 20:38

ipanfilo reviewed Jul 24, 2026

View reviewed changes

Comment thread setup.py

Addressed reviews

78deb67

ipanfilo approved these changes Jul 25, 2026

View reviewed changes

VeeraRajasekhar and others added 2 commits July 26, 2026 03:19

Set PYTHONSAFEPATH in CI to avoid source tree shadowing ROCm wheels

c544f05

VeeraRajasekhar force-pushed the veergopu/upgrade_ci_to_rock branch from 6680819 to 80bdb92 Compare July 27, 2026 03:45

ipanfilo reviewed Jul 27, 2026

View reviewed changes

ipanfilo requested changes Jul 27, 2026

View reviewed changes

Updated comments

f21c6ec

VeeraRajasekhar requested a review from ipanfilo July 27, 2026 17:05

Work around JAX 0.10.2 Triton FP8 GEMM failures in distributed layern…

98bb274

…orm MLP

VeeraRajasekhar added 2 commits July 28, 2026 17:57

Merge branch 'dev' into veergopu/upgrade_ci_to_rock

67a4b34

ipanfilo reviewed Jul 29, 2026

View reviewed changes

Addressed reviews

212a96a

VeeraRajasekhar requested a review from ipanfilo July 29, 2026 07:57

ipanfilo approved these changes Jul 29, 2026

View reviewed changes

VeeraRajasekhar merged commit 4d950f4 into dev Jul 29, 2026
6 of 11 checks passed

		AMDSMI_SRC = "/opt/rocm/share/amd_smi"


		def _install_amdsmi():

Uh oh!

Conversation

VeeraRajasekhar commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ipanfilo Jun 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ipanfilo Jun 30, 2026

Choose a reason for hiding this comment

Uh oh!

VeeraRajasekhar commented Jul 16, 2026

Uh oh!

Uh oh!

Uh oh!

ipanfilo left a comment

Choose a reason for hiding this comment

Uh oh!

ipanfilo left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ipanfilo Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

VeeraRajasekhar commented Jul 27, 2026

Uh oh!

VeeraRajasekhar commented Jul 27, 2026

Uh oh!

VeeraRajasekhar commented Jul 27, 2026

Uh oh!

ipanfilo Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

ipanfilo Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

VeeraRajasekhar commented May 29, 2026 •

edited

Loading