From c71ec89cfcf41c4f7e43a64bb4d27623aa87bbe1 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Fri, 7 Aug 2026 05:03:05 +0000 Subject: [PATCH] [ROCm] Do not use NCCL symmetric buffers on ROCm. 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 7a57c69ea6 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. --- xla/service/gpu/gpu_compiler.cc | 35 +++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/xla/service/gpu/gpu_compiler.cc b/xla/service/gpu/gpu_compiler.cc index 64a8c91a5b43a2..42330482f893f4 100644 --- a/xla/service/gpu/gpu_compiler.cc +++ b/xla/service/gpu/gpu_compiler.cc @@ -2504,6 +2504,36 @@ bool RequiresCollectiveOutput(const HloValue* value, const DebugOptions& opts) { return false; } +// Symmetric buffers require registering the collective buffers with +// ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC). RCCL declares that +// entry point in rccl.h and reports an NCCL version above the one that +// introduced it, so both the build-time and version-based checks pass, but it +// rejects every registration at runtime with ncclInvalidUsage. Enabling +// symmetric buffers therefore fails every multi-device collective on ROCm. +// +// This has to be decided from the device being compiled for, not from a +// build-time define: the defaults are set in +// DefaultDebugOptionsIgnoringFlags(), which is compiled into the +// backend-agnostic jaxlib, while only the plugin and PJRT artifacts are built +// with ROCm configured. +void DisableSymmetricBuffersIfUnsupported( + HloModule* module, const se::DeviceDescription& device_description) { + if (device_description.gpu_compute_capability().rocm_compute_capability() == + nullptr) { + return; + } + DebugOptions& opts = module->mutable_config().mutable_debug_options(); + if (opts.xla_enable_nccl_symmetric_buffers_for_collectives().empty() && + !opts.xla_gpu_experimental_enable_nccl_symmetric_buffers()) { + return; + } + LOG_FIRST_N(WARNING, 1) + << "Disabling NCCL symmetric buffers: RCCL does not implement collective " + "window registration (ncclCommWindowRegister)."; + opts.clear_xla_enable_nccl_symmetric_buffers_for_collectives(); + opts.set_xla_gpu_experimental_enable_nccl_symmetric_buffers(false); +} + void GpuCollectiveBufferAnalysis( HloModule* module, const HloAliasAnalysis& alias_analysis, std::function add_index_to_copy) { @@ -2960,6 +2990,11 @@ absl::StatusOr> GpuCompiler::RunBackend( const se::DeviceDescription& gpu_device_info = gpu_topology.gpu_target_config().device_description; + // Must happen before CompileToBackendResult(): buffer assignment colors + // collective buffers into the collective memory space and thunk emission + // selects the symmetric kernels, both from these debug options. + DisableSymmetricBuffersIfUnsupported(module.get(), gpu_device_info); + if (module->config().hlo_profiling_enabled() || VLOG_IS_ON(1)) { HloCostAnalysis::Options cost_analysis_options{ShapeSizeBytesFunction()}; cost_analysis_options.set_bytes_per_second(