[feat] make sure AdaScale works with Lightning trainer loop

## 🚀 Feature
as title

## Motivation

changing the LR schedule based on the effective step is not easy to do for the user, it would be nice to support it without user's change.

## Pitch

make sure adascale works without or with small user impacts

## Alternatives

not supporting it

## Additional context

surface from an internal discussion and use case