Add LARS to SGD and Momentum Optimizers (#6811)

[moting9]

ADD LARS to SGD and Momentum(no nesterov) optimizers for dense parameters update
@typhoonzero
LARSforFluid_ResNet50.pptx
related #6811

[typhoonzero]

Can put if (use\_lcoal\_lr): in the same line of else.

[moting9]

"param = param - learning_rate * velocity. " is called under "else."
if “if (use_lcoal_lr)” put under the same line of else , "param = param - learning_rate * velocity. " needs to write twice.

original logic:
else: \
if (use_lcoal_lr): \
learning_rate *= local_gw_ratio * sqrt(sumsq(param))
/ (sqrt(sumsq(gradient))+ weight_decay * sqrt(sumsq(param))) \
param = param - learning_rate * velocity. \

[typhoonzero]

Sorry I mean use else if {...}

[moting9]

Thanks.
change the comment from
$$
velocity = mu * velocity + gradient \
if (use_nesterov): \
param = param - gradient * learning_rate + mu * velocity * learning_rate \
else: \
param = param - learning_rate * velocity. \
$$
to
$$
velocity = mu * velocity + gradient \
if (use_nesterov): \
param = param - gradient * learning_rate + mu * velocity * learning_rate \
else if (use_lcoal_lr): \
learning_rate *= local_gw_ratio * sqrt(sumsq(param))
/ (sqrt(sumsq(gradient))+ weight_decay * sqrt(sumsq(param))) \
param = param - learning_rate * velocity. \
else: \
param = param - learning_rate * velocity. \
$$

[typhoonzero]

Please add the reference of LARS in the operator's doc and a bit more details.

[typhoonzero]

(float) weight decay => (float) LARS weight decay

[typhoonzero]

Same as above.

[typhoonzero]

Assigning to tensor may cause Eigen to do evaluation here, which may cause performance issue.

See: https://bitbucket.org/eigen/eigen/src/default/unsupported/Eigen/CXX11/src/Tensor/README.md?fileviewer=file-view-default#markdown-header-controlling-when-expression-are-evaluated

[moting9]

Thanks for the careful review! if "auto" is used instead of "Tensor", we can not perform p_norm > 0 and g_norm > 0 check due to lazy evaluation, there will be compilation error.

[typhoonzero]

Can add a epsilon to avoid zero-division.

[moting9]

It's weird if I use "auto" to calculate "param_out" as below, some elements in c++ output will be "-nan" and can't match python testcase output. yet the g_norm/p_norm/local_lr output are the same in both python and c++.

  if (use_local_lr) {
    auto p_norm = p.square().sum().sqrt();
    auto g_norm = g.square().sum().sqrt();
    auto local_lr = lr[0] * local_gw_ratio * (p_norm + epsilon) /
                    (g_norm + weight_decay * p_norm + epsilon);
    o = p - local_lr * g;
    std::cout << o;
  } else {
    o = p - lr[0] * g;
  }

here is python output for param_out
[0.69548035 0.28497267 0.22662182 0.5497418 0.718321 0.4220689 0.9801117 0.68425083]

here is c++ output for param_out
[ 0.69548 0.286139 -2.00666e+21 0.551315 -4.41094e+34 0.423106 0.980764 0.68483]

if I set the matrix size of weight and grad very small such as (4,5), the test_sgd_op.py unit test can pass,

I'm afraid it's due to the expression is a little long and lead to such error. if I use lr[0] to replace local_lr, such param_out discrepancy won't occur.

Any suggestion or comments? Thanks.

[moting9]

if I fallback to use tensor to calculate local_lr as below, param_out in both c++ and python can match and test case can pass.
T local_lr = lr[0];
if (use_local_lr) {
Eigen::TensorFixedSize<T, Eigen::Sizes<>, Eigen::RowMajor,
Eigen::DenseIndex> p_norm = p.square().sum().sqrt();
Eigen::TensorFixedSize<T, Eigen::Sizes<>, Eigen::RowMajor,
Eigen::DenseIndex> g_norm = g.square().sum().sqrt();
local_lr = lr[0] * local_gw_ratio * (p_norm(0) + epsilon) /
(g_norm(0) + weight_decay * p_norm(0) + epsilon);
}
o = p - local_lr * g;
std::cout << o;

[moting9]

If I use "auto" to calculate local_lr an param_out as below, local_lr can match in c++ print and python unit test case.
But for "o = p - local_lr * g", some elements in "o" are -nan in c++ print. I add eval for local_lr such as "o = p - local_lr.eval() * g", some element in "o" are very strange negative values as mentioned in comments above.
whether or not I add epsilon, the python test case " TestSGDLARSOp" and "TestMomentumOp3" failed.
So I fall back to use Eigen tensor to calculate local_lr and removed epsilon. python unit test can pass. Any comments or suggestion? Thanks!

if (use_local_lr) {
auto p_norm = p.square().sum().sqrt();
auto g_norm = g.square().sum().sqrt();
auto local_lr = lr[0] * local_gw_ratio * (p_norm + epsilon) /
(g_norm + weight_decay * p_norm + epsilon);
o = p - local_lr * g;
std::cout << o;
} else {
o = p - lr[0] * g;
}

[typhoonzero]

I didn't see the corresponding CUDA implement for this feature. So it's more proper if you can implement this as a "CPU_ONLY_KERNEL" or add CUDA support in both sgd and momentum op.

[moting9]

will implement it as a CPU_ONLY_KERNEL firstly. Thanks.

[typhoonzero]

CI fails because the CUDA implement is missing.

[luotao1]

close due to #10374