Veergopu/upgrade ci rock 714 by VeeraRajasekhar · Pull Request #690 · ROCm/TransformerEngine

VeeraRajasekhar · 2026-08-04T22:30:03Z

Description

Upgrade the TransformerEngine (ROCm fork) CI stack to ROCk 7.14 using the
new multi-arch TheRock wheelhouse. A single CI deps image now serves both
GPUs (gfx950;gfx942), with ROCm provided by the pip rocm-sdk
wheels (no /opt/rocm install). The version bump moves the stack to
torch 2.12 / triton 3.7.1 / jax 0.11 / FA 2.8.3, ROCm 7.14, which surfaced several
failures; this PR contains the upgrade plus the triage fixes and interim
workarounds needed to get the suite green.

Most of the newly-exposed failures are upstream torch 2.12 + ROCm regressions
(not TE bugs), each root-caused and reproduced with a TE-free reproducer; where
no upstream fix has landed yet, a narrowly-scoped, ROCm-gated interim workaround
is applied so CI stays green.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

CI / build infrastructure

Rebuild the CI deps image from the TheRock multi-arch wheelhouse (ROCm 7.14,
gfx950;gfx942); ROCm now comes from rocm-sdk wheels instead of /opt/rocm.
Resolve ROCm via rocm-sdk path --root in the core C++ suite (ci/_utils.sh,
transformer_engine/common/CMakeLists.txt, cuda_runtime.cpp) so it builds
and runs on an image with no /opt/rocm.
Update the default CI docker image (ci/ci_config.json).
benchmark_attention_rocm.py: use the rocm-sdk-core rocprofv3 to avoid a
duplicate librocprofiler-sdk.so.1 (core vs devel) SIGABRT.
test_sanity_import.py: adjust for the wheel-based ROCm layout
(test_lazy_init now passes).

Test fixes and interim workarounds

test_torch_compile::test_autocast_nested_custom — cherry-pick upstream
PR #3130 ([torch.compile] ...). Fixes a latent TE 2.17 bug (recipe reuse
in set_meta_tensor + a missing get_quantizer_roles() override) that only
runs under torch ≥ 2.11. NVIDIA/TransformerEngine@43093f16d6dc
test_gpt_cuda_graph — interim WAR: preferred_blas_library("cublas"),
IS_HIP_EXTENSION-gated. torch 2.12 creates a per-(device,stream) hipBLASLt
handle during graph capture → HIP 900. Upstream (PyTorch) fix pending.
triton_kernels/test_norms.py — interim WAR: new ROCm-gated conftest.py
that os._exit()s pytest's real status on pytest_sessionfinish. torch 2.12's
HIP teardown ordering trips a latent ROCr hsa_shut_down use-after-free at
process exit (segfault after all tests pass). Bypasses the buggy atexit path;
JUnit report and real exit status preserved.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

…_gpt_cuda_graph test_gpt_cuda_graph fails on ROCm with torch>=2.12 with: Hip error: 'operation not permitted when stream is capturing'(900) at hipblaslt/library/src/amd_detail/hipblaslt.cpp:164 then hangs. Root cause is upstream PyTorch, not TE. pytorch/pytorch#179053 (in 2.12) changed the ROCm hipBLASLt handle cache from per-device to per-(device, stream). The graph-capture stream is fresh, so the first matmul on it does a lazy hipblasLtCreate, whose internal ~25 MB hipMalloc is illegal during capture. The capture_begin pre-init added in pytorch/pytorch#180692 only covers the calling thread, so torch's own bmm in the captured backward (fp32 -> unfused attention) still creates a handle on the autograd thread during capture. Not reproducible on torch 2.10 (per-device handle) and not on CUDA (the per-stream path is #ifdef USE_ROCM). Route torch's own matmul/bmm off hipBLASLt for this test via torch.backends.cuda.preferred_blas_library(cublas), gated on IS_HIP_EXTENSION and restored in a finally. TE's own GEMMs are unaffected. Remove once the upstream capture-time pre-init covers all threads/streams (or hipblasLtCreate becomes capture-safe).

…ile (#3130) * [PyTorch] torch.compile: wrap pybind11 UB methods as compile-time constants; fix SP memory leak; test suite hook-up Wrap CommOverlapCore pybind11 methods that return compile-time constants so torch.compile(fullgraph=True) can trace through them without graph breaks: - `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py; `_ub_is_fp8()` in gemm.py - `with_cublasmp()` → `ub_is_cublasmp()` in base.py All callers in linear.py, layernorm_linear.py, layernorm_mlp.py, base.py, gemm.py, userbuffers_backward_linear.py and userbuffers_forward_linear.py updated. Fix quantized grad_output not being freed early for column-parallel SP backward. Row-parallel SP already called clear_tensor_data(grad_output) to release the gathered tensor; column-parallel SP quantizes grad_output to Float8TensorStorage but never freed it before returning. Under torch.compile reduce-overhead this leaves 3 live pool tensors at recording end and triggers "Detected 3 tensor(s) in the cudagraph pool not tracked as outputs". Extend the existing clear_tensor_data guard to cover both parallel modes. Fix custom-recipe quantizer state being re-initialised on every forward call even when the recipe object has not changed. The existing early-exit for CustomRecipeState was missing an identity check on the recipe object, so any repeated call with the same recipe would bypass the early-return and rebuild quantizers unnecessarily. Add `if recipe_state.recipe is recipe: return` to restore the intended caching behaviour. Add test_torch_compile.py to L0_pytorch_unittest so the autocast and existing compile tests run in CI. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Replace fp8_recipe in LinearBwdArgs with pre-resolved split-accumulator booleans LinearBwdArgs stored the entire FP8 recipe object so the backward could extract fp8_gemm_dgrad.use_split_accumulator and fp8_gemm_wgrad.use_split_accumulator at GEMM time. Recipe objects hold process-group references and are not serialisable as compile-time constants, making them incompatible with torch.compile custom-op paths. Replace fp8_recipe with two plain bool fields: - dgrad_use_split_accumulator (default _2X_ACC_DGRAD) - wgrad_use_split_accumulator (default _2X_ACC_WGRAD) These are resolved once in _linear_setup_ctx and passed into the args struct, so the backward consumes scalars instead of a live recipe object. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reset torch.compile state in destroy_ub to avoid stale assume_constant_result get_ub_is_fp8 bakes is_fp8_ubuf() as a compile-time constant; without a reset, destroy_ub + re-init with different FP8 settings would read stale values until recompile. Only affects in-memory caches, not disk. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Provide explicit QuantizerRoles in torch.compile custom-recipe test ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit the no-roles warning, which graph-breaks under fullgraph=True. qfactory dispatches on role.tensor_type instead of a pre-baked string key. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> (cherry picked from commit 43093f16d6dc3c0707c2b9ec3c3d07918fed598a)

…/test_norms

ipanfilo · 2026-08-05T03:43:26Z

Please update description, specifically list of changes

ipanfilo · 2026-08-05T03:31:17Z

+#
+# See LICENSE for license information.
+
+"""Work around a ROCm HSA-runtime teardown segfault after the triton_kernels tests.


It sounds like a serious ROCm bug. Is there ticket for that?

I will be syncing with Melantha to file ticket.

But do you have reproducer w/o TE?

ipanfilo · 2026-08-05T03:40:26Z

 #ifdef __HIP_PLATFORM_AMD__
    std::vector<std::pair<std::string, Path>> search_paths = {{"ROCM_PATH", ""},
                                                              {"HIP_PATH", ""},
+                                                              {"", string_path_rocm_root},


string_path_rocm_root is calculated build-time on build system. It cannot be used as runtime search path on test/run system. ROCM_PATH should be properly set if ROCm installed not as /opt/rocm

ipanfilo · 2026-08-05T03:53:22Z

+    # when stream is capturing"). The upstream capture_begin pre-init
+    # (pytorch/pytorch#180692) only covers the calling thread, not the autograd
+    # backward thread, so torch's own bmm in the captured backward still trips it.
+    # Route torch's matmul/bmm off hipBLASLt for this test until upstream fixes it.


Does upstream here stand for torch upstream? Is there ticket for that then?

ipanfilo · 2026-08-05T03:55:13Z

    """
-    import importlib.util, os
-    if importlib.util.find_spec("amdsmi") is not None or not os.path.isdir(AMDSMI_SRC):
+    import importlib.util


Does it continue properly working if ROCm is installed as OS package?

ipanfilo · 2026-08-05T14:04:58Z

+#ifdef __HIP_PLATFORM_AMD__
+namespace {
+
+using Path = std::filesystem::path;


What is a problem with ROCM_PATH that this fallback is needed? PR for ROCm 7.13 introduced rocm_init that sets it

The problem isn’t that _rocm_init is wrong — it’s that it never runs for the process that failed.

ci/core.sh → ctest → test_operator / test_util
This path:
Does not import transformer_engine
Does not run _rocm_init
Loads libtransformer_engine.so directly
Hits NVRTC → include_directory() → reads $ROCM_PATH → unset

So, to make it self-sufficient I added here.

I'm hesitate of making library spawning new process, moreover running applications related to python package. For direct transformer_engine.so import it might be OK to request user to set ROCM_PATH. For CI purposes it may be part of CI script

ipanfilo · 2026-08-05T14:05:54Z

+#
+# See LICENSE for license information.
+
+"""Work around a ROCm HSA-runtime teardown segfault after the triton_kernels tests.


But do you have reproducer w/o TE?

ipanfilo · 2026-08-05T14:25:36Z

+    """Path to the rocprofv3 shipped in rocm-sdk-core.
+
+    The `rocprofv3` on PATH is the rocm-sdk trampoline, which execs the copy in
+    rocm-sdk-devel. That copy loads librocprofiler-sdk.so.1 from _rocm_sdk_devel,


rocm_init uses get_devel_root() that should result in TE load devel libraries first, why does it load rocm_sdk_core variants?

VeeraRajasekhar and others added 7 commits August 4, 2026 16:57

Upgrade CI to TheRock 7.14

ee7f2f4

[ROCm] Work around torch 2.12 HSA shutdown segfault in triton_kernels…

e271bea

…/test_norms

Added comments

cadc944

Updated default docker image

9bf7e51

Added rocm_sdk support to core tests

b99d341

VeeraRajasekhar self-assigned this Aug 4, 2026

VeeraRajasekhar requested review from ipanfilo, wangye805 and wenchenvincent as code owners August 4, 2026 22:30

ipanfilo requested changes Aug 5, 2026

View reviewed changes

VeeraRajasekhar added 2 commits August 5, 2026 05:35

Updated docker image tag

a192833

Add runtime rocm_sdk resolution

32e2aca

ipanfilo reviewed Aug 5, 2026

View reviewed changes

Navigation Menu

Uh oh!

Veergopu/upgrade ci rock 714 - #690

Veergopu/upgrade ci rock 714#690
VeeraRajasekhar wants to merge 9 commits into
devfrom
veergopu/upgrade_ci_rock_714

VeeraRajasekhar commented Aug 4, 2026 •

edited

Loading

Uh oh!

ipanfilo commented Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

VeeraRajasekhar Aug 5, 2026 •

edited

Loading

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

VeeraRajasekhar Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

VeeraRajasekhar Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

ipanfilo Aug 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

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

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

ipanfilo commented Aug 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

VeeraRajasekhar commented Aug 4, 2026 •

edited

Loading

VeeraRajasekhar Aug 5, 2026 •

edited

Loading