Lighting wrapped optimizers do not allow return values

[thoglu]

This is both a bug report /feature request:

Any optimizer.step function might have a return value. However, NO return values are currently handled by the optimizer.LightningOptimizer class in the "step" function. Optimizers with return values in their step function therefore do not operate properly.

BugFix: Add return value to to the respective line, something like:

return_val=None
if make_optimizer_step:
      return_val=self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
....

return return_val

Also the __optimizer_step function must be able to pass on the return value

cc @Borda @rohitgr7

[github-actions]

Hi! thanks for your contribution!, great first issue!

[tchaton]

Dear @thoglu,

Thanks for reporting this bug :)
We will address it soon.

Best regards,
T.C

[ddrevicky]

Hi I wanted to take a look at this but have a couple of questions.

1. What would be the typical use case for returning a value from the step function? Basic torch optimizers return the result of computing the closure (if passed) e.g., Adam.

   • But in PL, for example in NativeMixedPrecisionPlugin closure is called before the optimizer.step and is not passed to step as arg (I think the reason is that torch.cuda.amp does not support passing closure to step). If one of the basic torch optimizers was used it would probably return None in this case.

2. Code-wise in PL it might be a little harder now to determine which value should be returned. Accelerator::optimizer_step has hooks for precision plugins for pre and post steps. In case of NativeMixedPrecisionPlugin the optimizer.step is actually called in the post_optimizer_step step while the Accelerator::run_optimizer_step is not executed.

[thoglu]

The use case (in my case, but I think more broadly) is for allowing to return optimizer-specific logging values for debug purposes. In Adam, for example, one could return first and second moments, and see how they evolve over time, as these are the central quantities of the optimizer. The logging can then happen in the main training loop function of PL.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!