diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc index ded60d2373..f7ab646650 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc @@ -230,13 +230,13 @@ Status launch_lamb_compute_direction( p_m2s[i], p_loss_scale, p_g_norm, - CudaT4(alphas[i]), - CudaT4(betas[i]), - CudaT2(lambdas[i]), - CudaT4(epsilons[i]), - CudaT2(max_norms[i]), - CudaT4(alpha_correction), - CudaT4(beta_correction), + alphas[i], + betas[i], + lambdas[i], + epsilons[i], + max_norms[i], + alpha_correction, + beta_correction, p_ds[i], p_m1_news[i], p_m2_news[i], @@ -276,7 +276,7 @@ Status launch_lamb_compute_direction( tensor_sizes_in_buckets[key], buckets[key], lamb_stage1, - p_loss_scale, p_g_norm, lambda, alpha, beta, epsilon, CudaT2(max_norm), alpha_correction, beta_correction); + p_loss_scale, p_g_norm, lambda, alpha, beta, epsilon, max_norm, alpha_correction, beta_correction); } return Status::OK(); diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu index 5ebc2fff49..f2138d82f7 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu @@ -20,41 +20,40 @@ __device__ __forceinline__ void _LambComputeDirectionRule( const T2& g, const T3& m1, const T3& m2, - const T3& alpha, - const T3& beta, - const T1& lambda, - const T3& epsilon, - const T3& alpha_correction, - const T3& beta_correction, + const float& alpha, + const float& beta, + const float& lambda, + const float& epsilon, + const float& alpha_correction, + const float& beta_correction, T2& d, T3& m1_new, T3& m2_new) { // Actual gradient. The scale is a product of loss' scale and // global gradient norm (if the norm > 1). - const T3 g_unscaled = T3(T1(g) / g_scale); + const T1 g_unscaled = T1(g) / g_scale; // A constant in Lamb's equation. - const T3 one = T3(1.0f); + const T1 one = T1(1.0f); // Update exponentially-averaged historical gradient - const T3 m1_new_tmp = alpha * m1 + (one - alpha) * g_unscaled; + const T1 m1_new_tmp = alpha * static_cast(m1) + (one - alpha) * g_unscaled; // Update exponentially-averaged historical squared gradient - const T3 m2_new_tmp = beta * m2 + (one - beta) * g_unscaled * g_unscaled; + const T1 m2_new_tmp = beta * static_cast(m2) + (one - beta) * g_unscaled * g_unscaled; // Compute unbiased 1st-order momentom. // The value alpha_correction is usually (1-alpha^t), // where t is the number of executed training iterations. - const T3 m1_new_tmp_corrected = m1_new_tmp / alpha_correction; + const T1 m1_new_tmp_corrected = m1_new_tmp / alpha_correction; // Compute unbiased 2nd-order momentom. // The value beta_correction is usually (1-beta^t), // where t is the number of executed training iterations. - const T3 m2_new_tmp_corrected = m2_new_tmp / beta_correction; + const T1 m2_new_tmp_corrected = m2_new_tmp / beta_correction; // Save regularized update direction to output. - const T2 d_tmp = lambda * w + - T1(m1_new_tmp_corrected / (_Sqrt(m2_new_tmp_corrected) + epsilon)); + const T1 d_tmp = lambda * w + m1_new_tmp_corrected / (_Sqrt(m2_new_tmp_corrected) + epsilon); // Things are updated only if the direction is finite. if (_IsFiniteScalar(d_tmp)) { @@ -76,13 +75,13 @@ __global__ void _LambComputeDirectionImpl( const T3* moment_2, const T1* loss_scale, const T_GRAD_NORM* g_norm, - T3 alpha, - T3 beta, - T1 lambda, - T3 epsilon, - T1 max_norm, - T3 alpha_correction, - T3 beta_correction, + 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, @@ -117,13 +116,13 @@ void LambComputeDirection( const T3* moment_2, const T1* loss_scale, const T_GRAD_NORM* grad_norm, - T3 alpha, - T3 beta, - T1 lambda, - T3 epsilon, - T1 max_norm, - T3 alpha_correction, - T3 beta_correction, + 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, @@ -160,13 +159,13 @@ void LambComputeDirection( const T3* moment_2, \ const T1* loss_scale, \ const T_GRAD_NORM* grad_norm, \ - T3 alpha, \ - T3 beta, \ - T1 lambda, \ - T3 epsilon, \ - T1 max_norm, \ - T3 alpha_correction, \ - T3 beta_correction, \ + float alpha, \ + float beta, \ + float lambda, \ + float epsilon, \ + float max_norm, \ + float alpha_correction, \ + float beta_correction, \ T2* weights_out, \ T3* moment_1_out, \ T3* moment_2_out, \ @@ -319,13 +318,13 @@ __global__ void LambMultiTensorComputeDirectionImpl( ChunkGroup<6> chunk_group, const T1* loss_scale, const T_GRAD_NORM* g_norm, - const T1 lambda, - const T3 alpha, - const T3 beta, - const T3 epsilon, + const float lambda, + const float alpha, + const float beta, + const float epsilon, const T1 max_norm, - const T3 alpha_correction, - const T3 beta_correction) { + const float alpha_correction, + const float beta_correction) { const int group_index = chunk_group.block_index_to_tensor_group_index[blockIdx.x]; const int tensor_size = chunk_group.tensor_sizes[group_index]; const int chunk_size = chunk_group.chunk_size; @@ -364,13 +363,13 @@ void LambMultiTensorComputeDirectionFunctor::operator() ChunkGroup<6> chunk_group, const T1* loss_scale, const T_GRAD_NORM* g_norm, - const T1 lambda, - const T3 alpha, - const T3 beta, - const T3 epsilon, - const T1 max_norm, - const T3 alpha_correction, - const T3 beta_correction) { + const float lambda, + const float alpha, + const float beta, + const float epsilon, + const float max_norm, + const float alpha_correction, + const float beta_correction) { const int thread_count = ChunkGroup<6>::thread_count_per_block; const int block_count = chunk_group.chunk_count; @@ -393,13 +392,13 @@ void LambMultiTensorComputeDirectionFunctor::operator() ChunkGroup<6> chunk_group, \ const T1* loss_scale, \ const T_GRAD_NORM* g_norm, \ - const T1 lambda, \ - const T3 alpha, \ - const T3 beta, \ - const T3 epsilon, \ - const T1 max_norm, \ - const T3 alpha_correction, \ - const T3 beta_correction); + const float lambda, \ + const float alpha, \ + const float beta, \ + const float epsilon, \ + const float max_norm, \ + const float alpha_correction, \ + const float beta_correction); INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(float, float, float, float) INSTANTIATE_LAMB_STAGE1_MULTI_TENSOR_FUNCTOR(double, double, double, double) diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h index d5bf742a1b..3be720fb8b 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h @@ -56,13 +56,13 @@ void LambComputeDirection( const T3* moment_2, const T1* loss_scale, const T_GRAD_NORM* grad_norm, - T3 alpha, - T3 beta, - T1 lambda, - T3 epsilon, - T1 max_norm, - T3 alpha_correction, - T3 beta_correction, + 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, @@ -112,13 +112,13 @@ struct LambMultiTensorComputeDirectionFunctor { ChunkGroup<6> chunk_group, const T1* loss_scale, const T_GRAD_NORM* grad_norm, - const T1 lambda, - const T3 alpha, - const T3 beta, - const T3 epsilon, - const T1 max_norm, - const T3 alpha_correction, - const T3 beta_correction); + 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 diff --git a/tools/ci_build/github/pai/pai-excluded-tests.txt b/tools/ci_build/github/pai/pai-excluded-tests.txt index bf99a01004..0564f28d22 100644 --- a/tools/ci_build/github/pai/pai-excluded-tests.txt +++ b/tools/ci_build/github/pai/pai-excluded-tests.txt @@ -1,10 +1,3 @@ -OptimizerTest.LambOptimizerTestExternalBaselineDouble -OptimizerTest.LambOptimizerTest5DTensorMixPrecision32_16 -OptimizerTest.LambOptimizerTestSimpleBaselineMixPrecision32_16 -OptimizerTest.LambOptimizerTestBaselineMixPrecision32_16 -OptimizerTest.LambOptimizerTestScalarMixPrecision32_16 -OptimizerTest.LambOptimizerTestScalarMixPrecision32_16_NoDefaultMaxNormClipping -OptimizerTest.LambOptimizerTestLarge CudaKernelTest.SparseSoftmaxCrossEntropy_LargeSizeTensor CudaKernelTest.NegativeLogLikelihoodLoss_TinySizeTensor CudaKernelTest.NegativeLogLikelihoodLoss_SmallSizeTensor