From 4e6cec4d09ca399c66541ee61109c3099af1a463 Mon Sep 17 00:00:00 2001 From: cloudhan Date: Tue, 22 Aug 2023 11:08:55 +0800 Subject: [PATCH] Update ck and enable test (#16383) Apply the fix in https://github.com/ROCmSoftwarePlatform/composable_kernel/issues/728 Introduce more kernel instances and allow the introduction of streamk and splitk. --- cmake/external/composable_kernel.cmake | 6 +++++- .../composable_kernel/Fix_Clang_Build.patch | 20 +++++-------------- .../impl.cuh | 10 ++++++---- onnxruntime/core/framework/tunable.h | 6 +++++- .../core/providers/rocm/math/softmax_ck.cuh | 5 +---- .../core/providers/rocm/tunable/gemm_ck.cuh | 11 ---------- .../kernel_explorer_interface.h | 16 +++++++++++++++ .../kernels/gemm_softmax_gemm_permute_test.py | 4 +++- .../kernels/rocm/elementwise.cu | 1 + .../kernels/rocm/gemm_fast_gelu_tunable.cu | 1 + .../kernels/rocm/gemm_tunable.cu | 3 +++ .../kernel_explorer/kernels/rocm/softmax.cu | 1 + .../multihead_attention_op_test.cc | 1 - .../orttraining-pai-ci-pipeline.yml | 13 ++++++------ 14 files changed, 54 insertions(+), 44 deletions(-) diff --git a/cmake/external/composable_kernel.cmake b/cmake/external/composable_kernel.cmake index fe57a5b532..d0b721c9a8 100644 --- a/cmake/external/composable_kernel.cmake +++ b/cmake/external/composable_kernel.cmake @@ -1,5 +1,5 @@ set(composable_kernel_URL https://github.com/ROCmSoftwarePlatform/composable_kernel.git) -set(composable_kernel_TAG ed3a2e52265e11daa366f47b082141a652b67c58) # 2023-04-10 21:02:17 +0800 +set(composable_kernel_TAG d52ec01652b7d620386251db92455968d8d90bdc) # 2023-08-18 11:14:59 +0800 set(PATCH ${PROJECT_SOURCE_DIR}/patches/composable_kernel/Fix_Clang_Build.patch) @@ -14,10 +14,14 @@ FetchContent_GetProperties(composable_kernel) if(NOT composable_kernel_POPULATED) FetchContent_Populate(composable_kernel) set(BUILD_DEV OFF CACHE BOOL "Disable -Weverything, otherwise, error: 'constexpr' specifier is incompatible with C++98 [-Werror,-Wc++98-compat]" FORCE) + # Exclude i8 device gemm instances due to excessive long compilation time and not being used + set(DTYPES fp32 fp16 bf16) + set(INSTANCES_ONLY ON) add_subdirectory(${composable_kernel_SOURCE_DIR} ${composable_kernel_BINARY_DIR} EXCLUDE_FROM_ALL) add_library(onnxruntime_composable_kernel_includes INTERFACE) target_include_directories(onnxruntime_composable_kernel_includes INTERFACE ${composable_kernel_SOURCE_DIR}/include ${composable_kernel_SOURCE_DIR}/library/include) + target_compile_definitions(onnxruntime_composable_kernel_includes INTERFACE __fp32__ __fp16__ __bf16__) endif() diff --git a/cmake/patches/composable_kernel/Fix_Clang_Build.patch b/cmake/patches/composable_kernel/Fix_Clang_Build.patch index ae2a015c1d..d564ffba91 100644 --- a/cmake/patches/composable_kernel/Fix_Clang_Build.patch +++ b/cmake/patches/composable_kernel/Fix_Clang_Build.patch @@ -1,5 +1,5 @@ diff --git a/CMakeLists.txt b/CMakeLists.txt -index f861e3020..f0b6bceae 100644 +index 514b98fde..59c8a568a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,7 +1,7 @@ @@ -11,7 +11,7 @@ index f861e3020..f0b6bceae 100644 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake") -@@ -41,27 +41,6 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) +@@ -94,27 +94,6 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CXX_EXTENSIONS OFF) message("CMAKE_CXX_COMPILER_ID: ${CMAKE_CXX_COMPILER_ID}") @@ -39,7 +39,7 @@ index f861e3020..f0b6bceae 100644 ## HIP find_package(HIP REQUIRED) # Override HIP version in config.h, if necessary. -@@ -83,8 +62,6 @@ if( DEFINED CK_OVERRIDE_HIP_VERSION_PATCH ) +@@ -136,8 +115,6 @@ if( DEFINED CK_OVERRIDE_HIP_VERSION_PATCH ) message(STATUS "CK_HIP_VERSION_PATCH overriden with ${CK_OVERRIDE_HIP_VERSION_PATCH}") endif() message(STATUS "Build with HIP ${HIP_VERSION}") @@ -48,17 +48,7 @@ index f861e3020..f0b6bceae 100644 ## tidy include(EnableCompilerWarnings) -@@ -273,9 +250,6 @@ rocm_package_setup_component(profiler - ) - - add_subdirectory(library) --add_subdirectory(example) --add_subdirectory(test) --add_subdirectory(profiler) - - #Create an interface target for the include only files and call it "composablekernels" - include(CMakePackageConfigHelpers) -@@ -301,11 +275,3 @@ rocm_install(FILES +@@ -391,11 +368,3 @@ rocm_install(FILES set(CPACK_RESOURCE_FILE_LICENSE "${CMAKE_CURRENT_SOURCE_DIR}/LICENSE") set(CPACK_RPM_PACKAGE_LICENSE "MIT") @@ -71,7 +61,7 @@ index f861e3020..f0b6bceae 100644 - HEADER_ONLY -) diff --git a/library/src/tensor_operation_instance/gpu/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/CMakeLists.txt -index c206c4dc0..b283eeb64 100644 +index 1d54a141b..4edd7dbfb 100644 --- a/library/src/tensor_operation_instance/gpu/CMakeLists.txt +++ b/library/src/tensor_operation_instance/gpu/CMakeLists.txt @@ -1,7 +1,13 @@ diff --git a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh index ae4aef5f99..0599318a40 100644 --- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh @@ -52,10 +52,11 @@ template , AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 64, 32, 128, 32, 8, 8, 2, 32, 32, 1, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec, 1>, DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 256, 128, 32, 64, 32, 8, 8, 2, 32, 32, 2, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>, DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 256, 128, 32, 128, 32, 8, 8, 2, 32, 32, 2, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>, #if ROCM_VERSION >= 50500 @@ -71,6 +72,7 @@ using device_batched_gemm_softmax_gemm_permute_instances = DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 64, 256, 64, 128, 32, 8, 8, 2, 16, 16, 1, 16, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 8, S<1, 16, 1,16>, 8, MaskingSpec>, DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 64, 256, 64, 64, 32, 8, 8, 2, 16, 16, 1, 16, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 4, S<1, 32, 1, 8>, 8, MaskingSpec>, // Padded fallback kernel + DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 128, 64, 128, 32, 8, 8, 2, 32, 32, 1, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec, 1>, DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 128, 64, 128, 32, 8, 8, 2, 32, 32, 1, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>, DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 64, 32, 128, 32, 8, 8, 2, 32, 32, 1, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec> // clang-format on diff --git a/onnxruntime/core/framework/tunable.h b/onnxruntime/core/framework/tunable.h index 91a7b2dfa5..96b4cc53a0 100644 --- a/onnxruntime/core/framework/tunable.h +++ b/onnxruntime/core/framework/tunable.h @@ -268,6 +268,10 @@ class TunableOp { WarmUp(candidate, params); auto approx_duration = Profile(candidate, params, approx_num_iter); + if (approx_duration > 2 * min_time) { + LOGS_DEFAULT(VERBOSE) << "FindFastestImpl skip slow instance " << op_sig << '(' << param_sig << ") id=" << i; + continue; + } int tuning_iter = std::max(1, int(std::min(double(max_tuning_iter), ctx->GetMaxTuningDurationMs() / approx_duration))); LOGS_DEFAULT(VERBOSE) << "FindFastestImpl run instance " << op_sig << '(' << param_sig << ") id=" << i << " " << tuning_iter << " times."; @@ -278,7 +282,7 @@ class TunableOp { id = static_cast(i); } } - ORT_ENFORCE(id >= 0, "Cannot found viable op"); + ORT_ENFORCE(id >= 0, "Could not find viable op"); LOGS_DEFAULT(VERBOSE) << "FindFastestImpl for " << op_sig << '(' << param_sig << ") found fastest with id=" << id; std::this_thread::sleep_for(std::chrono::milliseconds(50)); return id; diff --git a/onnxruntime/core/providers/rocm/math/softmax_ck.cuh b/onnxruntime/core/providers/rocm/math/softmax_ck.cuh index 502320fa5c..5830c9dd0b 100644 --- a/onnxruntime/core/providers/rocm/math/softmax_ck.cuh +++ b/onnxruntime/core/providers/rocm/math/softmax_ck.cuh @@ -34,7 +34,7 @@ auto GetCKSoftmaxTypeStringAndOps() { using OutDataType = typename CKDataTypeAdaptor::type; using AccDataType = typename CKDataTypeAdaptor::type; using DeviceSoftmax = ck::tensor_operation::device:: - DeviceSoftmax; + DeviceSoftmax; using InstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory; std::vector>>> ret; @@ -49,9 +49,6 @@ auto GetCKSoftmaxTypeStringAndOps() { TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( params->is_log_softmax, impl->GetTypeString(), " does not support log softmax"); - TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF( - impl->GetRank() != Rank || impl->GetNumReduceDim() != NumReduceDim, - impl->GetTypeString(), " does not support current Rank or NumReduceDim ", params->Signature()); std::vector in_lengths{1, 1, params->batch_count, params->softmax_elements}; std::vector in_strides{params->batch_count * params->input_stride, params->batch_count * params->input_stride, params->input_stride, 1}; diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh b/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh index 21edc20c0b..3e6f1612f2 100644 --- a/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh +++ b/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh @@ -46,17 +46,6 @@ auto GetCKGemmTypeStringAndOps() { std::vector>>> ret; for (auto&& impl : InstanceFactory::GetInstances()) { auto type_string = impl->GetTypeString(); - - // FIXME: ck upstream have bugs in some input shapes coupled with specific impls. The `IsSupportedArgument` is not - // sound, we exclude those implementation here for now. Check back later when AMD fixed them. - // - // The DeviceGemmXdl<256, 128, 144, 8, 8, 16, 16, 2, 9> and DeviceGemmXdl<256, 128, 144, 4, 8, 16, 16, 2, 9> only - // occurs in DeviceGemm for FP16. When k < 8, the result is wrong. - if (type_string == "DeviceGemmXdl<256, 128, 144, 8, 8, 16, 16, 2, 9>" || - type_string == "DeviceGemmXdl<256, 128, 144, 4, 8, 16, 16, 2, 9>") { - continue; - } - auto invoker = impl->MakeInvokerPointer(); auto ck_gemm_op = [impl = std::move(impl), invoker = std::move(invoker)](const GemmParams* params) -> Status { auto one = ToHipType::FromFloat(1.0f); diff --git a/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h b/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h index 0abc836b10..9eb0adcede 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h +++ b/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h @@ -94,6 +94,22 @@ class IKernelExplorer { int repeats_{100}; }; +class WithMaxTuningDurationMs { + public: + WithMaxTuningDurationMs(TuningContextT* ctx, int ms) : ctx_(ctx) { + original_tuning_duration_ = ctx_->GetMaxTuningDurationMs(); + ctx_->SetMaxTuningDurationMs(ms); + } + + ~WithMaxTuningDurationMs() { + ctx_->SetMaxTuningDurationMs(original_tuning_duration_); + } + + private: + TuningContextT* ctx_; + int original_tuning_duration_; +}; + pybind11::module GetKernelExplorerModule(); class KernelExplorerInit { diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py index 64c7c76a1a..6e1e431842 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py @@ -14,6 +14,8 @@ import numpy as np import pytest from utils import dtype_to_suffix, matmul, softmax +max_batch_size = int(os.environ.get("KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE", 64)) + def multinormal_distribution(num_distribution, num_element_per_dist): arrays = [] @@ -36,7 +38,7 @@ def get_ck_binding_name(dtype, biased: bool, masked: bool): dtypes = ["float16"] -batches = [1, 64] +batches = [1, max_batch_size] seqlens = [128, 512] total_seqlens = [128, 512] num_heads = [8, 12] diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/elementwise.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/elementwise.cu index cc76892c94..2151a8cc27 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/elementwise.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/elementwise.cu @@ -64,6 +64,7 @@ class ElementwiseTunable : public IKernelExplorer { } void Run() override { + WithMaxTuningDurationMs max_duration(TuningContext(), 250); ORT_THROW_IF_ERROR(op_(¶ms_)); } diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_tunable.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_tunable.cu index 6f0f8ca5e9..7ecb87828a 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_tunable.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_tunable.cu @@ -56,6 +56,7 @@ class GemmFastGeluTunable : public IKernelExplorer { } void Run() override { + WithMaxTuningDurationMs max_duration(TuningContext(), 250); ORT_THROW_IF_ERROR((op_(¶ms_))); } diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_tunable.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_tunable.cu index 0ca0079ab0..d1786f94b1 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_tunable.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_tunable.cu @@ -56,6 +56,7 @@ class GemmTunable : public IKernelExplorer { } void Run() override { + WithMaxTuningDurationMs max_duration(TuningContext(), 250); ORT_THROW_IF_ERROR(op_(¶ms_)); } @@ -117,6 +118,7 @@ class BatchedGemmTunable : public IBatchedGemmKernelExplorer { } void Run() override { + WithMaxTuningDurationMs max_duration(params_.TuningContext(), 250); ORT_THROW_IF_ERROR(op_(¶ms_)); } @@ -179,6 +181,7 @@ class StridedBatchedGemmTunable : public IKernelExplorer { } void Run() override { + WithMaxTuningDurationMs max_duration(params_.TuningContext(), 250); ORT_THROW_IF_ERROR(op_(¶ms_)); } diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu index a4f5329a31..6cc59ee257 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu @@ -111,6 +111,7 @@ class SoftmaxTunable : public IKernelExplorer { } void Run() override { + WithMaxTuningDurationMs max_duration(TuningContext(), 250); ORT_THROW_IF_ERROR(op_(¶ms_)); } diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc index c2230501b0..89b8c47fb4 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc @@ -478,7 +478,6 @@ TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize32_RightSidePadding_M } TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize32_RightSidePadding_Mask2D) { - ROCM_GTEST_SKIP("ROCm MHA expect failure due to ck bug"); AttentionTestData data; GetCrossAttentionData_Batch2_HeadSize32_RightSidePadding(data, false); RunMultiHeadAttentionTests(data, true); diff --git a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml index 81b8a980db..523390debc 100644 --- a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml @@ -186,13 +186,14 @@ jobs: --user onnxruntimedev \ --volume $(Build.SourcesDirectory):/onnxruntime_src \ --volume $(Build.BinariesDirectory):/build \ + -e OPENBLAS_NUM_THREADS=1 \ + -e OPENMP_NUM_THREADS=1 \ + -e MKL_NUM_THREADS=1 \ + -e KERNEL_EXPLORER_BUILD_DIR=/build/$(BuildConfig) \ + -e KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE=8 \ + -e KERNEL_EXPLORER_TEST_USE_CUPY=1 \ onnxruntimetrainingrocm-cibuild-rocm$(RocmVersion)-test \ - /bin/bash -c " - set -ex; \ - export KERNEL_EXPLORER_BUILD_DIR=/build/$(BuildConfig); \ - export KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE=8; \ - export KERNEL_EXPLORER_TEST_USE_CUPY=1; \ - pytest /onnxruntime_src/onnxruntime/python/tools/kernel_explorer/ -n 8 --reruns 1 --durations=100" + pytest /onnxruntime_src/onnxruntime/python/tools/kernel_explorer/ -n 4 --reruns 1 --durations=100 workingDirectory: $(Build.SourcesDirectory) displayName: 'Run kernel explorer tests' condition: succeededOrFailed()