Make training CUDA kernels to adhere established code structure patterns (#10735)

Current training optimizer kernels include CPU headers
  that affects changes that we can make in the CPU code with C++14 compiler and
  other refactoring efforts. Rearrange the kernel according to the established patterns
  and do not include headers that are not needed.

onnxruntime/contrib_ops/cuda/math/isfinite.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "contrib_ops/cuda/math/isfinite.h"
+#include "isfinite_impl.h"
 
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::common;

onnxruntime/contrib_ops/cuda/math/isfinite.cuh

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#pragma once
+
 #include <cuda_fp16.h>
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "contrib_ops/cuda/math/isfinite.h"

onnxruntime/contrib_ops/cuda/math/isfinite.h

@@ -4,7 +4,6 @@
 #pragma once
 #include "core/common/common.h"
 #include "core/providers/cuda/cuda_kernel.h"
-#include "core/providers/cuda/multi_tensor/common.cuh"
 
 namespace onnxruntime {
 namespace cuda {
@@ -31,10 +30,5 @@ class IsAllFiniteOp final : public CudaKernel {
   bool isinf_only_, isnan_only_;
 };
 
-template <typename T>
-struct IsAllFiniteFunctor {
-  void operator()(cudaStream_t stream, ChunkGroup<1> chunks, bool* output, const bool isinf_only, const bool isnan_only);
-};
-
 }  // namespace cuda
 }  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/math/isfinite_impl.cu

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include <cuda_fp16.h>
+#include "isfinite_impl.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "contrib_ops/cuda/math/isfinite.cuh"
 

onnxruntime/contrib_ops/cuda/math/isfinite_impl.h

@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include "core/providers/cuda/multi_tensor/common.cuh"
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T>
+struct IsAllFiniteFunctor {
+  void operator()(cudaStream_t stream, ChunkGroup<1> chunks, bool* output, const bool isinf_only, const bool isnan_only);
+};
+
+}
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/cuda_allocator.h

@@ -3,7 +3,8 @@
 
 #pragma once
 
-#include <unordered_set>
+
+#include "core/common/inlined_containers.h"
 #include "core/framework/allocator.h"
 #include "core/platform/ort_mutex.h"
 
@@ -47,7 +48,7 @@ class CUDAExternalAllocator : public CUDAAllocator {
   ExternalAlloc alloc_;
   ExternalFree free_;
   ExternalEmptyCache empty_cache_;
-  std::unordered_set<void*> reserved_;
+  InlinedHashSet<void*> reserved_;
 };
 
 //TODO: add a default constructor

onnxruntime/core/providers/cuda/multi_tensor/common.cuh

@@ -9,6 +9,8 @@
 #pragma once
 #include <vector>
 
+#include "core/common/common.h"
+
 namespace onnxruntime {
 namespace cuda {
 // initial reference from:
@@ -80,9 +82,9 @@ void launch_multi_tensor_functor(
     TMultiTensorFunctor multipleTensorKernel,
     TFunctorParams&&... kernelParams) {
   ORT_ENFORCE(tensor_sizes.size() > 0);
-  ORT_ENFORCE(tensor_sizes.size() < static_cast<size_t>(std::numeric_limits<int>::max()));
+  ORT_ENFORCE(tensor_sizes.size() < static_cast<size_t>(INT_MAX));
   ORT_ENFORCE(grouped_tensor_pointers.size() > 0);
-  ORT_ENFORCE(grouped_tensor_pointers.size() < static_cast<size_t>(std::numeric_limits<int>::max()));
+  ORT_ENFORCE(grouped_tensor_pointers.size() < static_cast<size_t>(INT_MAX));
   ORT_ENFORCE(chunk_size > 0);
   // Number of groups, for example, the number of updated weight tensors in Lamb optimizer.
   const int group_count = static_cast<int>(grouped_tensor_pointers.size());
@@ -92,7 +94,7 @@ void launch_multi_tensor_functor(
   int block_index = 0;
 
   // Check if 32-bit integer is enough.
-  ORT_ENFORCE(tensor_sizes.size() < static_cast<size_t>(std::numeric_limits<int>::max()));
+  ORT_ENFORCE(tensor_sizes.size() < static_cast<size_t>(INT_MAX));
   ORT_ENFORCE(grouped_tensor_pointers.size() == tensor_sizes.size());
   ORT_ENFORCE(group_size == ACTUAL_TENSOR_GROUP_SIZE[TensorGroupSize]);
   for (int i = 0; i < group_count; ++i) {

orttraining/orttraining/training_ops/cuda/math/isfinite.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "orttraining/training_ops/cuda/math/isfinite.h"
+#include "orttraining/training_ops/cuda/math/isfinite_impl.h"
 
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::common;

orttraining/orttraining/training_ops/cuda/math/isfinite.h

@@ -2,9 +2,9 @@
 // Licensed under the MIT License.
 
 #pragma once
+
 #include "core/common/common.h"
 #include "core/providers/cuda/cuda_kernel.h"
-#include "core/providers/cuda/multi_tensor/common.cuh"
 
 namespace onnxruntime {
 namespace cuda {
@@ -18,8 +18,5 @@ class IsFiniteOp final : public CudaKernel {
   Status ComputeInternal(OpKernelContext* context) const override;
 };
 
-template <typename TSrc>
-void IsFinite(cudaStream_t stream, const TSrc* input, bool* output, size_t N);
-
 }  // namespace cuda
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/math/isfinite_impl.cu

@@ -1,10 +1,12 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "isfinite_impl.h"
 #include <cuda_fp16.h>
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "contrib_ops/cuda/math/isfinite.cuh"
 
+
 namespace onnxruntime {
 namespace cuda {
 

orttraining/orttraining/training_ops/cuda/math/isfinite_impl.h

@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename TSrc>
+void IsFinite(cudaStream_t stream, const TSrc* input, bool* output, size_t N);
+
+}
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "mixed_precision_scale.h"
+#include "mixed_precision_scale_impl.h"
 
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::common;

orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale.h

@@ -8,14 +8,6 @@
 namespace onnxruntime {
 namespace cuda {
 
-template <typename SrcT, typename DstT>
-void Impl_MixedPrecisionScale(
-    cudaStream_t stream,
-    const SrcT* input_data,
-    const float* scale_data,
-    DstT* output_data,
-    size_t count);
-
 template <typename SrcT>
 class MixedPrecisionScale final : public CudaKernel {
  public:

orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale_impl.cu

@@ -5,9 +5,9 @@
 #pragma warning(disable : 4244)
 #endif
 
+#include "mixed_precision_scale_impl.h"
 #include <cuda_fp16.h>
 #include "core/providers/cuda/cu_inc/common.cuh"
-#include "mixed_precision_scale.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/math/mixed_precision_scale_impl.h

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+template <typename SrcT, typename DstT>
+void Impl_MixedPrecisionScale(
+    cudaStream_t stream,
+    const SrcT* input_data,
+    const float* scale_data,
+    DstT* output_data,
+    size_t count);
+}
+}  // namespace onnxruntime
+

orttraining/orttraining/training_ops/cuda/math/scale.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "orttraining/training_ops/cuda/math/scale.h"
+#include "orttraining/training_ops/cuda/math/scale_impl.h"
 
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::common;
@@ -25,6 +26,15 @@ namespace cuda {
           .InputMemoryType(OrtMemTypeCPUInput, 1),                             \
       Scale<T>);
 
+template <typename ScaleT>
+struct GetScaleValueImpl {
+  void operator()(const Tensor* scale, float& scale_value) const {
+    ORT_ENFORCE(scale->Shape().Size() == 1, "Scale input should have a single value.");
+    scale_value = static_cast<float>(*(scale->template Data<ScaleT>()));
+    ORT_ENFORCE(scale_value != 0.0f, "Scale value must not be 0.");
+  }
+};
+
 template <typename T>
 Scale<T>::Scale(const OpKernelInfo& info) : CudaKernel(info) {
   int64_t scale_down;

orttraining/orttraining/training_ops/cuda/math/scale.h

@@ -7,23 +7,6 @@
 namespace onnxruntime {
 namespace cuda {
 
-template <typename ScaleT>
-struct GetScaleValueImpl {
-  void operator()(const Tensor* scale, float& scale_value) const {
-    ORT_ENFORCE(scale->Shape().Size() == 1, "Scale input should have a single value.");
-    scale_value = static_cast<float>(*(scale->template Data<ScaleT>()));
-    ORT_ENFORCE(scale_value != 0.0f, "Scale value must not be 0.");
-  }
-};
-
-template <typename T>
-void Impl_Scale(
-    cudaStream_t stream,
-    const T* input_data,
-    const float scale_value,
-    T* output_data,
-    size_t count);
-
 template <typename T>
 class Scale final : public CudaKernel {
  public:

orttraining/orttraining/training_ops/cuda/math/scale_impl.cu

@@ -1,8 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "orttraining/training_ops/cuda/math/scale_impl.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
-#include "orttraining/training_ops/cuda/math/scale.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/math/scale_impl.h

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T>
+void Impl_Scale(
+    cudaStream_t stream,
+    const T* input_data,
+    const float scale_value,
+    T* output_data,
+    size_t count);
+ }
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/math/softmax_grad.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "orttraining/training_ops/cuda/math/softmax_grad.h"
+#include "orttraining/training_ops/cuda/math/softmax_grad_impl.h"
 
 #include "core/providers/common.h"
 #include "core/providers/cuda/cudnn_common.h"

orttraining/orttraining/training_ops/cuda/math/softmax_grad.h

@@ -8,9 +8,6 @@
 namespace onnxruntime {
 namespace cuda {
 
-template <typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
-void dispatch_softmax_backward(cudaStream_t stream, output_t* grad_input, const input_t* grad, const input_t* output, int softmax_elements, int softmax_elements_stride, int batch_count);
-
 template <typename T>
 class SoftmaxGrad final : public CudaKernel {
  public:

orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu

@@ -18,7 +18,7 @@
 
 // The code below is mostly copied from Pytorch PersistentSoftmax.cuh
 
-#include "orttraining/training_ops/cuda/math/softmax_grad.h"
+#include "orttraining/training_ops/cuda/math/softmax_grad_impl.h"
 
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "core/providers/cuda/math/softmax_warpwise_impl.cuh"

orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.h

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+template <typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+void dispatch_softmax_backward(cudaStream_t stream, output_t* grad_input, const input_t* grad, const input_t* output,
+                               int softmax_elements, int softmax_elements_stride, int batch_count);
+}
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/optimizer/adam.cc

@@ -1,11 +1,13 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "orttraining/training_ops/cuda/optimizer/adam.h"
+#include "orttraining/training_ops/cuda/optimizer/adam_impl.h"
+
 #include "core/providers/cuda/cuda_allocator.h"
 #include "core/providers/cuda/reduction/reduction_functions.h"
 #include "core/providers/cuda/math/binary_elementwise_ops.h"
 #include "orttraining/training_ops/cuda/optimizer/common.h"
-#include "orttraining/training_ops/cuda/optimizer/adam.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/adam.h

@@ -7,32 +7,6 @@
 
 namespace onnxruntime {
 namespace cuda {
-
-template <typename T1, typename T2, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM, typename T_MIXED_PRECISION_FP>
-void AdamOptimizerImpl(
-    cudaStream_t stream,
-    const T1* eta,
-    const T2 update_count,
-    const T3* weights,
-    const T_GRAD* grads,
-    const T4* moment_1,
-    const T4* moment_2,
-    const T3* loss_scale,
-    const T_GRAD_NORM* grad_norm,
-    const float alpha,
-    const float beta,
-    const float lambda,
-    const float epsilon,
-    const float max_norm,
-    const bool do_bias_correction,
-    const int64_t weight_decay_mode,
-    T4* moment_1_out,
-    T4* moment_2_out,
-    T3* weights_out,
-    T_GRAD* grads_out,
-    T_MIXED_PRECISION_FP* mixed_precision_weights_out,
-    size_t count);
-
 template <typename T1, typename T2, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM, typename T_MIXED_PRECISION_FP>
 class AdamOptimizer final : public CudaKernel {
  public:

orttraining/orttraining/training_ops/cuda/optimizer/adam_impl.cu

@@ -1,10 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "orttraining/training_ops/cuda/optimizer/adam_impl.h"
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "orttraining/training_ops/cuda/optimizer/common.cuh"
-#include "orttraining/training_ops/cuda/optimizer/adam.h"
 #include "orttraining/training_ops/cuda/optimizer/common.h"
 
 namespace onnxruntime {

orttraining/orttraining/training_ops/cuda/optimizer/adam_impl.h

@@ -0,0 +1,37 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cstdint>
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T1, typename T2, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM, typename T_MIXED_PRECISION_FP>
+void AdamOptimizerImpl(
+    cudaStream_t stream,
+    const T1* eta,
+    const T2 update_count,
+    const T3* weights,
+    const T_GRAD* grads,
+    const T4* moment_1,
+    const T4* moment_2,
+    const T3* loss_scale,
+    const T_GRAD_NORM* grad_norm,
+    const float alpha,
+    const float beta,
+    const float lambda,
+    const float epsilon,
+    const float max_norm,
+    const bool do_bias_correction,
+    const int64_t weight_decay_mode,
+    T4* moment_1_out,
+    T4* moment_2_out,
+    T3* weights_out,
+    T_GRAD* grads_out,
+    T_MIXED_PRECISION_FP* mixed_precision_weights_out,
+    size_t count);
+}
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc

@@ -1,11 +1,13 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "gradient_control.h"
+#include "gradient_control_impl.h"
+
 #include "core/providers/cuda/math/binary_elementwise_ops.h"
 #include "core/providers/cuda/reduction/reduction_functions.h"
 #include "core/providers/cuda/cuda_allocator.h"
 #include "common.h"
-#include "gradient_control.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.h

@@ -22,14 +22,5 @@ class InPlaceAccumulator final : public CudaKernel {
   Status ComputeInternal(OpKernelContext* context) const override;
 };
 
-// Implementation can be found in cuda file, optimizers_impl.cu
-template <typename T, typename T_GRAD>
-void InPlaceAccumulatorImpl(
-    cudaStream_t stream,
-    const T* gradient_buffer,
-    const T_GRAD* gradient,
-    T* accumulated_gradient,
-    size_t count);
-
 }  // namespace cuda
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/optimizer/gradient_control_impl.cu

@@ -2,10 +2,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "gradient_control_impl.h"
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "core/providers/cuda/atomic/common.cuh"
-#include "gradient_control.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/gradient_control_impl.h

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+// Implementation can be found in cuda file
+template <typename T, typename T_GRAD>
+void InPlaceAccumulatorImpl(
+    cudaStream_t stream,
+    const T* gradient_buffer,
+    const T_GRAD* gradient,
+    T* accumulated_gradient,
+    size_t count);
+}
+}  // namespace onnxruntime
+

orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc

@@ -1,12 +1,15 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#include <cmath>
-#include "core/providers/cuda/cuda_allocator.h"
+
+#include "orttraining/training_ops/cuda/optimizer/lamb.h"
+#include "orttraining/training_ops/cuda/optimizer/lamb_impl.h"
+
 #include "core/providers/cuda/reduction/reduction_functions.h"
 #include "core/providers/cuda/math/binary_elementwise_ops.h"
 #include "orttraining/training_ops/cuda/optimizer/common.h"
-#include "orttraining/training_ops/cuda/optimizer/lamb.h"
+
+#include <cmath>
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/lamb.h

@@ -4,7 +4,6 @@
 #pragma once
 #include "core/common/common.h"
 #include "core/providers/cuda/cuda_kernel.h"
-#include "core/providers/cuda/multi_tensor/common.cuh"
 
 namespace onnxruntime {
 namespace cuda {
@@ -43,156 +42,5 @@ class LambOptimizer final : public CudaKernel {
   bool do_bias_correction_;
 };
 
-// Implementation can be found in cuda file, optimizers_impl.cu
-// T1's precision should be higher than T2. It's used for
-// large tensors. Small tensors should use multi-tensor version
-// of this.
-template <typename T1, typename T2, typename T3, typename T_GRAD_NORM>
-void LambComputeDirection(
-    cudaStream_t stream,
-    const T1* weights,
-    const T2* grads,
-    const T3* moment_1,
-    const T3* moment_2,
-    const T1* loss_scale,
-    const T_GRAD_NORM* grad_norm,
-    float alpha,
-    float beta,
-    float lambda,
-    float epsilon,
-    float max_norm,
-    float alpha_correction,
-    float beta_correction,
-    T2* update_direction,
-    T3* moment_1_out,
-    T3* moment_2_out,
-    size_t count);
-
-// Implementation can be found in cuda file, optimizers_impl.cu
-// T2's precision should be higher than T1. It's used for
-// large tensors. Small tensors should use multi-tensor version
-// of this.
-template <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
-void LambUpdate(
-    cudaStream_t stream,
-    const T1* eta,
-    const float ratio_min,
-    const float ratio_max,
-    const T2* r_norm,
-    const T2* w_norm,
-    const T2* weights,
-    const T3* update_direction,
-    T2* weights_out,
-    T3* gradients_out,
-    T_MIXED_PRECISION_FP* mixed_precision_weights_out,
-    size_t count);
-
-// Lamb's stage 1 maps [w, g, m1, m2] to [d, m1_new, m2_new] where
-//  w: weight tensor
-//  g: gradient (reused to store update direction)
-//  m1: 1st momentum
-//  m2: 2nd momentum
-//  d: update direction
-//  m1_new: updated 1st momentum
-//  m2_new: updated 2nd momentum
-// Because we reuse g to store d, there are only 6 tensors in total and
-// therefore the type of chunk_group is ChunkGroup<6>.
-//
-// Tensor pointers associated with the i-th tensor in this chunk:
-//  w: chunk_group.tensor_ptrs[0][i]
-//  g (or d): chunk_group.tensor_ptrs[1][i]
-//  m1: chunk_group.tensor_ptrs[2][i]
-//  m2: chunk_group.tensor_ptrs[3][i]
-//  m1_new: chunk_group.tensor_ptrs[4][i]
-//  m2_new: chunk_group.tensor_ptrs[5][i]
-template <typename T1, typename T2, typename T3, typename T_GRAD_NORM>
-struct LambMultiTensorComputeDirectionFunctor {
-  void operator()(
-      cudaStream_t stream,
-      ChunkGroup<6> chunk_group,
-      const T1* loss_scale,
-      const T_GRAD_NORM* grad_norm,
-      const float lambda,
-      const float alpha,
-      const float beta,
-      const float epsilon,
-      const float max_norm,
-      const float alpha_correction,
-      const float beta_correction);
-};
-
-// Lamb's reduction maps [w, d] to [w_norm, d_norm] where
-//  w: weight tensor
-//  d: update direction
-//  w_norm: norm of w
-//  d_norm: norm of d
-// There are 4 distinct tensors in total and therefore the
-// type of chunk_group is ChunkGroup<4>.
-//
-// Tensor pointers associated with the i-th tensor in this chunk:
-//  w: chunk_group.tensor_ptrs[0][i]
-//  d: chunk_group.tensor_ptrs[1][i]
-//  w_norm: chunk_group.tensor_ptrs[2][i]
-//  d_norm: chunk_group.tensor_ptrs[3][i]
-template <typename TIn1, typename TIn2, typename TOut1, typename TOut2, typename TBuf>
-struct LambMultiTensorReductionFunctor {
-  void operator()(
-      cudaStream_t stream,
-      ChunkGroup<4> chunk_group,
-      const CudaKernel& kernel,
-      void* reduction_buffer,
-      size_t reduction_buffer_size);
-};
-
-// Lamb's reduction mapping [w, d] to [w_norm, d_norm] spans multiples thread blocks
-//
-// This includes any block-index for which it holds
-//    i-th tensor-index == chunk_group.block_index_to_tensor_group_index[ block-index ]
-// and where i-th tensor-index corresponds to tensor group w(i), d(i), w_norm(i), d_norm(i)
-// (see above)
-//
-// The above span of blocks corresponding i-th tensor will be contiguous.
-// To perform an ORDERED reduction across the thread blocks for i-th tensor,
-//   the following struct is passed for every tensor.
-// It consists of fields:
-//   'leading_block' := lowest block-index corresponding i-th tensor
-//   'number_blocks' := number block-index "
-//   'completed_blocks' := initialized to zero (for internal use)
-// Note 'completed_blocks' prevents inter-block reduction until intra-block reduction is complete.
-struct LambMultiTensorSyncRangeAndLock {
-  int leading_block;
-  int number_blocks;
-  int completed_blocks;
-};
-
-// Lamb's stage 2 maps [w_norm, w_norm, w, d] to [w_new, g_new, w_mixed_precision_new] where
-//  w_norm: norm of w
-//  d_norm: norm of d
-//  w: weight tensor
-//  d: update direction
-//  w_new: updated weight tensor
-//  g_new: updated gradient tensor
-//  w_mixed_precision_new: updated weight tensor of mixed-precision type
-// There are 7 distinct tensors in total and therefore the
-// type of chunk_group is ChunkGroup<7>.
-//
-// Tensor pointers associated with the i-th tensor in this chunk:
-//  w_norm: chunk_group.tensor_ptrs[0][i]
-//  d_norm: chunk_group.tensor_ptrs[1][i]
-//  w: chunk_group.tensor_ptrs[2][i]
-//  d: chunk_group.tensor_ptrs[3][i]
-//  w_new: chunk_group.tensor_ptrs[4][i]
-//  g_new: chunk_group.tensor_ptrs[5][i]
-//  w_mixed_precision_new: chunk_group.tensor_ptrs[6][i]
-template <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
-struct LambMultiTensorUpdateFunctor {
-  void operator()(
-      cudaStream_t stream,
-      ChunkGroup<7> chunk_group,
-      const T1* eta,
-      const float ratio_min,
-      const float ratio_max);
-};
-
 }  // namespace cuda
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/optimizer/lamb_impl.cu

@@ -1,15 +1,13 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "orttraining/training_ops/cuda/optimizer/lamb_impl.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
-#include "core/providers/cuda/cuda_allocator.h"
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/atomic/common.cuh"
 #include "core/providers/cuda/reduction/reduction_utils.cuh"
 #include "contrib_ops/cuda/math/isfinite.cuh"
-#include "orttraining/training_ops/cuda/optimizer/common.h"
 #include "orttraining/training_ops/cuda/optimizer/common.cuh"
-#include "orttraining/training_ops/cuda/optimizer/lamb.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/lamb_impl.h

@@ -0,0 +1,164 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include "core/providers/cuda/cuda_kernel.h"
+#include "core/providers/cuda/multi_tensor/common.cuh"
+
+namespace onnxruntime {
+namespace cuda {
+// Implementation can be found in cuda file.
+// T1's precision should be higher than T2. It's used for
+// large tensors. Small tensors should use multi-tensor version
+// of this.
+template <typename T1, typename T2, typename T3, typename T_GRAD_NORM>
+void LambComputeDirection(
+    cudaStream_t stream,
+    const T1* weights,
+    const T2* grads,
+    const T3* moment_1,
+    const T3* moment_2,
+    const T1* loss_scale,
+    const T_GRAD_NORM* grad_norm,
+    float alpha,
+    float beta,
+    float lambda,
+    float epsilon,
+    float max_norm,
+    float alpha_correction,
+    float beta_correction,
+    T2* update_direction,
+    T3* moment_1_out,
+    T3* moment_2_out,
+    size_t count);
+
+// Implementation can be found in cuda file.
+// T2's precision should be higher than T1. It's used for
+// large tensors. Small tensors should use multi-tensor version
+// of this.
+template <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
+void LambUpdate(
+    cudaStream_t stream,
+    const T1* eta,
+    const float ratio_min,
+    const float ratio_max,
+    const T2* r_norm,
+    const T2* w_norm,
+    const T2* weights,
+    const T3* update_direction,
+    T2* weights_out,
+    T3* gradients_out,
+    T_MIXED_PRECISION_FP* mixed_precision_weights_out,
+    size_t count);
+
+// Lamb's stage 1 maps [w, g, m1, m2] to [d, m1_new, m2_new] where
+//  w: weight tensor
+//  g: gradient (reused to store update direction)
+//  m1: 1st momentum
+//  m2: 2nd momentum
+//  d: update direction
+//  m1_new: updated 1st momentum
+//  m2_new: updated 2nd momentum
+// Because we reuse g to store d, there are only 6 tensors in total and
+// therefore the type of chunk_group is ChunkGroup<6>.
+//
+// Tensor pointers associated with the i-th tensor in this chunk:
+//  w: chunk_group.tensor_ptrs[0][i]
+//  g (or d): chunk_group.tensor_ptrs[1][i]
+//  m1: chunk_group.tensor_ptrs[2][i]
+//  m2: chunk_group.tensor_ptrs[3][i]
+//  m1_new: chunk_group.tensor_ptrs[4][i]
+//  m2_new: chunk_group.tensor_ptrs[5][i]
+template <typename T1, typename T2, typename T3, typename T_GRAD_NORM>
+struct LambMultiTensorComputeDirectionFunctor {
+  void operator()(
+      cudaStream_t stream,
+      ChunkGroup<6> chunk_group,
+      const T1* loss_scale,
+      const T_GRAD_NORM* grad_norm,
+      const float lambda,
+      const float alpha,
+      const float beta,
+      const float epsilon,
+      const float max_norm,
+      const float alpha_correction,
+      const float beta_correction);
+};
+
+// Lamb's reduction maps [w, d] to [w_norm, d_norm] where
+//  w: weight tensor
+//  d: update direction
+//  w_norm: norm of w
+//  d_norm: norm of d
+// There are 4 distinct tensors in total and therefore the
+// type of chunk_group is ChunkGroup<4>.
+//
+// Tensor pointers associated with the i-th tensor in this chunk:
+//  w: chunk_group.tensor_ptrs[0][i]
+//  d: chunk_group.tensor_ptrs[1][i]
+//  w_norm: chunk_group.tensor_ptrs[2][i]
+//  d_norm: chunk_group.tensor_ptrs[3][i]
+template <typename TIn1, typename TIn2, typename TOut1, typename TOut2, typename TBuf>
+struct LambMultiTensorReductionFunctor {
+  void operator()(
+      cudaStream_t stream,
+      ChunkGroup<4> chunk_group,
+      const CudaKernel& kernel,
+      void* reduction_buffer,
+      size_t reduction_buffer_size);
+};
+
+// Lamb's reduction mapping [w, d] to [w_norm, d_norm] spans multiples thread blocks
+//
+// This includes any block-index for which it holds
+//    i-th tensor-index == chunk_group.block_index_to_tensor_group_index[ block-index ]
+// and where i-th tensor-index corresponds to tensor group w(i), d(i), w_norm(i), d_norm(i)
+// (see above)
+//
+// The above span of blocks corresponding i-th tensor will be contiguous.
+// To perform an ORDERED reduction across the thread blocks for i-th tensor,
+//   the following struct is passed for every tensor.
+// It consists of fields:
+//   'leading_block' := lowest block-index corresponding i-th tensor
+//   'number_blocks' := number block-index "
+//   'completed_blocks' := initialized to zero (for internal use)
+// Note 'completed_blocks' prevents inter-block reduction until intra-block reduction is complete.
+struct LambMultiTensorSyncRangeAndLock {
+  int leading_block;
+  int number_blocks;
+  int completed_blocks;
+};
+
+// Lamb's stage 2 maps [w_norm, w_norm, w, d] to [w_new, g_new, w_mixed_precision_new] where
+//  w_norm: norm of w
+//  d_norm: norm of d
+//  w: weight tensor
+//  d: update direction
+//  w_new: updated weight tensor
+//  g_new: updated gradient tensor
+//  w_mixed_precision_new: updated weight tensor of mixed-precision type
+// There are 7 distinct tensors in total and therefore the
+// type of chunk_group is ChunkGroup<7>.
+//
+// Tensor pointers associated with the i-th tensor in this chunk:
+//  w_norm: chunk_group.tensor_ptrs[0][i]
+//  d_norm: chunk_group.tensor_ptrs[1][i]
+//  w: chunk_group.tensor_ptrs[2][i]
+//  d: chunk_group.tensor_ptrs[3][i]
+//  w_new: chunk_group.tensor_ptrs[4][i]
+//  g_new: chunk_group.tensor_ptrs[5][i]
+//  w_mixed_precision_new: chunk_group.tensor_ptrs[6][i]
+template <typename T1, typename T2, typename T3, typename T_MIXED_PRECISION_FP>
+struct LambMultiTensorUpdateFunctor {
+  void operator()(
+      cudaStream_t stream,
+      ChunkGroup<7> chunk_group,
+      const T1* eta,
+      const float ratio_min,
+      const float ratio_max);
+};
+
+}
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/optimizer/sg.cc

@@ -1,10 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#include "core/providers/cuda/cuda_allocator.h"
+#include "sg.h"
+#include "sg_impl.h"
+
 #include "core/providers/cuda/reduction/reduction_functions.h"
 #include "core/providers/cuda/math/binary_elementwise_ops.h"
-#include "sg.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/sg.h

@@ -8,16 +8,6 @@
 namespace onnxruntime {
 namespace cuda {
 
-template <typename T>
-void SGDOptimizerImpl(
-    cudaStream_t stream,
-    const T* eta,
-    const T* weights,
-    const T* gradients,
-    T* weight_out,
-    T* gradients_out,
-    size_t count);
-
 class SGDOptimizer final : public CudaKernel {
  public:
   SGDOptimizer(const OpKernelInfo& info) : CudaKernel(info) {}

orttraining/orttraining/training_ops/cuda/optimizer/sg_impl.cu

@@ -1,10 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "sg_impl.h"
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "core/providers/cuda/atomic/common.cuh"
-#include "sg.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/sg_impl.h

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+template <typename T>
+void SGDOptimizerImpl(
+    cudaStream_t stream,
+    const T* eta,
+    const T* weights,
+    const T* gradients,
+    T* weight_out,
+    T* gradients_out,
+    size_t count);
+}
+}  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/reduction/all.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "orttraining/training_ops/cuda/reduction/all.h"
+#include "orttraining/training_ops/cuda/reduction/all_impl.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/reduction/all.h

@@ -15,8 +15,5 @@ class All final : public CudaKernel {
   Status ComputeInternal(OpKernelContext* context) const override;
 };
 
-template<typename T>
-void LaunchAllKernel(cudaStream_t stream, const T* data, const int size, bool* output);
-
 }  // namespace cuda
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/reduction/all_impl.cu

@@ -1,11 +1,12 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#include "orttraining/training_ops/cuda/reduction/all.h"
+#include "orttraining/training_ops/cuda/reduction/all_impl.h"
 
 #include <thrust/logical.h>
 #include <thrust/functional.h>
 #include <thrust/execution_policy.h>
+
 #ifdef _WIN32
 #pragma warning(disable : 4244)
 #endif

orttraining/orttraining/training_ops/cuda/reduction/all_impl.h

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace onnxruntime {
+namespace cuda {
+template <typename T>
+void LaunchAllKernel(cudaStream_t stream, const T* data, const int size, bool* output);
+}
+}  // namespace onnxruntime
+

orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "orttraining/training_ops/cuda/reduction/reduction_all.h"
+#include "orttraining/training_ops/cuda/reduction/reduction_all_impl.h"
 
 #include "core/providers/cuda/reduction/reduction_functions.h"
 #include "core/providers/cuda/shared_inc/accumulation_type.h"

orttraining/orttraining/training_ops/cuda/reduction/reduction_all.h

@@ -3,7 +3,6 @@
 
 #pragma once
 #include "core/providers/cuda/cuda_kernel.h"
-#include "core/providers/cuda/multi_tensor/common.cuh"
 
 namespace onnxruntime {
 namespace cuda {
@@ -16,13 +15,5 @@ class ReduceAllL2 final : public CudaKernel {
   Status ComputeInternal(OpKernelContext* context) const override;
 };
 
-template <typename TIn, typename TOut>
-struct MultiTensorReduceL2 {
-  void operator()(cudaStream_t stream, ChunkGroup<1> chunk_group, TOut* output);
-};
-
-template<typename Tin, typename Tout>
-void ScalarSqrt(cudaStream_t stream, Tin* input, Tout* output);
-
 }  // namespace cuda
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/reduction/reduction_all_impl.cu

@@ -1,10 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#include "orttraining/training_ops/cuda/reduction/reduction_all.h"
+#include "reduction_all_impl.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/atomic/common.cuh"
+#include "core/providers/cuda/multi_tensor/common.cuh"
 #include "core/providers/cuda/reduction/reduction_utils.cuh"
 #include "core/providers/cuda/shared_inc/accumulation_type.h"
 

orttraining/orttraining/training_ops/cuda/reduction/reduction_all_impl.h

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include "core/providers/cuda/multi_tensor/common.cuh"
+
+namespace onnxruntime {
+namespace cuda {
+template <typename TIn, typename TOut>
+struct MultiTensorReduceL2 {
+  void operator()(cudaStream_t stream, ChunkGroup<1> chunk_group, TOut* output);
+};
+
+template <typename Tin, typename Tout>
+void ScalarSqrt(cudaStream_t stream, Tin* input, Tout* output);
+}
+}  // namespace onnxruntime
+

orttraining/orttraining/training_ops/cuda/tensor/gather_nd_grad_impl.h

@@ -2,7 +2,9 @@
 // Licensed under the MIT License.
 
 #pragma once
-#include "core/providers/cuda/shared_inc/cuda_utils.h"
+
+#include <stdint.h>
+#include <cuda_runtime.h>
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/rocm/math/softmax_grad.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "orttraining/training_ops/rocm/math/softmax_grad.h"
+#include "orttraining/training_ops/rocm/math/softmax_grad_impl.h"
 
 #include "core/providers/common.h"
 #include "core/providers/rocm/miopen_common.h"

orttraining/orttraining/training_ops/rocm/math/softmax_grad_impl.cu

@@ -19,7 +19,7 @@
 // The code below is mostly copied from Pytorch PersistentSoftmax.cuh
 #include "hip/hip_runtime.h"
 
-#include "orttraining/training_ops/rocm/math/softmax_grad.h"
+#include "orttraining/training_ops/rocm/math/softmax_grad_impl.h"
 
 #include "core/providers/rocm/cu_inc/common.cuh"
 #include "core/providers/rocm/math/softmax_warpwise_impl.cuh"

orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "orttraining/training_ops/rocm/reduction/reduction_all.h"
+#include "orttraining/training_ops/rocm/reduction/reduction_all_impl.h"
 
 #include "core/providers/rocm/reduction/reduction_functions.h"
 #include "core/providers/rocm/shared_inc/accumulation_type.h"