Veergopu/upgrade ci rock 714 - #690
Open
VeeraRajasekhar wants to merge 9 commits into
Open
Conversation
…_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)
VeeraRajasekhar
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
August 4, 2026 22:30
Collaborator
|
Please update description, specifically list of changes |
ipanfilo
requested changes
Aug 5, 2026
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| """Work around a ROCm HSA-runtime teardown segfault after the triton_kernels tests. |
Collaborator
There was a problem hiding this comment.
It sounds like a serious ROCm bug. Is there ticket for that?
Contributor
Author
There was a problem hiding this comment.
I will be syncing with Melantha to file ticket.
Collaborator
There was a problem hiding this comment.
But do you have reproducer w/o TE?
| #ifdef __HIP_PLATFORM_AMD__ | ||
| std::vector<std::pair<std::string, Path>> search_paths = {{"ROCM_PATH", ""}, | ||
| {"HIP_PATH", ""}, | ||
| {"", string_path_rocm_root}, |
Collaborator
There was a problem hiding this comment.
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
| # 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. |
Collaborator
There was a problem hiding this comment.
Does upstream here stand for torch upstream? Is there ticket for that then?
| """ | ||
| import importlib.util, os | ||
| if importlib.util.find_spec("amdsmi") is not None or not os.path.isdir(AMDSMI_SRC): | ||
| import importlib.util |
Collaborator
There was a problem hiding this comment.
Does it continue properly working if ROCm is installed as OS package?
ipanfilo
reviewed
Aug 5, 2026
| #ifdef __HIP_PLATFORM_AMD__ | ||
| namespace { | ||
|
|
||
| using Path = std::filesystem::path; |
Collaborator
There was a problem hiding this comment.
What is a problem with ROCM_PATH that this fallback is needed? PR for ROCm 7.13 introduced rocm_init that sets it
Contributor
Author
There was a problem hiding this comment.
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.
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| """Work around a ROCm HSA-runtime teardown segfault after the triton_kernels tests. |
Collaborator
There was a problem hiding this comment.
But do you have reproducer w/o TE?
| """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, |
Collaborator
There was a problem hiding this comment.
rocm_init uses get_devel_root() that should result in TE load devel libraries first, why does it load rocm_sdk_core variants?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 piprocm-sdkwheels (no
/opt/rocminstall). The version bump moves the stack totorch 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
Changes
CI / build infrastructure
gfx950;gfx942); ROCm now comes fromrocm-sdkwheels instead of/opt/rocm.rocm-sdk path --rootin the core C++ suite (ci/_utils.sh,transformer_engine/common/CMakeLists.txt,cuda_runtime.cpp) so it buildsand runs on an image with no
/opt/rocm.ci/ci_config.json).benchmark_attention_rocm.py: use the rocm-sdk-corerocprofv3to avoid aduplicate
librocprofiler-sdk.so.1(core vs devel) SIGABRT.test_sanity_import.py: adjust for the wheel-based ROCm layout(
test_lazy_initnow passes).Test fixes and interim workarounds
test_torch_compile::test_autocast_nested_custom— cherry-pick upstreamPR #3130 (
[torch.compile] ...). Fixes a latent TE 2.17 bug (recipe reusein
set_meta_tensor+ a missingget_quantizer_roles()override) that onlyruns 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) hipBLASLthandle during graph capture → HIP 900. Upstream (PyTorch) fix pending.
triton_kernels/test_norms.py— interim WAR: new ROCm-gatedconftest.pythat
os._exit()s pytest's real status onpytest_sessionfinish. torch 2.12'sHIP teardown ordering trips a latent ROCr
hsa_shut_downuse-after-free atprocess exit (segfault after all tests pass). Bypasses the buggy atexit path;
JUnit report and real exit status preserved.
Checklist: