Repository navigation
Conversation
Symmetric buffers make XLA register collective buffers with ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC). RCCL declares that entry point in rccl.h and reports NCCL_VERSION_CODE 23004 -- above the version that introduced symmetric memory -- so both the link-time and any version-based check pass. It then rejects every registration at runtime with ncclInvalidUsage. Since 7a57c69 turned these filters on by default for ALLREDUCE, ALLGATHER and REDUCESCATTER up to 8 MB, every multi-device collective on ROCm fails with: INTERNAL: RCCL operation ncclCommWindowRegister( comm, addr.opaque(), addr.size(), &win, 0x01) failed: invalid usage Verified against RCCL 2.30.4 on 8x MI355X with a standalone RCCL reproducer: registration fails for every buffer size and for both NCCL_WIN_COLL_SYMMETRIC and NCCL_WIN_DEFAULT, on ncclMemAlloc'd buffers. So this is neither an alignment problem nor specific to the symmetric flag -- window registration is simply not implemented. Clear the filters during HLO compilation, keyed on the compute capability of the device being compiled for. A build-time guard would not work: the defaults are set in DefaultDebugOptionsIgnoringFlags(), which is compiled into the backend-agnostic jaxlib, whereas only the plugin and PJRT artifacts are built with ROCm configured, so TENSORFLOW_USE_ROCM would never be defined there in a released build. Keying on the device also keeps a build that configures both CUDA and ROCm correct. RunBackend() is the single choke point: buffer assignment colors the collective buffers and thunk emission selects the symmetric kernels, and both run under CompileToBackendResult(). The HLO passes do not read these filters -- CollectiveBackendAssigner is driven by the separate xla_gpu_experimental_collectives_mode flag. Validated on 8x MI355X (gfx950) with jax //tests:pmap_test_gpu, forcing the filters on via --xla_enable_nccl_symmetric_buffers_for_collectives=allreduce:8388608, allgather:8388608,reducescatter:8388608 so the result does not depend on the compiled-in defaults: without this change: FAILED, 40 errors, all ncclCommWindowRegister "invalid usage", on all 3 attempts with this change: PASSED, 288 tests, 0 RCCL failures, and the warning fires once from RunBackend() Note there is no runtime escape hatch: the flag setter for xla_enable_nccl_symmetric_buffers_for_collectives only appends to the list and never clears it, so the defaults cannot be turned off via XLA_FLAGS.
mminutoli
marked this pull request as ready for review
August 7, 2026 06:33
mminutoli
added a commit
to ROCm/jax
that referenced
this pull request
Aug 7, 2026
XLA 7a57c69ea6 enabled xla_enable_nccl_symmetric_buffers_for_collectives
by default for ALLREDUCE, ALLGATHER and REDUCESCATTER up to 8 MB. That
makes XLA register collective buffers with
ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC), which RCCL declares
but rejects at runtime with ncclInvalidUsage, so every multi-GPU
collective on ROCm fails:
INTERNAL: RCCL operation ncclCommWindowRegister(
comm, addr.opaque(), addr.size(), &win, 0x01) failed: invalid usage
Pin XLA to ROCm/xla branch rocm-jaxlib-0.11.1-alpha, which is the
previously pinned upstream revision
1c32308ad0599517cd4dd2cfae136851b0712bb5 plus a single cherry-picked
commit that disables symmetric buffers when compiling for a ROCm device.
The download URL has to move to the fork as well, since the commit does
not exist in openxla/xla. tf_mirror_urls() is kept because
tf_http_archive() requires a mirror as the first URL; that mirror does
not carry fork archives, so it 404s and Bazel falls back to GitHub.
Revert both files to openxla/xla once
openxla/xla#46878 lands upstream.
Validated on 8x MI355X (gfx950) with //tests:pmap_test_gpu, forcing the
filters on so the result does not depend on compiled-in defaults:
without the fix 40 ncclCommWindowRegister failures on all 3 attempts,
with it 288 tests pass and the gate logs once from RunBackend().
| // DefaultDebugOptionsIgnoringFlags(), which is compiled into the | ||
| // backend-agnostic jaxlib, while only the plugin and PJRT artifacts are built | ||
| // with ROCm configured. | ||
| void DisableSymmetricBuffersIfUnsupported( |
Member
There was a problem hiding this comment.
Shouldn't this be part of amdgpu_compiler?
Member
|
Hi @sohaibiftikhar, please have a look at this hot-fix instead: #46895 |
mminutoli
added a commit
to ROCm/jax
that referenced
this pull request
Aug 7, 2026
XLA 7a57c69ea6 enabled xla_enable_nccl_symmetric_buffers_for_collectives
by default for ALLREDUCE, ALLGATHER and REDUCESCATTER up to 8 MB. That
makes XLA register collective buffers with
ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC), which RCCL declares
but rejects at runtime with ncclInvalidUsage, so every multi-GPU
collective on ROCm fails:
INTERNAL: RCCL operation ncclCommWindowRegister(
comm, addr.opaque(), addr.size(), &win, 0x01) failed: invalid usage
Ragged-all-to-all fails the same way through a separate route: the
one-shot NCCL barrier is on by default since 26072180fa, and
RaggedAllToAllThunk::PrepareCollective requests a symmetric allocation
for its output buffer regardless of the symmetric-buffers filters.
Pin XLA to ROCm/xla branch rocm-jaxlib-0.11.1-alpha, which is the
previously pinned upstream revision
1c32308ad0599517cd4dd2cfae136851b0712bb5 plus a single commit that gates
both routes behind a runtime ROCm platform probe.
The download URL has to move to the fork as well, since the commit does
not exist in openxla/xla. tf_mirror_urls() is kept because
tf_http_archive() requires a mirror as the first URL; that mirror does
not carry fork archives, so it 404s and Bazel falls back to GitHub.
Validated on 8x MI355X: //tests:ragged_collective_test_gpu passes 12
tests and //tests:pmap_test_gpu passes 288 (11 skipped) with the
symmetric-buffer filters forced on the command line. As negative
controls, forcing either escape hatch back on reproduces the original
failure exactly -- 107 of 288 pmap errors with
xla_gpu_experimental_enable_nccl_symmetric_buffers=true, and 10 of 12
ragged errors with xla_gpu_ragged_all_to_all_mode=symmetric -- all
ncclCommWindowRegister invalid usage.
Revert both files to openxla/xla once
openxla/xla#46878 lands upstream.
mminutoli
added a commit
to ROCm/jax
that referenced
this pull request
Aug 7, 2026
XLA 7a57c69ea6 enabled xla_enable_nccl_symmetric_buffers_for_collectives
by default for ALLREDUCE, ALLGATHER and REDUCESCATTER up to 8 MB. That
makes XLA register collective buffers with
ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC), which RCCL declares
but rejects at runtime with ncclInvalidUsage, so every multi-GPU
collective on ROCm fails:
INTERNAL: RCCL operation ncclCommWindowRegister(
comm, addr.opaque(), addr.size(), &win, 0x01) failed: invalid usage
Ragged-all-to-all fails the same way through a separate route: the
one-shot NCCL barrier is on by default since 26072180fa, and
RaggedAllToAllThunk::PrepareCollective requests a symmetric allocation
for its output buffer regardless of the symmetric-buffers filters.
Pin XLA to ROCm/xla branch rocm-jaxlib-0.11.1-alpha, which is the
previously pinned upstream revision
1c32308ad0599517cd4dd2cfae136851b0712bb5 plus a single commit that gates
both routes behind a runtime ROCm platform probe.
The download URL has to move to the fork as well, since the commit does
not exist in openxla/xla. tf_mirror_urls() is kept because
tf_http_archive() requires a mirror as the first URL; that mirror does
not carry fork archives, so it 404s and Bazel falls back to GitHub.
Validated on 8x MI355X: //tests:ragged_collective_test_gpu passes 12
tests and //tests:pmap_test_gpu passes 288 (11 skipped) with the
symmetric-buffer filters forced on the command line. As negative
controls, forcing either escape hatch back on reproduces the original
failure exactly -- 107 of 288 pmap errors with
xla_gpu_experimental_enable_nccl_symmetric_buffers=true, and 10 of 12
ragged errors with xla_gpu_ragged_all_to_all_mode=symmetric -- all
ncclCommWindowRegister invalid usage.
Revert both files to openxla/xla once
openxla/xla#46878 lands upstream.
mminutoli
added a commit
to ROCm/jax
that referenced
this pull request
Aug 10, 2026
XLA 7a57c69ea6 enabled xla_enable_nccl_symmetric_buffers_for_collectives
by default for ALLREDUCE, ALLGATHER and REDUCESCATTER up to 8 MB. That
makes XLA register collective buffers with
ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC), which RCCL declares
but rejects at runtime with ncclInvalidUsage, so every multi-GPU
collective on ROCm fails:
INTERNAL: RCCL operation ncclCommWindowRegister(
comm, addr.opaque(), addr.size(), &win, 0x01) failed: invalid usage
Ragged-all-to-all fails the same way through a separate route: the
one-shot NCCL barrier is on by default since 26072180fa, and
RaggedAllToAllThunk::PrepareCollective requests a symmetric allocation
for its output buffer regardless of the symmetric-buffers filters.
Pin XLA to ROCm/xla branch rocm-jaxlib-0.11.1-alpha, which is the
previously pinned upstream revision
1c32308ad0599517cd4dd2cfae136851b0712bb5 plus a single commit that gates
both routes behind a runtime ROCm platform probe.
The download URL has to move to the fork as well, since the commit does
not exist in openxla/xla. tf_mirror_urls() is kept because
tf_http_archive() requires a mirror as the first URL; that mirror does
not carry fork archives, so it 404s and Bazel falls back to GitHub.
Validated on 8x MI355X: //tests:ragged_collective_test_gpu passes 12
tests and //tests:pmap_test_gpu passes 288 (11 skipped) with the
symmetric-buffer filters forced on the command line. As negative
controls, forcing either escape hatch back on reproduces the original
failure exactly -- 107 of 288 pmap errors with
xla_gpu_experimental_enable_nccl_symmetric_buffers=true, and 10 of 12
ragged errors with xla_gpu_ragged_all_to_all_mode=symmetric -- all
ncclCommWindowRegister invalid usage.
Revert both files to openxla/xla once
openxla/xla#46878 lands upstream.
This was referenced Aug 11, 2026
copybara-service Bot
pushed a commit
that referenced
this pull request
Aug 11, 2026
Imported from GitHub PR #46895 For now, we disable symmetric memory support on ROCM platform which was enabled by default in 7a57c69 This is a hot-fix, later on we can add an RCCL version check for that. Still, symmetric memory can be enforced on ROCM with **--xla_gpu_experimental_enable_nccl_symmetric_buffers=true** XLA flag. This can be used instead of this fix: #46878 Copybara import of the project: -- 3e23a78 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: disable rccl symm memory runtime check skipped symm memory test for rocm Revert "skipped symm memory test for rocm" This reverts commit 30d805a. disabled commwindowregister call -- ae6d9dc by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert -- c682659 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert Merging this change closes #46895 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46895 from ROCm:pemeliya/rccl_disable_symm_memory c682659 PiperOrigin-RevId: 962837161
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Aug 11, 2026
Imported from GitHub PR openxla/xla#46895 For now, we disable symmetric memory support on ROCM platform which was enabled by default in openxla/xla@7a57c69 This is a hot-fix, later on we can add an RCCL version check for that. Still, symmetric memory can be enforced on ROCM with **--xla_gpu_experimental_enable_nccl_symmetric_buffers=true** XLA flag. This can be used instead of this fix: openxla/xla#46878 Copybara import of the project: -- 3e23a7877e70c81fa7cf997a66e4526d2c539ccf by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: disable rccl symm memory runtime check skipped symm memory test for rocm Revert "skipped symm memory test for rocm" This reverts commit 30d805ac6fbf9e83c217936fde63973d68f829db. disabled commwindowregister call -- ae6d9dc8b2db762cde70bac3596e20144e34187f by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert -- c682659d74d43cfb9cfa82aa935dc7d0f00c6f5f by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert Merging this change closes #46895 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#46895 from ROCm:pemeliya/rccl_disable_symm_memory c682659d74d43cfb9cfa82aa935dc7d0f00c6f5f PiperOrigin-RevId: 962837161
copybara-service Bot
pushed a commit
that referenced
this pull request
Aug 11, 2026
Imported from GitHub PR #46895 For now, we disable symmetric memory support on ROCM platform which was enabled by default in 7a57c69 This is a hot-fix, later on we can add an RCCL version check for that. Still, symmetric memory can be enforced on ROCM with **--xla_gpu_experimental_enable_nccl_symmetric_buffers=true** XLA flag. This can be used instead of this fix: #46878 Copybara import of the project: -- 3e23a78 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: disable rccl symm memory runtime check skipped symm memory test for rocm Revert "skipped symm memory test for rocm" This reverts commit 30d805a. disabled commwindowregister call -- ae6d9dc by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert -- c682659 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert Merging this change closes #46895 COPYBARA_INTEGRATE_REVIEW=#46895 from ROCm:pemeliya/rccl_disable_symm_memory c682659 PiperOrigin-RevId: 962907538
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Aug 11, 2026
Imported from GitHub PR openxla/xla#46895 For now, we disable symmetric memory support on ROCM platform which was enabled by default in openxla/xla@7a57c69 This is a hot-fix, later on we can add an RCCL version check for that. Still, symmetric memory can be enforced on ROCM with **--xla_gpu_experimental_enable_nccl_symmetric_buffers=true** XLA flag. This can be used instead of this fix: openxla/xla#46878 Copybara import of the project: -- 3e23a7877e70c81fa7cf997a66e4526d2c539ccf by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: disable rccl symm memory runtime check skipped symm memory test for rocm Revert "skipped symm memory test for rocm" This reverts commit 30d805ac6fbf9e83c217936fde63973d68f829db. disabled commwindowregister call -- ae6d9dc8b2db762cde70bac3596e20144e34187f by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert -- c682659d74d43cfb9cfa82aa935dc7d0f00c6f5f by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>: revert Merging this change closes #46895 PiperOrigin-RevId: 962907538
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
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.
Problem
7a57c69ea6("Enable NCCL symmetric kernels for collectives up to 8 MB") turned onxla_enable_nccl_symmetric_buffers_for_collectivesby default for ALLREDUCE, ALLGATHER and REDUCESCATTER up to 8 MB. This breaks every multi-device collective on ROCm:Symmetric buffers require registering collective buffers via
ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC). RCCL declares that entry point inrccl.h, so the build and link both succeed, and it reportsNCCL_VERSION_CODE 23004— above the NCCL version that introduced symmetric memory — so version-based feature detection passes too. It then rejects every registration at runtime.Root cause
The general issue is that RCCL's version code aliases into NCCL's numbering, so any
NCCL_VERSION_CODE-based capability check silently misfires for RCCL. This particular feature is one instance of that.Fix
Clear the filters during HLO compilation, keyed on the compute capability of the device being compiled for.
Why not a build-time guard. The defaults are set in
DefaultDebugOptionsIgnoringFlags(), which compiles into a backend-agnosticjaxlib. Interjecting the issue here would required a customjaxlibfor ROCm.Why
RunBackend(). It is the single choke point: buffer assignment colors the collective buffers and thunk emission selects the symmetric kernels, and both run underCompileToBackendResult().Validation
8× MI355X (gfx950), JAX
//tests:pmap_test_gpu. The filters were forced on explicitly so the result does not depend on compiled-in defaults:ncclCommWindowRegister ... invalid usage, on all 3 attemptsRunBackend()Notes for reviewers
add_...and never clears, so--xla_enable_nccl_symmetric_buffers_for_collectives=appends to the defaults rather than replacing them, and there is no way to turn the feature off viaXLA_FLAGSon any platform. Deliberately not changed here to keep this focused, but worth fixing. It is what turned this regression into a hard outage with no mitigation.