From 6fc5237d9e79dac76886368d4b5c7e3d81d5a8e2 Mon Sep 17 00:00:00 2001 From: Cian Hayes Date: Mon, 1 Feb 2021 21:54:16 -0800 Subject: [PATCH] Introduce --enable_training_ops build flag (#6523) * minimal_build with training ops * Removing redundant comment from an earlier attempt at a fix * Fixing a bad merge conflict resolution * Responding to PR feedback * tweaking the makefiles based on feedback * combining two enable_training blocks in CMakeLists.txt --- cmake/CMakeLists.txt | 6 ++++ cmake/onnxruntime_providers.cmake | 30 ++++++++++++++++--- .../providers/cpu/cpu_execution_provider.cc | 4 +-- .../core/providers/cpu/tensor/scatter.cc | 4 +-- .../training_ops/cpu/cpu_training_kernels.cc | 6 ++++ .../training_ops/cpu/nn/dropout_7.cc | 2 +- tools/ci_build/build.py | 4 +++ 7 files changed, 47 insertions(+), 9 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 61c85324b3..c031fc6830 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -142,6 +142,7 @@ option(onnxruntime_FUZZ_TEST "Enable Fuzz testing" OFF) option(onnxruntime_ENABLE_NVTX_PROFILE "Enable NVTX profile." OFF) option(onnxruntime_ENABLE_MEMORY_PROFILE "Enable memory profile." OFF) option(onnxruntime_ENABLE_TRAINING "Enable training functionality." OFF) +option(onnxruntime_ENABLE_TRAINING_OPS "Include training operators but no training session support." OFF) option(onnxruntime_ENABLE_TRAINING_E2E_TESTS "Enable training end-to-end tests." OFF) option(onnxruntime_ENABLE_CPU_FP16_OPS "Build with advanced instruction sets" ON) option(onnxruntime_USE_NCCL "Build with NCCL support" OFF) @@ -1238,8 +1239,13 @@ if (onnxruntime_USE_DML) include(dml) endif() +if (onnxruntime_ENABLE_TRAINING_OPS) + add_compile_definitions(ENABLE_TRAINING_OPS) +endif() + if (onnxruntime_ENABLE_TRAINING) add_compile_definitions(ENABLE_TRAINING) + add_compile_definitions(ENABLE_TRAINING_OPS) if (UNIX) # Find MPI diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index 26d4ed0629..3c6cd8fbe3 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -129,6 +129,25 @@ if (onnxruntime_USE_FEATURIZERS) list(APPEND onnxruntime_providers_src ${onnxruntime_cpu_featurizers_cc_srcs}) endif() +if (onnxruntime_ENABLE_TRAINING_OPS) + file(GLOB_RECURSE onnxruntime_cpu_training_ops_srcs CONFIGURE_DEPENDS + "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/*.h" + "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/*.cc" + ) + + source_group(TREE ${ORTTRAINING_ROOT}/ FILES ${onnxruntime_cpu_training_ops_srcs}) + list(APPEND onnxruntime_providers_src ${onnxruntime_cpu_training_ops_srcs}) + + file(GLOB_RECURSE onnxruntime_cpu_full_training_only_srcs + "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/gist/*.h" + "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/gist/*.cc" + "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/tensorboard/*.h" + "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/tensorboard/*.cc" + ) + + list(REMOVE_ITEM onnxruntime_providers_src ${onnxruntime_cpu_full_training_only_srcs}) +endif() + if (onnxruntime_ENABLE_TRAINING) file(GLOB_RECURSE onnxruntime_cpu_training_ops_srcs CONFIGURE_DEPENDS "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/*.h" @@ -174,8 +193,11 @@ target_include_directories(onnxruntime_providers PRIVATE ${ONNXRUNTIME_ROOT} ${e add_dependencies(onnxruntime_providers onnx ${onnxruntime_EXTERNAL_DEPENDENCIES}) -if (onnxruntime_ENABLE_TRAINING) +if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS) target_include_directories(onnxruntime_providers PRIVATE ${ORTTRAINING_ROOT}) +endif() + +if (onnxruntime_ENABLE_TRAINING) add_dependencies(onnxruntime_providers tensorboard) onnxruntime_add_include_to_target(onnxruntime_providers tensorboard) @@ -208,7 +230,7 @@ if (onnxruntime_USE_CUDA) list(APPEND onnxruntime_providers_cuda_src ${onnxruntime_cuda_contrib_ops_cc_srcs} ${onnxruntime_cuda_contrib_ops_cu_srcs}) endif() - if (onnxruntime_ENABLE_TRAINING) + if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS) file(GLOB_RECURSE onnxruntime_cuda_training_ops_cc_srcs CONFIGURE_DEPENDS "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/*.h" "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/*.cc" @@ -253,7 +275,7 @@ if (onnxruntime_USE_CUDA) target_compile_options(onnxruntime_providers_cuda PRIVATE "$<$:SHELL:-Xcompiler /wd4127>") endif() onnxruntime_add_include_to_target(onnxruntime_providers_cuda onnxruntime_common onnxruntime_framework onnx onnx_proto protobuf::libprotobuf flatbuffers) - if (onnxruntime_ENABLE_TRAINING) + if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS) onnxruntime_add_include_to_target(onnxruntime_providers_cuda onnxruntime_training) target_link_libraries(onnxruntime_providers_cuda PRIVATE onnxruntime_training) endif() @@ -267,7 +289,7 @@ if (onnxruntime_USE_CUDA) target_include_directories(onnxruntime_providers_cuda PRIVATE ${PROJECT_SOURCE_DIR}/external/cub) endif() - if (onnxruntime_ENABLE_TRAINING) + if (onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_TRAINING_OPS) target_include_directories(onnxruntime_providers_cuda PRIVATE ${ORTTRAINING_ROOT} ${MPI_INCLUDE_DIRS}) if (onnxruntime_USE_NCCL) diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index 15d98894f9..872ba119c7 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -13,7 +13,7 @@ #include "featurizers_ops/cpu/cpu_featurizers_kernels.h" #endif -#ifdef ENABLE_TRAINING +#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS) #include "orttraining/training_ops/cpu/cpu_training_kernels.h" #endif @@ -1840,7 +1840,7 @@ static Status RegisterCPUKernels(KernelRegistry& kernel_registry) { #ifdef ML_FEATURIZERS ORT_RETURN_IF_ERROR(::onnxruntime::featurizers::RegisterCpuMSFeaturizersKernels(kernel_registry)); #endif -#ifdef ENABLE_TRAINING +#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS) ORT_RETURN_IF_ERROR(::onnxruntime::contrib::RegisterCpuTrainingKernels(kernel_registry)); #endif return Status::OK(); diff --git a/onnxruntime/core/providers/cpu/tensor/scatter.cc b/onnxruntime/core/providers/cpu/tensor/scatter.cc index 2c047676cc..dbb3b14e2c 100644 --- a/onnxruntime/core/providers/cpu/tensor/scatter.cc +++ b/onnxruntime/core/providers/cpu/tensor/scatter.cc @@ -5,7 +5,7 @@ #include "core/common/common.h" #include "core/framework/op_kernel.h" #include "core/providers/common.h" -#ifdef ENABLE_TRAINING +#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS) #include "orttraining/training_ops/cpu/tensor/gather_elements_grad_impl.h" #endif @@ -255,7 +255,7 @@ Status Scatter::Compute(OpKernelContext* context) const { return status; } -#ifdef ENABLE_TRAINING +#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS) namespace contrib { diff --git a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc index bd666a3402..7d74ac12b8 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc @@ -58,12 +58,14 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_float, DropoutGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_double, DropoutGrad); +#ifdef ENABLE_TRAINING class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryScalar); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryHistogram); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryMerge); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, SummaryText); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GistBinarizeEncoder); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GistBinarizeDecoder); +#endif class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, LayerNormalizationGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SimplifiedLayerNormalizationGrad); @@ -147,18 +149,22 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, +#ifdef ENABLE_TRAINING BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, +#endif BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, +#ifdef ENABLE_TRAINING BuildKernelCreateInfo, BuildKernelCreateInfo, +#endif BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cpu/nn/dropout_7.cc b/orttraining/orttraining/training_ops/cpu/nn/dropout_7.cc index 1e94114267..ed126eb260 100644 --- a/orttraining/orttraining/training_ops/cpu/nn/dropout_7.cc +++ b/orttraining/orttraining/training_ops/cpu/nn/dropout_7.cc @@ -6,7 +6,7 @@ namespace onnxruntime { -#ifdef ENABLE_TRAINING +#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS) ONNX_CPU_OPERATOR_KERNEL( Dropout, 7, diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 04eefbfce6..61da3c719c 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -159,6 +159,8 @@ def parse_arguments(): "--enable_memory_profile", action='store_true', help="Enable memory profile in ORT.") parser.add_argument( "--enable_training", action='store_true', help="Enable training in ORT.") + parser.add_argument( + "--enable_training_ops", action='store_true', help="Enable training ops in inference graph.") parser.add_argument( "--disable_nccl", action='store_true', help="Disable Nccl.") parser.add_argument( @@ -711,6 +713,8 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home "ON" if args.enable_nvtx_profile else "OFF"), "-Donnxruntime_ENABLE_TRAINING=" + ( "ON" if args.enable_training else "OFF"), + "-Donnxruntime_ENABLE_TRAINING_OPS=" + ( + "ON" if args.enable_training_ops else "OFF"), # Enable advanced computations such as AVX for some traininig related ops. "-Donnxruntime_ENABLE_CPU_FP16_OPS=" + ( "ON" if args.enable_training else "OFF"),