Add NCCL_MAX_NCHANNELS env variable to multi gpu tests

ROCm · Nov 28, 2024 · ee307e3 · ee307e3
1 parent 1d70f15
commit ee307e3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/build_tools/rocm/run_xla_multi_gpu.sh b/build_tools/rocm/run_xla_multi_gpu.sh
@@ -71,6 +71,7 @@ bazel \
     --test_env=TF_GPU_COUNT=$TF_GPU_COUNT \
     --action_env=XLA_FLAGS=--xla_gpu_force_compilation_parallelism=16 \
     --action_env=XLA_FLAGS=--xla_gpu_enable_llvm_module_compilation_parallelism=true \
+    --action_env=NCCL_MAX_NCHANNELS=1 \
     -- //xla/tests:collective_ops_test_e2e_gpu_amd_any \
        //xla/tests:collective_ops_test_gpu_amd_any \
        //xla/tests:replicated_io_feed_test_gpu_amd_any \