Subtle difference with Pytorch AdamW?

[kyleliang919]

transformers/src/transformers/optimization.py

Line 648 in e5fd865

denom = exp_avg_sq.sqrt().add_(group["eps"])

It does correction after epsilon is added, whereas pytorch：
https://pytorch.org/docs/stable/_modules/torch/optim/adamw.html#AdamW

step = _get_value(step_t)

bias_correction1 = 1 - beta1**step
bias_correction2 = 1 - beta2**step

step_size = lr / bias_correction1

bias_correction2_sqrt = bias_correction2**0.5

if amsgrad:
    # Maintains the maximum of all 2nd moment running avg. till now
    torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])

    # Use the max. for normalizing running avg. of gradient
    denom = (max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt).add_(eps)
else:
    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)

param.addcdiv_(exp_avg, denom, value=-step_size)

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[Rocketknight1]

cc @ArthurZucker - is there a reason for us to maintain our own AdamW? In future, should we just wrap Torch's optimizer for this class?

[ArthurZucker]

Yeah no reason anymore AFAIK!

[Rocketknight1]

Opened a PR at #36177 to do that!

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.