squad multi-lamb
diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py index 0297a1d..f358635 100644 --- a/python/mxnet/ndarray/contrib.py +++ b/python/mxnet/ndarray/contrib.py
@@ -605,3 +605,27 @@ etas=etas, name=name, **kwargs) + +def multi_lamb_update(weights, grads, mean, var, temp_g, step_count, + out=None, name=None, num_tensors=0, **kwargs): + if not num_tensors: + num_tensors = len(weights) + temp_list = _flatten_list(zip(weights, grads, mean, var, temp_g)) + return ndarray._internal._multi_lamb_update(*temp_list, + out=out, + num_tensors=num_tensors, + step_count=step_count, + name=name, + **kwargs) + +def multi_mp_lamb_update(weights, grads, mean, var, temp_g, weights32, step_count, + out=None, name=None, num_tensors=0, **kwargs): + if not num_tensors: + num_tensors = len(weights) + temp_list = _flatten_list(zip(weights, grads, mean, var, temp_g, weights32)) + return ndarray._internal._multi_mp_lamb_update(*temp_list, + out=out, + num_tensors=num_tensors, + step_count=step_count, + name=name, + **kwargs)
diff --git a/python/mxnet/optimizer/optimizer.py b/python/mxnet/optimizer/optimizer.py index 42d418e..7f455f2 100644 --- a/python/mxnet/optimizer/optimizer.py +++ b/python/mxnet/optimizer/optimizer.py
@@ -36,13 +36,14 @@ preloaded_multi_sgd_mom_update, preloaded_multi_mp_sgd_update, preloaded_multi_mp_sgd_mom_update, lamb_update_phase1, lamb_update_phase2, mp_lamb_update_phase1, mp_lamb_update_phase2) +from ..ndarray.contrib import (multi_lamb_update, multi_mp_lamb_update) from ..ndarray import sparse from ..random import normal from ..util import is_np_array __all__ = [ 'AdaDelta', 'AdaGrad', 'Adam', 'Adamax', 'DCASGD', 'FTML', 'Ftrl', 'LARS', 'LBSGD', - 'NAG', 'NDabs', 'Nadam', 'Optimizer', 'RMSProp', 'SGD', 'SGLD', 'Signum', 'LAMB', + 'MultiLAMB', 'NAG', 'NDabs', 'Nadam', 'Optimizer', 'RMSProp', 'SGD', 'SGLD', 'Signum', 'LAMB', 'Test', 'Updater', 'ccSGD', 'create', 'get_updater', 'register' ] @@ -1044,6 +1045,95 @@ self._update_impl(index, weight, grad, state, multi_precision=use_multi_precision) +@register +class MultiLAMB(Optimizer): + """multiLAMB optimizer. + """ + def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-6, + lower_bound=1e-3, upper_bound=10.0, bias_correction=False, **kwargs): + super(MultiLAMB, self).__init__(learning_rate=learning_rate, **kwargs) + self.beta1 = beta1 + self.beta2 = beta2 + self.epsilon = epsilon + self.lower_bound = lower_bound + self.upper_bound = upper_bound + self.bias_correction = bias_correction + self.aggregate_num = max(1, min(50, int(os.getenv('MXNET_OPTIMIZER_AGGREGATION_SIZE', "50")))) + + def create_state(self, index, weight): + stype = weight.stype + dtype = weight.dtype + return (zeros(weight.shape, weight.context, dtype=dtype, stype=stype), # mean + zeros(weight.shape, weight.context, dtype=dtype, stype=stype), # variance + zeros(weight.shape, weight.context, dtype=dtype, stype=stype)) # temp_g + + def _update_impl(self, index, weights, grads, states, multi_precision=False): + step_count = [] + if not isinstance(index, (tuple, list)): + weights = [weights] + grads = [grads] + states = [states] + self._update_count(index) + step_count.append(self._index_update_count[index]) + lr = self._get_lr(index) + wd = self._get_wd(index) + else: + for i, (weight, grad) in enumerate(zip(weights, grads)): + assert(isinstance(weight, NDArray)) + assert(isinstance(grad, NDArray)) + self._update_count(i) + step_count.append(self._index_update_count[i]) + lr = self._get_lr(index[0]) + wd = self._get_wd(index[0]) + + kwargs = {'learning_rate': lr, 'beta1': self.beta1, 'beta2': self.beta2, + 'epsilon': self.epsilon, 'wd': wd, + 'lower_bound': self.lower_bound, 'upper_bound': self.upper_bound, + 'bias_correction': self.bias_correction, + 'rescale_grad': self.rescale_grad} + + if self.clip_gradient: + kwargs['clip_gradient'] = self.clip_gradient + + updated_tensors = 0 + while updated_tensors < len(weights): + sidx = updated_tensors + eidx = min(updated_tensors + self.aggregate_num, len(weights)) + if not multi_precision: + mean, var, temp_g = list(zip(*states[sidx:eidx])) + multi_lamb_update(weights[sidx:eidx], + grads[sidx:eidx], + mean, var, temp_g, + out=weights[sidx:eidx], + step_count=step_count[sidx:eidx], + **kwargs) + else: + mean_var_g = list(zip(*states[sidx:eidx]))[1] + temp = list(zip(*mean_var_g)) + mean = temp[0] + var = temp[1] + temp_g = temp[2] + multi_mp_lamb_update(weights[sidx:eidx], + grads[sidx:eidx], + mean, var, temp_g, + list(zip(*states[sidx:eidx]))[0], + out=weights[sidx:eidx], + step_count=step_count[sidx:eidx], + **kwargs) + + updated_tensors += self.aggregate_num + + def update(self, index, weight, grad, state): + self._update_impl(index, weight, grad, state, multi_precision=False) + + def update_multi_precision(self, index, weight, grad, state): + if not isinstance(index, (tuple, list)): + use_multi_precision = self.multi_precision and weight.dtype == numpy.float16 + else: + use_multi_precision = self.multi_precision and weight[0].dtype == numpy.float16 + self._update_impl(index, weight, grad, state, + multi_precision=use_multi_precision) + # @register class LBSGD(Optimizer):
diff --git a/python/mxnet/test_utils.py b/python/mxnet/test_utils.py index abdf570..6eaec47 100644 --- a/python/mxnet/test_utils.py +++ b/python/mxnet/test_utils.py
@@ -2141,36 +2141,60 @@ def compare_optimizer(opt1, opt2, shape, dtype, w_stype='default', g_stype='default', - rtol=1e-4, atol=1e-5, compare_states=True): + rtol=1e-4, atol=1e-5, compare_states=True, ntensors=1): """Compare opt1 and opt2.""" - if w_stype == 'default': - w2 = mx.random.uniform(shape=shape, ctx=default_context(), dtype=dtype) - w1 = w2.copyto(default_context()) - elif w_stype in ('row_sparse', 'csr'): - w2 = rand_ndarray(shape, w_stype, density=1, dtype=dtype) - w1 = w2.copyto(default_context()).tostype('default') + if ntensors == 1: + if w_stype == 'default': + w2 = mx.random.uniform(shape=shape, ctx=default_context(), dtype=dtype) + w1 = w2.copyto(default_context()) + elif w_stype in ('row_sparse', 'csr'): + w2 = rand_ndarray(shape, w_stype, density=1, dtype=dtype) + w1 = w2.copyto(default_context()).tostype('default') + else: + raise Exception("type not supported yet") + if g_stype == 'default': + g2 = mx.random.uniform(shape=shape, ctx=default_context(), dtype=dtype) + g1 = g2.copyto(default_context()) + elif g_stype in ('row_sparse', 'csr'): + g2 = rand_ndarray(shape, g_stype, dtype=dtype) + g1 = g2.copyto(default_context()).tostype('default') + else: + raise Exception("type not supported yet") + + state1 = opt1.create_state_multi_precision(0, w1) + state2 = opt2.create_state_multi_precision(0, w2) + if compare_states: + compare_ndarray_tuple(state1, state2) + + opt1.update_multi_precision(0, w1, g1, state1) + opt2.update_multi_precision(0, w2, g2, state2) + if compare_states: + compare_ndarray_tuple(state1, state2, rtol=rtol, atol=atol) + assert_almost_equal(w1, w2, rtol=rtol, atol=atol) else: - raise Exception("type not supported yet") - if g_stype == 'default': - g2 = mx.random.uniform(shape=shape, ctx=default_context(), dtype=dtype) - g1 = g2.copyto(default_context()) - elif g_stype in ('row_sparse', 'csr'): - g2 = rand_ndarray(shape, g_stype, dtype=dtype) - g1 = g2.copyto(default_context()).tostype('default') - else: - raise Exception("type not supported yet") + # test multi-tensor: Opt1 single-tensor, Opt2 multi-tensor + from copy import deepcopy + if not isinstance(shape, list): + w1 = [mx.random.uniform(shape=shape, ctx=default_context(), dtype=dtype) for i in range(ntensors)] + g1 = [mx.random.uniform(shape=shape, ctx=default_context(), dtype=dtype) for i in range(ntensors)] + else: + w1, g1 = [], [] + for s in shape: + w1.append(mx.random.uniform(shape=s, ctx=default_context(), dtype=dtype)) + g1.append(mx.random.uniform(shape=s, ctx=default_context(), dtype=dtype)) + w1 = tuple(w1) + w2 = deepcopy(w1) + g1 = tuple(g1) + g2 = deepcopy(g1) + state2 = [opt2.create_state_multi_precision(0, w2[i]) for i in range(ntensors)] - state1 = opt1.create_state_multi_precision(0, w1) - state2 = opt2.create_state_multi_precision(0, w2) - if compare_states: - compare_ndarray_tuple(state1, state2) - - opt1.update_multi_precision(0, w1, g1, state1) - opt2.update_multi_precision(0, w2, g2, state2) - if compare_states: - compare_ndarray_tuple(state1, state2, rtol=rtol, atol=atol) - assert_almost_equal(w1.asnumpy(), w2.asnumpy(), rtol=rtol, atol=atol) - + opt2.update_multi_precision(list(range(ntensors)), w2, g2, state2) + for i in range(ntensors): + state1 = opt1.create_state_multi_precision(i, w1[i]) + opt1.update_multi_precision(i, w1[i], g1[i], state1) + if compare_states: + compare_ndarray_tuple(state1, state2[i], rtol, atol) + compare_ndarray_tuple(w1[i], w2[i], rtol, atol) def same_symbol_structure(sym1, sym2): """Compare two symbols to check if they have the same computation graph structure.
diff --git a/src/operator/contrib/multi_lamb-inl.h b/src/operator/contrib/multi_lamb-inl.h new file mode 100644 index 0000000..107a4c9 --- /dev/null +++ b/src/operator/contrib/multi_lamb-inl.h
@@ -0,0 +1,332 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * Copyright (c) 2019 by Contributors + * \file multi_lamb-inl.h + * \brief vectorized lars coefficient computed from sums of squared weights and grads + * \author Moises Hernandez + */ +#ifndef MXNET_OPERATOR_CONTRIB_MULTI_LAMB_INL_H_ +#define MXNET_OPERATOR_CONTRIB_MULTI_LAMB_INL_H_ + +#include <dmlc/parameter.h> +#include <mxnet/operator.h> +#include <mxnet/operator_util.h> +#include <mxnet/op_attr_types.h> +#include <mshadow/base.h> +#include <nnvm/op.h> +#include <nnvm/op_attr_types.h> +#include <vector> +#include "../operator_common.h" +#include "../mshadow_op.h" +#include "../mxnet_op.h" +#include "../tensor/init_op.h" +#include "../tensor/util/tensor_util-inl.h" +#include "multi_sum_sq-inl.h" + +namespace mxnet { +namespace op { + +namespace multilamb { +enum MultiLambUpdateResource {kTempSpace}; +} // namespace multilamb + +struct MultiLAMBParam : public dmlc::Parameter<MultiLAMBParam> { + float learning_rate; + float beta1; + float beta2; + float epsilon; + float wd; + float rescale_grad; + float lower_bound; + float upper_bound; + float clip_gradient; + bool bias_correction; + int num_tensors; + mxnet::Tuple<int> step_count; + + DMLC_DECLARE_PARAMETER(MultiLAMBParam) { + DMLC_DECLARE_FIELD(learning_rate) + .set_default(0.001f) + .describe("Learning rate"); + DMLC_DECLARE_FIELD(beta1) + .set_default(0.9f) + .describe("Exponential decay rate for the first moment estimates."); + DMLC_DECLARE_FIELD(beta2) + .set_default(0.999f) + .describe("Exponential decay rate for the second moment estimates."); + DMLC_DECLARE_FIELD(epsilon) + .set_default(1e-6f) + .describe("Small value to avoid division by 0."); + DMLC_DECLARE_FIELD(wd) + .set_default(0.0f) + .describe("Weight decay augments the objective function with a " + "regularization term that penalizes large weights. " + "The penalty scales with the square of the magnitude of each weight."); + DMLC_DECLARE_FIELD(rescale_grad) + .set_default(1.0f) + .describe("Gradient rescaling factor"); + DMLC_DECLARE_FIELD(lower_bound) + .set_default(1e-3f) + .describe("Lower limit of norm of weight."); + DMLC_DECLARE_FIELD(upper_bound) + .set_default(10.0f) + .describe("Upper limit of norm of weight."); + DMLC_DECLARE_FIELD(clip_gradient) + .set_default(-1.0f) + .describe("Clip gradient to the range of [-clip_gradient, clip_gradient] " + "If clip_gradient <= 0, gradient clipping is turned off. " + "grad = max(min(grad, clip_gradient), -clip_gradient)."); + DMLC_DECLARE_FIELD(bias_correction) + .set_default(false) + .describe("Whether to use bias correction."); + DMLC_DECLARE_FIELD(step_count) + .describe("Step count for each tensor"); + DMLC_DECLARE_FIELD(num_tensors) + .set_default(1) + .describe("Number of tensors"); + } +}; + +template<typename ParamType, int input_stride> +inline bool MultiLAMB_InferShape(const nnvm::NodeAttrs& attrs, + mxnet::ShapeVector *in_attrs, + mxnet::ShapeVector *out_attrs) { + const ParamType& param = dmlc::get<ParamType>(attrs.parsed); + CHECK_EQ(in_attrs->size(), input_stride * param.num_tensors); + CHECK_EQ(out_attrs->size(), param.num_tensors); + + bool all_inferred = true; + auto& input_shapes = *in_attrs; + auto& output_shapes = *out_attrs; + + CHECK_EQ(param.step_count.ndim(), param.num_tensors) + << "Number of step counts is inconsistent with num_weights." + << "Expected number of step counts: " + << param.num_tensors << ", and got " << param.step_count.ndim(); + + // Weights, gradients, mean and variance + for (int i = 0; i < param.num_tensors; ++i) { + mxnet::ShapeVector input_vec; + mxnet::ShapeVector output_vec({output_shapes[i]}); + for (int j = 0; j < input_stride; ++j) { + input_vec.push_back(input_shapes[i * input_stride + j]); + } + all_inferred = all_inferred && ElemwiseShape<input_stride, 1>(attrs, &input_vec, &output_vec); + } + return all_inferred; +} + +template <typename ParamType, int input_stride> +inline bool MP_MultiLAMB_InferType(const nnvm::NodeAttrs& attrs, + std::vector<int> *in_attrs, + std::vector<int> *out_attrs) { + const ParamType& param = dmlc::get<ParamType>(attrs.parsed); + CHECK_EQ(in_attrs->size(), input_stride * param.num_tensors); + CHECK_EQ(out_attrs->size(), param.num_tensors); + + bool all_inferred = true; + auto& input_types = *in_attrs; + auto& output_types = *out_attrs; + + // weights, gradients + for (int i = 0; i < param.num_tensors; ++i) { + std::vector<int> input_vec; + std::vector<int> output_vec({output_types[i]}); + for (int j = 0; j < 2; ++j) { + input_vec.push_back(input_types[i * input_stride + j]); + } + all_inferred = all_inferred && + ElemwiseType<2, 1>(attrs, &input_vec, &output_vec); + } + + // mean, var, temp_g, weights32 (master copies of weights) + for (int i = 0; i < param.num_tensors; ++i) { + TYPE_ASSIGN_CHECK(input_types, input_stride * i + 2, mshadow::kFloat32); + TYPE_ASSIGN_CHECK(input_types, input_stride * i + 3, mshadow::kFloat32); + TYPE_ASSIGN_CHECK(input_types, input_stride * i + 4, mshadow::kFloat32); + TYPE_ASSIGN_CHECK(input_types, input_stride * i + input_stride - 1, mshadow::kFloat32); + } + return all_inferred; +} + +template<typename T> +class LAMB_type_identity { + public: + using type = T; +}; + +template<typename T> +class LAMB_single_precision { + public: + using type = float; +}; + +template<typename DType, typename MPDType> +struct MultiLAMBKernelParam { + static const int N = 50; + size_t ntensors; + size_t max_size; + size_t total_size; + size_t sizes[N]; + DType* weights[N]; + DType* grads[N]; + MPDType* mean[N]; + MPDType* var[N]; + MPDType* temp_g[N]; + MPDType* weights32[N]; + DType* out_data[N]; + int step_count[N]; + + // gpu + int chunk_size = 65536; + int nchunks; +}; + +template<typename xpu, + typename DType, + typename MPDType, + typename ParamType = MultiLAMBParam, + int input_stride = 5> +void FillMultiLAMBKernelParam(const nnvm::NodeAttrs& attrs, + const OpContext &ctx, + const std::vector<TBlob> &inputs, + const std::vector<TBlob> &outputs, + MultiLAMBKernelParam<DType, MPDType> *multi_param) { + const ParamType& p = nnvm::get<ParamType>(attrs.parsed); + mxnet_op::Stream<xpu>* s = ctx.get_stream<xpu>(); + + multi_param->ntensors = p.num_tensors; + multi_param->total_size = 0; + multi_param->max_size = 0; + multi_param->nchunks = 0; + + constexpr bool isSame = std::is_same<DType, MPDType>::value; + for (size_t i = 0; i < multi_param->ntensors; ++i) { + const auto idx = i * input_stride; + multi_param->sizes[i] = inputs[idx].shape_.Size(); + multi_param->total_size += multi_param->sizes[i]; + if (multi_param->max_size < multi_param->sizes[i]) + multi_param->max_size = multi_param->sizes[i]; + + multi_param->weights[i] = inputs[idx].FlatTo2D<xpu, DType>(s).dptr_; + multi_param->grads[i] = inputs[idx + 1].FlatTo2D<xpu, DType>(s).dptr_; + multi_param->mean[i] = inputs[idx + 2].FlatTo2D<xpu, MPDType>(s).dptr_; + multi_param->var[i] = inputs[idx + 3].FlatTo2D<xpu, MPDType>(s).dptr_; + multi_param->temp_g[i] = inputs[idx + 4].FlatTo2D<xpu, MPDType>(s).dptr_; + // if mixed precision, then the last input in a set + // is 32-bit master copy of the weights + if (!isSame) + multi_param->weights32[i] = inputs[idx + input_stride - 1].FlatTo2D<xpu, MPDType>(s).dptr_; + multi_param->out_data[i] = outputs[i].FlatTo2D<xpu, DType>(s).dptr_; + multi_param->nchunks += (multi_param->sizes[i] + multi_param->chunk_size - 1) + /multi_param->chunk_size; + } + memcpy(multi_param->step_count, p.step_count.begin(), multi_param->ntensors * sizeof(int)); +} + +using namespace mxnet_op; +template<typename MPDType, typename DType> +void call_kernel1(Stream<cpu>* s); +template<typename MPDType, typename DType> +void call_kernel1(Stream<gpu>* s); + +template<typename MPDType, typename DType> +void call_kernel2(Stream<cpu>* s); +template<typename MPDType, typename DType> +void call_kernel2(Stream<gpu>* s); + +template<typename xpu, template<typename> class MPTypeChooser, int input_stride> +inline void multiLAMB(const nnvm::NodeAttrs& attrs, + const OpContext &ctx, + const std::vector<TBlob> &inputs, + const std::vector<OpReqType> &req, + const std::vector<TBlob> &outputs) { + auto param = nnvm::get<MultiLAMBParam>(attrs.parsed); + Stream<xpu>* s = ctx.get_stream<xpu>(); + + MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, { + using MPDType = typename MPTypeChooser<DType>::type; + MultiLAMBKernelParam<DType, MPDType> kernel_params; + FillMultiLAMBKernelParam<xpu, DType, MPDType, MultiLAMBParam, input_stride> + (attrs, ctx, inputs, outputs, &kernel_params); + + // create vector of TBlob with all the weights contiguous + std::vector<TBlob> weights; + for (size_t index = 0; index < kernel_params.ntensors; ++index) { + weights.emplace_back(inputs[index*input_stride]); + } + // create vector of TBlob with all the temp_g contiguous + std::vector<TBlob> temp_g; + for (size_t index = 0; index < kernel_params.ntensors; ++index) { + temp_g.emplace_back(inputs[index*input_stride+4]); + } + + // Calculate amount of temporary storage + size_t workspace_size = 2 * kernel_params.ntensors * sizeof(float) + + 2 * kernel_params.nchunks * sizeof(int); + + // Request temporary storage + Tensor<xpu, 1, char> workspace = + ctx.requested[multilamb::kTempSpace].get_space_typed<xpu, 1, char>( + Shape1(workspace_size), s); + + // Create tensors + size_t pos_wspace = 0; + Tensor<xpu, 1, float> r1(reinterpret_cast<float*>(&workspace[pos_wspace]), + Shape1(kernel_params.ntensors), s); + pos_wspace += kernel_params.ntensors * sizeof(float); + Tensor<xpu, 1, float> r2(reinterpret_cast<float*>(&workspace[pos_wspace]), + Shape1(kernel_params.ntensors), s); + pos_wspace += kernel_params.ntensors * sizeof(float); + Tensor<xpu, 1, int> block_to_tensor(reinterpret_cast<int*>(&workspace[pos_wspace]), + Shape1(kernel_params.nchunks), s); + pos_wspace += kernel_params.nchunks * sizeof(int); + Tensor<xpu, 1, int> block_to_chunk(reinterpret_cast<int*>(&workspace[pos_wspace]), + Shape1(kernel_params.nchunks), s); + + MultiSumSqRun<xpu>(weights, kernel_params.ntensors, r1.dptr_, s); + call_kernel1<MPDType, DType>(s, kernel_params, param, block_to_tensor.dptr_, + block_to_chunk.dptr_); + MultiSumSqRun<xpu>(temp_g, kernel_params.ntensors, r2.dptr_, s); + call_kernel2<MPDType, DType>(s, kernel_params, param, r1.dptr_, r2.dptr_, + block_to_tensor.dptr_, block_to_chunk.dptr_, + req[0]); + }); +} + +template<typename xpu, bool MP> +inline void multiLAMBUpdate(const nnvm::NodeAttrs& attrs, + const OpContext &ctx, + const std::vector<TBlob> &inputs, + const std::vector<OpReqType> &req, + const std::vector<TBlob> &outputs) { + if (!MP) { + multiLAMB<xpu, LAMB_type_identity, 5> + (attrs, ctx, inputs, req, outputs); + } else { + multiLAMB<xpu, LAMB_single_precision, 6> + (attrs, ctx, inputs, req, outputs); + } +} + +} // namespace op +} // namespace mxnet +#endif // MXNET_OPERATOR_CONTRIB_MULTI_LAMB_INL_H_
diff --git a/src/operator/contrib/multi_lamb.cc b/src/operator/contrib/multi_lamb.cc new file mode 100644 index 0000000..ebfa7a2 --- /dev/null +++ b/src/operator/contrib/multi_lamb.cc
@@ -0,0 +1,245 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * Copyright (c) 2019 by Contributors + * \file multi_lamb.cc + * \brief vectorized LAMB coefficient computed from sums of squared weights and grads + * \author Moises Hernandez + */ + +#include "./multi_lamb-inl.h" +#include "../elemwise_op_common.h" + +namespace mxnet { +namespace op { + +template<typename MPDType, bool has_mixed_precision> +struct MultiLAMB_step1_kernel { + template<typename DType> + MSHADOW_XINLINE static void Map(int i, + const MultiLAMBKernelParam<DType, MPDType>& kernel_params, + const float learning_rate, + const float beta1, const float beta2, + const float epsilon, + const float wd, + const float clip_gradient, + const bool bias_correction, + const float rescale_grad) { + using namespace mshadow_op; + for (size_t index = 0; index < kernel_params.ntensors; ++index) { + if ((size_t)i < kernel_params.sizes[index]) { + MPDType w = has_mixed_precision ? kernel_params.weights32[index][i]: + MPDType(kernel_params.weights[index][i]); + MPDType scaled_grad = static_cast<MPDType>(kernel_params.grads[index][i])*rescale_grad; + if (clip_gradient >= 0.0f) + scaled_grad = mshadow_op::clip::Map(scaled_grad, static_cast<MPDType>(clip_gradient)); + MPDType mean = static_cast<MPDType>(beta1) * kernel_params.mean[index][i] + + (static_cast<MPDType>(1.0f) - static_cast<MPDType>(beta1)) * scaled_grad; + MPDType var = static_cast<MPDType>(beta2) * kernel_params.var[index][i] + + (static_cast<MPDType>(1.0f) - static_cast<MPDType>(beta2)) * scaled_grad * scaled_grad; + kernel_params.mean[index][i] = mean; + kernel_params.var[index][i] = var; + + MPDType g; + if (bias_correction) { + MPDType mean_hat = mean / (static_cast<MPDType>(1.0f) - + power::Map(static_cast<MPDType>(beta1), + static_cast<MPDType>(kernel_params.step_count[index]))); + MPDType var_hat = var / (static_cast<MPDType>(1.0f) - + power::Map(static_cast<MPDType>(beta2), + static_cast<MPDType>(kernel_params.step_count[index]))); + g = mean_hat / (sqrt(var_hat) + epsilon) + wd * w; + } else { + g = mean / (sqrt(var) + epsilon) + wd * w; + } + kernel_params.temp_g[index][i] = g; + } + } + } +}; + +template<typename MPDType, bool has_mixed_precision> +struct MultiLAMB_step2_kernel { + template<typename DType> + MSHADOW_XINLINE static void Map(int i, + const MultiLAMBKernelParam<DType, MPDType>& kernel_params, + const float* sumSqWeigths, + const float* sumSqtemp_g, + const float learning_rate, + const float lower_bound, + const float upper_bound, + const OpReqType req) { + for (size_t index = 0; index < kernel_params.ntensors; ++index) { + if ((size_t)i < kernel_params.sizes[index]) { + MPDType w = has_mixed_precision ? kernel_params.weights32[index][i]: + MPDType(kernel_params.weights[index][i]); + float r1 = sqrt(sumSqWeigths[index]); + float r2 = sqrt(sumSqtemp_g[index]); + r1 = std::min(std::max(r1, lower_bound), upper_bound); + + // calculate lamb_trust_ratio + MPDType r; + if (r1 == 0.0f || r2 == 0.0f) + r = 1.0f; + else + r = r1/r2; + + MPDType lr_adjusted = learning_rate * r; + w -= lr_adjusted * kernel_params.temp_g[index][i]; + + // update weights + if (has_mixed_precision) + kernel_params.weights32[index][i] = w; + KERNEL_ASSIGN(kernel_params.out_data[index][i], req, w); + } + } + } +}; + +template<typename MPDType, typename DType> +void call_kernel1(Stream<cpu>* s, + const MultiLAMBKernelParam<DType, MPDType>& kernel_params, + const MultiLAMBParam ¶m, + int* block_to_tensor, + int* block_to_chunk) { + Kernel<MultiLAMB_step1_kernel<MPDType, !std::is_same<DType, MPDType>::value>, cpu>:: + Launch(s, kernel_params.max_size, + kernel_params, + param.learning_rate, + param.beta1, param.beta2, + param.epsilon, + param.wd, + param.clip_gradient, + param.bias_correction, + param.rescale_grad); +} + +template<typename MPDType, typename DType> +void call_kernel2(Stream<cpu>* s, + const MultiLAMBKernelParam<DType, MPDType>& kernel_params, + const MultiLAMBParam ¶m, + float* r1, float* r2, + int* block_to_tensor, + int* block_to_chunk, + const OpReqType req) { + Kernel<MultiLAMB_step2_kernel<MPDType, !std::is_same<DType, MPDType>::value>, cpu>:: + Launch(s, kernel_params.max_size, + kernel_params, + r1, r2, + param.learning_rate, + param.lower_bound, param.upper_bound, + req); +} + +DMLC_REGISTER_PARAMETER(MultiLAMBParam); + +std::vector<std::string> LAMBParamToVector(uint32_t num_args, const char *pName[], size_t nParams) { + std::vector<std::string> ret; + for (uint32_t i = 0; i < num_args; ++i) { + const auto idx = std::to_string(i); + for (size_t j = 0; j < nParams; ++j) + ret.push_back(std::string(pName[i]) + idx); + } + + return ret; +} + +inline uint32_t num_tensors(const nnvm::NodeAttrs& attrs) { + return static_cast<uint32_t>(dmlc::get<MultiLAMBParam>(attrs.parsed).num_tensors); +} + +NNVM_REGISTER_OP(_multi_lamb_update) +.describe(R"code(Compute the LAMB coefficients of multiple weights and grads" +)code" ADD_FILELINE) +.set_num_inputs([](const nnvm::NodeAttrs& attrs) { + return num_tensors(attrs) * 5; + }) +.set_num_outputs([](const nnvm::NodeAttrs& attrs) { + return num_tensors(attrs); + }) +.set_attr_parser(ParamParser<MultiLAMBParam>) +.set_attr<mxnet::FInferShape>("FInferShape", MultiLAMB_InferShape<MultiLAMBParam, 5>) +.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<-1, -1>) +.set_attr<nnvm::FListInputNames>("FListInputNames", + [](const NodeAttrs& attrs) { + const char *paramName[] = {"weight_", "grad_", "mean_", "var_", "temp_g"}; + return LAMBParamToVector(num_tensors(attrs), paramName, sizeof(paramName)/sizeof(paramName[0])); + }) +// mutable: mean, var, temp_g, +.set_attr<nnvm::FMutateInputs>("FMutateInputs", + [](const nnvm::NodeAttrs& attrs) { + std::vector<uint32_t> ret; + const auto iMax = num_tensors(attrs); + for (size_t i = 0; i < iMax; ++i) { + ret.push_back(i * 5 + 2); + ret.push_back(i * 5 + 3); + ret.push_back(i * 5 + 4); + } + return ret; + }) +.set_attr<FResourceRequest>("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; + }) +.set_attr<FCompute>("FCompute<cpu>", multiLAMBUpdate<cpu, false>) +.add_argument("data", "NDArray-or-Symbol[]", "data") +.add_arguments(MultiLAMBParam::__FIELDS__()); + + +NNVM_REGISTER_OP(_multi_mp_lamb_update) +.describe(R"code(Compute the LAMB coefficients of multiple weights and grads with Mix Precision" +)code" ADD_FILELINE) +.set_num_inputs([](const nnvm::NodeAttrs& attrs) { + return num_tensors(attrs) * 6; + }) +.set_num_outputs([](const nnvm::NodeAttrs& attrs) { + return num_tensors(attrs); + }) +.set_attr_parser(ParamParser<MultiLAMBParam>) +.set_attr<mxnet::FInferShape>("FInferShape", MultiLAMB_InferShape<MultiLAMBParam, 6>) +.set_attr<nnvm::FInferType>("FInferType", MP_MultiLAMB_InferType<MultiLAMBParam, 6>) +.set_attr<nnvm::FListInputNames>("FListInputNames", + [](const NodeAttrs& attrs) { + const char *paramName[] = {"weight_", "grad_", "mean_", "var_", "temp_g", "weight32_"}; + return LAMBParamToVector(num_tensors(attrs), paramName, sizeof(paramName)/sizeof(paramName[0])); + }) +// mutable: mean, var, temp_g, weights32 +.set_attr<nnvm::FMutateInputs>("FMutateInputs", + [](const nnvm::NodeAttrs& attrs) { + std::vector<uint32_t> ret; + const auto iMax = num_tensors(attrs); + for (size_t i = 0; i < iMax; ++i) { + ret.push_back(i * 6 + 2); + ret.push_back(i * 6 + 3); + ret.push_back(i * 6 + 4); + ret.push_back(i * 6 + 5); + } + return ret; + }) +.set_attr<FResourceRequest>("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; + }) +.set_attr<FCompute>("FCompute<cpu>", multiLAMBUpdate<cpu, true>) +.add_argument("data", "NDArray-or-Symbol[]", "data") +.add_arguments(MultiLAMBParam::__FIELDS__()); + +} // namespace op +} // namespace mxnet
diff --git a/src/operator/contrib/multi_lamb.cu b/src/operator/contrib/multi_lamb.cu new file mode 100644 index 0000000..ed357d8 --- /dev/null +++ b/src/operator/contrib/multi_lamb.cu
@@ -0,0 +1,254 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * Copyright (c) 2019 by Contributors + * \file multi_lamb.cu + * \brief vectorized lamb coefficient computed from sums of squared weights and grads + * \author Moises Hernandez + */ + +#include "./multi_lamb-inl.h" + +namespace mxnet { +namespace op { + +#define BLOCK_SIZE_LAMB 512 +#define ILP_LAMB 4 + +template<bool has_mixed_precision, typename MPDType, typename DType> +__global__ void kernel_step1(const MultiLAMBKernelParam<DType, MPDType> kernel_params, + const float learning_rate, + const float beta1, const float beta2, + const MPDType beta3, const MPDType beta4, + const float epsilon, + const float wd, + const float clip_gradient, + const bool bias_correction, + const float rescale_grad, + int* block_to_tensor, + int* block_to_chunk) { + const int tensorID = block_to_tensor[blockIdx.x]; + const int chunckID = block_to_chunk[blockIdx.x]; + const int startPos = chunckID * kernel_params.chunk_size + threadIdx.x; + const int stopPos = chunckID * kernel_params.chunk_size + kernel_params.chunk_size; + + MPDType biascorrection1, biascorrection2; + if (bias_correction) { + biascorrection1 = 1.0f - std::pow(beta1, kernel_params.step_count[tensorID]); + biascorrection2 = 1.0f - std::pow(beta2, kernel_params.step_count[tensorID]); + } else { + biascorrection1 = 1.0f; + biascorrection2 = 1.0f; + } + + MPDType r_weight[ILP_LAMB]; + MPDType r_grad[ILP_LAMB]; + MPDType r_mean[ILP_LAMB]; + MPDType r_var[ILP_LAMB]; + MPDType r_g[ILP_LAMB]; + + for (size_t i=startPos; i < stopPos && i < kernel_params.sizes[tensorID]; + i+= blockDim.x*ILP_LAMB) { +#pragma unroll + for (int ii = 0; ii < ILP_LAMB; ii++) { + int load_pos = i + ii*blockDim.x; + if (load_pos < stopPos && load_pos < kernel_params.sizes[tensorID]) { + r_weight[ii] = has_mixed_precision ? kernel_params.weights32[tensorID][load_pos]: + static_cast<MPDType>(kernel_params.weights[tensorID][load_pos]); + r_grad[ii] = static_cast<MPDType>(kernel_params.grads[tensorID][load_pos]); + r_mean[ii] = kernel_params.mean[tensorID][load_pos]; + r_var[ii] = kernel_params.var[tensorID][load_pos]; + } else { + r_weight[ii] = static_cast<MPDType>(0); + r_grad[ii] = static_cast<MPDType>(0); + r_mean[ii] = static_cast<MPDType>(0); + r_var[ii] = static_cast<MPDType>(0); + } + } +#pragma unroll + for (int ii = 0; ii < ILP_LAMB; ii++) { + r_grad[ii] = r_grad[ii] * rescale_grad; + if (clip_gradient >= 0.0f) + r_grad[ii] = max(min(r_grad[ii], clip_gradient), -clip_gradient); + r_mean[ii] = static_cast<MPDType>(beta1) * r_mean[ii] + beta3 * r_grad[ii]; + r_var[ii] = static_cast<MPDType>(beta2) * r_var[ii] + beta4 * r_grad[ii] * r_grad[ii]; + r_g[ii] = (r_mean[ii] / biascorrection1) / (sqrtf(r_var[ii] / biascorrection2) + epsilon) + + wd * r_weight[ii]; + } +#pragma unroll + for (int ii = 0; ii < ILP_LAMB; ii++) { + int store_pos = i + ii*blockDim.x; + if (store_pos < stopPos && store_pos < kernel_params.sizes[tensorID]) { + kernel_params.mean[tensorID][store_pos] = r_mean[ii]; + kernel_params.var[tensorID][store_pos] = r_var[ii]; + kernel_params.temp_g[tensorID][store_pos] = r_g[ii]; + } + } + } +} + +template<bool has_mixed_precision, typename MPDType, typename DType> +__global__ void kernel_step2(const MultiLAMBKernelParam<DType, MPDType> kernel_params, + const float* sumSqWeigths, + const float* sumSqtemp_g, + const float learning_rate, + const float lower_bound, + const float upper_bound, + int* block_to_tensor, + int* block_to_chunk, + const OpReqType req) { + const int tensorID = block_to_tensor[blockIdx.x]; + const int chunckID = block_to_chunk[blockIdx.x]; + const int startPos = chunckID * kernel_params.chunk_size + threadIdx.x; + const int stopPos = chunckID * kernel_params.chunk_size + kernel_params.chunk_size; + + MPDType r1 = sqrtf(sumSqWeigths[tensorID]); + MPDType r2 = sqrtf(sumSqtemp_g[tensorID]); + r1 = min(max(r1, lower_bound), upper_bound); + + MPDType lr_adjusted; + if (r1 == 0.0f || r2 == 0.0f) + lr_adjusted = learning_rate; + else + lr_adjusted = learning_rate * r1/r2; + + MPDType r_weight[ILP_LAMB]; + MPDType r_g[ILP_LAMB]; + + for (size_t i=startPos; i < stopPos && i < kernel_params.sizes[tensorID]; + i+= blockDim.x*ILP_LAMB) { +#pragma unroll + for (int ii = 0; ii < ILP_LAMB; ii++) { + int load_pos = i + ii*blockDim.x; + if (load_pos < stopPos&& load_pos < kernel_params.sizes[tensorID]) { + r_weight[ii] = has_mixed_precision ? kernel_params.weights32[tensorID][load_pos]: + static_cast<MPDType>(kernel_params.weights[tensorID][load_pos]); + r_g[ii] = kernel_params.temp_g[tensorID][load_pos]; + } + } +#pragma unroll + for (int ii = 0; ii < ILP_LAMB; ii++) { + r_weight[ii] -= lr_adjusted * r_g[ii]; + } +#pragma unroll + for (int ii = 0; ii < ILP_LAMB; ii++) { + int store_pos = i + ii*blockDim.x; + if (store_pos < stopPos && store_pos < kernel_params.sizes[tensorID]) { + if (has_mixed_precision) + kernel_params.weights32[tensorID][store_pos] = r_weight[ii]; + KERNEL_ASSIGN(kernel_params.out_data[tensorID][store_pos], req, r_weight[ii]); + } + } + } +} + +template<typename MPDType, typename DType> +void call_kernel1(Stream<gpu>* s, + const MultiLAMBKernelParam<DType, MPDType>& kernel_params, + const MultiLAMBParam ¶m, + int* block_to_tensor, + int* block_to_chunk) { + int nblocks = kernel_params.nchunks; + int* host_block2tensor = reinterpret_cast<int*>(malloc(kernel_params.nchunks*sizeof(int))); + int* host_block2chunk = reinterpret_cast<int*>(malloc(kernel_params.nchunks*sizeof(int))); + int chunkID = 0; + for (size_t index = 0; index < kernel_params.ntensors; ++index) { + int current_chunk = 0; + for (size_t j = 0; j < kernel_params.sizes[index]; j+=kernel_params.chunk_size) { + host_block2tensor[chunkID] = index; + host_block2chunk[chunkID] = current_chunk; + current_chunk++; + chunkID++; + } + } + cudaMemcpyAsync(block_to_tensor, host_block2tensor, kernel_params.nchunks*sizeof(int), + cudaMemcpyHostToDevice, Stream<gpu>::GetStream(s)); + cudaMemcpyAsync(block_to_chunk, host_block2chunk, kernel_params.nchunks*sizeof(int), + cudaMemcpyHostToDevice, Stream<gpu>::GetStream(s)); + + bool has_mixed_precision = !std::is_same<DType, MPDType>::value; + MPDType beta3 = 1.0 - param.beta1; + MPDType beta4 = 1.0 - param.beta2; + + if (has_mixed_precision) + kernel_step1<true><<<nblocks, BLOCK_SIZE_LAMB, 0, Stream<gpu>::GetStream(s)>>>( + kernel_params, + param.learning_rate, + param.beta1, param.beta2, + beta3, beta4, + param.epsilon, param.wd, + param.clip_gradient, + param.bias_correction, + param.rescale_grad, + block_to_tensor, + block_to_chunk); + else + kernel_step1<false><<<nblocks, BLOCK_SIZE_LAMB, 0, Stream<gpu>::GetStream(s)>>>( + kernel_params, + param.learning_rate, + param.beta1, param.beta2, + beta3, beta4, + param.epsilon, param.wd, + param.clip_gradient, + param.bias_correction, + param.rescale_grad, + block_to_tensor, + block_to_chunk); + } + +template<typename MPDType, typename DType> +void call_kernel2(Stream<gpu>* s, + const MultiLAMBKernelParam<DType, MPDType>& kernel_params, + const MultiLAMBParam ¶m, + float* r1, float* r2, + int* block_to_tensor, + int* block_to_chunk, + const OpReqType req) { + size_t nblocks = kernel_params.nchunks; + bool has_mixed_precision = !std::is_same<DType, MPDType>::value; + if (has_mixed_precision) + kernel_step2<true><<<nblocks, BLOCK_SIZE_LAMB, 0, Stream<gpu>::GetStream(s)>>>( + kernel_params, + r1, r2, + param.learning_rate, + param.lower_bound, param.upper_bound, + block_to_tensor, + block_to_chunk, + req); + else + kernel_step2<false><<<nblocks, BLOCK_SIZE_LAMB, 0, Stream<gpu>::GetStream(s)>>>( + kernel_params, + r1, r2, + param.learning_rate, + param.lower_bound, param.upper_bound, + block_to_tensor, + block_to_chunk, + req); +} + + +NNVM_REGISTER_OP(_multi_lamb_update) +.set_attr<FCompute>("FCompute<gpu>", multiLAMBUpdate<gpu, false>); + +NNVM_REGISTER_OP(_multi_mp_lamb_update) +.set_attr<FCompute>("FCompute<gpu>", multiLAMBUpdate<gpu, true>); + +} // namespace op +} // namespace mxnet
diff --git a/tests/python/unittest/test_optimizer.py b/tests/python/unittest/test_optimizer.py index 8f86bdf..44487d9 100644 --- a/tests/python/unittest/test_optimizer.py +++ b/tests/python/unittest/test_optimizer.py
@@ -775,6 +775,104 @@ dtype, w_stype='default', g_stype='row_sparse', rtol=1e-4, atol=2e-5) +# MultiLAMB +class PyMultiLAMB(mx.optimizer.Optimizer): + """python reference implemenation of lamb""" + def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-6, + lower_bound=1e-3, upper_bound=10.0, bias_correction=False, + multi_precision=False, clip_gradient=-1, **kwargs): + super(PyMultiLAMB, self).__init__(learning_rate=learning_rate, **kwargs) + self.beta1 = beta1 + self.beta2 = beta2 + self.epsilon = epsilon + self.lower_bound = lower_bound + self.upper_bound = upper_bound + self.bias_correction = bias_correction + self.multi_precision = multi_precision + self.clip_gradient = clip_gradient + + def create_state(self, index, weight): + return (mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype), + mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype), + mx.nd.zeros(weight.shape, weight.context, dtype=weight.dtype)) + + def update(self, index, weight, grad, state): + self._update_count(index) + lr = self._get_lr(index) + wd = self._get_wd(index) + index_update_count = self._index_update_count[index] + + grad *= self.rescale_grad + if self.clip_gradient >= 0: + grad = mx.nd.clip(grad, -self.clip_gradient, self.clip_gradient) + + mean, var, temp_g = state + mean[:] = self.beta1 * mean + (1. - self.beta1) * grad.astype(mean.dtype) + var[:] = self.beta2 * var + (1. - self.beta2) * mx.nd.square(grad.astype(mean.dtype)) + + r1 = weight.norm() + if self.lower_bound: + r1 = mx.nd.maximum(r1, self.lower_bound) + if self.upper_bound: + r1 = mx.nd.minimum(r1, self.upper_bound) + + if self.bias_correction: + mean_hat = mean / (1. - mx.nd.power(self.beta1, index_update_count)) + var_hat = var / (1. - mx.nd.power(self.beta2, index_update_count)) + else: + mean_hat = mean + var_hat = var + + temp_g[:] = mean_hat / (mx.nd.sqrt(var_hat) + self.epsilon) + wd * weight + r2 = temp_g.norm() + # calculate lamb_trust_ratio + r = 1. if r1 == 0. or r2 == 0. else r1 / r2 + lr *= r + weight[:] -= lr * temp_g + +@with_seed() +def test_multilamb(): + opt1 = PyMultiLAMB + opt2 = mx.optimizer.MultiLAMB + #set_default_context(mx.gpu(0)) + + # shapes as Bert-large + dims_x = [1024, 4096, 1024, 1024] + dims_y = [1, 1, 1024, 4096] + dims_occurrences = [9, 1, 4, 2] + nlayers = 4 # 24 + extra_dims_x=[30522, 512, 30522] + extra_dims_y=[1, 1024, 1024] + shapes=[] + for l in range(nlayers): + for i, (dx,dy) in enumerate(zip(dims_x, dims_y)): + for j in range(dims_occurrences[i]): + shapes.append((dx,dy)) + for dx,dy in zip(extra_dims_x, extra_dims_y): + shapes.append((dx,dy)) + + cg_options = [{}, {'clip_gradient': 0.4}, {'clip_gradient': 0.5}] + rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}] + wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}] + bias_options = [{'bias_correction': False}, {'bias_correction': True}] + + for dtype in [np.float16, np.float32, np.float64]: + for cg_option in cg_options: + for rg_option in rg_options: + for wd_option in wd_options: + for bias_option in bias_options: + kwarg = {} + kwarg.update(cg_option) + kwarg.update(rg_option) + kwarg.update(wd_option) + kwarg.update(bias_option) + if (dtype == np.float16): + kwarg.update({'multi_precision': True}) + atol = 1e-3 + rtol = 1e-6 + compare_optimizer(opt1(**kwarg), opt2(**kwarg), shapes, dtype, + rtol=rtol, atol=atol, ntensors=len(shapes)) + # AdaMax class PyAdamax(mx.optimizer.Optimizer): """The python reference of AdaMax optimizer.