Add NCCL_MAX_NCHANNELS env variable to multi gpu tests

ROCm · Nov 28, 2024 · 2aa5cd2 · 2aa5cd2
1 parent d529b07
commit 2aa5cd2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/build_tools/rocm/run_xla_multi_gpu.sh b/build_tools/rocm/run_xla_multi_gpu.sh
@@ -85,7 +85,8 @@ 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 \
-    -- //xla/tests:collective_ops_e2e_test_gpu_amd_any \
+    --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 \
        //xla/tools/multihost_hlo_runner:functional_hlo_runner_test_gpu_amd_any \