diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.cc index 475bb0ee7f8..1cf6363cfc9 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.cc @@ -24,8 +24,8 @@ namespace { static constexpr size_t kSizeFloat32 = sizeof(float); static constexpr size_t kSizeFloat16 = sizeof(float16); static constexpr size_t kScalarIndex = 0; -static constexpr size_t kFusedAdaFactorInputNum = 12; -static constexpr size_t kFusedAdaFactorWorkSpaceNum = 3; +static constexpr size_t kStandardInputNum = 12; +static constexpr size_t kWorkSpaceNum = 3; static constexpr size_t kBatchSize = 10000; static auto constexpr kEnableScaleParameter = "enable_scale_parameter"; static auto constexpr kEnableFirstMoment = "enable_first_moment"; @@ -37,15 +37,9 @@ static constexpr float kEps = 1e-30; void FusedAdaFactorCPUKernel::InitInputOutputSize(const CNodePtr &kernel_node) { CPUKernel::InitInputOutputSize(kernel_node); - if (param_dtype_ == kNumberTypeFloat16) { - (void)workspace_size_list_.emplace_back(elem_num_ * kSizeFloat16); - (void)workspace_size_list_.emplace_back(elem_num_ / last_row_dim_size_ * kSizeFloat16); - (void)workspace_size_list_.emplace_back(elem_num_ / last_col_dim_size_ * kSizeFloat16); - } else { - (void)workspace_size_list_.emplace_back(elem_num_ * kSizeFloat32); - (void)workspace_size_list_.emplace_back(elem_num_ / last_row_dim_size_ * kSizeFloat32); - (void)workspace_size_list_.emplace_back(elem_num_ / last_col_dim_size_ * kSizeFloat32); - } + (void)workspace_size_list_.emplace_back(elem_num_ * kSizeFloat32); + (void)workspace_size_list_.emplace_back(elem_num_ / last_row_dim_size_ * kSizeFloat32); + (void)workspace_size_list_.emplace_back(elem_num_ / last_col_dim_size_ * kSizeFloat32); } void FusedAdaFactorCPUKernel::InitKernel(const CNodePtr &kernel_node) { @@ -93,14 +87,14 @@ float FusedAdaFactorCPUKernel::CalcRMS(T *input, size_t elem_num) { } template -void FusedAdaFactorCPUKernel::FactorUpdate(T *update, const std::vector &inputs, +void FusedAdaFactorCPUKernel::FactorUpdate(float *update, const std::vector &inputs, const std::vector &workspaces) { auto beta2t = reinterpret_cast(inputs[BETA2T]->addr)[kScalarIndex]; auto grad = reinterpret_cast(inputs[GRAD]->addr); auto exp_avg_sq_row = reinterpret_cast(inputs[EXP_AVG_SQ_ROW]->addr); auto exp_avg_sq_col = reinterpret_cast(inputs[EXP_AVG_SQ_COL]->addr); - auto r_factor = reinterpret_cast(workspaces[R_FACTOR]->addr); - auto c_factor = reinterpret_cast(workspaces[C_FACTOR]->addr); + auto r_factor = reinterpret_cast(workspaces[R_FACTOR]->addr); + auto c_factor = reinterpret_cast(workspaces[C_FACTOR]->addr); auto one_minus_beta2t = 1 - beta2t; std::function task; @@ -115,7 +109,7 @@ void FusedAdaFactorCPUKernel::FactorUpdate(T *update, const std::vector(update[reduce_start + j]); + row_reduce += update[reduce_start + j]; } row_reduce = row_reduce / row_dim_size; auto tmp = static_cast(exp_avg_sq_row[i]) * beta2t + row_reduce * one_minus_beta2t; @@ -135,8 +129,7 @@ void FusedAdaFactorCPUKernel::FactorUpdate(T *update, const std::vector(exp_avg_sq_row[reduce_start + j]) / col_reduce); - r_factor[reduce_start + j] = static_cast(tmp); + r_factor[reduce_start + j] = std::sqrt(static_cast(exp_avg_sq_row[reduce_start + j]) / col_reduce); } } }; @@ -149,13 +142,13 @@ void FusedAdaFactorCPUKernel::FactorUpdate(T *update, const std::vector(update[reduce_start + j * row_dim_size]); + row_reduce += update[reduce_start + j * row_dim_size]; } row_reduce = row_reduce / col_dim_size; auto tmp = static_cast(exp_avg_sq_col[i]) * beta2t + row_reduce * one_minus_beta2t; tmp = std::max(tmp, kEps); exp_avg_sq_col[i] = static_cast(tmp); - c_factor[i] = static_cast(std::sqrt(tmp)); + c_factor[i] = std::sqrt(tmp); } }; CPUKernelUtils::ParallelFor(task, exp_avg_sq_col_elem_num, kBatchSize); @@ -166,12 +159,8 @@ void FusedAdaFactorCPUKernel::FactorUpdate(T *update, const std::vector(r_factor[slice * col_dim_size + col_i]); - auto right = static_cast(c_factor[slice * row_dim_size + row_i]); - auto norm = left * right; - norm = std::max(norm, kEps); - auto tmp = static_cast(grad[i]) / norm; - update[i] = static_cast(tmp); + auto norm = r_factor[slice * col_dim_size + col_i] * c_factor[slice * row_dim_size + row_i]; + update[i] = static_cast(grad[i]) * global_norm_reciprocal_ / std::max(norm, kEps); } }; CPUKernelUtils::ParallelFor(task, elem_num_, kBatchSize); @@ -190,7 +179,7 @@ void FusedAdaFactorCPUKernel::LaunchKernel(const std::vector &inputs auto param = reinterpret_cast(inputs[PARAM]->addr); auto exp_avg = reinterpret_cast(inputs[EXP_AVG]->addr); auto exp_avg_sq = reinterpret_cast(inputs[EXP_AVG_SQ]->addr); - auto update = reinterpret_cast(workspaces[UPDATE]->addr); + auto update = reinterpret_cast(workspaces[UPDATE]->addr); auto one_minus_beta1 = 1 - beta1; auto one_minus_beta2t = 1 - beta2t; if (clip_threshold <= 0) { @@ -211,23 +200,22 @@ void FusedAdaFactorCPUKernel::LaunchKernel(const std::vector &inputs // update = grad * grad + eps[0] task = [&](size_t start, size_t end) { for (size_t i = start; i < end; ++i) { - auto tmp = static_cast(grad[i]); - update[i] = static_cast(tmp * tmp + epsilon[0]); + auto tmp = static_cast(grad[i]) * global_norm_reciprocal_; + update[i] = tmp * tmp + epsilon[0]; } }; CPUKernelUtils::ParallelFor(task, elem_num_, kBatchSize); if (need_factor_) { - FactorUpdate(update, inputs, workspaces); + FactorUpdate(update, inputs, workspaces); } else { // no factor task = [&](size_t start, size_t end) { for (size_t i = start; i < end; ++i) { - auto tmp = static_cast(exp_avg_sq[i]) * beta2t + static_cast(update[i]) * one_minus_beta2t; + auto tmp = static_cast(exp_avg_sq[i]) * beta2t + update[i] * one_minus_beta2t; tmp = std::max(tmp, kEps); exp_avg_sq[i] = static_cast(tmp); - tmp = static_cast(grad[i]) / std::sqrt(static_cast(exp_avg_sq[i])); - update[i] = static_cast(tmp); + update[i] = static_cast(grad[i]) * global_norm_reciprocal_ / std::sqrt(tmp); } }; CPUKernelUtils::ParallelFor(task, elem_num_, kBatchSize); @@ -244,18 +232,16 @@ void FusedAdaFactorCPUKernel::LaunchKernel(const std::vector &inputs auto update_coff = learning_rate / std::max(update_rms_thres, 1.0f); task = [&](size_t start, size_t end) { for (size_t i = start; i < end; ++i) { - auto tmp = static_cast(update[i]) * update_coff; - update[i] = static_cast(tmp); + update[i] = update[i] * update_coff; if (enable_first_moment_) { - tmp = static_cast(exp_avg[i]) * beta1 + static_cast(update[i]) * one_minus_beta1; - exp_avg[i] = static_cast(tmp); - update[i] = exp_avg[i]; + update[i] = static_cast(exp_avg[i]) * beta1 + update[i] * one_minus_beta1; + exp_avg[i] = static_cast(update[i]); } if (enable_weight_decay_) { - tmp = static_cast(param[i]) * weight_decay * learning_rate; - param[i] = param[i] - update[i] - static_cast(tmp); + auto tmp = static_cast(param[i]) * weight_decay * learning_rate; + param[i] = static_cast(static_cast(param[i]) - update[i] - tmp); } else { - param[i] = param[i] - update[i]; + param[i] = static_cast(static_cast(param[i]) - update[i]); } } }; @@ -265,15 +251,17 @@ void FusedAdaFactorCPUKernel::LaunchKernel(const std::vector &inputs bool FusedAdaFactorCPUKernel::Launch(const std::vector &inputs, const std::vector &workspaces, const std::vector &outputs) { - if (inputs.size() != kFusedAdaFactorInputNum) { - MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of inputs should be " << kFusedAdaFactorInputNum - << ", but got: " << inputs.size(); + if (inputs.size() == kStandardInputNum + 1) { + auto global_norm = reinterpret_cast(inputs[GLOBAL_NORM]->addr)[kScalarIndex]; + if (global_norm < kEps) { + global_norm_reciprocal_ = 1.0f; + } else { + global_norm_reciprocal_ = 1.0f / global_norm; + } } - if (workspaces.size() != kFusedAdaFactorWorkSpaceNum) { - MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of workspaces should be " - << kFusedAdaFactorWorkSpaceNum << ", but got: " << workspaces.size(); - } - CheckParam(inputs, workspaces, outputs); + + CheckInputAddresses(inputs); + CheckWorkspaceAddresses(workspaces); if (param_dtype_ == kNumberTypeFloat16) { LaunchKernel(inputs, workspaces, outputs); } else { @@ -282,9 +270,12 @@ bool FusedAdaFactorCPUKernel::Launch(const std::vector &inpu return true; } -void FusedAdaFactorCPUKernel::CheckParam(const std::vector &inputs, - const std::vector &workspaces, - const std::vector &) const { +void FusedAdaFactorCPUKernel::CheckInputAddresses(const std::vector &inputs) const { + if (inputs.size() < kStandardInputNum) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of inputs should be at least " << kStandardInputNum + << ", but got: " << inputs.size(); + } + if (inputs[EPSILON]->size != kSizeFloat32 << 1) { MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'epsilon' should be " << (kSizeFloat32 << 1) << ", but got " << inputs[EPSILON]->size; @@ -293,7 +284,6 @@ void FusedAdaFactorCPUKernel::CheckParam(const std::vector & MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'beta1' should be " << kSizeFloat32 << ", but got " << inputs[BETA1]->size; } - if (inputs[BETA1]->size != kSizeFloat32) { MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'beta1' should be " << kSizeFloat32 << ", but got " << inputs[BETA1]->size; @@ -320,10 +310,6 @@ void FusedAdaFactorCPUKernel::CheckParam(const std::vector & MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'gradient' should be " << param_size << ", but got " << inputs[GRAD]->size; } - if (workspaces[UPDATE]->size != param_size) { - MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'update ' should be " << param_size - << ", but got " << workspaces[0]->size; - } if (enable_first_moment_ && inputs[EXP_AVG]->size != param_size) { MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'exp_avg' should be " << param_size @@ -347,5 +333,28 @@ void FusedAdaFactorCPUKernel::CheckParam(const std::vector & << param_size / last_col_dim_size_ << ", but got " << inputs[EXP_AVG_SQ_COL]->size; } } + +void FusedAdaFactorCPUKernel::CheckWorkspaceAddresses(const std::vector &workspaces) const { + if (workspaces.size() != kWorkSpaceNum) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of workspaces should be " << kWorkSpaceNum + << ", but got: " << workspaces.size(); + } + + size_t update_size = elem_num_ * kSizeFloat32; + + if (workspaces[UPDATE]->size != elem_num_ * kSizeFloat32) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'update ' should be " << update_size + << ", but got " << workspaces[0]->size; + } + + if (workspaces[R_FACTOR]->size != update_size / last_row_dim_size_) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'r_factor' should be " + << update_size / last_row_dim_size_ << ", but got " << workspaces[R_FACTOR]->size; + } + if (workspaces[C_FACTOR]->size != update_size / last_col_dim_size_) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the address size of 'c_factor' should be " + << update_size / last_col_dim_size_ << ", but got " << workspaces[C_FACTOR]->size; + } +} } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h index c9430aade41..7514a49cd39 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/fused_ada_factor_cpu_kernel.h @@ -33,8 +33,9 @@ class FusedAdaFactorCPUKernel : public CPUKernel { const std::vector &outputs) override; private: - void CheckParam(const std::vector &inputs, const std::vector &workspaces, - const std::vector &outputs) const; + void CheckInputAddresses(const std::vector &inputs) const; + void CheckWorkspaceAddresses(const std::vector &workspaces) const; + template void LaunchKernel(const std::vector &inputs, const std::vector &workspaces, const std::vector &outputs); @@ -43,7 +44,7 @@ class FusedAdaFactorCPUKernel : public CPUKernel { float CalcRMS(T *input, size_t elem_num); template - void FactorUpdate(T *update, const std::vector &inputs, const std::vector &workspaces); + void FactorUpdate(float *update, const std::vector &inputs, const std::vector &workspaces); bool enable_scale_parameter_{false}; bool enable_first_moment_{false}; @@ -53,6 +54,7 @@ class FusedAdaFactorCPUKernel : public CPUKernel { size_t last_row_dim_size_{1}; size_t last_col_dim_size_{1}; TypeId param_dtype_{kTypeUnknown}; + float global_norm_reciprocal_{1.0f}; enum InputEnum { EPSILON, @@ -66,7 +68,8 @@ class FusedAdaFactorCPUKernel : public CPUKernel { EXP_AVG, EXP_AVG_SQ_ROW, EXP_AVG_SQ_COL, - EXP_AVG_SQ + EXP_AVG_SQ, + GLOBAL_NORM }; enum WorkspaceEnum { UPDATE, R_FACTOR, C_FACTOR }; @@ -105,6 +108,42 @@ MS_REG_CPU_KERNEL(FusedAdaFactor, .AddInputAttr(kNumberTypeFloat16) .AddOutputAttr(kNumberTypeFloat16), FusedAdaFactorCPUKernel) + +MS_REG_CPU_KERNEL(FusedAdaFactorWithGlobalNorm, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + FusedAdaFactorCPUKernel) + +MS_REG_CPU_KERNEL(FusedAdaFactorWithGlobalNorm, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat16) + .AddInputAttr(kNumberTypeFloat16) + .AddInputAttr(kNumberTypeFloat16) + .AddInputAttr(kNumberTypeFloat16) + .AddInputAttr(kNumberTypeFloat16) + .AddInputAttr(kNumberTypeFloat16) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat16), + FusedAdaFactorCPUKernel) } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/utils/utils.h b/mindspore/ccsrc/utils/utils.h index 7c6264a3ecf..19e840d617b 100644 --- a/mindspore/ccsrc/utils/utils.h +++ b/mindspore/ccsrc/utils/utils.h @@ -215,6 +215,7 @@ constexpr auto kAdamWeightDecayName = "AdamWeightDecay"; constexpr auto kFusedCastAdamWeightDecayName = "FusedCastAdamWeightDecay"; constexpr auto kFusedAdamName = "FusedAdam"; constexpr auto kFusedAdaFactorName = "FusedAdaFactor"; +constexpr auto kFusedAdaFactorWithGlobalNormName = "FusedAdaFactorWithGlobalNorm"; constexpr auto kFusedSparseAdamName = "FusedSparseAdam"; constexpr auto kFusedMatMulBiasAddName = "FusedMatMulBiasAdd"; constexpr auto kDeadNodeName = "DeadNode"; @@ -693,6 +694,7 @@ const std::set kOptOperatorSet = {kMomentumOpName, kFusedCastAdamWeightDecayName, kFusedAdamName, kFusedAdaFactorName, + kFusedAdaFactorWithGlobalNormName, kFusedSparseAdamName, kFusedMulApplyMomentumOpName, kFusedWeightScaleApplyMomentum, diff --git a/mindspore/python/mindspore/ops/operations/__init__.py b/mindspore/python/mindspore/ops/operations/__init__.py index b60a1477f02..0ed361accbe 100644 --- a/mindspore/python/mindspore/ops/operations/__init__.py +++ b/mindspore/python/mindspore/ops/operations/__init__.py @@ -35,7 +35,8 @@ from .array_ops import (Argmax, Argmin, Cast, Concat, Pack, Stack, Unpack, Unsta BatchToSpace, SpaceToBatchND, BatchToSpaceND, BroadcastTo, InplaceUpdate, ReverseSequence, EmbeddingLookup, Unique, GatherD, Identity, Range, MaskedFill, MaskedSelect, SearchSorted, TensorScatterMax, TensorScatterMin, TensorScatterSub, ScatterElements, ExtractVolumePatches) -from .comm_ops import (AllGather, AllReduce, NeighborExchange, NeighborExchangeV2, AlltoAll, _AllSwap, ReduceScatter, Broadcast, +from .comm_ops import (AllGather, AllReduce, NeighborExchange, NeighborExchangeV2, AlltoAll, _AllSwap, ReduceScatter, + Broadcast, _MirrorOperator, _MirrorMiniStepOperator, _MiniStepAllGather, ReduceOp, _VirtualDataset, _VirtualOutput, _VirtualDiv, _GetTensorSlice, _VirtualAdd, _VirtualAssignAdd, _VirtualAccuGrad, _HostAllGather, _HostReduceScatter, _MirrorMicroStepOperator, _MicroStepAllGather) @@ -44,7 +45,8 @@ from .debug_ops import (ImageSummary, InsertGradientOf, HookBackward, ScalarSumm from .control_ops import GeSwitch, Merge from .inner_ops import (ScalarCast, Randperm, NoRepeatNGram, LambApplyOptimizerAssign, LambApplyWeightAssign, MakeRefKey, - FusedWeightScaleApplyMomentum, FusedCastAdamWeightDecay, FusedAdaFactor) + FusedWeightScaleApplyMomentum, FusedCastAdamWeightDecay, FusedAdaFactor, + FusedAdaFactorWithGlobalNorm) from .math_ops import (Abs, ACos, Asin, Asinh, AddN, AccumulateNV2, AssignAdd, AssignSub, Atan2, BatchMatMul, BitwiseAnd, BitwiseOr, Ger, @@ -178,6 +180,7 @@ __all__ = [ 'FusedSparseLazyAdam', 'AdamNoUpdateParam', 'FusedAdaFactor', + 'FusedAdaFactorWithGlobalNorm', 'Softplus', 'Softmax', 'Softsign', diff --git a/mindspore/python/mindspore/ops/operations/inner_ops.py b/mindspore/python/mindspore/ops/operations/inner_ops.py index 326bc92e270..329b4daf2a0 100755 --- a/mindspore/python/mindspore/ops/operations/inner_ops.py +++ b/mindspore/python/mindspore/ops/operations/inner_ops.py @@ -689,3 +689,25 @@ class FusedAdaFactor(PrimitiveWithInfer): learning_rate_type, grad_type, param_type, exp_avg_type, exp_avg_sq_row_type, exp_avg_sq_col_type, exp_avg_sq_type): return param_type + + +class FusedAdaFactorWithGlobalNorm(FusedAdaFactor): + r""" + Divide global norm for gradient in FusedAdaFactor, and refer to super class for FusedAdaFactor details + """ + + @prim_attr_register + def __init__(self, enable_scale_parameter=False, enable_first_moment=False, enable_weight_decay=False): + super(FusedAdaFactorWithGlobalNorm, self).__init__(enable_scale_parameter, enable_first_moment, + enable_weight_decay) + + def infer_shape(self, epsilon_shape, clip_threshold_shape, beta1_shape, beta2t_shape, weight_decay_shape, + learning_rate_shape, grad_shape, param_shape, exp_avg_shape, exp_avg_sq_row_shape, + exp_avg_sq_col_shape, exp_avg_sq_shape, global_norm_shape): + validator.check("grad_shape", grad_shape, "param_shape", param_shape, Rel.EQ, self.name) + return param_shape + + def infer_dtype(self, epsilon_type, clip_threshold_type, beta1_type, beta2t_type, weight_decay_type, + learning_rate_type, grad_type, param_type, exp_avg_type, exp_avg_sq_row_type, + exp_avg_sq_col_type, exp_avg_sq_type, global_norm_type): + return param_type diff --git a/tests/st/ops/cpu/test_fused_ada_factor_op.py b/tests/st/ops/cpu/test_fused_ada_factor_op.py index 7f8ef4297d1..51b2737642a 100644 --- a/tests/st/ops/cpu/test_fused_ada_factor_op.py +++ b/tests/st/ops/cpu/test_fused_ada_factor_op.py @@ -39,6 +39,17 @@ class Net(nn.Cell): return out +class NetWithGlobalNorm(Net): + def __init__(self): + super(NetWithGlobalNorm, self).__init__() + self.opt = ops.FusedAdaFactorWithGlobalNorm() + + def construct(self, epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad, global_norm): + out = self.opt(epsilon, clip_threshold, beta1, beta2, weight_decay, lr, grad, self.param, self.exp_avg, + self.exp_avg_sq_row, self.exp_avg_sq_col, self.exp_avg_sq, global_norm) + return out + + @pytest.mark.level0 @pytest.mark.platform_x86_cpu @pytest.mark.env_onecard @@ -54,3 +65,20 @@ def test_adafactor(): net((1e-30, 1e-3), 1.0, 0.9, 0.8, 1e-2, 0.03, gradient) diff = net.param.asnumpy() - np.ones(param_shape) * 0.97 assert np.all(diff < 1e-3) + + +@pytest.mark.level0 +@pytest.mark.platform_x86_cpu +@pytest.mark.env_onecard +def test_adafactor_with_global_norm(): + ''' + Feature: AdaFactor + Description: Test AdaFactor + Expectation: Run success + ''' + context.set_context(mode=context.GRAPH_MODE, device_target="CPU") + net = NetWithGlobalNorm() + gradient = Tensor(np.ones(param_shape), mstype.float32) + net((1e-30, 1e-3), 1.0, 0.9, 0.8, 1e-2, 0.03, gradient, 10.0) + diff = net.param.asnumpy() - np.ones(param_shape) * 0.97 + assert np.all(diff < 1e-3) diff --git a/tests/ut/cpp/kernel/cpu/fused_ada_factor_cpu_kernel_test.cc b/tests/ut/cpp/kernel/cpu/fused_ada_factor_cpu_kernel_test.cc index bb37cb14f9e..10cb9393548 100644 --- a/tests/ut/cpp/kernel/cpu/fused_ada_factor_cpu_kernel_test.cc +++ b/tests/ut/cpp/kernel/cpu/fused_ada_factor_cpu_kernel_test.cc @@ -57,7 +57,7 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { grad_.resize(elem_num_); exp_avg_.resize(elem_num_); exp_avg_sq_.resize(elem_num_); - update_.resize(elem_num_); + update_.resize(elem_num_, 0.0f); for (size_t i = 0; i < elem_num_; ++i) { auto ptr = (float16 *)param_.data(); ptr[i] = static_cast(1.0f); @@ -67,8 +67,6 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { ptr[i] = static_cast(0.0f); ptr = (float16 *)exp_avg_sq_.data(); ptr[i] = static_cast(0.0f); - ptr = (float16 *)update_.data(); - ptr[i] = static_cast(0.0f); } auto r_factor_num = elem_num_ / last_row_dim_size_; @@ -77,8 +75,6 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { for (size_t i = 0; i < r_factor_num; ++i) { auto ptr = (float16 *)exp_avg_sq_row_.data(); ptr[i] = static_cast(0.0f); - ptr = (float16 *)r_factor_.data(); - ptr[i] = static_cast(0.0f); } auto c_factor_num = elem_num_ / last_col_dim_size_; @@ -87,8 +83,6 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { for (size_t i = 0; i < c_factor_num; ++i) { auto ptr = (float16 *)exp_avg_sq_col_.data(); ptr[i] = static_cast(0.0f); - ptr = (float16 *)c_factor_.data(); - ptr[i] = static_cast(0.0f); } } @@ -99,7 +93,7 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { return kernel_addr; } - void CreateAddress() { + void CreateAddress(bool enable_global_norm) { constexpr size_t eps_num = 2; inputs_.push_back(CreateKernelAddress(epsilon_.data(), eps_num, kSizeFloat32)); inputs_.push_back(CreateKernelAddress(&clip_threshold_, 1, kSizeFloat32)); @@ -113,36 +107,37 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { inputs_.push_back(CreateKernelAddress(exp_avg_sq_row_.data(), elem_num_ / last_row_dim_size_, type_size_)); inputs_.push_back(CreateKernelAddress(exp_avg_sq_col_.data(), elem_num_ / last_col_dim_size_, type_size_)); inputs_.push_back(CreateKernelAddress(exp_avg_sq_.data(), elem_num_, type_size_)); - workspace_.push_back(CreateKernelAddress(update_.data(), elem_num_, type_size_)); - workspace_.push_back(CreateKernelAddress(r_factor_.data(), elem_num_ / last_row_dim_size_, type_size_)); - workspace_.push_back(CreateKernelAddress(c_factor_.data(), elem_num_ / last_col_dim_size_, type_size_)); + workspace_.push_back(CreateKernelAddress(update_.data(), elem_num_, kSizeFloat32)); + workspace_.push_back(CreateKernelAddress(r_factor_.data(), elem_num_ / last_row_dim_size_, kSizeFloat32)); + workspace_.push_back(CreateKernelAddress(c_factor_.data(), elem_num_ / last_col_dim_size_, kSizeFloat32)); + if (enable_global_norm) { + inputs_.push_back(CreateKernelAddress(&global_norm_, 1, kSizeFloat32)); + } } - void ComputeFp32() { + void ComputeFp32(bool enable_global_norm) { ada_factor_->param_dtype_ = kNumberTypeFloat32; type_size_ = sizeof(float); InitDataFp32(); - CreateAddress(); + CreateAddress(enable_global_norm); ada_factor_->Launch(inputs_, workspace_, outputs_); - constexpr float result = 0.97; for (size_t i = 0; i < elem_num_; ++i) { - EXPECT_TRUE(std::fabs(param_[i] - result) < 1e-6); + EXPECT_TRUE(std::fabs(param_[i] - result_) < 1e-6); } } - void ComputeFp16() { + void ComputeFp16(bool enable_global_norm) { ada_factor_->param_dtype_ = kNumberTypeFloat16; type_size_ = sizeof(float16); InitDataFp16(); - CreateAddress(); + CreateAddress(enable_global_norm); ada_factor_->Launch(inputs_, workspace_, outputs_); - constexpr float result = 0.97; auto ptr = (float16 *)param_.data(); for (size_t i = 0; i < elem_num_; ++i) { - EXPECT_TRUE(std::fabs(static_cast(ptr[i]) - result) < 1e-3); + EXPECT_TRUE(std::fabs(static_cast(ptr[i]) - result_) < 1e-3); } } @@ -152,6 +147,8 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { float beta1_ = 0.9; float beta2t_ = 0.8; float weight_decay_ = 1e-2; + float global_norm_ = 10.0f; + float result_ = 0.97; std::vector param_; std::vector grad_; std::vector exp_avg_; @@ -179,7 +176,7 @@ class FusedAdaFactorCpuKernelTest : public UT::Common { /// Expectation: pass TEST_F(FusedAdaFactorCpuKernelTest, compute_fp32_factor) { ada_factor_->need_factor_ = true; - ComputeFp32(); + ComputeFp32(false); } /// Feature: FusedAdaFactor @@ -187,7 +184,23 @@ TEST_F(FusedAdaFactorCpuKernelTest, compute_fp32_factor) { /// Expectation: pass TEST_F(FusedAdaFactorCpuKernelTest, compute_fp32_no_factor) { ada_factor_->need_factor_ = false; - ComputeFp32(); + ComputeFp32(false); +} + +/// Feature: FusedAdaFactor +/// Description: Run FusedAdaFactor that needs factor state with fp32 data inputs and global norm +/// Expectation: pass +TEST_F(FusedAdaFactorCpuKernelTest, compute_fp32_factor_global_norm) { + ada_factor_->need_factor_ = true; + ComputeFp32(true); +} + +/// Feature: FusedAdaFactor +/// Description: Run FusedAdaFactor that doesn't need factor state with fp32 data inputs and global norm +/// Expectation: pass +TEST_F(FusedAdaFactorCpuKernelTest, compute_fp32_no_factor_global_norm) { + ada_factor_->need_factor_ = false; + ComputeFp32(true); } /// Feature: FusedAdaFactor @@ -195,7 +208,7 @@ TEST_F(FusedAdaFactorCpuKernelTest, compute_fp32_no_factor) { /// Expectation: pass TEST_F(FusedAdaFactorCpuKernelTest, compute_fp16_factor) { ada_factor_->need_factor_ = true; - ComputeFp16(); + ComputeFp16(false); } /// Feature: FusedAdaFactor @@ -203,7 +216,23 @@ TEST_F(FusedAdaFactorCpuKernelTest, compute_fp16_factor) { /// Expectation: pass TEST_F(FusedAdaFactorCpuKernelTest, compute_fp16_no_factor) { ada_factor_->need_factor_ = false; - ComputeFp16(); + ComputeFp16(false); +} + +/// Feature: FusedAdaFactor +/// Description: Run FusedAdaFactor that needs factor state with fp16 data inputs and global norm +/// Expectation: pass +TEST_F(FusedAdaFactorCpuKernelTest, compute_fp16_factor_global_norm) { + ada_factor_->need_factor_ = true; + ComputeFp16(true); +} + +/// Feature: FusedAdaFactor +/// Description: Run FusedAdaFactor that doesn't need factor state with fp16 data inputs and global norm +/// Expectation: pass +TEST_F(FusedAdaFactorCpuKernelTest, compute_fp16_no_factor_global_norm) { + ada_factor_->need_factor_ = false; + ComputeFp16(true); } } // namespace kernel } // namespace mindspore