fix: fix qadam NAN problem

[wangraying]

fix #485

[wangraying]

This pr is ready, please take a look.

[ganshaoduo]

I think maybe it's better to keep the 'weight_decay' and set it to 0 by default, for consistency with the original Adam optimizer.

[wangraying]

I think maybe it's better to keep the 'weight_decay' and set it to 0 by default, for consistency with the original Adam optimizer.

Thanks @ganshaoduo , wondering how do we use weight_decay when calculating momentum, the relative code is here. The original paper does not discuss the situation with weight_decay, so I am not sure whether this will still converge in theory. In fact, I did some simple tests, seems the performance is not as good (I did not tune the parameters much).

[ganshaoduo]

Well spotted! Actually, according to the current implementation, we only apply the 'weight_decay' in the warm-up phase but not in the compression phase. This change of grad is only counted if step_id < self.warmup_steps. However, when we calculate the momentum here, the 'weight_decay' has been ignored.

About the necessity of "weight_decay", I would say it does not have a direct impact on the convergence. All it does is to add a small value to the gradient, which won't hurt the theoretical analysis of Qadam in my view.

Therefore, the question here is whether we add the 'weight_decay' operation here such that we use the same weight decay during both warm-up and compression phases, OR we delete the weight_decay operation here such that we do not consider the weight decay at all. I would prefer the former.

[wangraying]

Thanks for your detailed explanations, Shaoduo. Yes, I added weight_decay using the first way you mentioned, but the performance drops on the synthetic benchmark. Since this pr is for fixing NAN problem for QAdam, how about we removing weight_decay temporarily here but adding an issue about this. We could add it back when it is totally prepared after more validations.

[wangraying]

The current implementation of weight_decay is added by me last time during refactoring. Our version #0 does not have this weight_decay. And we have not used it in our benchmarks and published results. Or maybe keeping current implementation unchanged is also a good idea?

Any opinions from other reviewers?

[wangraying]

After discussions offline, we decided to rollback the modifications for weight_decay.