fix the timeout bug of some communication api on A100 (#54513)

PaddlePaddle · Jun 13, 2023 · 60e3e35 · 60e3e35
1 parent 9857bfc
commit 60e3e35
Showing 1 changed file with 21 additions and 6 deletions.
diff --git a/test/collective/CMakeLists.txt b/test/collective/CMakeLists.txt
@@ -79,8 +79,13 @@ if((WITH_GPU OR WITH_ROCM) AND (LINUX))
   py_test_modules(
     test_collective_alltoall_api MODULES test_collective_alltoall_api ENVS
     "http_proxy=;https_proxy=;PYTHONPATH=..:${PADDLE_BINARY_DIR}/python")
-  set_tests_properties(test_collective_alltoall_api
-                       PROPERTIES TIMEOUT "120" LABELS "RUN_TYPE=DIST")
+  if(${CUDA_ARCH_NAME} STREQUAL "Ampere")
+    set_tests_properties(test_collective_alltoall_api
+                         PROPERTIES TIMEOUT "160" LABELS "RUN_TYPE=DIST")
+  else()
+    set_tests_properties(test_collective_alltoall_api
+                         PROPERTIES TIMEOUT "120" LABELS "RUN_TYPE=DIST")
+  endif()
 endif()
 if((WITH_GPU OR WITH_ROCM) AND (LINUX))
   bash_test_modules(
@@ -130,8 +135,13 @@ if((WITH_GPU OR WITH_ROCM) AND (LINUX))
   py_test_modules(
     test_collective_broadcast_api MODULES test_collective_broadcast_api ENVS
     "http_proxy=;https_proxy=;PYTHONPATH=..:${PADDLE_BINARY_DIR}/python")
-  set_tests_properties(test_collective_broadcast_api
-                       PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=DIST")
+  if(${CUDA_ARCH_NAME} STREQUAL "Ampere")
+    set_tests_properties(test_collective_broadcast_api
+                         PROPERTIES TIMEOUT "360" LABELS "RUN_TYPE=DIST")
+  else()
+    set_tests_properties(test_collective_broadcast_api
+                         PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=DIST")
+  endif()
 endif()
 if((WITH_GPU OR WITH_ROCM) AND (LINUX))
   py_test_modules(
@@ -228,8 +238,13 @@ if((WITH_GPU OR WITH_ROCM) AND (LINUX))
     test_collective_reduce_scatter_api MODULES
     test_collective_reduce_scatter_api ENVS
     "http_proxy=;https_proxy=;PYTHONPATH=..:${PADDLE_BINARY_DIR}/python")
-  set_tests_properties(test_collective_reduce_scatter_api
-                       PROPERTIES TIMEOUT "150" LABELS "RUN_TYPE=DIST")
+  if(${CUDA_ARCH_NAME} STREQUAL "Ampere")
+    set_tests_properties(test_collective_reduce_scatter_api
+                         PROPERTIES TIMEOUT "210" LABELS "RUN_TYPE=DIST")
+  else()
+    set_tests_properties(test_collective_reduce_scatter_api
+                         PROPERTIES TIMEOUT "150" LABELS "RUN_TYPE=DIST")
+  endif()
 endif()
 if((WITH_GPU OR WITH_ROCM) AND (LINUX))
   py_test_modules(